################################################################################
## THIS IS AN UNTESTED WORK IN PROGRESS DON'T USE YET - Travis
################################################################################
from functools import lru_cache
from typing import Optional, List, Dict, Any, Sequence, Set
import logging
import uuid
from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.providers.amazon.aws.hooks.athena import AthenaHook
from airflow.exceptions import AirflowException
from config import Config
log = logging.getLogger(__name__)
####################################################################################################
## Athena to PostgreSQL Type Mapping
####################################################################################################
ATHENA_TO_PG_TYPE_MAP = {
# Numeric types
"tinyint": "smallint",
"smallint": "smallint",
"int": "integer",
"integer": "integer",
"bigint": "bigint",
"float": "real",
"double": "double precision",
"decimal": "numeric",
# String types
"string": "text",
"char": "text",
"varchar": "text",
# Boolean
"boolean": "boolean",
# Date/time types
"date": "date",
"timestamp": "timestamp",
# Binary
"binary": "bytea",
"varbinary": "bytea",
# Complex types (stored as text, user can override)
"array": "text",
"map": "jsonb",
"struct": "jsonb",
"row": "jsonb",
# JSON
"json": "jsonb",
}
def _map_athena_type_to_pg(athena_type: str) -> str:
"""
Map an Athena data type to PostgreSQL equivalent.
Handles parameterized types like varchar(255), decimal(10,2), array<string>.
"""
athena_type = athena_type.lower().strip()
# Handle parameterized types
base_type = athena_type.split("(")[0].split("<")[0]
# Special handling for decimal with precision
if base_type == "decimal" and "(" in athena_type:
# decimal(10,2) -> numeric(10,2)
precision = athena_type.split("(")[1].rstrip(")")
return f"numeric({precision})"
# Special handling for varchar with length
if base_type in ("varchar", "char") and "(" in athena_type:
length = athena_type.split("(")[1].rstrip(")")
return f"varchar({length})"
return ATHENA_TO_PG_TYPE_MAP.get(base_type, "text")
####################################################################################################
## Lazy Load Imports for Heavy Libraries
####################################################################################################
@lru_cache(maxsize=None)
def _import_boto3():
"""
Lazy import of boto3.
Ensures it's only loaded when the operator executes, not when the DAG parses.
"""
try:
import importlib
return importlib.import_module("boto3")
except ModuleNotFoundError as e:
raise AirflowException(
f"Required AWS library missing: {e.name}. "
"Ensure boto3 is installed in the Airflow environment."
)
####################################################################################################
## Athena to Postgres Operator
####################################################################################################
[docs]
class AthenaToPostgresOperator(BaseOperator):
"""
Execute an Athena query and load the results into a PostgreSQL table.
Uses the ``aws_s3`` PostgreSQL extension for efficient loading.
Falls back to streaming COPY if configured (TODO).
Workflow:
1. Check if target table exists; if not and create_table_if_not_exists=True, continue.
2. Run Athena Query (SELECT).
3. If table doesn't exist, create it from Athena schema (with column_type_overrides).
4. Create staging table (with TEXT columns for special types).
5. Import S3 CSV into staging via aws_s3 extension.
6. INSERT from staging to target with type transformations.
7. Drop staging table and cleanup S3 output.
Array Handling:
Athena arrays are exported as ``[val1, val2]`` strings in CSV. This operator
auto-detects PostgreSQL array columns and transforms them during INSERT.
Use ``column_transforms`` to override auto-detection or handle complex types.
Note: The default array transform assumes comma-space (``, ``) delimiter between
elements. If your Athena data uses a different delimiter, provide a custom
transform via ``column_transforms``.
Geometry/Geography Handling:
PostGIS geometry and geography columns are auto-detected from the target table
schema. Values from Athena (WKT or hex-encoded EWKB) are transformed using
native ``::geometry`` / ``::geography`` casts with proper SRID handling:
- Use ``geometry_srid`` to set a default SRID for all geometry columns (e.g., 4326).
- Use ``geometry_srids`` dict for per-column SRID overrides.
- NULL/empty values are handled gracefully (no error on empty WKT).
- When using ``create_table_if_not_exists``, geometry columns MUST be specified
in ``column_type_overrides`` (e.g., {"geom": "geometry(Point, 4326)"}) since
Athena exports geometry as strings.
:param query: The Athena SQL query to execute.
:type query: str
:param target_table: Name of the PostgreSQL target table.
:type target_table: str
:param target_schema: PostgreSQL schema. Defaults to "public".
:type target_schema: str
:param postgres_conn_id: Airflow connection ID for PostgreSQL. Defaults to "postgres_default".
:type postgres_conn_id: str
:param aws_conn_id: Airflow connection ID for AWS. Defaults to "aws_default".
:type aws_conn_id: str
:param s3_bucket: S3 bucket for Athena query output.
:type s3_bucket: str
:param s3_prefix: S3 prefix for Athena query output.
:type s3_prefix: str
:param if_exists: How to handle existing data - "append", "replace", or "upsert". Defaults to "append".
:type if_exists: str
:param upsert_keys: Columns for upsert matching (required if if_exists="upsert").
:type upsert_keys: Optional[List[str]]
:param column_transforms: Dict of {column_name: SQL_expression}. Use {col} placeholder
for the column reference. Overrides auto-detected array transforms.
Example: {"tags": "string_to_array({col}, ',')::text[]"}
:type column_transforms: Optional[Dict[str, str]]
:param create_table_if_not_exists: If True, create target table from Athena query schema
when it doesn't exist. Defaults to False.
:type create_table_if_not_exists: bool
:param column_type_overrides: Dict of {column_name: postgresql_type}. Override inferred
PostgreSQL types when auto-creating table. Example: {"id": "uuid", "geom": "geometry"}.
Note: Geometry/geography columns MUST be specified here since Athena stores them as
strings; without override they become text and the geometry cast won't be applied.
:type column_type_overrides: Optional[Dict[str, str]]
:param geometry_srid: Default SRID for all geometry columns. Common values: 4326 (WGS84),
3857 (Web Mercator). If None, no SRID is specified in ST_GeomFromText(). Defaults to None.
:type geometry_srid: Optional[int]
:param geometry_srids: Dict of {column_name: srid} for per-column SRID overrides.
Takes precedence over geometry_srid. Example: {"location": 4326, "boundary": 3857}.
:type geometry_srids: Optional[Dict[str, int]]
:param treat_empty_as_null: If True, treat empty CSV fields as NULL during import.
Matches Athena's NULL representation in CSV output. Defaults to True.
:type treat_empty_as_null: bool
:param use_aws_s3_extension: Use aws_s3 extension. Defaults to True.
:type use_aws_s3_extension: bool
:param cleanup_s3: Delete S3 output after load. Defaults to True.
:type cleanup_s3: bool
:param region_name: AWS region. Defaults to "us-east-1".
:type region_name: str
:param workgroup: Athena workgroup. Defaults to "primary".
:type workgroup: str
:param dry_run: If 'true', no changes will be made to the target database.
Defaults to 'false'. We are using str here to allow templating.
:type dry_run: str
:param include_columns: List of column names to include in the target table insert.
If specified, only these columns will be inserted, allowing Prisma/DB defaults
to handle auto-generated columns (id, createdAt, etc.). The Athena query must
return columns in the same order as this list. If None, all columns are included.
:type include_columns: Optional[List[str]]
"""
template_fields: Sequence[str] = (
"query",
"target_table",
"target_schema",
"s3_bucket",
"s3_prefix",
"postgres_conn_id",
"aws_conn_id",
"dry_run",
"include_columns",
)
ui_color = "#C0FFEE"
def __init__(
self,
*,
query: str,
target_table: str,
target_schema: str = "public",
postgres_conn_id: str = "postgres_default",
aws_conn_id: str = "aws_default",
s3_bucket: str,
s3_prefix: str,
if_exists: str = "append", # append, replace, upsert
upsert_keys: Optional[List[str]] = None, # columns to check for uniqueness
column_transforms: Optional[
Dict[str, str]
] = None, # custom SQL transforms per column
create_table_if_not_exists: bool = False,
column_type_overrides: Optional[
Dict[str, str]
] = None, # override inferred PG types when creating table
geometry_srid: Optional[int] = None, # default SRID for geometry columns
geometry_srids: Optional[
Dict[str, int]
] = None, # per-column SRID overrides
treat_empty_as_null: bool = True,
use_aws_s3_extension: bool = True,
cleanup_s3: bool = True,
region_name: str = "us-east-1",
workgroup: str = "primary",
dry_run: str = "false",
include_columns: Optional[List[str]] = None,
skip_if_source_missing: bool = False,
**kwargs,
):
super().__init__(**kwargs)
self.query = query
self.target_table = target_table
self.target_schema = target_schema
self.postgres_conn_id = postgres_conn_id
self.aws_conn_id = aws_conn_id
self.s3_bucket = s3_bucket
self.s3_prefix = s3_prefix
self.if_exists = if_exists
self.upsert_keys = upsert_keys
self.column_transforms = column_transforms or {}
self.create_table_if_not_exists = create_table_if_not_exists
self.column_type_overrides = column_type_overrides or {}
self.geometry_srid = geometry_srid
self.geometry_srids = geometry_srids or {}
self.treat_empty_as_null = treat_empty_as_null
self.use_aws_s3_extension = use_aws_s3_extension
self.cleanup_s3 = cleanup_s3
self.region_name = region_name
self.workgroup = workgroup
self.dry_run = dry_run
self.include_columns = include_columns
self.skip_if_source_missing = skip_if_source_missing
################################################################################################
## Athena Query Execution
################################################################################################
def _run_athena_query(self) -> str:
"""
Execute the Athena query and return the Query Execution ID.
"""
hook = AthenaHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
# Ensure prefix ends with /
prefix = (
self.s3_prefix if self.s3_prefix.endswith("/") else f"{self.s3_prefix}/"
)
output_location = f"s3://{self.s3_bucket}/{prefix}"
log.info(f"Executing Athena query. Output location: {output_location}")
query_id = hook.run_query(
self.query,
query_context={"Database": Config.DATA_LAKE_DB}, # Context is flexible
result_configuration={"OutputLocation": output_location},
workgroup=self.workgroup,
)
# Wait for completion and verify success
status = hook.poll_query_status(query_id)
if status not in ("SUCCEEDED",):
# Fetch error detail from Athena
boto3 = _import_boto3()
client = boto3.client("athena", region_name=self.region_name)
exec_resp = client.get_query_execution(QueryExecutionId=query_id)
reason = (
exec_resp.get("QueryExecution", {})
.get("Status", {})
.get("StateChangeReason", "Unknown")
)
if self.skip_if_source_missing and "TABLE_NOT_FOUND" in reason.upper():
log.info(
f"Athena source table not found (query {query_id}). "
f"Skipping due to skip_if_source_missing=True. Reason: {reason}"
)
return None
raise AirflowException(
f"Athena query {query_id} failed with status: {status}. Reason: {reason}"
)
return query_id
def _get_athena_query_schema(self, query_id: str) -> List[Dict[str, str]]:
"""
Fetch column names and types from Athena query results metadata.
Returns list of {"name": col_name, "athena_type": type_string}.
"""
boto3 = _import_boto3()
client = boto3.client("athena", region_name=self.region_name)
response = client.get_query_results(
QueryExecutionId=query_id,
MaxResults=1, # We only need metadata, not data
)
columns = []
for col_info in response["ResultSet"]["ResultSetMetadata"]["ColumnInfo"]:
columns.append(
{
"name": col_info["Name"],
"athena_type": col_info["Type"],
}
)
return columns
def _create_target_table(
self,
pg_hook: PostgresHook,
athena_schema: List[Dict[str, str]],
) -> None:
"""
Create the target table based on Athena query schema.
Applies column_type_overrides for user-specified type mappings.
"""
col_defs = []
for col in athena_schema:
col_name = col["name"]
quoted_col = f'"{col_name}"'
if col_name in self.column_type_overrides:
pg_type = self.column_type_overrides[col_name]
else:
pg_type = _map_athena_type_to_pg(col["athena_type"])
col_defs.append(f"{quoted_col} {pg_type}")
col_list = ", ".join(col_defs)
target_fqn = f'"{self.target_schema}"."{self.target_table}"'
log.info(f"Creating target table {target_fqn}...")
pg_hook.run(f"CREATE TABLE {target_fqn} ({col_list});")
################################################################################################
## Schema Introspection & Column Transforms
################################################################################################
def _load_column_metadata(self, pg_hook: PostgresHook) -> None:
"""
Load all column metadata in a single query and cache results.
Populates: _columns, _array_columns, _enum_columns, _geometry_columns,
_geometry_column_srids, _special_columns.
"""
sql = """
SELECT
c.column_name,
c.data_type,
c.udt_name,
COALESCE(t.typtype, '') AS typtype
FROM information_schema.columns c
LEFT JOIN pg_type t ON t.typname = c.udt_name
WHERE c.table_schema = %s AND c.table_name = %s
ORDER BY c.ordinal_position;
"""
results = pg_hook.get_records(sql, (self.target_schema, self.target_table))
self._columns: List[Dict[str, Any]] = []
self._array_columns: Dict[str, str] = {}
self._enum_columns: Dict[str, str] = {}
self._geometry_columns: Dict[str, str] = {}
self._geometry_column_srids: Dict[str, int] = {}
self._json_columns: Set[str] = set()
for row in results:
col_name, data_type, udt_name, typtype = row
if any([c['name'] == col_name for c in self._columns]):
continue # skip duplicate columns (enums can appear once per value)
self._columns.append(
{
"name": col_name,
"data_type": data_type,
"udt_name": udt_name,
}
)
if data_type == "ARRAY":
self._array_columns[col_name] = udt_name.lstrip("_")
elif typtype == "e":
self._enum_columns[col_name] = udt_name
elif udt_name in ("geometry", "geography"):
self._geometry_columns[col_name] = udt_name
elif udt_name in ("json", "jsonb"):
self._json_columns.add(col_name)
# Collect all enum type names in the schema for array cast qualification
enum_results = pg_hook.get_records(
"SELECT typname FROM pg_type t JOIN pg_namespace n ON t.typnamespace = n.oid "
"WHERE n.nspname = %s AND t.typtype = 'e';",
(self.target_schema,),
)
self._enum_type_names: Set[str] = {r[0] for r in enum_results} if enum_results else set()
if self._geometry_columns:
self._load_geometry_srids(pg_hook)
self._special_columns: Set[str] = set()
self._special_columns.update(self._array_columns.keys())
self._special_columns.update(self._enum_columns.keys())
self._special_columns.update(self._geometry_columns.keys())
self._special_columns.update(self._json_columns)
self._special_columns.update(self.column_transforms.keys())
log.info(f"Special columns: {list(self._special_columns)}")
def _load_geometry_srids(self, pg_hook: PostgresHook) -> None:
"""
Query geometry_columns view to get SRID for each geometry/geography column.
Falls back to 0 (unknown) if not found. Only called if geometry columns exist.
"""
sql = """
SELECT f_geometry_column, srid
FROM geometry_columns
WHERE f_table_schema = %s AND f_table_name = %s;
"""
try:
results = pg_hook.get_records(sql, (self.target_schema, self.target_table))
for col_name, srid in results:
if srid and srid > 0:
self._geometry_column_srids[col_name] = srid
except Exception as e:
log.warning(f"Could not query geometry_columns for SRID info: {e}")
def _build_select_columns(self) -> str:
"""
Build SELECT column list with transformations for special column types.
Priority: user-defined column_transforms > arrays > enums > geometry > passthrough.
Requires _load_column_metadata() to have been called first.
If include_columns is set, only those columns are included.
"""
select_parts = []
working_columns = self._get_working_columns()
for col in working_columns:
col_name = col["name"]
quoted_col = f'"{col_name}"'
if col_name in self.column_transforms:
transform = self.column_transforms[col_name].format(col=quoted_col)
select_parts.append(f"{transform} AS {quoted_col}")
elif col_name in self._array_columns:
element_type = self._array_columns[col_name]
# Schema-qualify custom enum types; built-in types resolve without qualification
if element_type in self._enum_type_names:
type_cast = f'"{self.target_schema}"."{element_type}"[]'
else:
type_cast = f"{element_type}[]"
transform = (
f"CASE WHEN {quoted_col} IS NULL OR NULLIF(TRIM({quoted_col}), '') IS NULL THEN NULL "
f"WHEN TRIM({quoted_col}) IN ('[]', '{{}}') THEN ARRAY[]::{type_cast} "
f"ELSE regexp_split_to_array(trim(both '[]{{}}' from {quoted_col}), '\\s*,\\s*')::{type_cast} "
f"END"
)
select_parts.append(f"{transform} AS {quoted_col}")
elif col_name in self._enum_columns:
enum_type = self._enum_columns[col_name]
transform = (
f"CASE WHEN NULLIF(TRIM({quoted_col}), '') IS NOT NULL "
f'THEN {quoted_col}::"{self.target_schema}"."{enum_type}" ELSE NULL END'
)
select_parts.append(f"{transform} AS {quoted_col}")
elif col_name in self._json_columns:
json_type = next(
c["udt_name"] for c in self._columns if c["name"] == col_name
)
transform = (
f"CASE WHEN NULLIF(TRIM({quoted_col}), '') IS NOT NULL "
f"THEN {quoted_col}::{json_type} ELSE NULL END"
)
select_parts.append(f"{transform} AS {quoted_col}")
elif col_name in self._geometry_columns:
geom_type = self._geometry_columns[col_name]
srid = self.geometry_srids.get(
col_name,
self._geometry_column_srids.get(col_name, self.geometry_srid),
)
# Use native casts (::geometry / ::geography) instead of
# ST_GeomFromText so that both WKT *and* hex-encoded
# EWKB (the default output of SELECT *) are accepted.
if geom_type == "geography":
transform = (
f"CASE WHEN NULLIF(TRIM({quoted_col}), '') IS NOT NULL "
f"THEN {quoted_col}::geography ELSE NULL END"
)
elif srid is not None:
transform = (
f"CASE WHEN NULLIF(TRIM({quoted_col}), '') IS NOT NULL "
f"THEN ST_SetSRID({quoted_col}::geometry, {srid}) ELSE NULL END"
)
else:
transform = (
f"CASE WHEN NULLIF(TRIM({quoted_col}), '') IS NOT NULL "
f"THEN {quoted_col}::geometry ELSE NULL END"
)
select_parts.append(f"{transform} AS {quoted_col}")
else:
# When include_columns is set, staging uses all TEXT columns,
# so we need to cast back to the target type
if self.include_columns and col["data_type"] not in ("text", "character varying"):
target_type = col["data_type"]
select_parts.append(f"{quoted_col}::{target_type} AS {quoted_col}")
else:
select_parts.append(quoted_col)
return ", ".join(select_parts)
def _get_working_columns(self) -> List[Dict[str, Any]]:
"""
Return the columns to use for staging/insert operations.
If include_columns is set, filter to only those columns and preserve
the order specified in include_columns (must match Athena query order
for correct CSV positional import).
"""
if not self.include_columns:
return self._columns
# Preserve include_columns order (must match Athena query order)
col_map = {col["name"]: col for col in self._columns}
return [col_map[name] for name in self.include_columns if name in col_map]
def _create_staging_table(
self,
pg_hook: PostgresHook,
staging_table: str,
target_fqn: str,
):
"""
Create staging table with TEXT columns for special types that need transformation.
Arrays, enums, geometry, and user-defined transforms all use TEXT in staging.
Requires _load_column_metadata() to have been called first.
If include_columns is set, staging table only contains those columns.
"""
working_columns = self._get_working_columns()
# If include_columns is set, always create explicit column list
# Otherwise, use LIKE for simple case with no special columns
if not self.include_columns and not self._special_columns:
pg_hook.run(
f"CREATE TABLE {staging_table} (LIKE {target_fqn} INCLUDING DEFAULTS);"
)
return
col_defs = []
for col in working_columns:
col_name = col["name"]
log.info(f"Creating staging column: {col_name}")
quoted_col = f'"{col_name}"'
# When using include_columns, use TEXT for all columns for simplicity
# since we're importing from CSV and will cast during INSERT
if self.include_columns or col_name in self._special_columns:
col_defs.append(f"{quoted_col} TEXT")
else:
col_defs.append(f"{quoted_col} {col['data_type']}")
col_list = ", ".join(col_defs)
pg_hook.run(f"CREATE TABLE {staging_table} ({col_list});")
################################################################################################
## S3 Import via aws_s3 Extension
################################################################################################
def _get_s3_file_key(self, query_id: str) -> str:
"""
Construct the S3 key for the query output CSV.
Athena standard output: s3://bucket/prefix/query_id.csv
"""
prefix = (
self.s3_prefix if self.s3_prefix.endswith("/") else f"{self.s3_prefix}/"
)
# Athena output naming convention
return f"{prefix}{query_id}.csv"
def _import_via_aws_s3(
self, pg_hook: PostgresHook, staging_table: str, s3_key: str
):
"""
Run aws_s3.table_import_from_s3 on the Postgres DB.
"""
# Options for Athena CSV output:
# Athena CSVs are quoted, comma separated, with header.
# Postgres COPY options: format csv, header true, quote '"'
options = "format csv, header true, quote '\"'"
if self.treat_empty_as_null:
options += ", null ''"
# Check if we are using IAM role based access (preferred) or credentials
# If using role-based access, the RDS instance must have an IAM role attached
# with permission to read from the S3 bucket.
# We prefer the 5-arg signature (without creds) if available,
# implying the RDS instance has an attached IAM role.
sql = f"""
SELECT aws_s3.table_import_from_s3(
'{staging_table}',
'', -- empty column list = all columns in order
$$({options})$$,
'{self.s3_bucket}',
'{s3_key}',
'{self.region_name}'
);
"""
log.info(
f"Importing from s3://{self.s3_bucket}/{s3_key} into {staging_table} via aws_s3 extension..."
)
pg_hook.run(sql)
################################################################################################
## S3 Cleanup
################################################################################################
def _cleanup_s3_file(self, s3_key: str):
"""Delete the Athena output CSV and metadata file from S3."""
boto3 = _import_boto3()
s3 = boto3.client("s3", region_name=self.region_name)
log.info(f"Cleaning up s3://{self.s3_bucket}/{s3_key}")
try:
s3.delete_object(Bucket=self.s3_bucket, Key=s3_key)
s3.delete_object(Bucket=self.s3_bucket, Key=s3_key + ".metadata")
except Exception as e:
log.warning(f"Failed to cleanup S3 file: {e}")
################################################################################################
## Postgres Load
################################################################################################
def _postgres_load(self, hook: PostgresHook, query: str) -> None:
"""Run a SQL command on Postgres."""
if self.dry_run:
log.info(f"Dry run enabled - skipping Postgres command:\n\n{query}")
return
hook.run(query)
################################################################################################
## Main Execute
################################################################################################
[docs]
def execute(self, context) -> Dict[str, str]:
self.dry_run = self.dry_run.strip().lower() == 'true'
valid_if_exists = ("append", "replace", "upsert")
if self.if_exists not in valid_if_exists:
raise ValueError(
f"if_exists must be one of {valid_if_exists}, got: {self.if_exists!r}"
)
if not self.use_aws_s3_extension:
raise NotImplementedError(
"Streaming COPY not yet implemented. Enable use_aws_s3_extension."
)
if self.if_exists == "upsert" and not self.upsert_keys:
raise ValueError("upsert_keys must be provided for if_exists='upsert'")
pg_hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
# 1. Check if target table exists
self._load_column_metadata(pg_hook)
table_exists = bool(self._columns)
if not table_exists and not self.create_table_if_not_exists:
raise AirflowException(
f"Target table not found: {self.target_schema}.{self.target_table}. "
"Set create_table_if_not_exists=True to auto-create."
)
# 2. Run Athena Query
query_id = self._run_athena_query()
if query_id is None:
log.info("Skipping — no source data to restore.")
return {"status": "skipped", "reason": "source_table_not_found"}
s3_key = self._get_s3_file_key(query_id)
# 3. Create target table if needed (after Athena query so we have schema)
if not table_exists:
athena_schema = self._get_athena_query_schema(query_id)
self._create_target_table(pg_hook, athena_schema)
# Reload column metadata now that table exists
self._load_column_metadata(pg_hook)
# 4. Validate column references
column_names = {col["name"] for col in self._columns}
if self.include_columns:
invalid_includes = set(self.include_columns) - column_names
if invalid_includes:
raise ValueError(
f"include_columns contain columns not in target table: {invalid_includes}. "
f"Available columns: {column_names}"
)
log.info(f"Using include_columns: {self.include_columns}")
if self.column_transforms:
invalid_transforms = set(self.column_transforms.keys()) - column_names
if invalid_transforms:
raise ValueError(
f"column_transforms contain columns not in target table: {invalid_transforms}"
)
# Warn if transforms reference columns not in include_columns
if self.include_columns:
unused_transforms = set(self.column_transforms.keys()) - set(self.include_columns)
if unused_transforms:
log.warning(
f"column_transforms {unused_transforms} will be ignored because they are "
f"not in include_columns: {self.include_columns}"
)
if self.if_exists == "upsert":
invalid_keys = set(self.upsert_keys) - column_names
if invalid_keys:
raise ValueError(
f"upsert_keys contain columns not in target table: {invalid_keys}"
)
# Ensure upsert_keys are included when using include_columns
if self.include_columns:
missing_upsert_keys = set(self.upsert_keys) - set(self.include_columns)
if missing_upsert_keys:
raise ValueError(
f"upsert_keys {missing_upsert_keys} must be included in include_columns "
f"for upsert to work correctly. Current include_columns: {self.include_columns}"
)
stg_name = f"__stg_{self.target_table}_{str(uuid.uuid4())[:8]}"
staging_table = f'"{self.target_schema}"."{stg_name}"'
target_fqn = f'"{self.target_schema}"."{self.target_table}"'
try:
# 3. Create Staging Table (with TEXT columns for special types)
log.info(f"Creating staging table {staging_table}...")
self._create_staging_table(pg_hook, staging_table, target_fqn)
# 4. Import S3 -> Staging
self._import_via_aws_s3(pg_hook, staging_table, s3_key)
# Log row count
count_result = pg_hook.get_first(f"SELECT COUNT(*) FROM {staging_table};")
row_count = count_result[0] if count_result else 0
log.info(f"Imported {row_count} rows from S3 into staging.")
# 5. Build SELECT with array/custom transformations
select_columns = self._build_select_columns()
# Build column list for INSERT (needed when include_columns is set)
working_columns = self._get_working_columns()
insert_column_list = ", ".join([f'"{col["name"]}"' for col in working_columns])
# 6. Sync Staging -> Target
if self.if_exists == "replace":
log.info(f"Truncating {target_fqn} and inserting from staging...")
self._postgres_load(
pg_hook,
f"""
TRUNCATE TABLE {target_fqn};
INSERT INTO {target_fqn} ({insert_column_list}) SELECT {select_columns} FROM {staging_table};
""",
)
elif self.if_exists == "append":
log.info(f"Appending from staging to {target_fqn}...")
self._postgres_load(
pg_hook,
f"INSERT INTO {target_fqn} ({insert_column_list}) SELECT {select_columns} FROM {staging_table};",
)
elif self.if_exists == "upsert":
match_condition = " AND ".join(
[
f'{target_fqn}."{k}" = {staging_table}."{k}"'
for k in self.upsert_keys
]
)
log.info(
f"Upserting into {target_fqn} using keys: {self.upsert_keys}..."
)
self._postgres_load(
pg_hook,
f"""
DELETE FROM {target_fqn}
USING {staging_table}
WHERE {match_condition};
INSERT INTO {target_fqn} ({insert_column_list}) SELECT {select_columns} FROM {staging_table};
""",
)
log.info("Load complete.")
finally:
# Cleanup Staging
log.info(f"Dropping staging table {staging_table}...")
pg_hook.run(f"DROP TABLE IF EXISTS {staging_table};")
# Cleanup S3
if self.cleanup_s3:
self._cleanup_s3_file(s3_key)
return {
"query_id": query_id,
"s3_key": s3_key,
"target_table": f"{self.target_schema}.{self.target_table}",
}