from config import Config
from util import Util
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException
from typing import Dict, Optional, List, Tuple
from collections import defaultdict
[docs]
class GCSToS3RenameOperator(BaseOperator):
"""
Transfer files from Google Cloud Storage to Amazon S3 with the ability to rename files.
This operator supports copying a single GCS file to multiple S3 destinations.
:param file_mappings: List of tuples mapping GCS object paths to S3 destination keys
e.g., [
("gs://my-bucket/path/to/file.txt", "s3://dest-bucket1/new/path1.txt"),
("gs://my-bucket/path/to/file.txt", "s3://dest-bucket2/new/path2.txt"),
("gs://other-bucket/data.csv", "s3://dest-bucket3/data.csv")
]
Note: GCS paths should include the bucket name with "gs://" prefix
:type file_mappings: List[Tuple[str, str]]
:param gcp_project: The GCP project ID
:type gcp_project: str
:param replace: Whether to replace existing files in S3
If True, will upload all files replacing existing ones in the destination bucket.
If False, will upload only files that don't exist in the destination bucket.
:type replace: bool
"""
template_fields = ["file_mappings"]
ui_color = "#E0FFFF" # Light Cyan
@apply_defaults
def __init__(
self,
file_mappings: List[Tuple[str, str]],
gcp_project: Optional[str] = None,
replace: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.file_mappings = file_mappings
self.gcp_project = gcp_project
self.replace = replace
def _ensure_s3_path(self, path: str) -> str:
"""
Ensure the path has the s3:// prefix.
:param path: S3 path
:return: Path with s3:// prefix
"""
if not path.startswith("s3://"):
return f"s3://{path}"
return path
def _parse_gcs_path(self, gcs_path: str) -> tuple:
"""
Parse a GCS path into bucket name and object path.
:param gcs_path: Full GCS path (gs://bucket-name/path/to/object)
:return: Tuple of (bucket_name, object_path)
"""
# Handle paths with or without gs:// prefix
if gcs_path.startswith("gs://"):
path_without_prefix = gcs_path[5:]
else:
path_without_prefix = gcs_path
# Split into bucket and object path
parts = path_without_prefix.split("/", 1)
bucket_name = parts[0]
object_path = parts[1] if len(parts) > 1 else ""
return bucket_name, object_path
[docs]
def execute(self, context: Dict) -> None:
try:
from google.cloud import storage
import tempfile
import awswrangler as wr
import boto3
# Initialize GCS client
credentials = Util.get_gcp_credential(
service_account_conn_id=Config.GCP_SERVICE_ACCOUNT,
scopes=None,
)
storage_client = storage.Client(
project=self.gcp_project, credentials=credentials
)
# Track successful transfers for reporting
successful_transfers = []
failed_transfers = []
# Create boto3 s3 client for checking if files exist
s3_client = boto3.client("s3")
# Group files by GCS bucket and object for efficiency
# This allows us to download each unique GCS file only once
bucket_grouped_files = defaultdict(lambda: defaultdict(list))
for gcs_full_path, dest_s3_key in self.file_mappings:
gcs_bucket, gcs_object_path = self._parse_gcs_path(gcs_full_path)
bucket_grouped_files[gcs_bucket][gcs_object_path].append(
(gcs_full_path, dest_s3_key)
)
# Process each bucket and its files
for gcs_bucket, object_destinations in bucket_grouped_files.items():
self.log.info(f"Processing files from GCS bucket {gcs_bucket}")
# Get the bucket once for all files
bucket = storage_client.bucket(gcs_bucket)
# For each unique object in this bucket
for gcs_object_path, destinations in object_destinations.items():
# Get the blob once for all destinations
blob = bucket.blob(gcs_object_path)
try:
# Check if the blob exists
blob.reload()
except Exception as e:
gcs_full_path = f"gs://{gcs_bucket}/{gcs_object_path}"
self.log.error(
f"File not found in GCS: {gcs_full_path} - {str(e)}"
)
for gcs_path, _ in destinations:
failed_transfers.append(
(gcs_path, f"File not found: {str(e)}")
)
continue
# Download the file once
with tempfile.NamedTemporaryFile() as temp_file:
self.log.info(
f"Downloading {gcs_object_path} from GCS bucket {gcs_bucket}"
)
blob.download_to_filename(temp_file.name)
# Upload to each S3 destination
for gcs_full_path, dest_s3_key in destinations:
try:
# Ensure proper S3 path format
dest_s3_key = self._ensure_s3_path(dest_s3_key)
# Parse S3 bucket and key from the destination path
s3_path_parts = dest_s3_key.replace("s3://", "").split(
"/", 1
)
s3_bucket = s3_path_parts[0]
s3_key = (
s3_path_parts[1] if len(s3_path_parts) > 1 else ""
)
# Check if file exists in S3 and skip if not replacing
if not self.replace:
try:
s3_client.head_object(
Bucket=s3_bucket, Key=s3_key
)
self.log.info(
f"File already exists in S3 at {dest_s3_key} and replace=False, skipping."
)
continue
except Exception:
# File doesn't exist, proceed with upload
pass
# Upload to S3
self.log.info(f"Uploading to S3 at {dest_s3_key}")
wr.s3.upload(
local_file=temp_file.name, path=dest_s3_key
)
successful_transfers.append(
(gcs_full_path, dest_s3_key)
)
self.log.info(
f"✓ Transferred {gcs_full_path} to S3 at {dest_s3_key}"
)
except Exception as e:
failed_transfers.append((gcs_full_path, str(e)))
self.log.error(
f"Failed to transfer {gcs_full_path} to {dest_s3_key}: {str(e)}"
)
# Continue with other destinations
continue
# Report summary
total_mappings = len(self.file_mappings)
unique_sources = sum(
len(objects) for objects in bucket_grouped_files.values()
)
self.log.info(
f"Summary: Transferred {len(successful_transfers)} of {total_mappings} file mappings "
f"({unique_sources} unique source files from {len(bucket_grouped_files)} GCS buckets)"
)
if failed_transfers:
self.log.warning(f"Failed to transfer {len(failed_transfers)} files")
for name, error in failed_transfers:
self.log.warning(f" - {name}: {error}")
if not successful_transfers and self.file_mappings:
raise AirflowException("No files were successfully transferred")
except Exception as ex:
self.log.error(f"Error details: {str(ex)}")
raise AirflowException(f"Failed to transfer files from GCS to S3") from ex