Source code for operators.postgres.utils.get_latest_fabric_fixed_availability_version

from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowException
from typing import TypedDict

from config import Config


class FabricVersion(TypedDict):
    version_date: str
    update_date: str
    version: int
    process_uuid: str




class GetLatestFabricFixedAvailabilityVersionOperator(BaseOperator):
    """
    Operator to retrieve the latest fabric fixed availability (FFA) version from the fabric_versions table.

    Returns a FabricVersion dict with version, version_date, update_date, and process_uuid.

    :param postgres_conn_id: Airflow connection id for the Postgres database.
    :param schema: Schema containing the fabric_versions table.
    :param table: Name of the fabric_versions table.
    :param status: Optional status filter (e.g., 'ingested').
    :param fa_ver: Optional fabric address version filter.
    :param fs_ver: Optional fabric fixed availability version filter (takes precedence over fa_ver).
    """

    template_fields = ("postgres_conn_id", "status", "fa_ver", "fs_ver")
    ui_color = "#4D6A42"

    VERSION_DATE_DICT = {
        1: "2023-06-30",
        2: "2022-12-31",
        3: "2023-06-30",
        4: "2023-12-31",
        5: "2024-06-30",
        6: "2024-12-31",
        7: "2025-06-30",
    }

    def __init__(
        self,
        postgres_conn_id: str = Config.GAME_PG,
        schema: str = "data",
        table: str = "fabric_versions",
        status: str | None = None,
        fa_ver: int | None = None,
        fs_ver: int | None = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.postgres_conn_id = postgres_conn_id
        self.schema = schema
        self.table = table
        self.status = status
        self.fa_ver = fa_ver
        self.fs_ver = fs_ver



    def execute(self, context) -> FabricVersion:
        """Retrieve the latest fabric version from Postgres."""
        from util.format import FormatUtil

        where_clauses = []

        if self.fa_ver:
            version_date = self.VERSION_DATE_DICT.get(int(self.fa_ver))
            if version_date:
                where_clauses.append(f"version_date = '{version_date}'")

        if self.fs_ver:
            where_clauses = [f"version = {int(self.fs_ver)}"]

        if self.status:
            where_clauses.append(f"status = '{self.status}'")

        where_cond = f" WHERE {' AND '.join(where_clauses)}" if where_clauses else ""

        sql = f"""
            SELECT version_date, update_date, version, process_uuid
            FROM {self.schema}.{self.table}
            {where_cond}
            ORDER BY version DESC
            LIMIT 1
        """

        self.log.info(
            "Querying fabric_versions with conn_id='%s', status='%s'",
            self.postgres_conn_id,
            self.status,
        )

        try:
            hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
            result = hook.get_first(sql)
        except Exception as e:
            self.log.exception(
                "Error querying fabric_versions with conn_id '%s': %s",
                self.postgres_conn_id,
                e,
            )
            raise AirflowException(
                f"Failed to query fabric_versions (conn_id='{self.postgres_conn_id}')."
            ) from e

        if not result:
            raise AirflowException(
                f"No fabric version found matching criteria (status='{self.status}')."
            )

        fabric_version: FabricVersion = {
            "version_date": FormatUtil.to_date_string(result[0]),
            "update_date": FormatUtil.to_date_string(result[1]),
            "version": int(result[2]),
            "process_uuid": str(result[3]),
        }

        self.log.info(
            "Found fabric version: %d (version_date=%s, status=%s)",
            fabric_version["version"],
            fabric_version["version_date"],
            self.status,
        )

        return fabric_version