"""
Note that this is part of a refactor of the flat file exporter.
This first step strictly copies the functionality of the existing dag.
It will be adapted to use ECS instead of local processing as another phase
"""
from config import Config
from util import Util
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook
[docs]
class QueryToGDriveCSVOperator(BaseOperator):
""" """
ui_color = "#420a69"
@apply_defaults
def __init__(
self,
source_postgres_db: str,
source_datalake_db: str,
source_sql: str,
custom_columns: dict,
target_file_name: str,
gdrive_parent_folder_id: str,
pg_flag: bool = False,
gdrive_overwrite_existing: bool = False,
*args,
**kwargs,
):
super(QueryToGDriveCSVOperator, self).__init__(*args, **kwargs)
self.pg_flag = pg_flag
self.source_postgres_db = source_postgres_db
self.source_datalake_db = source_datalake_db
self.source_sql = source_sql
self.custom_columns = custom_columns
self.target_file_name = target_file_name
self.gdrive_parent_folder_id = gdrive_parent_folder_id
self.gdrive_overwrite_existing = gdrive_overwrite_existing
[docs]
def execute(self, context):
try:
import awswrangler as wr
from tempfile import NamedTemporaryFile
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
# Create gdrive service
credentials = Util.get_gcp_credential()
gdrive_service = build(
"drive",
"v3",
credentials=credentials,
)
# Ensure correct file name handling
filename = self.target_file_name
print(filename)
print(self.source_sql)
# Select execution connection
if self.pg_flag == True:
src_hook = PostgresHook(postgres_conn_id=self.source_postgres_db)
df_source = src_hook.get_pandas_df(self.source_sql)
else:
df_source = wr.athena.read_sql_query(
sql=self.source_sql,
database=self.source_datalake_db,
)
# Rename columns if custom_columns is provided
if self.custom_columns:
df_source = df_source.rename(
columns=self.custom_columns, errors="raise", inplace=True
)
# Write the DataFrame to a temporary CSV file
with NamedTemporaryFile(suffix=".csv") as csv_fh:
df_source.to_csv(csv_fh.name, index=False, header=True)
tmp_file_name = csv_fh.name
print(
f"""
- Now uploading file: {filename} to Google Drive.
- Parent folder ID: {self.gdrive_parent_folder_id}
- Overwrite: {self.gdrive_overwrite_existing}
"""
)
media = MediaFileUpload(tmp_file_name, mimetype="text/csv")
file_metadata = {
"name": filename,
"parents": (
[self.gdrive_parent_folder_id]
if self.gdrive_parent_folder_id
else []
),
}
#########
### Overwrite Case
#########
if self.gdrive_overwrite_existing:
items = (
gdrive_service.files()
.list(
q=(
f"'{self.gdrive_parent_folder_id}' in parents and trashed=false"
if self.gdrive_parent_folder_id
else None
),
pageSize=1000, # default is 100
fields="nextPageToken, files(id, name, mimeType)",
supportsAllDrives=True,
corpora="allDrives",
includeItemsFromAllDrives=True,
)
.execute()
).get("files", [])
names_existing = [item["name"] for item in items]
if filename in names_existing:
print(f"File {filename} already exists in the parent folder")
file_id = items[names_existing.index(filename)]["id"]
# STEP 2. Update the existing file
update_file = (
gdrive_service.files()
.update(
fileId=file_id,
body={"name": filename},
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(
"File with same name exists. Updated file for the ID: ",
update_file["id"],
)
print(
f'https://drive.google.com/file/d/{update_file["id"]}/view'
)
return update_file["id"]
else:
print(
"File with same name does not exist. Creating a new file."
)
created_file = (
gdrive_service.files()
.create(
body=file_metadata,
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(f'Created a new File ID: {created_file["id"]}')
print(
f'https://drive.google.com/file/d/{created_file["id"]}/view'
)
return created_file["id"]
#########
### New File Case
#########
else:
print(
"Developer does not want to overwrite the file. Creating a new file."
)
created_file = (
gdrive_service.files()
.create(
body=file_metadata,
media_body=media,
fields="id",
supportsAllDrives=True,
)
.execute()
)
print(f'Created a new File ID: {created_file["id"]}')
return created_file["id"]
except Exception as e:
print(e)
raise e