"""Operator that transfers files from Google Drive (exporting Docs-editor files) to an S3 bucket."""
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from util import Util
from io import BytesIO
from config import Config
DEFAULT_S3_BUCKET = Config.DATA_LAKE_BUCKET
# Map Google Docs Editor types to (export MIME type, file extension)
GOOGLE_EXPORT_MAP = {
"application/vnd.google-apps.document": ("application/pdf", ".pdf"),
"application/vnd.google-apps.spreadsheet": (
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
".xlsx",
),
"application/vnd.google-apps.presentation": ("application/pdf", ".pdf"),
}
[docs]
class GoogleDriveToS3Operator(BaseOperator):
"""
Operator to transfer files from Google Drive to an S3 bucket.
:param google_drive_id: The Google Drive folder ID to transfer.
:type google_drive_id: str
:param s3_key: The S3 key (path) where the file will be stored.
:type s3_key: str
:param s3_bucket_name: The S3 bucket name where the file will be stored.
:type s3_bucket_name: str
:param remove_spaces_from_filename: If True, spaces in filenames will be replaced with underscores.
:type remove_spaces_from_filename: bool
:param overwrite_file: If True, existing files in S3 will be overwritten.
:type overwrite_file: bool
:param recursive: If True, recursively process subfolders when a folder ID is provided.
:type recursive: bool
"""
template_fields = (
"google_drive_file_id",
"s3_bucket_name",
"s3_key",
"overwrite_file",
"recursive",
)
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
google_drive_id: str,
s3_key: str,
s3_bucket_name: str = DEFAULT_S3_BUCKET,
remove_spaces_from_filename: bool = False,
overwrite_file: bool = True,
skip_spreadsheets: bool = True,
recursive: bool = True,
*args,
**kwargs,
):
super(GoogleDriveToS3Operator, self).__init__(*args, **kwargs)
self.google_drive_file_id = google_drive_id
self.s3_bucket_name = s3_bucket_name
self.s3_key = s3_key
self.remove_spaces_from_filename = remove_spaces_from_filename # this is only relevant when transferring many files
# s3 doesn't lke spaces when it's serving files back to the browser
self.overwrite_file = overwrite_file
self.skip_spreadsheets = skip_spreadsheets
self.recursive = recursive
self.downloaded_files = []
[docs]
def execute(self, context):
from googleapiclient.http import MediaIoBaseDownload
import boto3
s3_client = boto3.client("s3")
service = Util.construct_gdrive_service()
print("DriveToS3: Initialized services")
s3_files = Util.get_existing_s3_files(
bucket=self.s3_bucket_name,
prefix=self.s3_key,
file_type="",
full_key=True,
)
if Util.is_drive_folder(
service,
self.google_drive_file_id,
):
print("DriveToS3: Detected folder")
# Download all files in the folder - memory risk if it gets massive tho
file_list = Util.get_all_gdrive_files(
service, self.google_drive_file_id, recursive=self.recursive
)
for each_file in file_list:
print(f"DriveToS3: Handling {each_file['name']}")
file_name = (
each_file["name"]
if not self.remove_spaces_from_filename
else each_file["name"].replace(" ", "_")
)
print(each_file)
if (
each_file["mimeType"] == "application/vnd.google-apps.spreadsheet"
and self.skip_spreadsheets == True
):
print("DriveToS3: Skipping spreadsheet")
continue
target_s3_key = (
f"{self.s3_key}/{file_name}" if self.s3_key else file_name
)
if target_s3_key in s3_files:
warning = f"File already exists at {self.s3_bucket_name}/{target_s3_key}"
if self.overwrite_file:
print(f"{warning}. Overwriting file.")
else:
print(
f"{warning}. Skipping file as overwrite flag is {self.overwrite_file}"
)
continue
self.download_file(
service,
s3_client,
MediaIoBaseDownload,
google_drive_file_id=each_file["id"],
s3_key=target_s3_key,
)
else:
file = (
service.files()
.get(
fileId=self.google_drive_file_id,
supportsAllDrives=True,
supportsTeamDrives=True,
)
.execute()
)
file_name = file.get("name")
if self.remove_spaces_from_filename:
file_name = file_name.replace(" ", "_")
target_s3_key = f"{self.s3_key}/{file_name}" if self.s3_key else file_name
print(f"{target_s3_key=}")
if target_s3_key in s3_files:
warning = f"File already exists at {self.s3_bucket_name}/{target_s3_key}"
if self.overwrite_file:
print(f"{warning}. Overwriting file.")
self.download_file(
service,
s3_client,
MediaIoBaseDownload,
s3_key=target_s3_key,
)
else:
print(
f"{warning}. Skipping file as overwrite flag is {self.overwrite_file}"
)
else:
self.download_file(
service,
s3_client,
MediaIoBaseDownload,
s3_key=target_s3_key,
)
return self.downloaded_files
[docs]
def download_file(
self,
service,
s3,
downloader_class,
google_drive_file_id=None,
s3_key=None,
):
target_file_id = (
self.google_drive_file_id
if not google_drive_file_id
else google_drive_file_id
)
print(f"DriveToS3: Downloading {target_file_id}")
target_s3_key = self.s3_key if not s3_key else s3_key
try:
mime_type = Util.get_mimetype(
service,
target_file_id,
)
if mime_type in GOOGLE_EXPORT_MAP:
export_mime, extension = GOOGLE_EXPORT_MAP[mime_type]
request = service.files().export_media(
fileId=target_file_id, mimeType=export_mime
)
if not target_s3_key.endswith(extension):
target_s3_key = f"{target_s3_key}{extension}"
print(f"DriveToS3: Exporting Google file as {extension}")
else:
request = service.files().get_media(
fileId=target_file_id,
supportsTeamDrives=True,
supportsAllDrives=True,
)
file_data = BytesIO()
downloader = downloader_class(
file_data,
request,
)
# Download the file from Google Drive
done = False
while not done:
status, done = downloader.next_chunk()
self.log.info("Download %d%%." % int(status.progress() * 100))
print("Download %d%%." % int(status.progress() * 100))
file_data.seek(0)
# S3 authentication and upload
s3.put_object(Body=file_data, Bucket=self.s3_bucket_name, Key=target_s3_key)
print(
f"File {target_file_id} has been uploaded to {self.s3_bucket_name}/{target_s3_key}"
)
self.downloaded_files.append(
{
"bucket": self.s3_bucket_name,
"key": target_s3_key,
"google_drive_id": target_file_id,
"mime_type": mime_type,
}
)
except Exception as e:
print(f"Error downloading file: {e}")