Source code for operators.lake_pit_table

from config import Config

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




class LakePitTableOperator(BaseOperator):
    """
    Create/insert data into a point in time table in the data lake.

    :param pit_table_name: the name of the point in table that will be created/inserted into
    :type table_name: str
    :param source_table_name: the name of the base table that will be read from
    :type table_name: str
    :param version: the version of the data, gets inserted into the version column in the pit table
    :type version: str
    :param database: the database the table is in, defaults to Config.DATA_LAKE_DB
    :type database: str
    :param version_column_name: the name of the version column in the pit table, defaults to "version_date"
    :type version_column_name: str

    :raises AirflowException: _description_
    """

    ui_color = "#9842f5"
    template_fields = ("pit_table_name", "source_table_name", "version")

    @apply_defaults
    def __init__(
        self: object,
        pit_table_name: str,
        source_table_name: str,
        version: str,
        database: str = Config.DATA_LAKE_DB,
        version_column_name: str = "version_date",
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

        self.pit_table_name = pit_table_name
        self.source_table_name = source_table_name
        self.version = version
        self.database = database
        self.version_column_name = version_column_name



    def execute(
        self: object,
        context: dict,
    ) -> None:
        import awswrangler as wr
        from util.backup_table import backup_table

        if wr.catalog.does_table_exist(
            database=Config.DATA_LAKE_DB, table=self.pit_table_name
        ):
            # Get the column names of the tables
            df_og_col_names = wr.athena.read_sql_query(
                sql=f"SELECT * FROM {self.pit_table_name} LIMIT 1",
                database=Config.DATA_LAKE_DB,
            )
            og_col_names = [
                col for col in df_og_col_names.columns.tolist() if col != self.version_column_name
            ]

            df_new_col_names = wr.athena.read_sql_query(
                sql=f"SELECT * FROM {self.source_table_name} LIMIT 1", database=Config.DATA_LAKE_DB
            )
            new_col_names = df_new_col_names.columns.tolist()

            if og_col_names == new_col_names:
                # If the columns are the same, just insert the data
                insert_query = f"INSERT INTO {self.pit_table_name} SELECT *, '{self.version}' AS {self.version_column_name} FROM {self.source_table_name}"
                print(insert_query)
                wr.athena.start_query_execution(insert_query, Config.DATA_LAKE_DB, wait=True)
            else:
                # If the columns are different, create a new table containing all columns from both tables, and mark the missing columns as NULL
                selection_query = "SELECT "
                for col in new_col_names:
                    if col in og_col_names:
                        selection_query += f"{col}, "
                    else:
                        selection_query += f"NULL AS {col}, "

                for col in og_col_names:
                    if col not in new_col_names:
                        selection_query += f"{col}, "

                selection_query = (
                    selection_query[:-2] + f" FROM {self.pit_table_name} UNION ALL SELECT "
                )

                for col in new_col_names:
                    selection_query += f"{col}, "

                for col in og_col_names:
                    if col not in new_col_names:
                        selection_query += f"NULL AS {col}, "

                selection_query = selection_query[:-2] + f" FROM {self.source_table_name}"

                print(
                    f"Detected column mismatch between {self.source_table_name} and {self.pit_table_name}. Creating new backup table with correct columns\n\n{selection_query}\n"
                )

                backup_table(
                    source_table_name=self.source_table_name,
                    backup_table_name=self.pit_table_name,
                    compare_checksum=False,
                    verbose=True,
                    new_partitions=[self.version_column_name],
                    custom_selection=selection_query,
                    overwrite=True
                )
        else:
            print(f"Creating new table {self.pit_table_name}")

            backup_table(
                source_table_name=self.source_table_name,
                backup_table_name=self.pit_table_name,
                compare_checksum=False,
                verbose=True,
                new_partitions=[self.version_column_name],
                custom_selection=f"SELECT *, '{self.version}' AS {self.version_column_name} FROM {self.source_table_name}",
            )