from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.hooks.base import BaseHook
from airflow.exceptions import AirflowException
from config import Config
[docs]
class GdriveZipToS3Operator(BaseOperator):
"""
This operator will zip files from a Google Drive folder and upload the zip file to an S3 bucket.
The current implementation will raise an exception if the total size of the files exceeds a threshold.
A future implementation will trigger an ECS execution to zip the files if the threshold is exceeded.
The default upload location is the BOSS private bucket.
:return: None
:rtype: None
:param output_zip_file_name: The name of the output zip file.
:type output_zip_file_name: str
:param source_gdrive_folder_url: The URL of the Google Drive folder to zip files from.
:type source_gdrive_folder_url: str
:param s3_key: The S3 key where the zip file will be uploaded.
:type s3_key: str
:param s3_bucket_name: The name of the S3 bucket where the zip file will be uploaded.
:type s3_bucket_name: str
:param remove_spaces_from_filename: Whether to remove spaces from the filenames in the zip file.
:type remove_spaces_from_filename: bool
:param little_zip_threshold_mb: The threshold in MB for the total size of files to be zipped locally.
:type little_zip_threshold_mb: int
"""
ui_color = "#A69696"
template_fields = [
"source_gdrive_folder_url",
"output_zip_file_name",
"s3_key",
"generate_presigned_url",
"presigned_url_expiry",
"s3_bucket_name",
"remove_spaces_from_filename",
]
@apply_defaults
def __init__(
self,
output_zip_file_name: str,
source_gdrive_folder_url: str,
s3_key: str = "",
s3_bucket_name: str = Config.BOSS_PRIVATE_BUCKET,
remove_spaces_from_filename: bool = False,
little_zip_threshold_mb: int = 500,
generate_presigned_url: bool = True,
presigned_url_expiry: int = 3600, # 1 hour
*args,
**kwargs,
):
super(GdriveZipToS3Operator, self).__init__(*args, **kwargs)
self.source_gdrive_folder_url = source_gdrive_folder_url
self.s3_bucket_name = s3_bucket_name
self.s3_key = s3_key
self.remove_spaces_from_filename = remove_spaces_from_filename
self.output_zip_file_name = output_zip_file_name
self.little_zip_threshold_mb = little_zip_threshold_mb
self.generate_presigned_url = generate_presigned_url
self.presigned_url_expiry = presigned_url_expiry
@staticmethod
def _extract_folder_id(url: str) -> str:
# Extract the folder ID from the Google Drive folder URL
import re
# handles URLs like
# https://drive.google.com/drive/folders/<ID>
# https://drive.google.com/open?id=<ID>
patterns = [r"/folders/([a-zA-Z0-9_-]+)", r"id=([a-zA-Z0-9_-]+)"]
for p in patterns:
m = re.search(p, url)
if m:
return m.group(1)
raise AirflowException(f"Could not parse folder ID from URL: {url}")
def _construct_gdrive_service(self):
from google.oauth2 import service_account
from googleapiclient.discovery import build
import json
# Define scope of access
scopes = [
"https://www.googleapis.com/auth/drive",
"https://www.googleapis.com/auth/spreadsheets",
]
service_acccount_conn_id = Config.GDRIVE_SERVICE_ACCOUNT
gdrive_hook = BaseHook.get_connection(service_acccount_conn_id)
gcp_json_credentials = json.loads(gdrive_hook.extra)
credentials = service_account.Credentials.from_service_account_info(
gcp_json_credentials, scopes=scopes
)
gdrive_service = build("drive", "v3", credentials=credentials)
return gdrive_service
[docs]
def execute(self, context):
import io
import zipfile
import boto3
from util import Util
source_folder_id = self._extract_folder_id(self.source_gdrive_folder_url)
self.log.info(f"Extracted folder ID: {source_folder_id}")
gdrive_service = self._construct_gdrive_service()
self.log.info("Constructed Google Drive service")
# List files in the source folder
results = (
gdrive_service.files()
.list(
q=(
f"'{source_folder_id}' in parents and trashed=false"
if source_folder_id
else None
),
pageSize=1000,
fields="nextPageToken, files(id, name, size, mimeType)",
supportsAllDrives=True,
corpora="allDrives",
includeItemsFromAllDrives=True,
)
.execute()
)
# check if the directory is null
files = results.get("files", [])
if not files:
raise AirflowException(
"No files found in the specified Google Drive folder."
)
# Calculate total size of source files
total_bytes = sum(int(file["size"]) for file in files if "size" in file)
total_mb = total_bytes / (1024 * 1024)
self.log.info(f"Total size of source files: {total_mb:.2f} MB")
# Check if the total size exceeds the threshold to be zipped locally
# For now, raise an exception if the total size exceeds the threshold
# This will later trigger an ecs execution to zip the files
if total_mb > self.little_zip_threshold_mb:
raise AirflowException(
f"Total size of source files ({total_mb:.2f} MB) exceeds the threshold of {self.little_zip_threshold_mb} MB. "
"Please use a different operator to zip the files."
)
# Create a zip file in memory
zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
for file in files:
file_id = file["id"]
file_name = file["name"]
# Remove spaces from filename if required
if self.remove_spaces_from_filename:
file_name = file_name.replace(" ", "_")
self.log.info(f"Downloading file: {file_name}")
# Download the file content
request = gdrive_service.files().get_media(fileId=file_id)
file_content = request.execute()
# Write the file content to the zip
zip_file.writestr(file_name, file_content)
# Format S3 key
if not self.s3_key or self.s3_key == "" or self.s3_key == "null":
output_key = self.output_zip_file_name
else:
output_key = f"{self.s3_key}/{self.output_zip_file_name}"
# Multipart upload to S3
s3_client = boto3.client("s3")
zip_buffer.seek(0)
try:
s3_client.upload_fileobj(
zip_buffer,
self.s3_bucket_name,
output_key,
ExtraArgs={"ContentType": "application/zip"},
)
self.log.info(f"Successfully uploaded zip file to S3: {output_key}")
except Exception as e:
raise AirflowException(f"Failed to upload zip file to S3: {str(e)}")
# Generate presigned URL if required. Signed by Util.get_presigned_url with the
# automation user rather than the task role, so the expiry actually holds; the
# upload above stays on the task role
if self.generate_presigned_url:
try:
presigned_url = Util.get_presigned_url(
bucket=self.s3_bucket_name,
key=output_key,
expiration=int(self.presigned_url_expiry),
)
self.log.info(f"Generated presigned URL: {presigned_url}")
return presigned_url
except Exception as e:
raise AirflowException(f"Failed to generate presigned URL: {str(e)}")
return None