"""
BuildRedactedTileSqlOperator
Builds a SELECT query from a lake table with PII/sensitive columns excluded
for public-facing tile generation.
Usage:
from operators.tile.build_redacted_tile_sql import BuildRedactedTileSqlOperator
build_sql = BuildRedactedTileSqlOperator(
task_id="build_public_tile_sql",
lake_table="c_cpc_merge_tile_ca",
)
# Use output in CreateTileOperator
create_tile = CreateTileOperator(
task_id="create_public_tile",
source_sql='{{ ti.xcom_pull(task_ids="build_public_tile_sql") }}',
...
)
"""
from typing import Dict, List, Optional
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
# Default columns to exclude from public tiles (PII and sensitive data)
DEFAULT_EXCLUDE_COLUMNS = [
"location_id",
"address_primary",
"city",
"state",
"zip",
"zip_suffix",
"unit_count",
"building_type_code",
"land_use_code",
"address_confidence_code",
"county_geoid",
"block_geoid",
"h3_9",
"pt_lon",
"pt_lat",
# "cai_category",
# "speedrank_no_cai",
]
# Default column renames for public tiles
# bsl_flag → is_cab: Security by obscurity to protect CQA's data in public tiles.
# See: https://readynet.slack.com/archives/C067JGS04FM/p1718292720928659?thread_ts=1718147815.565329&cid=C067JGS04FM
DEFAULT_COLUMN_RENAMES = {
"bsl_flag": "is_cab",
}
[docs]
class BuildRedactedTileSqlOperator(BaseOperator):
"""
Builds a SELECT query from a lake table with PII/sensitive columns excluded (current implementation EXCLUDE PII columns rather than REDACT text).
Queries the Athena information_schema to get column names, filters out
excluded columns, applies any column renames, and returns a SELECT statement
suitable for use with CreateTileOperator.
:param lake_table: Source lake table name (e.g., "c_cpc_merge_tile_ca")
:type lake_table: str
:param database: Athena database name. Defaults to Config.DATA_LAKE_DB
:type database: str
:param exclude_columns: List of column names to exclude from the SELECT.
Defaults to DEFAULT_EXCLUDE_COLUMNS (PII columns).
:type exclude_columns: List[str]
:param column_renames: Dict mapping original column names to new names.
Defaults to DEFAULT_COLUMN_RENAMES ({"bsl_flag": "is_cab"}).
:type column_renames: Dict[str, str]
:param additional_exclude_columns: Additional columns to exclude beyond the defaults.
Merged with exclude_columns.
:type additional_exclude_columns: List[str]
:return: SELECT SQL string via XCom
"""
ui_color = "#ed17ed"
ui_fgcolor = "#FFFFFF"
template_fields = ("lake_table", "database")
@apply_defaults
def __init__(
self,
lake_table: str,
database: str = None,
exclude_columns: List[str] = None,
column_renames: Dict[str, str] = None,
additional_exclude_columns: List[str] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.lake_table = lake_table
self.database = database
self.exclude_columns = (
exclude_columns if exclude_columns is not None else DEFAULT_EXCLUDE_COLUMNS
)
self.column_renames = (
column_renames if column_renames is not None else DEFAULT_COLUMN_RENAMES
)
self.additional_exclude_columns = additional_exclude_columns or []
[docs]
def execute(self, context) -> str:
import awswrangler as wr
import sqlparse
database = self.database or Config.DATA_LAKE_DB
# Combine exclude lists
all_excluded = set(self.exclude_columns + self.additional_exclude_columns)
# Query information_schema for column names
sql_get_columns = f"""
SELECT column_name
FROM information_schema.columns
WHERE table_schema = '{database.lower()}'
AND table_name = '{self.lake_table}'
"""
self.log.info(f"Querying columns for {database}.{self.lake_table}")
df_columns = wr.athena.read_sql_query(
sql=sql_get_columns,
database=database,
)
if df_columns.empty:
raise ValueError(
f"No columns found for table {database}.{self.lake_table}. "
"Verify the table exists."
)
# Filter and transform columns
columns = []
for col in df_columns["column_name"].tolist():
if col in all_excluded:
continue
if col in self.column_renames:
columns.append(f"{col} as {self.column_renames[col]}")
else:
columns.append(col)
if not columns:
raise ValueError(
f"All columns were excluded from {self.lake_table}. "
"Check exclude_columns configuration."
)
# Build SELECT statement
select_sql = f"SELECT {', '.join(columns)} FROM {self.lake_table}"
self.log.info(f"Built redacted SQL with {len(columns)} columns:\n{columns}")
formatted_sql = sqlparse.format(select_sql, reindent=True, keyword_case="upper")
self.log.info(f"SQL:\n{formatted_sql}")
return select_sql