"""
Load Fabric Address Full Text Search Operator
This module provides an operator for loading FCC Fabric Address data into a
Postgres full text search table. It handles the complete workflow including
tenant resolution, fabric version lookup, and dynamic SQL selection.
The operator extends TableToPGOperator to leverage existing transfer and
hot-swap functionality while encapsulating domain-specific logic for:
1. State abbreviation to tenant resolution
2. Fabric version lookup (bead profile, ULFW override, or ARC override)
3. ECS instance sizing based on state data volume
4. First-time vs update SQL selection based on table existence
Usage:
from operators.postgres.load_fabric_address_full_text_search.operator import LoadFabricAddressFullTextSearchOperator
load_fabric_address = LoadFabricAddressFullTextSearchOperator(
task_id="load_fabric_address_fts",
state_abbr="CA",
dest_db="boss_grapi_stage",
ulfw_override=True,
)
"""
from __future__ import annotations
from pathlib import Path
from typing import Any, Dict, Literal, Sequence
from jinja2 import Template
from config import Config
from util import Util
from operators.datalake_to_pg import TableToPGOperator
################################################################################
## Constants
################################################################################
SQL_DIR = Path(__file__).parent / "sql"
TGT_TABLE = "pit_fcc_addresses_full_text_search"
LARGE_STATES_ABBR = ["FL", "TX", "OH", "IN", "NY", "PA"]
HUGE_STATES_ABBR = ["CA", "VA", "TX"]
CPC_FABRIC_VERSION = {"VA": "2", "MI": "3"}
ULFW_FABRIC_VERSION = {"CA": "6", "OK": "6", "TX": "6", "MS": "6", "AZ": "6"}
LANGUAGE_CONFIG = {"GU": "simple"}
SQL_TEMPLATE_DIR = "operators/postgres/load_fabric_address_full_text_search/sql"
POST_TRANSFER_SQL = """
CREATE INDEX idx_{{uuid}}_address_street_address ON {{schema}}.{{stage_table_name}} USING gin (to_tsvector('{{language}}', address_primary));
CREATE INDEX idx_{{uuid}}_address_full_addr ON {{schema}}.{{stage_table_name}} USING gin (full_address);
CREATE INDEX idx_{{uuid}}_address_city ON {{schema}}.{{stage_table_name}} (city);
CREATE INDEX idx_{{uuid}}_address_zip_code ON {{schema}}.{{stage_table_name}} (zip_code);
CREATE INDEX idx_{{uuid}}_address_location_id ON {{schema}}.{{stage_table_name}} (location_id);
"""
################################################################################
## Operator
################################################################################
[docs]
class LoadFabricAddressFullTextSearchOperator(TableToPGOperator):
"""
Load FCC Fabric Address data into a Postgres full text search table.
Extends TableToPGOperator with domain-specific logic for fabric address
data. Automatically resolves tenant from state abbreviation, determines
appropriate fabric version, and selects ECS instance size based on state
data volume.
:param state_abbr: Two-letter state abbreviation (e.g., "CA", "TX").
:param dest_db: Destination Postgres connection ID.
:param ulfw_override: If True, use ULFW fabric version (6) for applicable states.
:param arc_override: If True, source fabric version from fabric_version_arc in bead profile.
:param dry_run: If True, operator exits before doing any work.
:param tenant_override: If non-empty, use this tenant schema instead of
deriving from state_abbr. Required for program-specific schemas
(e.g., CASF's `california_casf` where state=CA but tenant != "california").
:param program_id_override: If non-empty, use this program_id in the
FTS SQL filter instead of the hardcoded "Bead". Required for non-BEAD
programs (e.g., "BIG" for CASF).
:param program_stage_id_override: If non-empty, use this program_stage_id
in the FTS SQL filter instead of the value derived from override flags
(ARC / BOB / ACCEPTED). Required for non-BEAD stages (e.g., "BIG_2026").
:param fabric_version_override: If non-empty, use this fabric version in
the FTS SQL filter instead of looking it up via the bead profile.
Required when the FTS run should match a fabric version that doesn't
correspond to the state's bead profile entry (e.g., CASF on v6).
:param source_table_override: If non-empty, use this Athena table as the
source for FTS curation instead of the default `c_fabric_sbo_v{version}`.
Table must expose the columns curate_data*.sql selects (`location_id`,
`address_primary`, `city`, `state`, `zip`, `latitude`, `longitude`,
`stateabbr`). Needed for CASF, where FTS should read from the CASF
merge tile so `location_id` carries the swapped `casf_id` value
(DATASME-1908).
"""
template_fields: Sequence[str] = (
"state_abbr",
"dest_db",
"ulfw_override",
"arc_override",
"dry_run",
"tenant_override",
"program_id_override",
"program_stage_id_override",
"fabric_version_override",
"source_table_override",
)
def __init__(
self,
state_abbr: str,
dest_db: str,
ulfw_override: bool = False,
arc_override: bool = False,
dry_run: bool = False,
tenant_override: str = "",
program_id_override: str = "",
program_stage_id_override: str = "",
fabric_version_override: str = "",
source_table_override: str = "",
*args,
**kwargs,
):
self.state_abbr = state_abbr
self.ulfw_override = ulfw_override
self.arc_override = arc_override
self.dry_run = dry_run
self.tenant_override = tenant_override
self.program_id_override = program_id_override
self.program_stage_id_override = program_stage_id_override
self.fabric_version_override = fabric_version_override
self.source_table_override = source_table_override
super().__init__(
dest_db=dest_db,
dest_table=TGT_TABLE,
copy_table_sql="",
schema="public",
create_table_sql="create_table.sql",
post_transfer_sql=POST_TRANSFER_SQL,
template=SQL_TEMPLATE_DIR,
*args,
**kwargs,
)
############################################################################
## Overrides
############################################################################
[docs]
def read_sql_template(self, sql: str, params: Dict[str, str] = {}) -> str:
"""Override to use SQL files from operator directory instead of DAG_ROOT."""
if sql.lower().endswith(".sql"):
sql_file_path = SQL_DIR / sql
self.log.info(f"Opening SQL file {sql_file_path}")
sql = sql_file_path.read_text()
sql_jinja_template = Template(sql)
sql_to_run = sql_jinja_template.render(params)
self.log.info(f"{sql_to_run=}")
return sql_to_run
############################################################################
## Helpers
############################################################################
def _resolve_tenant(self, state_abbr: str) -> str:
"""Resolve state abbreviation to tenant name."""
tenant_statefips = Util.tenant_to_statefips()
for tenant, fips in tenant_statefips.items():
state_for_tenant = next(
(
abbr
for abbr, state_fips in Config.ALL_STATE_FIPS.items()
if state_fips == fips
),
None,
)
if state_for_tenant == state_abbr:
return tenant
raise ValueError(f"Could not resolve tenant for state: {state_abbr}")
def _get_fabric_version(self, state_abbr: str) -> str:
"""Determine fabric version based on overrides and bead profile."""
state_profile = Util.bead_profile(state_abbr, check_pre=False)
if self.arc_override:
fabric_version = state_profile.get("fabric_version_arc")
if fabric_version is None:
raise ValueError(
f"fabric_version_arc not found for state {state_abbr}. Check state profile."
)
elif self.ulfw_override and ULFW_FABRIC_VERSION.get(state_abbr.upper()):
fabric_version = ULFW_FABRIC_VERSION.get(state_abbr.upper())
else:
fabric_version = state_profile.get(
"fabric_version", CPC_FABRIC_VERSION.get(state_abbr.upper())
)
if fabric_version is None:
raise ValueError(
f"Fabric version not found for state {state_abbr}. Check state profile."
)
return fabric_version
def _get_ecs_size(self, state_abbr: str) -> Literal["", "large", "xlarge"]:
"""Determine ECS instance size based on state data volume."""
if state_abbr in HUGE_STATES_ABBR:
return "xlarge"
elif state_abbr in LARGE_STATES_ABBR:
return "large"
return ""
def _get_grapi_athena_alias(self) -> str:
"""Determine the appropriate GRAPI Athena federated connection alias."""
if "stage" in self.dest_db.lower():
return Config.GRAPI_ATHENA_STAGE
return Config.GRAPI_ATHENA
def _get_grapi_pg_connection(self) -> str:
"""Determine the appropriate GRAPI PostgreSQL connection ID."""
if "stage" in self.dest_db.lower():
return Config.GRAPI_PG_STAGE
return Config.GRAPI_PG
def _table_exists_in_tenant(self, tenant: str, grapi_db: str) -> bool:
"""Check if the target table already exists in the tenant schema."""
import awswrangler as wr
from botocore.exceptions import ClientError
sql = f"SELECT * FROM {grapi_db}.{tenant}.{TGT_TABLE} LIMIT 1"
try:
wr.athena.read_sql_query(sql, database=Config.DATA_LAKE_DB)
return True
except ClientError as e:
msg = str(e)
if "does not exist" in msg or "Table not found" in msg:
return False
self.log.error(
"Error checking table existence for %s.%s: %s", tenant, TGT_TABLE, e
)
raise
except Exception as e:
error_msg = str(e)
if "does not exist" in error_msg or "Table not found" in error_msg:
return False
self.log.error("Unexpected error checking table existence: %s", e)
raise
def _get_program_stage_id(self) -> str:
"""Determine program stage ID based on overrides."""
if self.arc_override:
return "ARC"
elif self.ulfw_override:
return "BOB"
return "ACCEPTED"
############################################################################
## Execute
############################################################################
[docs]
def execute(self, context: Any) -> Any:
if self.ulfw_override and self.arc_override:
raise ValueError("Cannot set both ULFW and ARC overrides at the same time")
for variable in ["partitions", "copy_partitions", "dest_db"]:
self.resolve_self_x_com_attributes(variable, context=context)
if self.dry_run:
self.log.info("Dry run enabled - skipping execution")
return None
state_abbr = self.state_abbr.upper()
# Override-first lookups — empty overrides fall back to existing
# state-derived defaults (no behavior change for non-overridden flows).
# `str(...)` coerces because Airflow Jinja may pass typed Param values
# (e.g. int 6 instead of "6") through to template_fields unchanged.
tenant = str(self.tenant_override or "").strip() or self._resolve_tenant(state_abbr)
fabric_version = str(self.fabric_version_override or "").strip() or self._get_fabric_version(state_abbr)
program_id = str(self.program_id_override or "").strip() or "Bead"
program_stage_id = str(self.program_stage_id_override or "").strip() or self._get_program_stage_id()
source_table = str(self.source_table_override or "").strip() or f"c_fabric_sbo_v{fabric_version}"
grapi_athena_alias = self._get_grapi_athena_alias()
grapi_pg_conn = self._get_grapi_pg_connection()
table_exists = self._table_exists_in_tenant(tenant, grapi_athena_alias)
copy_sql_file = (
"curate_data.sql" if table_exists else "curate_data_first_time.sql"
)
copy_sql_file = "curate_data_demo_tenant.sql" if tenant == "demo_base" else copy_sql_file
language = LANGUAGE_CONFIG.get(state_abbr, "english")
self.log.info(
"Loading Fabric Address FTS for state=%s tenant=%s fabric_version=%s table_exists=%s",
state_abbr,
tenant,
fabric_version,
table_exists,
)
self.schema = tenant
self.dest_db = grapi_pg_conn
self.copy_table_sql = copy_sql_file
self.task_suffix = self._get_ecs_size(state_abbr)
self.stage_table_name = f"{self.dest_table}_stage"
self.map_name = f"{grapi_pg_conn}.{self.schema}.{self.dest_table}"
self.query_params.update(
{
"grapi_link": grapi_athena_alias,
"tenant": tenant,
"state_abbr": state_abbr,
"fabric_version": fabric_version,
"program_id": program_id,
"program_stage_id": program_stage_id,
"language": language,
"schema": tenant,
"stage_table_name": self.stage_table_name,
"source_table": source_table,
}
)
return super().execute(context)