import io
import zipstream
from typing import List, Dict
from airflow.models.baseoperator import BaseOperator
from airflow.utils.context import Context
from airflow.exceptions import AirflowException
COMPRESSED_EXT = {
".zip",
".gz",
".xz",
".bz2",
".7z",
".rar",
".jpg",
".jpeg",
".png",
".gif",
".pdf",
".mp4",
".mov",
".mkv",
".mp3",
".aac",
}
[docs]
class ZipS3Operator(BaseOperator):
"""
Stream-zips files grouped by `zip_group` (from Athena metadata) into ≤ N GB archives
and uploads the ZIPs back to S3 without calculating sizes or re-grouping.
:param files: list of dicts with keys:
{"bucket": str, "key": str, "arcname": str, "file_size_byte": int, "zip_group": int}
:param dest_bucket: bucket where ZIP(s) will be stored
:param dest_prefix: folder/prefix inside dest_bucket
:param email_recipients: list of e-mail addresses (only passed through XCom so the e-mail task can reuse)
:param max_zip_gb: nominal size of each ZIP in GiB (for naming only)
:param overwrite: if True, delete existing ZIPs under dest_prefix before uploading new ones
:param zip_name: archive file name to use instead of the default
``zip_part_{group}.zip``. When the operator holds more than one
group, ``_part{group}`` is appended to keep the keys unique.
"""
template_fields = (
"files",
"dest_bucket",
"dest_prefix",
"email_recipients",
"zip_name",
)
def __init__(
self,
*,
zip_group: int,
files: List[Dict[str, str]],
dest_bucket: str,
dest_prefix: str,
email_recipients: List[str],
max_zip_gb: int = 3,
overwrite: bool = False,
zip_name: str = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
if not files:
raise AirflowException("files list must not be empty")
self.files = files
self.dest_bucket = dest_bucket
self.dest_prefix = dest_prefix.rstrip("/")
self.email_recipients = email_recipients
self.max_zip_gb = max_zip_gb
self.overwrite = overwrite
self.zip_group = int(zip_group)
self.zip_name = zip_name
def _zip_key(self, group: int, single_group: bool) -> str:
"""Object key for one zip group. Defaults to the legacy zip_part_ naming."""
if not self.zip_name:
return f"{self.dest_prefix}/zip_part_{group}.zip"
stem = self.zip_name
if stem.lower().endswith(".zip"):
stem = stem[: -len(".zip")]
suffix = "" if single_group else f"_part{group}"
return f"{self.dest_prefix}/{stem}{suffix}.zip"
def _delete_existing_zips(self, s3):
"""Remove existing zip_part_*.zip under dest_prefix."""
paginator = s3.get_paginator("list_objects_v2")
to_delete = []
prefix = f"{self.dest_prefix}/zip_part_"
for page in paginator.paginate(Bucket=self.dest_bucket, Prefix=prefix):
for obj in page.get("Contents", []):
to_delete.append({"Key": obj["Key"]})
if to_delete:
for i in range(0, len(to_delete), 1000):
s3.delete_objects(
Bucket=self.dest_bucket,
Delete={"Objects": to_delete[i : i + 1000]},
)
self.log.info(
"Removed %d old ZIP(s) from %s/%s",
len(to_delete),
self.dest_bucket,
self.dest_prefix,
)
else:
self.log.info(
"No existing ZIPs to delete under %s/%s",
self.dest_bucket,
self.dest_prefix,
)
def _upload_group(self, s3, group, zip_key, part_size=8 * 1024 * 1024):
mp = s3.create_multipart_upload(Bucket=self.dest_bucket, Key=zip_key)
upload_id = mp["UploadId"]
parts, part_number = [], 1
buffer = bytearray()
def chunked_reader(bucket, key, chunk_size):
resp = s3.get_object(Bucket=bucket, Key=key)
body = resp["Body"]
try:
for chunk in body.iter_chunks(chunk_size=chunk_size):
yield chunk
finally:
body.close()
try:
zs = zipstream.ZipFile(mode="w", allowZip64=True)
for f in group:
ext = f["key"].rsplit(".", 1)[-1].lower()
compress = (
zipstream.ZIP_STORED
if ext in COMPRESSED_EXT
else zipstream.ZIP_DEFLATED
)
zs.write_iter(
f["arcname"],
chunked_reader(f["bucket"], f["key"], part_size),
compress_type=compress,
)
for chunk in zs:
buffer.extend(chunk)
while len(buffer) >= part_size:
part = bytes(buffer[:part_size])
buffer = buffer[part_size:]
resp = s3.upload_part(
Bucket=self.dest_bucket,
Key=zip_key,
PartNumber=part_number,
UploadId=upload_id,
Body=part,
)
parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
part_number += 1
if buffer:
resp = s3.upload_part(
Bucket=self.dest_bucket,
Key=zip_key,
PartNumber=part_number,
UploadId=upload_id,
Body=bytes(buffer),
)
parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
s3.complete_multipart_upload(
Bucket=self.dest_bucket,
Key=zip_key,
UploadId=upload_id,
MultipartUpload={"Parts": parts},
)
except:
s3.abort_multipart_upload(
Bucket=self.dest_bucket, Key=zip_key, UploadId=upload_id
)
raise
[docs]
def execute(self, context: Context):
import boto3
from botocore.client import Config as BotoConfig
from util import Util
# Initialize S3 client. Reads/writes stay on the task role; only the presigned
# URLs are signed by Util.get_presigned_url with the automation user, since
# task role credentials are temporary and cap a url's life at the session token
bootstrap = boto3.client("s3")
region = (
bootstrap.get_bucket_location(Bucket=self.dest_bucket)["LocationConstraint"]
or "us-east-1"
)
s3 = boto3.client(
"s3",
region_name=region,
config=BotoConfig(signature_version="s3v4"),
)
# Group files by zip_group field from metadata
groups = {}
for f in self.files:
grp = f.get("zip_group", 1)
groups.setdefault(grp, []).append(f)
# Prepare presigned URL expiry
raw_days = context.get("params", {}).get("expiration_in_days", 7)
try:
days_int = int(
raw_days[0] if isinstance(raw_days, (list, tuple)) else raw_days
)
except Exception:
days_int = 7
# 7 days is AWS's hard ceiling for SigV4 presigned URLs.
expiry_seconds = min(days_int, 7) * 24 * 3600
self.log.info("Presigned-URL expiry set to %s s", expiry_seconds)
presigned_urls = []
objects = []
single_group = len(groups) == 1
for grp, files in sorted(groups.items()):
zip_key = self._zip_key(grp, single_group)
self._upload_group(s3, files, zip_key)
url = Util.get_presigned_url(
bucket=self.dest_bucket,
key=zip_key,
expiration=expiry_seconds,
region_name=region,
)
presigned_urls.append(url)
try:
size = s3.head_object(Bucket=self.dest_bucket, Key=zip_key)[
"ContentLength"
]
except Exception as exc: # size is cosmetic; never fail the upload for it
self.log.warning("Could not size %s: %s", zip_key, exc)
size = None
objects.append(
{
"key": zip_key,
"name": zip_key.rsplit("/", 1)[-1],
"bytes": size,
"url": url,
"file_count": len(files),
}
)
return {
"download_urls": presigned_urls,
"recipients": self.email_recipients,
"objects": objects,
}