Source code for util.metadata_pg

# metadata_pg.py
from config import Config
from airflow.providers.postgres.hooks.postgres import PostgresHook
from typing import Dict, Any, List, Union, Optional
from datetime import datetime
import json

FAFS_METADATA_KEY = "fafs_vers"




class MetadataUtil:
    """Manages metadata in PostgreSQL for Airflow DAGs."""

    _conn_id = Config.GAME_PG
    _schema = "data"
    _table_name = "datateam_meta"

    @staticmethod
    def _get_hook() -> PostgresHook:
        return PostgresHook(postgres_conn_id=MetadataUtil._conn_id)

    @staticmethod
    def _ensure_table_exists() -> None:
        hook = MetadataUtil._get_hook()
        hook.run(
            f"""
            CREATE TABLE IF NOT EXISTS {MetadataUtil._schema}.{MetadataUtil._table_name} (
                id SERIAL PRIMARY KEY,
                dag_id VARCHAR(203) NOT NULL,
                metadata_key VARCHAR(304) NOT NULL,
                metadata_value VARCHAR(405),
                execution_date TIMESTAMP WITH TIME ZONE,
                json_data JSONB
            );
            CREATE INDEX IF NOT EXISTS idx_{MetadataUtil._table_name}_dag_id_metadata_key
                ON {MetadataUtil._schema}.{MetadataUtil._table_name} (dag_id, metadata_key);
            CREATE INDEX IF NOT EXISTS idx_{MetadataUtil._table_name}_execution_date
                ON {MetadataUtil._schema}.{MetadataUtil._table_name} (execution_date DESC);
        """
        )



    @staticmethod
    def store_pgmeta(
        dag_id: str,
        metadata_key: str,
        metadata_value: Any,
        execution_date: Union[datetime, str],
        json_data: Optional[Dict[str, Any]] = None,
    ) -> None:
        MetadataUtil._ensure_table_exists()
        hook = MetadataUtil._get_hook()
        execution_date = (
            execution_date.isoformat()
            if isinstance(execution_date, datetime)
            else execution_date
        )
        try:
            hook.run(
                f"""
                INSERT INTO {MetadataUtil._schema}.{MetadataUtil._table_name}
                (dag_id, metadata_key, metadata_value, execution_date, json_data)
                VALUES (%s, %s, %s, %s, %s)
            """,
                parameters=(
                    dag_id,
                    metadata_key,
                    str(metadata_value),
                    execution_date,
                    json.dumps(json_data) if json_data else None,
                ),
            )
        except Exception as e:
            print(f"Error storing metadata: {e}")




    @staticmethod
    def get_pgmeta(
        dag_id: str,
        metadata_key: str,
        mode: str = "latest",
        month: Optional[str] = None,
        limit: Optional[int] = None,
        parname: Optional[str] = None,
    ) -> Union[Optional[str], Dict[str, Any], List[Dict[str, Any]]]:
        """
        Retrieves metadata based on specified parameters and mode.

        Modes:
        - 'latest': Latest entry
        - 'max': Maximum value
        - 'all': All entries
        - 'fafs_month': FAFS data for a specific month
        - 'fafs_max': Maximum FAFS versions
        - 'fafs_latest': Latest FAFS versions
        """
        MetadataUtil._ensure_table_exists()
        hook = MetadataUtil._get_hook()
        query, params = MetadataUtil._build_pgmeta_query(
            dag_id, metadata_key, mode, month, limit, parname
        )
        try:
            if mode in ["max", "latest", "fafs_month"]:
                result = hook.get_first(query, parameters=params)
            elif mode in ["fafs_max", "fafs_latest", "all"]:
                result = hook.get_records(query, parameters=params)
            else:
                raise ValueError(f"Invalid mode: {mode}")

            return MetadataUtil._process_pgmeta_result(result, mode, parname)
        except Exception as e:
            print(f"Error retrieving metadata: {e}")
            print(f"Query: {query}")
            print(f"Parameters: {params}")
            return None if mode in ["max", "latest", "fafs_month"] else []


    @staticmethod
    def _build_pgmeta_query(
        dag_id: str,
        metadata_key: str,
        mode: str,
        month: Optional[str] = None,
        limit: Optional[int] = None,
        parname: Optional[str] = None,
    ) -> tuple:
        """Builds SQL query based on retrieval mode."""
        base_query = f"FROM {MetadataUtil._schema}.{MetadataUtil._table_name} WHERE dag_id = %s AND metadata_key = %s"
        params = [dag_id, metadata_key]

        if mode == "max":
            query = f"SELECT MAX(metadata_value::text) {base_query}"
        elif mode == "latest":
            query = f"SELECT id, dag_id, metadata_key, metadata_value, execution_date, json_data {base_query} ORDER BY execution_date DESC LIMIT 1"
        elif mode == "all":
            query = f"SELECT id, dag_id, metadata_key, metadata_value, execution_date, json_data {base_query} ORDER BY execution_date DESC"
            if limit:
                query += f" LIMIT {limit}"
        elif mode == "fafs_month":
            query = f"SELECT metadata_value, json_data->'{FAFS_METADATA_KEY}' as versions {base_query}"
            if month:
                query += " AND metadata_value = %s"
                params.append(month)
            query += " ORDER BY execution_date DESC LIMIT 1"
        elif mode == "fafs_max":
            query = f"""
                WITH versions_data AS (
                    SELECT jsonb_array_elements(json_data->'{FAFS_METADATA_KEY}') as version 
                    FROM {MetadataUtil._schema}.{MetadataUtil._table_name} 
                    WHERE dag_id = %s AND metadata_key = %s
                )
                SELECT 
                    version->>'parname' as parname,
                    MAX((version->>'fa_ver')::int)::text as fa_ver,
                    MAX((version->>'fs_ver')::int)::text as fs_ver
                FROM versions_data
            """
            if parname:
                query += " WHERE version->>'parname' = %s"
                params.append(parname)
            query += " GROUP BY version->>'parname'"
        elif mode == "fafs_latest":
            query = f"""
                WITH latest_record AS (
                    SELECT json_data->'{FAFS_METADATA_KEY}' as versions
                    FROM {MetadataUtil._schema}.{MetadataUtil._table_name}
                    WHERE dag_id = %s AND metadata_key = %s
                    ORDER BY execution_date DESC
                    LIMIT 1
                ), unnested_versions AS (
                    SELECT jsonb_array_elements(versions) as version
                    FROM latest_record
                )
                SELECT version
                FROM unnested_versions
            """
            if parname:
                query += " WHERE version->>'parname' = %s"
                params.append(parname)
        else:
            raise ValueError(f"Invalid mode: {mode}")

        return query, params

    @staticmethod
    def _process_pgmeta_result(result, mode, parname=None):
        if mode == "max":
            return result[0] if result else None
        elif mode == "latest":
            return MetadataUtil._format_pgmeta_result(result) if result else None
        elif mode == "fafs_month":
            return {result[0]: result[1]} if result else None
        elif mode == "fafs_max":
            if result:
                versions = {
                    row[0]: {"fa_ver": row[1], "fs_ver": row[2]}
                    for row in result
                    if row[0]
                }
                return versions.get(parname) if parname else versions
            return None if parname else {}
        elif mode == "fafs_latest":
            if result:
                versions = [dict(version[0]) for version in result if version[0]]
                latest_versions = {
                    v["parname"]: {"fa_ver": v["fa_ver"], "fs_ver": v["fs_ver"]}
                    for v in versions
                }
                return latest_versions.get(parname) if parname else latest_versions
            return None if parname else {}
        else:
            return result

    @staticmethod
    def _format_pgmeta_result(row: tuple) -> Dict[str, Any]:
        """Formats database row into a dictionary."""
        result = dict(
            zip(
                [
                    "id",
                    "dag_id",
                    "metadata_key",
                    "metadata_value",
                    "execution_date",
                    "json_data",
                ],
                row,
            )
        )
        if result.get("execution_date"):
            result["execution_date"] = result["execution_date"].isoformat()
        if result.get("json_data"):
            result["json_data"] = json.loads(result["json_data"])
        return result

    # Simplified interface for common operations
    # users could fall back to the more general store_pgmeta and get_pgmeta methods for more complex queries if needed.


    @staticmethod
    def store_pgmeta_version(
        dag_id: str,
        version_key: str,
        version_value: str,
        execution_date: Union[datetime, str],
    ):
        """Stores version metadata for DAGs."""
        MetadataUtil.store_pgmeta(
            dag_id=dag_id,
            metadata_key=version_key,
            metadata_value=version_value,
            execution_date=execution_date,
        )




    @staticmethod
    def get_pgmeta_latest_version(dag_id: str, version_key: str) -> Optional[str]:
        """Retrieves the latest version for a DAG."""
        result = MetadataUtil.get_pgmeta(dag_id, version_key, mode="latest")
        return result["metadata_value"] if result else None




    @staticmethod
    def get_pgmeta_max_version(dag_id: str, version_key: str) -> Optional[str]:
        """Retrieves the maximum version for a DAG."""
        return MetadataUtil.get_pgmeta(dag_id, version_key, mode="max")




    @staticmethod
    def get_pgmeta_all_versions(
        dag_id: str, version_key: str, limit: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Retrieves all versions for a DAG, with optional limit."""
        return MetadataUtil.get_pgmeta(dag_id, version_key, mode="all", limit=limit)


    ## FAFS intefaces for common operations


    @staticmethod
    def store_pgmeta_fafs(
        dag_id: str,
        month: str,
        version_info: List[Dict[str, Union[str, int]]],
        execution_date: Union[datetime, str],
    ):
        """Stores FAFS metadata for a given month."""
        MetadataUtil.store_pgmeta(
            dag_id=dag_id,
            metadata_key=FAFS_METADATA_KEY,
            metadata_value=month,
            execution_date=execution_date,
            json_data={FAFS_METADATA_KEY: version_info},
        )




    @staticmethod
    def get_pgmeta_fafs_versions(
        dag_id: str,
        month: Optional[str] = None,
    ) -> Dict[str, List[Dict[str, Union[str, int]]]]:
        """Retrieves FAFS versions for a specific month or the latest."""
        return MetadataUtil.get_pgmeta(
            dag_id, FAFS_METADATA_KEY, mode="fafs_month", month=month
        )




    @staticmethod
    def get_pgmeta_fafs_max_versions(
        dag_id: str,
        parname: Optional[str] = None,
    ) -> Union[Dict[str, Dict[str, Union[str, int]]], Dict[str, Union[str, int]]]:
        """Retrieves maximum FAFS versions, optionally for a specific parname."""
        return MetadataUtil.get_pgmeta(
            dag_id, FAFS_METADATA_KEY, mode="fafs_max", parname=parname
        )




    @staticmethod
    def get_pgmeta_fafs_latest_versions(
        dag_id: str,
        parname: Optional[str] = None,
    ) -> Union[Dict[str, Dict[str, Union[str, int]]], Dict[str, Union[str, int]]]:
        """Retrieves latest FAFS versions from the most recent record, optionally for a specific parname."""
        return MetadataUtil.get_pgmeta(
            dag_id, FAFS_METADATA_KEY, mode="fafs_latest", parname=parname
        )