Source code for operators.athena_to_s3

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




class AthenaToS3Operator(BaseOperator):
    """
    An operator that will execute a query against Athena and write the results to S3.

    The operator wraps awswrangler's athena.unload functionality to execute a query and then write
    the results to S3. It handles the configuration of the S3 bucket, prefix, and other parameters.

    Parameters:
    :param source_db: The source database to run the query against.
    :type source_db: str

    :param sql: The SQL query to execute in Athena.
    :type prefix: str

    :param s3_bucket: The S3 bucket where the results will be stored.
    :type s3_bucket: str

    :param s3_prefix: The S3 prefix where the results will be stored.
    :type s3_prefix: str

    :param file_format: The file format for the output files (e.g., 'parquet', 'csv').
    :type file_format: str

    :param compression: The compression type for the output files (e.g., 'gzip', 'snappy').
    :type compression: str

    :param delimiter: The delimiter for CSV files.
    :type delimiter: str
    """

    ui_color = "#71D1E5"

    template_fields = (
        "sql",
        "s3_bucket",
        "s3_prefix",
        "s3_filename",
        "file_format",
        "compression",
        "delimiter",
        "data_source",
        "source_db",
    )

    @apply_defaults
    def __init__(
        self,
        sql: str,
        source_db: str,
        s3_bucket: str,
        s3_prefix: str,
        s3_filename: str,
        file_format: str,
        compression: str = "None",
        delimiter: str = "None",
        data_source: str = "AwsDataCatalog",
        *args,
        **kwargs,
    ):
        super(AthenaToS3Operator, self).__init__(*args, **kwargs)
        self.sql = sql
        self.source_db = source_db
        self.data_source = data_source
        self.s3_filename = s3_filename
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_prefix
        self.file_format = file_format
        self.compression = compression
        self.delimiter = delimiter



    def execute(self, context):
        import awswrangler as wr

        try:
            wr.athena.unload(
                sql=self.sql,
                path=f"s3://{self.s3_bucket}/{self.s3_prefix}/{self.s3_filename}",
                database=self.source_db,
                data_source=self.data_source,
                file_format=self.file_format,
                compression=self.compression,
                field_delimiter=self.delimiter,
            )
            self.log.info("Data successfully unloaded to S3.")
        except Exception as e:
            raise AirflowException(f"Failed to unload data to S3: {str(e)}")