Source code for operators.iceberg.remove_iceberg_table

from airflow.exceptions import AirflowException
from airflow.utils.decorators import apply_defaults
from typing import Dict, Optional, Any
from operators.iceberg.base_iceberg_table import BaseIcebergTableOperator
from config import Config




class RemoveIcebergTableOperator(BaseIcebergTableOperator):
    """
    Completely remove an Iceberg table by deleting all data and then dropping the table.

    This operator first deletes all data in the Iceberg table and then drops the table
    using AWS Athena.

    :param table_name: The name of the Iceberg table to remove
    :type table_name: str
    :param database: The name of the database containing the Iceberg table
    :type database: str
    :param skip_if_not_exists: Whether to skip without error if table doesn't exist
    :type skip_if_not_exists: bool
    """

    template_fields = BaseIcebergTableOperator.template_fields + ["skip_if_not_exists"]

    ui_color = "#E8DAEF"  # Light Purple (matching base operator)

    @apply_defaults
    def __init__(
        self,
        table_name: str,
        database: Optional[str] = None,
        skip_if_not_exists: bool = True,
        *args,
        **kwargs,
    ):
        # Don't verify table exists in base operator - we'll handle it specially
        super().__init__(
            table_name=table_name,
            database=database,
            verify_exists=False,
            *args,
            **kwargs,
        )
        self.skip_if_not_exists = skip_if_not_exists



    def execute(self, context: Dict) -> Dict[str, Any]:
        """
        Execute the table removal operation.
        """
        import awswrangler as wr

        # Check if table exists - handle according to skip_if_not_exists setting
        table_exists = wr.catalog.does_table_exist(
            database=self.database,
            table=self.table_name,
        )

        if not table_exists:
            message = f"Table {self.database}.{self.table_name} does not exist"
            if self.skip_if_not_exists:
                self.log.info(f"❌ {message}, skipping remove")
                return
            raise AirflowException(f"{message}")

        self.log.info(
            f"🧹 Cleaning up: Deleting data + Dropping table {self.database}.{self.table_name}"
        )

        try:
            # First delete all data
            self._execute_athena_query(
                sql=f"DELETE FROM {self.database}.{self.table_name}",
            )

            # Then drop the table
            self._execute_athena_query(
                sql=f"DROP TABLE {self.database}.{self.table_name}",
            )

            self.log.info(
                f"✅ Table {self.database}.{self.table_name} removed successfully"
            )

        except Exception as e:
            self.log.error(f"Error details: {str(e)}")
            raise AirflowException(
                f"Failed to remove Iceberg table {self.database}.{self.table_name}"
            ) from e