"""
Point-in-Time Ingestion Operators
This module provides operators for ingesting point-in-time (PIT) data from the
Athena data lake into PostgreSQL (Boss Platform API). Each program stage has
its own operator that handles stage-specific transformations while sharing
common helper functions for the core ETL logic.
Usage:
from operators.point_in_time.ingest import PitIngestArcOperator
ingest_arc = PitIngestArcOperator(
task_id="ingest_arc_pit",
dest_db="BOSS_PLATFORM_DB",
schema="virginia",
lake_table="c_arc_merge_tile_va_mission_control",
program_id="Bead",
program_stage="ARC",
fabric_version=6,
fixed_availability_version=99,
reference_date="2025-01-15",
)
"""
import logging
from abc import abstractmethod
from pathlib import Path
from typing import List, Literal
from airflow.exceptions import AirflowConfigException
from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.utils.context import Context
from jinja2 import Template
from psycopg2 import sql as psql
from config import Config
from util import Util
log = logging.getLogger(__name__)
################################################################################
## Constants
################################################################################
SQL_DIR = Path(__file__).parent / "SQL" / "ingest"
COMMON_SQL_DIR = SQL_DIR / "common"
STAGES_SQL_DIR = SQL_DIR / "stages"
TEMP_TABLE = "temp_pit_address_metadata"
DEST_TABLE_PIT = "pit_address_metadata"
DEST_TABLE_ADDRESS = "pit_fabric_addresses"
DEFAULT_LARGE_STATES = ["california", "new_york", "texas", "ohio"]
################################################################################
## Helper Functions - SQL Rendering
################################################################################
def _render_sql(sql_path: Path, variables: dict) -> str:
"""Render a Jinja SQL template with the provided variables."""
return Template(sql_path.read_text()).render(variables)
################################################################################
## Helper Functions - Temp Table Management
################################################################################
def _drop_temp_table(hook: PostgresHook, schema: str, table: str = TEMP_TABLE) -> None:
"""Drop a temporary table if it exists."""
query = psql.SQL("DROP TABLE IF EXISTS {}.{}").format(
psql.Identifier(schema), psql.Identifier(table)
)
hook.run(query.as_string(hook.get_conn()), autocommit=True)
def _temp_table_exists(
hook: PostgresHook, schema: str, table: str = TEMP_TABLE
) -> bool:
"""Check if a temporary table exists in the given schema."""
query = """
SELECT EXISTS (
SELECT 1 FROM information_schema.tables
WHERE table_schema = %s AND table_name = %s
)
"""
result = hook.get_first(query, parameters=(schema, table))
return result[0] if result else False
def _add_attributes_column_and_index(
hook: PostgresHook, schema: str, table: str = TEMP_TABLE
) -> None:
"""
Convert attributes_text (VARCHAR) to attributes (JSON) column and add index.
Required because Athena exports JSON as text which must be cast in PG.
"""
query = psql.SQL(
"""
ALTER TABLE {schema}.{table} ADD COLUMN IF NOT EXISTS attributes json;
UPDATE {schema}.{table} SET attributes=cast(attributes_text AS json);
ALTER TABLE {schema}.{table} DROP COLUMN attributes_text;
CREATE INDEX ON {schema}.{table} (md5_pkey_hash);
"""
).format(schema=psql.Identifier(schema), table=psql.Identifier(table))
hook.run(query.as_string(hook.get_conn()), autocommit=True)
################################################################################
## Helper Functions - Data Transfer
################################################################################
def _transfer_to_temp_table(
query: str,
dest_db: str,
schema: str,
table: str = TEMP_TABLE,
task_suffix: Literal["", "large", "xlarge"] = "",
) -> None:
"""
Transfer data from Athena to a PostgreSQL temp table via S3.
Uses Util.transfer_datalake_db which handles the ECS task for large transfers.
"""
temp_s3_path = f"temp/{table}"
Util.transfer_datalake_db(
table_name=f"{schema}.{table}",
temp_key=temp_s3_path,
tgt_query=query,
dest_db=dest_db,
drop_table=False,
task_suffix=task_suffix,
use_pg_copy=True,
)
################################################################################
## Helper Functions - PIT Metadata Upsert
################################################################################
def _upsert_pit_metadata(hook: PostgresHook, schema: str) -> None:
"""
Upsert records from temp table into pit_address_metadata.
Uses ON CONFLICT to update existing records by (md5_pkey_hash, md5_pit_hash).
"""
variables = {
"schema_name": schema,
"temp_table": TEMP_TABLE,
"dest_table": DEST_TABLE_PIT,
"address_table": DEST_TABLE_ADDRESS,
}
sql = _render_sql(COMMON_SQL_DIR / "upsert_pit_metadata.sql", variables)
hook.run(sql, autocommit=True)
################################################################################
## Helper Functions - Address Sync
################################################################################
def _sync_addresses(
hook: PostgresHook,
dest_db: str,
schema: str,
lake_table: str,
athena_alias: str,
task_suffix: Literal["", "large", "xlarge"] = "",
) -> None:
"""
Sync new/changed addresses from lake table to pit_fabric_addresses.
Workflow:
1. Extract addresses from lake that don't exist or differ in PG
2. Load to temp table via S3
3. Upsert into pit_fabric_addresses
4. Drop temp table
"""
variables = {
"schema_name": schema,
"lake_table": lake_table,
"athena_alias": athena_alias,
"address_table": DEST_TABLE_ADDRESS,
}
extract_sql = _render_sql(COMMON_SQL_DIR / "extract_address_updates.sql", variables)
_drop_temp_table(hook, schema)
_transfer_to_temp_table(
query=extract_sql,
dest_db=dest_db,
schema=schema,
task_suffix=task_suffix,
)
if not _temp_table_exists(hook, schema):
log.info("No address updates to sync - temp table was not created")
return
index_query = psql.SQL("CREATE INDEX ON {}.{} (md5_pkey_hash)").format(
psql.Identifier(schema), psql.Identifier(TEMP_TABLE)
)
hook.run(index_query.as_string(hook.get_conn()), autocommit=True)
upsert_variables = {
"schema_name": schema,
"temp_table": TEMP_TABLE,
"dest_table": DEST_TABLE_ADDRESS,
}
upsert_sql = _render_sql(COMMON_SQL_DIR / "upsert_address.sql", upsert_variables)
hook.run(upsert_sql, autocommit=True)
_drop_temp_table(hook, schema)
################################################################################
## Helper Functions - Cleanup & Stats
################################################################################
def _cleanup_mismatched_versions(
hook: PostgresHook,
schema: str,
program_id: str,
program_stage: str,
fabric_version: int,
fixed_availability_version: int,
) -> None:
"""
Remove the PIT records this load supersedes: same program, program stage and
fabric version, at a different FFA version.
A permutation is (program_id, program_stage_id, fabric_version,
fixed_availability_version), and several can coexist per program and stage.
The app picks one by program, stage and fabric version but not FFA version,
so two FFA versions under one fabric version would show up as duplicate
locations. Other fabric versions and other programs are separate datasets
and are kept (DATASME-2150: MS has Bead/ARC at fabric 6 and fabric 9).
"""
query = psql.SQL(
"""
DELETE FROM {schema}.{table}
WHERE program_id = %s
AND program_stage_id = %s
AND fabric_version = %s
AND fixed_availability_version != %s
"""
).format(schema=psql.Identifier(schema), table=psql.Identifier(DEST_TABLE_PIT))
hook.run(
query.as_string(hook.get_conn()),
parameters=(
program_id,
program_stage,
fabric_version,
fixed_availability_version,
),
autocommit=True,
)
def _update_table_stats(hook: PostgresHook, schema: str) -> None:
"""Run ANALYZE on PIT tables to update query planner statistics."""
for table in [DEST_TABLE_PIT, DEST_TABLE_ADDRESS]:
query = psql.SQL("ANALYZE {}.{}").format(
psql.Identifier(schema), psql.Identifier(table)
)
hook.run(query.as_string(hook.get_conn()), autocommit=True)
def _drop_program_stage_records(
hook: PostgresHook, schema: str, program_id: str, program_stage: str
) -> None:
"""
Delete all existing records for a program/stage combination.
Used when drop_existing_stage=True to fully replace stage data.
"""
query = psql.SQL(
"""
DELETE FROM {schema}.{table}
WHERE program_id = %s AND program_stage_id = %s
"""
).format(schema=psql.Identifier(schema), table=psql.Identifier(DEST_TABLE_PIT))
hook.run(
query.as_string(hook.get_conn()),
parameters=(program_id, program_stage),
autocommit=True,
)
################################################################################
## Base Operator
################################################################################
[docs]
class PitIngestBaseOperator(BaseOperator):
"""
Base operator for Point-in-Time ingestion. Handles common ETL logic while
delegating stage-specific SQL generation to subclasses.
Subclasses must implement `get_extract_sql_path()` to return the path to
their stage-specific extraction SQL template.
Workflow:
1. Optionally drop existing program stage records
2. Sync addresses from lake to pit_fabric_addresses (optional)
3. Extract PIT data from lake using stage-specific SQL
4. Load to temp table, convert attributes column
5. Upsert into pit_address_metadata
6. Cleanup superseded FFA versions of this program/stage/fabric (optional)
7. Update table statistics
:param dest_db: Config attribute name for destination DB (e.g., "BOSS_PLATFORM_DB")
:type dest_db: str
:param schema: Tenant schema name (e.g., "virginia", "arkansas")
:type schema: str
:param lake_table: Source merge tile table in Athena (e.g., "c_arc_merge_tile_va")
:type lake_table: str
:param program_id: Program identifier (e.g., "Bead", "Baseline")
:type program_id: str
:param program_stage: Program stage (e.g., "ARC", "BOB", "ACCEPTED")
:type program_stage: str
:param fabric_version: FCC fabric version used in source data
:type fabric_version: int
:param fixed_availability_version: Fixed availability version used in source data
:type fixed_availability_version: int
:param athena_alias: Cross-DB alias for Athena queries. Defaults to Config.ATHENA_GRAPI_DB_ALIAS
:type athena_alias: str
:param sub_program_id: Sub-program identifier if applicable
:type sub_program_id: str
:param reference_date: Program reference date (YYYY-MM-DD format). Required.
:type reference_date: str
:param sync_addresses: Whether to sync pit_fabric_addresses. Defaults to True
:type sync_addresses: bool
:param drop_existing_stage: If True, delete existing stage records before insert. Defaults to False
:type drop_existing_stage: bool
:param cleanup_mismatched_versions: If True, remove records of this program, stage and fabric version at a different FFA version. Other fabric versions are kept. Defaults to True
:type cleanup_mismatched_versions: bool
:param extract_sql_override: Path to custom extraction SQL template. Overrides stage default
:type extract_sql_override: str
:param task_suffix: ECS task size suffix ("", "large", "xlarge"). Auto-detected if not set
:type task_suffix: str
:param large_states: List of state schemas that require "large" task suffix. Defaults to DEFAULT_LARGE_STATES
:type large_states: List[str]
"""
ui_color = "#8B5CF6"
ui_fgcolor = "#FFFFFF"
template_fields = (
"dest_db",
"schema",
"lake_table",
"program_id",
"sub_program_id",
"program_stage",
"reference_date",
"athena_alias",
"fabric_version",
"fixed_availability_version",
"drop_existing_stage",
"task_suffix",
"attributes_sql",
"pre_ineligible_sql",
)
def __init__(
self,
dest_db: str,
schema: str,
lake_table: str,
program_id: str,
program_stage: str,
fabric_version: int,
fixed_availability_version: int,
reference_date: str,
athena_alias: str = None,
sub_program_id: str = None,
sync_addresses: bool = True,
drop_existing_stage: bool = False,
cleanup_mismatched_versions: bool = True,
extract_sql_override: str = None,
attributes_sql: str = "",
pre_ineligible_sql: str = "False",
task_suffix: Literal["", "large", "xlarge"] = "",
large_states: List[str] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.dest_db = dest_db
self.schema = schema
self.lake_table = lake_table
self.program_id = program_id
self.program_stage = program_stage
self.fabric_version = fabric_version
self.fixed_availability_version = fixed_availability_version
self.reference_date = reference_date
self.athena_alias = athena_alias
self.sub_program_id = sub_program_id or ""
self.sync_addresses = sync_addresses
self.drop_existing_stage = drop_existing_stage
self.cleanup_mismatched_versions = cleanup_mismatched_versions
self.extract_sql_override = extract_sql_override
# SQL expression rendered into the extract template as the value of
# `attributes_text`. Empty / unset falls back to the JSON-empty-object
# literal `'{}'`, preserving the operator's original behavior. Callers
# can pass a one-line SQL fragment (e.g. CASF concats casf_id + cbg_mhi)
# to write per-row attributes without forking the stage extract SQL.
self.attributes_sql = attributes_sql
# SQL expression rendered into the extract template's `is_pre_ineligible`
# CASE. Default `"False"` preserves the legacy behavior (flag never fires
# for tenants whose lake table doesn't carry the column). Tenants whose
# merge tile projects `is_pre_ineligible` (e.g. CASF post-DATASME-1893,
# CPC's STATE_WITH_PRE_ELIGIBLE) pass `"is_pre_ineligible"` to read the
# lake column through directly. Any boolean SQL expression is valid.
self.pre_ineligible_sql = pre_ineligible_sql
self.task_suffix = task_suffix
self.large_states = (
large_states if large_states is not None else DEFAULT_LARGE_STATES
)
############################################################################
## Abstract Methods - Subclasses Must Implement
############################################################################
############################################################################
## Template Variables
############################################################################
[docs]
def get_task_suffix(self) -> str:
"""Determine ECS task size. Large states get 'large' suffix by default."""
if self.task_suffix:
return self.task_suffix
return "large" if self.schema in self.large_states else ""
############################################################################
## Execute
############################################################################
[docs]
def execute(self, context: Context) -> dict:
# Coerce templated strings to proper types (Jinja renders everything as strings)
if isinstance(self.fabric_version, str):
self.fabric_version = int(self.fabric_version)
if isinstance(self.fixed_availability_version, str):
self.fixed_availability_version = int(self.fixed_availability_version)
if isinstance(self.drop_existing_stage, str):
self.drop_existing_stage = self.drop_existing_stage.lower() in (
"true",
"1",
"yes",
)
# Validate dest_db after template rendering
if not hasattr(Config, self.dest_db):
raise AirflowConfigException(
f"dest_db '{self.dest_db}' not found in Config. "
f"Valid options include: BOSS_PLATFORM_DB, BOSS_PLATFORM_STAGE_DB"
)
dest_db_conn = getattr(Config, self.dest_db)
hook = PostgresHook(postgres_conn_id=dest_db_conn)
task_suffix = self.get_task_suffix()
if self.athena_alias is None:
self.athena_alias = Config.ATHENA_GRAPI_DB_ALIAS
# Validate SQL path before starting
sql_path = (
Path(self.extract_sql_override)
if self.extract_sql_override
else self.get_extract_sql_path()
)
if not sql_path.exists():
raise FileNotFoundError(
f"Extract SQL template not found: {sql_path}. "
f"Ensure the file exists or check extract_sql_override parameter."
)
try:
# Step 1: Drop existing stage records if requested
if self.drop_existing_stage:
self.log.info("Dropping existing program stage records")
_drop_program_stage_records(
hook, self.schema, self.program_id, self.program_stage
)
# Step 2: Sync addresses from lake to PG
if self.sync_addresses:
self.log.info("Syncing addresses from lake to PG")
_sync_addresses(
hook=hook,
dest_db=dest_db_conn,
schema=self.schema,
lake_table=self.lake_table,
athena_alias=self.athena_alias,
task_suffix=task_suffix,
)
# Step 3: Extract PIT data from lake
self.log.info("Extracting PIT data from lake")
_drop_temp_table(hook, self.schema)
extract_sql = _render_sql(sql_path, self.get_extract_variables())
self.log.info(f"Extract SQL: {extract_sql}")
_transfer_to_temp_table(
query=extract_sql,
dest_db=dest_db_conn,
schema=self.schema,
task_suffix=task_suffix,
)
# Step 4: Convert attributes column and add index
self.log.info("Converting attributes column and adding index")
_add_attributes_column_and_index(hook, self.schema)
# Step 5: Upsert into pit_address_metadata
self.log.info("Upserting PIT metadata")
_upsert_pit_metadata(hook, self.schema)
# Step 6: Cleanup mismatched versions
if self.cleanup_mismatched_versions:
self.log.info("Cleaning up mismatched version records")
_cleanup_mismatched_versions(
hook=hook,
schema=self.schema,
program_id=self.program_id,
program_stage=self.program_stage,
fabric_version=self.fabric_version,
fixed_availability_version=self.fixed_availability_version,
)
# Step 7: Update table statistics
self.log.info("Updating table statistics")
_update_table_stats(hook, self.schema)
self.log.info("PIT ingestion complete")
finally:
# Ensure temp table is cleaned up even on failure
self.log.info("Cleaning up temp table")
_drop_temp_table(hook, self.schema)
return {
"schema": self.schema,
"lake_table": self.lake_table,
"program_id": self.program_id,
"program_stage": self.program_stage,
"fabric_version": self.fabric_version,
"fixed_availability_version": self.fixed_availability_version,
"reference_date": self.reference_date,
"dest_table_pit": f"{self.schema}.{DEST_TABLE_PIT}",
"dest_table_address": f"{self.schema}.{DEST_TABLE_ADDRESS}",
}
################################################################################
## ARC Operator
################################################################################
[docs]
class PitIngestArcOperator(PitIngestBaseOperator):
"""
Operator for ingesting ARC program stage data into Point-in-Time tables.
Uses extraction SQL at: SQL/ingest/stages/arc/extract_pit.sql
"""
custom_operator_name = "PitIngestArcOperator"