Source code for operators.iceberg.base_iceberg_table

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException
from config import Config
from typing import Dict, Any




class BaseIcebergTableOperator(BaseOperator):
    """
    Base operator for working with Iceberg tables in AWS Athena.

    This abstract class provides common functionality for all Iceberg table operators
    including connection handling, validation, and utility methods.

    :param database: The database containing the Iceberg table
    :type database: str
    :param table_name: The name of the Iceberg table to operate on
    :type table_name: str
    :param verify_exists: Whether to verify the table exists before operations
    :type verify_exists: bool
    """

    template_fields = ["database", "table_name"]

    ui_color = "#E8DAEF"  # Light Purple base color for all Iceberg operators

    @apply_defaults
    def __init__(
        self,
        table_name: str,
        database: str = None,
        verify_exists: bool = True,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.table_name = table_name
        self.database = database or Config.DATA_LAKE_DB
        self.verify_exists = verify_exists



    def pre_execute(self, context: Dict) -> None:
        """
        Hook called before execution. https://airflow.apache.org/docs/apache-airflow/1.10.12/_api/airflow/operators/index.html#airflow.operators.BaseOperator.pre_execute
        Verifies table exists if verify_exists is True.
        """
        if self.verify_exists:
            self._verify_table_exists()


    def _verify_table_exists(self) -> bool:
        """
        Verify the specified Iceberg table exists.

        :return: True if table exists, raises exception if not
        :rtype: bool
        :raises: AirflowException if table does not exist
        """
        import awswrangler as wr

        if not wr.catalog.does_table_exist(
            database=self.database,
            table=self.table_name,
        ):
            raise AirflowException(
                f"Table {self.database}.{self.table_name} does not exist"
            )

        return True

    def _execute_athena_query(self, sql: str, wait: bool = True) -> Dict[str, Any]:
        """
        Execute a query in Athena with proper error handling.

        :param sql: SQL query to execute
        :type sql: str
        :param wait: Whether to wait for query completion
        :type wait: bool
        :return: Query execution details
        :rtype: Dict[str, Any]
        """
        import awswrangler as wr

        try:
            self.log.info(f"Executing query: {sql}")
            response = wr.athena.start_query_execution(
                sql=sql,
                database=self.database,
                wait=wait,
            )
            self.log.info(f"Query execution ID: {response['QueryExecutionId']}")
        except Exception as e:
            self.log.error(f"Failed to execute Athena query: {sql}")
            self.log.error(f"Error details: {str(e)}")
            raise AirflowException(f"Athena query execution failed: {str(e)}") from e

    def _get_table_statistics(self) -> Dict[str, Any]:
        """
        Get statistics for the Iceberg table.

        :return: Dictionary containing file count and size statistics
        :rtype: Dict[str, Any]
        """
        import awswrangler as wr
        import pandas as pd

        try:
            file_stats = wr.athena.read_sql_query(
                f"SELECT count(*) as file_count, "
                f"sum(file_size_in_bytes)/1024/1024 as total_size_mb "
                f'FROM "{self.database}"."{self.table_name}$files"',
                database=self.database,
            )

            # Handle case where there are no files
            if file_stats.empty or pd.isna(file_stats["total_size_mb"].iloc[0]):
                return {"file_count": 0, "total_size_mb": 0}

            return file_stats.iloc[0].to_dict()
        except Exception as e:
            self.log.warning(f"Failed to get table statistics: {str(e)}")
            return {"file_count": "unknown", "total_size_mb": "unknown"}