Source code for operators.s3_unzip

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class UnzipOperator(BaseOperator):
    """
    Unzip files from S3 to S3.

    :param src_bucket: source S3 bucket name
    :type src_bucket: str
    :param src_key: source S3 object key of the zip file
    :type src_key: str
    :param dest_key: destination S3 prefix where unzipped files will be stored
    :type dest_key: str
    """

    ui_color = "#FF69B4"

    @apply_defaults
    def __init__(
        self: object,
        src_bucket: str,
        src_key: str,
        dest_key: str,
        *args: list,
        **kwargs: dict,
    ):
        super().__init__(
            *args,
            **kwargs,
        )
        self.src_bucket = src_bucket
        self.src_key = src_key
        self.dest_key = dest_key



    def execute(
        self: object,
        context: dict,
    ) -> None:
        try:
            import boto3
            import zipfile
            from io import BytesIO
            import os

            s3_resource = boto3.resource("s3")

            # Get the destination key
            unzip_key = (
                self.dest_key
                if self.dest_key is not None
                else os.path.splitext(self.src_key.replace("zipped", "unzipped"))[0]
            )

            self.log.info(f"Source: s3://{self.src_bucket}/{self.src_key}")
            self.log.info(f"Destination: s3://{self.src_bucket}/{unzip_key}")

            # Get the zip file from S3
            zip_obj = s3_resource.Object(
                bucket_name=self.src_bucket,
                key=self.src_key,
            )
            buffer = BytesIO(zip_obj.get()["Body"].read())

            # Unzip and extract files
            file_count = 0
            skipped_count = 0
            z = zipfile.ZipFile(buffer)

            # Log total number of files in the zip
            total_files = len(z.namelist())
            self.log.info(f"Found {total_files} entries in zip file")

            for filename in z.namelist():
                # Skip directories/folders
                if filename.endswith("/"):
                    skipped_count += 1
                    continue

                # Construct the full destination key path, preserving directory structure
                dest_file_key = os.path.join(unzip_key, filename)

                # Upload the file to S3
                s3_resource.meta.client.upload_fileobj(
                    z.open(filename),
                    Bucket=self.src_bucket,
                    Key=dest_file_key,
                )
                file_count += 1

                # Only log every 10 files to reduce log spam for large zips
                if file_count % 10 == 0 or file_count == 1:
                    self.log.info(
                        f"Extracted {file_count}/{total_files-skipped_count} files"
                    )

            self.log.info(f"Successfully extracted {file_count} files from zip")
            self.log.info(f"Skipped {skipped_count} directory entries")
            self.log.info(f"Files extracted to s3://{self.src_bucket}/{unzip_key}")

        except Exception as e:
            self.log.error(f"Error unzipping file: {e}")
            raise e