"""This operator is used to transfer files from S3 to Google Drive."""
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from util import Util
# Above this size the Drive upload switches to the resumable protocol. A
# non-resumable upload sends the whole file as one request body — that means the
# entire file in worker memory and no recovery from a mid-flight failure, which
# multi-GiB ZIPs cannot afford. Small attachments stay non-resumable: one round
# trip beats the resumable handshake.
GDRIVE_RESUMABLE_THRESHOLD_BYTES = 8 * 1024**2
[docs]
class S3ToGoogleDriveOperator(BaseOperator):
"""
Operator to transfer files from S3 to a Google Drive folder.
:param s3_paths: List of S3 keys/paths for the source files.
:type s3_paths: list[str]
:param dest_folder_id: The destination Google Drive folder ID.
:type dest_folder_id: str
:param s3_bucket: The S3 bucket name where the files are stored.
:type s3_bucket: str
:param overwrite: Flag to indicate whether to overwrite existing files in Google Drive.
:type overwrite: bool
:param dest_file_name: The destination file name in Google Drive. If not provided, the original file name will be used.
:type dest_file_name: str
:param google_service_account_file: Path to the Google service account JSON file.
:type google_service_account_file: str
"""
template_fields = ("s3_paths", "dest_folder_id", "dest_file_name")
@apply_defaults
def __init__(
self,
s3_paths,
dest_folder_id,
s3_bucket,
overwrite,
dest_file_name=None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.s3_paths = s3_paths
self.dest_folder_id = dest_folder_id
self.s3_bucket = s3_bucket
self.overwrite = overwrite
self.dest_file_name = dest_file_name
[docs]
def execute(self, context):
from googleapiclient.http import MediaFileUpload
import boto3
import os
import tempfile
import mimetypes
# Initialize S3 client
s3_client = boto3.client("s3")
# Set up credentials and initialize the Google Drive API client
# Create gdrive service
gdrive_service = Util.construct_gdrive_service()
for s3_path in self.s3_paths:
# use provided dest_file_name or fall back to the file name in s3_path
original_file_name = os.path.basename(s3_path)
file_name = self.dest_file_name or original_file_name
self.log.info(
f"Downloading file {file_name} from s3://{self.s3_bucket}/{s3_path}"
)
# Create a temporary file to hold the downloaded file
with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
local_path = tmp_file.name
# Download file from S3
s3_client.download_file(
self.s3_bucket,
s3_path,
local_path,
)
# Guess the MIME type based on the file name
mime_type, _ = mimetypes.guess_type(file_name)
if mime_type is None:
mime_type = "application/octet-stream"
# Prepare metadata for the Google Drive file
file_metadata = {
"name": file_name,
"parents": [self.dest_folder_id],
}
# execute() drives the resumable protocol itself (it loops next_chunk
# internally), so the create/update calls below need no special casing.
local_bytes = os.path.getsize(local_path)
resumable = local_bytes > GDRIVE_RESUMABLE_THRESHOLD_BYTES
self.log.info(
"Uploading %s (%.2f MiB) to Drive, resumable=%s",
file_name,
local_bytes / 1024**2,
resumable,
)
media = MediaFileUpload(
local_path,
mimetype=mime_type,
resumable=resumable,
)
# Check if a file with the same name already exists in the destination folder
query = f"'{self.dest_folder_id}' in parents and name='{file_name}' and trashed=false"
response = (
gdrive_service.files()
.list(
q=query,
pageSize=1000, # default is 100. Doesn't work with >1000, will fail with returned "Invalid value '2000'. Values must be within the range: [value: 1, value: 1000]
fields="nextPageToken, files(id, name, mimeType)",
supportsAllDrives=True,
corpora="allDrives",
includeItemsFromAllDrives=True,
)
.execute()
)
existing_files = response.get("files", [])
if existing_files:
if self.overwrite:
# Overwrite by updating the first matching file
file_id = existing_files[0]["id"]
self.log.info(
f"File {file_name} exists. Overwriting it (updating file ID {file_id})."
)
gdrive_service.files().update(
fileId=file_id,
body={"name": file_name},
media_body=media,
fields="id",
supportsAllDrives=True,
).execute()
else:
self.log.info(
f"File {file_name} already exists in Google Drive. Skipping upload due to overwrite=False."
)
os.remove(local_path)
continue
else:
self.log.info(
f"Uploading file {file_name} to Google Drive folder {self.dest_folder_id}"
)
gdrive_service.files().create(
body=file_metadata,
media_body=media,
fields="id",
supportsAllDrives=True,
).execute()
# Remove the temporary file after upload/update
os.remove(local_path)
self.log.info(f"File {file_name} processed successfully.")
self.log.info("All files have been processed successfully.")
return True