Source code for operators.lake_msck

from config import Config

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException




class LakeMsckOperator(BaseOperator):
    """
    Run MSCK REPAIR TABLE on a table in the data lake.
    Running twice to ensure all partitions are picked up.

    :param table_name: the name of the table to repair
    :type table_name: str
    :param database: the database the table is in, defaults to Config.DATA_LAKE_DB
    :type database

    :raises AirflowException: _description_
    """

    ui_color = "#FF69B4"

    template_fields = [
        "table_name",
        "database",
        "workgroup",
    ]

    @apply_defaults
    def __init__(
        self: object,
        table_name: str,
        database: str = Config.DATA_LAKE_DB,
        workgroup: str = Config.DEFAULT_WORKGROUP,
        *args,
        **kwargs,
    ):

        super().__init__(*args, **kwargs)
        self.table_name = table_name
        self.database = database
        self.workgroup = workgroup



    def execute(
        self: object,
        context: dict,
    ) -> None:
        try:
            import awswrangler as wr

            for _ in range(2):
                wr.athena.start_query_execution(
                    sql=f"MSCK REPAIR TABLE `{self.table_name}`;",
                    database=self.database,
                    workgroup=self.workgroup,
                    wait=True,
                )
        except Exception as ex:
            raise AirflowException(
                f"Failed to run MSCK REPAIR TABLE on {self.table_name} in {self.database}"
            ) from ex
        else:
            self.log.info(
                f"MSCK REPAIR TABLE completed successfully on {self.table_name} in {self.database}"
            )