Source code for operators.iceberg.maintain_iceberg_table

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




class MaintainIcebergTableOperator(BaseIcebergTableOperator):
    """
    Performs maintenance operations on Iceberg tables in AWS Athena.

    This operator combines small files, cleans up old snapshots, and provides
    before/after statistics on the table.

    :param table_name: The name of the Iceberg table to maintain
    :type table_name: str
    :param database: The database containing the Iceberg table
    :type database: str
    :param optimize: Whether to run OPTIMIZE to combine small files
    :type optimize: bool
    :param vacuum: Whether to run VACUUM to clean up old snapshots
    :type vacuum: bool
    :param optimize_options: Additional options for the OPTIMIZE command
    :type optimize_options: str
    """

    template_fields = BaseIcebergTableOperator.template_fields + ["optimize_options"]

    ui_color = "#D6EAF8"  # Light Blue shade

    @apply_defaults
    def __init__(
        self,
        table_name: str,
        database: Optional[str] = None,
        optimize: bool = True,
        vacuum: bool = True,
        optimize_options: str = "REWRITE DATA USING BIN_PACK",
        *args,
        **kwargs,
    ):
        super().__init__(table_name=table_name, database=database, *args, **kwargs)
        self.optimize = optimize
        self.vacuum = vacuum
        self.optimize_options = optimize_options



    def execute(self, context: Dict) -> Dict[str, Any]:
        """
        Execute the maintenance operations on the Iceberg table.

        :param context: Airflow context
        :type context: Dict
        """
        # Get statistics before maintenance
        stats_before = self._get_table_statistics()
        self.log.info(f"BEFORE MAINTENANCE, table statistics: {stats_before}")

        # Run optimize if requested
        if self.optimize:
            optimize_query = (
                f"OPTIMIZE {self.database}.{self.table_name} {self.optimize_options}"
            )
            self._execute_athena_query(sql=optimize_query)

        # Run vacuum if requested
        if self.vacuum:
            vacuum_query = f"VACUUM {self.database}.{self.table_name}"
            self._execute_athena_query(sql=vacuum_query)

        # Get statistics after maintenance
        stats_after = self._get_table_statistics()
        self.log.info(f"AFTER MAINTENANCE, table statistics: {stats_after}")

        # Calculate improvements
        file_count_reduction = 0
        size_reduction_mb = 0

        if isinstance(stats_before["file_count"], (int, float)) and isinstance(
            stats_after["file_count"], (int, float)
        ):
            file_count_reduction = (
                stats_before["file_count"] - stats_after["file_count"]
            )

        if isinstance(stats_before["total_size_mb"], (int, float)) and isinstance(
            stats_after["total_size_mb"], (int, float)
        ):
            size_reduction_mb = (
                stats_before["total_size_mb"] - stats_after["total_size_mb"]
            )

        # Log a summary of the operation
        self.log.info(
            f"✅ Maintenance completed for {self.database}.{self.table_name}. "
            f"Files: {stats_before['file_count']} → {stats_after['file_count']} "
            f"({file_count_reduction:+.0f}), "
            f"Size: {stats_before['total_size_mb']:.2f}MB → {stats_after['total_size_mb']:.2f}MB "
            f"({size_reduction_mb:+.2f}MB)"
        )