Source code for operators.postgres.utils.backup_table

"""
Backup Table Operator

This module provides an operator for backing up a Postgres table to a temporary
table within the same schema. Useful for preserving data before destructive
operations like schema syncs.

Usage:
    from operators.postgres.utils.backup_table import BackupTableOperator

    backup_users = BackupTableOperator(
        task_id="backup_users",
        postgres_conn_id="boss_grapi_stage",
        tenant="arkansas",
        source_table="User",
        backup_suffix="_backup",
    )
"""

from __future__ import annotations

from typing import Sequence

from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.utils.context import Context


################################################################################
## Constants
################################################################################

DEFAULT_BACKUP_SUFFIX = "_backup"
STAGE_CONN_PATTERN = "stage"

SQL_DROP_BACKUP = """
DROP TABLE IF EXISTS {tenant}."{table}{suffix}";
"""

SQL_CREATE_BACKUP = """
CREATE TABLE {tenant}."{table}{suffix}" AS
SELECT * FROM {tenant}."{table}";
"""

SQL_CREATE_BACKUP_CAST = """
CREATE TABLE {tenant}."{table}{suffix}" AS
SELECT {columns} FROM {tenant}."{table}";
"""

SQL_GET_ENUM_COLUMNS = """
SELECT c.column_name
FROM information_schema.columns AS c
JOIN pg_type AS t ON t.typname = c.udt_name
WHERE c.table_schema = '{tenant}'
    AND c.table_name = '{table}'
    AND t.typtype = 'e';
"""


################################################################################
## Operator
################################################################################




class BackupTableOperator(BaseOperator):
    """
    Backup a Postgres table to a temporary table in the same schema.

    Creates a copy of the source table with a suffix appended to the name.
    Drops any existing backup table with the same name first.

    :param postgres_conn_id: Airflow connection ID for the Postgres database.
    :param tenant: The tenant schema name.
    :param source_table: The table name to backup.
    :param backup_suffix: Suffix to append to backup table name (default: "_backup").
    :param cast_enums_to_text: If True, cast ENUM columns to TEXT in the backup
        to avoid dependencies on ENUM types (default: False).
    """

    ui_color = "#BADA55"
    template_fields: Sequence[str] = ("postgres_conn_id", "tenant", "source_table")

    def __init__(
        self,
        *,
        postgres_conn_id: str,
        tenant: str,
        source_table: str,
        backup_suffix: str = DEFAULT_BACKUP_SUFFIX,
        cast_enums_to_text: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.postgres_conn_id = postgres_conn_id
        self.tenant = tenant
        self.source_table = source_table
        self.backup_suffix = backup_suffix
        self.cast_enums_to_text = cast_enums_to_text

    ############################################################################
    ## Execute
    ############################################################################



    def execute(self, context: Context) -> str:
        if STAGE_CONN_PATTERN not in self.postgres_conn_id.lower():
            raise ValueError(
                f"Safety check failed: '{self.postgres_conn_id}' does not appear to be a stage connection. "
                "This operator should only modify stage databases."
            )

        hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)

        backup_table = f"{self.source_table}{self.backup_suffix}"

        drop_sql = SQL_DROP_BACKUP.format(
            tenant=self.tenant,
            table=self.source_table,
            suffix=self.backup_suffix,
        )

        if self.cast_enums_to_text:
            enum_sql = SQL_GET_ENUM_COLUMNS.format(
                tenant=self.tenant,
                table=self.source_table,
            )
            enum_cols = {
                row[0] for row in hook.get_records(enum_sql)
            }
            self.log.info("Enum columns to cast as TEXT: %s", enum_cols)

            all_cols_sql = (
                f"SELECT column_name FROM information_schema.columns "
                f"WHERE table_schema = '{self.tenant}' "
                f"AND table_name = '{self.source_table}' "
                f"ORDER BY ordinal_position"
            )
            all_cols = [row[0] for row in hook.get_records(all_cols_sql)]

            columns = ", ".join(
                f'"{c}"::TEXT AS "{c}"' if c in enum_cols else f'"{c}"'
                for c in all_cols
            )

            create_sql = SQL_CREATE_BACKUP_CAST.format(
                tenant=self.tenant,
                table=self.source_table,
                suffix=self.backup_suffix,
                columns=columns,
            )
        else:
            create_sql = SQL_CREATE_BACKUP.format(
                tenant=self.tenant,
                table=self.source_table,
                suffix=self.backup_suffix,
            )

        self.log.info(
            "Backing up %s.%s to %s.%s",
            self.tenant,
            self.source_table,
            self.tenant,
            backup_table,
        )

        hook.run(drop_sql)
        hook.run(create_sql)

        self.log.info("Backup complete: %s.%s", self.tenant, backup_table)
        return backup_table