Source code for operators.iceberg.drop_iceberg_partition

from airflow.utils.decorators import apply_defaults
from typing import Dict, Optional, Any, Union
from operators.iceberg.base_iceberg_table import BaseIcebergTableOperator
from operators.iceberg.maintain_iceberg_table import MaintainIcebergTableOperator




class DropIcebergPartitionOperator(BaseIcebergTableOperator):
    """
    Drops a partition from an Iceberg table and optionally runs maintenance afterward.

    This operator deletes data from a specific partition using a WHERE clause and
    can trigger maintenance operations after the deletion.

    The DropIcebergPartitionOperator inherits from BaseIcebergTableOperator, so by default it will already verify the table exists during the pre_execute phase. Since your verify_exists parameter is set to True by default in the base operator, all derived operators will perform this check unless explicitly disabled.

    :param table_name: The Iceberg table name
    :type table_name: str
    :param partition_col: The partition column name
    :type partition_col: str
    :param partition_value: The partition value to drop
    :type partition_value: str, int, or float
    :param database: The database containing the table
    :type database: str
    :param additional_filter: Additional WHERE conditions (without 'AND')
    :type additional_filter: str
    :param run_maintenance: Whether to run maintenance after dropping the partition
    :type run_maintenance: bool
    :param maintenance_options: Optional parameters for maintenance operation
    :type maintenance_options: dict
    """

    template_fields = BaseIcebergTableOperator.template_fields + [
        "partition_col",
        "partition_value",
        "additional_filter",
    ]

    ui_color = "#FADBD8"  # Light Red/Pink shade

    @apply_defaults
    def __init__(
        self,
        table_name: str,
        partition_col: str,
        partition_value: Union[str, int, float],
        database: Optional[str] = None,
        additional_filter: Optional[str] = None,
        run_maintenance: bool = True,
        maintenance_options: Optional[Dict[str, Any]] = None,
        *args,
        **kwargs,
    ):
        super().__init__(table_name=table_name, database=database, *args, **kwargs)
        self.partition_col = partition_col
        self.partition_value = partition_value
        self.additional_filter = additional_filter
        self.run_maintenance = run_maintenance
        self.maintenance_options = maintenance_options or {}



    def execute(self, context: Dict) -> Dict[str, Any]:
        """
        Execute the partition drop operation.
        :param context: Airflow context
        :type context: Dict
        """

        # Build the WHERE clause
        where_clause = self._build_partition_filter()

        # Execute the DELETE operation
        sql = f"DELETE FROM {self.database}.{self.table_name} WHERE {where_clause}"
        self.log.info(f"Dropping partition data with: {sql}")

        self._execute_athena_query(sql=sql)

        self.log.info(f"Successfully dropped data from partition where {where_clause}")

        # Run maintenance if requested
        if self.run_maintenance:
            self.log.info(f"Running maintenance on table after partition drop")
            maintenance_op = MaintainIcebergTableOperator(
                task_id=f"maintenance_after_partition_drop",
                table_name=self.table_name,
                database=self.database,
                **self.maintenance_options,
            )

            # Execute the maintenance operation directly
            maintenance_op.execute(context)


    def _build_partition_filter(self) -> str:
        """
        Build the WHERE clause filter for the partition drop operation.

        :return: SQL WHERE clause condition (without the WHERE keyword)
        :rtype: str
        """
        # Handle different value types appropriately
        if isinstance(self.partition_value, str):
            partition_condition = f"{self.partition_col}='{self.partition_value}'"
        else:
            partition_condition = f"{self.partition_col}={self.partition_value}"

        # Add any additional WHERE conditions if provided
        if self.additional_filter:
            return f"({partition_condition}) AND ({self.additional_filter})"

        return partition_condition