from typing import Union, List, Optional, Tuple, Dict, Literal
import logging
import uuid
import time
from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
from airflow.exceptions import AirflowException
from config import Config
####################################################################################################
## Constants
####################################################################################################
_SIMPLE_DEFAULTS_MAP = {
"now()": "GETDATE()",
"current_timestamp": "GETDATE()",
"CURRENT_TIMESTAMP": "GETDATE()",
"true": "TRUE",
"false": "FALSE",
}
# Cap VARCHAR to Redshift max
_MAX_VARCHAR = 65535
# Redshift's hard per-value size limit for GEOMETRY/GEOGRAPHY (bytes of EWKB).
# Values larger than this cannot be loaded; we null them out so the row still loads.
_MAX_GEOMETRY_BYTES = 1048447
# Functions that are not supported in Redshift and should be dropped from defaults
# The intended use of this operator is to replace data as-is, so we drop these functions
# since we will not be adding new data to a replicated table
_UNSUPPORTED_FUNC_DEFAULTS = {
# Postgres-only UUID generators
"gen_random_uuid()",
"uuid_generate_v4()",
"uuid_generate_v1()",
"uuid_nil()",
# we may need to add some timestamp stuff here later
}
####################################################################################################
## Utility Functions
####################################################################################################
def _cap_varchar(n: Optional[int]) -> int:
"""Cap VARCHAR length to Redshift maximum."""
if not n or n <= 0:
return _MAX_VARCHAR
return min(n, _MAX_VARCHAR)
def _normalize_default(defval: Optional[str]) -> Optional[str]:
"""
Normalize simple default values from Postgres to Redshift equivalents.
Returns None if default should be dropped (unsupported).
"""
import re
if not defval:
return None
dv = defval.strip()
# Drop sequence-driven defaults
if dv.lower().startswith("nextval("):
return None
# Strip PG casts ::type
if "::" in dv:
dv = dv.split("::", 1)[0].strip()
# Map simple literals/funcs we explicitly support (case-insensitive)
key = dv if dv in _SIMPLE_DEFAULTS_MAP else dv.lower()
if key in _SIMPLE_DEFAULTS_MAP:
return _SIMPLE_DEFAULTS_MAP[key]
# Keep NULL
if dv.upper() == "NULL":
return "NULL"
# Keep numeric literals
if re.fullmatch(r"[-+]?\d+(\.\d+)?", dv):
return dv
if len(dv) >= 2 and dv[0] == dv[-1] == "'":
return dv
# Drop known PG-only function defaults
func_call_re = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*\s*\(.*\)$")
if dv.lower() in _UNSUPPORTED_FUNC_DEFAULTS or func_call_re.match(dv):
logging.info("Dropping unsupported default %r (function-like).", dv)
return None
# default case
return None
def _as_bool(v: Union[str, bool, None]) -> bool:
"""Convert input to boolean. Accepts bool, None, or string representations.
None -> False"""
if isinstance(v, bool):
return v
if v is None:
return False
return str(v).strip().lower() in {"1", "true", "t", "yes", "y", "on"}
def _norm_sortkey(val: Union[str, List[str], None]) -> Union[str, List[str]]:
"""Normalize sortkey input into either "AUTO" or list of columns."""
import json, ast
if val is None:
return "AUTO"
if isinstance(val, (list, tuple)):
xs = [str(c).strip().strip('"').strip("'") for c in val if str(c).strip()]
return "AUTO" if len(xs) == 1 and xs[0].upper() == "AUTO" else xs
s = str(val).strip()
if s.upper() == "AUTO":
return "AUTO"
if s.startswith("[") and s.endswith("]"):
try:
parsed = json.loads(s)
except Exception:
parsed = ast.literal_eval(s)
return _norm_sortkey(parsed)
if "," in s:
return [c.strip().strip('"').strip("'") for c in s.split(",") if c.strip()]
return [s.strip().strip('"').strip("'")]
def _case_prefix(flag: Union[str, bool]) -> str:
return "SET enable_case_sensitive_identifier TO true;\n" if _as_bool(flag) else ""
def _looks_templated(v: Optional[str]) -> bool:
return isinstance(v, str) and ("{{" in v or "}}" in v)
####################################################################################################
## Type Conversion Logic
####################################################################################################
def _pg_to_rs_type(
row: dict,
logger: Optional[logging.Logger] = None,
identity_policy: str = "strip",
) -> str:
"""
Map Postgres column (from information_schema) to a Redshift Serverless type.
Expects keys: data_type, udt_name, character_maximum_length, numeric_precision,
numeric_scale, is_serial (bool), domain_name (optional).
See here for Redshift/Postgres type mapping:
https://docs.aws.amazon.com/redshift/latest/dg/federated-data-types.html
"""
udt = (row.get("udt_name") or "").lower()
data_type = (row.get("data_type") or "").lower()
char_len = row.get("character_maximum_length")
num_prec = row.get("numeric_precision")
num_scale = row.get("numeric_scale")
################################################################################################
# Helpers
################################################################################################
def _int4() -> str:
if not is_id:
return "INTEGER"
if identity_policy == "preserve":
return "INTEGER IDENTITY(1,1)"
if identity_policy == "as_varchar":
return f"VARCHAR(20)" # fits signed BIGINT text; adjust if you like
return "INTEGER" # identity_policy == "strip"
def _int8() -> str:
if not is_id:
return "BIGINT"
if identity_policy == "preserve":
return "BIGINT IDENTITY(1,1)"
if identity_policy == "as_varchar":
return f"VARCHAR(20)"
return "BIGINT"
################################################################################################
# Integers & identity
################################################################################################
# Integers
# 2-byte
if udt == "int2":
return "SMALLINT"
# is_id will be True if the column is a serial/identity column
is_id = row.get("is_serial") or str(row.get("is_identity", "")).upper() == "YES"
# 4-byte
if udt == "int4":
return _int4()
# 8-byte
if udt == "int8":
return _int8()
# Numerics
if data_type in ("numeric", "decimal"):
# if num_prec and num_scale are both set, use them
if num_prec is not None and num_scale is not None:
return f"DECIMAL({num_prec},{num_scale})"
# Redshift supports up to precision 38
return "DECIMAL(38,18)"
# Floating
if udt == "float4":
return "REAL"
if udt == "float8":
return "DOUBLE PRECISION"
# Boolean
if data_type == "boolean":
return "BOOLEAN"
# Character
if data_type in ("character varying", "varchar"):
return f"VARCHAR({_cap_varchar(char_len)})"
if data_type in ("character", "char"):
# CHAR must have a length; default to 1 if absent
return f"CHAR({char_len if char_len and char_len > 0 else 1})"
if data_type == "text":
return f"VARCHAR({_MAX_VARCHAR})"
# Datetime / time
if data_type in ("timestamp without time zone", "timestamp"):
return "TIMESTAMP"
if data_type in ("timestamp with time zone", "timestamptz"):
# Redshift supports TIMESTAMPTZ and stores as UTC
return "TIMESTAMPTZ"
if data_type == "date":
return "DATE"
if data_type == "time without time zone":
return "TIME"
if data_type == "time with time zone":
return "TIME"
# Semi-structured
if data_type in ("json", "jsonb"):
return "SUPER"
# Arrays
# information_schema reports data_type='ARRAY' and udt_name like '_int4'
if data_type == "array" or udt.startswith("_"):
# Prefer SUPER to preserve structure
return "SUPER"
# Binary
if data_type == "bytea":
return "VARBYTE"
# UUID (no native type in Redshift)
if udt == "uuid":
return "VARCHAR(36)"
# Enums & user-defined
# Spatial (PostGIS in PG reports udt_name geometry/geography)
if udt == "geometry":
return "GEOMETRY"
if udt == "geography":
return "GEOGRAPHY"
# user-defined / enums
if data_type == "user-defined":
if logger:
logger.warning(
"Mapping USER-DEFINED type to VARCHAR(255): %s.%s (%s/%s)",
row.get("table_schema", "?"),
row.get("column_name", "?"),
data_type,
udt,
)
else:
logging.warning(
"Mapping USER-DEFINED type to VARCHAR(255): %s.%s (%s/%s)",
row.get("table_schema", "?"),
row.get("column_name", "?"),
data_type,
udt,
)
return "VARCHAR(255)"
# fallback with explicit warning
msg = (
f"Falling back to VARCHAR(255) for "
f"{row.get('table_schema','?')}.{row.get('column_name','?')} "
f"({data_type}/{udt})"
)
if logger:
logger.warning(msg)
else:
logging.warning(msg)
return "VARCHAR(255)"
####################################################################################################
## Redshift Create Table Operator from Postgres Source
####################################################################################################
[docs]
class PostgresToRedshiftDDLOperator(BaseOperator):
"""
Generate (and optionally execute) a Redshift Serverless CREATE TABLE
equivalent to a source Postgres table.
- Pulls column metadata from Postgres (information_schema for PK).
- Maps types to Redshift equivalents (TIMESTAMPTZ, SUPER, VARBYTE, GEOMETRY/GEOGRAPHY, etc).
- Preserves NOT NULL, simple DEFAULTs; PK is added as informational metadata.
- Appends DISTSTYLE AUTO and AUTO SORTKEY (override via args).
- Please Note that Redshift does not enforce primary key constraints, but this operator will include the PK in the DDL for documentation purposes.
:param source_table: Name of the source Postgres table (CaSe SeNsItIvE).
:type source_table: str
:param target_table: Name of the target Redshift table (CaSe SeNsItIvE). Typically the same as source_table.
:type target_table: str
:param pg_conn_id: Airflow connection ID for the source Postgres database.
:type pg_conn_id: str
:param rs_database: Name of the Redshift database to connect to (typically "dev").
:type rs_database: str
:param schema: Name of the Redshift schema to create the table in (CaSe SeNsItIvE).
:type schema: str
:param region: AWS region of the Redshift Serverless instance, defaults to "us-east-1".
:type region: str
:param workgroup_name: Redshift Serverless workgroup name, defaults to Config.REDSHIFT_WORKGROUP.
:type workgroup_name: str
:param dry_run_mode: If True, do not execute the DDL, just return it. Defaults to False. Can also be the string "true"/"false" (case-insensitive).
:type dry_run_mode: Union[str, bool]
:param diststyle: Redshift distribution style, one of "AUTO", "EVEN", "ALL", "KEY". Defaults to "AUTO".
:type diststyle: str
:param sortkey: Redshift sort key, either "AUTO" or a list of column names. Defaults to "AUTO".
:type sortkey: Union[str, List[str], None]
:param distkey: Redshift distribution key column name, required if diststyle is "KEY". Defaults to None.
:type distkey: Optional[str]
:param include_pk_metadata: If True, include PRIMARY KEY metadata in the DDL. Defaults to True. Can also be the string "true"/"false" (case-insensitive).
:type include_pk_metadata: Union[str, bool]
:param create_schema_if_needed: If True, includes a "CREATE SCHEMA IF NOT EXISTS" statement before the CREATE TABLE. Defaults to True.
:type create_schema_if_needed: bool
:param aws_conn_id: Airflow connection ID for AWS credentials. Defaults to "aws_default".
:type aws_conn_id: str
:param case_sensitive_identifiers: If True, sets Redshift to treat identifiers as case-sensitive for the session. Defaults to True. Can also be the string "true"/"false" (case-insensitive).
:type case_sensitive_identifiers: Union[str, bool]
:param identity_policy: How to handle identity/serial columns from Postgres. One of:
- "strip" (default): create as normal integer columns (no identity).
- "preserve": create as IDENTITY columns (starts at 1, increments by 1).
- "as_varchar": create as VARCHAR(20) to hold textual representation of the value.
:type identity_policy: Literal["preserve", "strip", "as_varchar"]
:param drop_if_exists: If True, includes a "DROP TABLE IF EXISTS" statement before the CREATE TABLE. Defaults to False. Can also be the string "true"/"false" (case-insensitive).
:type drop_if_exists: Union[str, bool]
:param nullable_override: If True, drops all NOT NULL constraints from the generated DDL (except on PRIMARY KEY columns). Useful when the source Postgres schema has NOT NULL but the data contains NULLs. Defaults to False. Can also be the string "true"/"false" (case-insensitive).
:type nullable_override: Union[str, bool]
"""
template_fields = (
"source_table",
"target_table",
"rs_database",
"schema",
"diststyle",
"sortkey",
"distkey",
"include_pk_metadata",
"pg_conn_id",
"dry_run_mode",
"create_schema_if_needed",
"aws_conn_id",
"region",
"workgroup_name",
"case_sensitive_identifiers",
"identity_policy",
"drop_if_exists",
"nullable_override",
)
ui_color = "#c2acee"
def __init__(
self,
*,
source_table: str,
target_table: str,
pg_conn_id: str,
rs_database: str,
schema: str,
region: str = "us-east-1",
workgroup_name: str = Config.REDSHIFT_WORKGROUP,
dry_run_mode: Union[str, bool] = False,
diststyle: str = "AUTO", # AUTO | EVEN | ALL | KEY
sortkey: Union[str, List[str], None] = None, # "AUTO" | list[str] | None
distkey: Optional[str] = None, # required if diststyle == KEY
include_pk_metadata: Union[str, bool] = True,
create_schema_if_needed: bool = True,
aws_conn_id: str = "aws_default",
case_sensitive_identifiers: Union[str, bool] = True,
identity_policy: Literal["preserve", "strip", "as_varchar"] = "strip",
drop_if_exists: Union[str, bool] = False,
nullable_override: Union[str, bool] = False,
**kwargs,
):
super().__init__(**kwargs)
self.source_table = source_table
self.target_table = target_table
self.pg_conn_id = pg_conn_id
self.rs_database = rs_database
self.schema = schema
self.region = region
self.workgroup_name = workgroup_name
self.dry_run_mode = dry_run_mode
self.diststyle = diststyle
self.sortkey = sortkey
self.distkey = distkey
self.include_pk_metadata = include_pk_metadata
self.create_schema_if_needed = create_schema_if_needed
self.aws_conn_id = aws_conn_id
self.case_sensitive_identifiers = case_sensitive_identifiers
self.identity_policy = identity_policy
self.drop_if_exists = drop_if_exists
self.nullable_override = nullable_override
################################################################################################
## Handle SUPER defaults
################################################################################################
@staticmethod
def _json_default_for_super(raw: str) -> Optional[str]:
"""
Convert common PG defaults to valid Redshift SUPER defaults.
Returns a SQL fragment (without the word DEFAULT) or None if we can't convert.
"""
if raw is None:
return None
s = str(raw).strip()
# PG array literal with no elements (ARRAY[]) -> JSON empty array
if s.upper().replace(" ", "") == "ARRAY[]":
return "JSON_PARSE('[]')"
# Bare JSON-looking defaults like [] or {}
if (s.startswith("[") and s.endswith("]")) or (
s.startswith("{") and s.endswith("}")
):
# escape single quotes in case (rare for []/{})
j = s.replace("'", "''")
return f"JSON_PARSE('{j}')"
# Empty string as a SUPER default -> treat like empty JSON object
if s == "''":
return "JSON_PARSE('{}')"
return None
################################################################################################
## Create Hooks
################################################################################################
def _get_pg_hook(self) -> PostgresHook:
"""Returns a PostgresHook instance."""
return PostgresHook(postgres_conn_id=self.pg_conn_id)
def _get_rs_hook(self) -> RedshiftDataHook:
"""Returns a RedshiftDataHook instance."""
return RedshiftDataHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
################################################################################################
## Fetch Metadata From Postgres
################################################################################################
def _fetch_pg_columns_and_pk(self) -> Tuple[List[dict], List[str]]:
"""
Return (columns, pk_cols). Each column dict contains:
column_name, data_type, udt_name, character_maximum_length,
numeric_precision, numeric_scale, is_nullable, column_default, is_serial.
"""
sql_cols = """
SELECT
c.table_schema,
c.table_name,
c.ordinal_position,
c.column_name,
c.data_type,
c.udt_name,
c.character_maximum_length,
c.numeric_precision,
c.numeric_scale,
c.is_nullable,
c.column_default,
(position('nextval(' in coalesce(c.column_default, '')) = 1) AS is_serial,
c.is_identity
FROM information_schema.columns c
WHERE c.table_schema = %s
AND c.table_name = %s
ORDER BY c.ordinal_position;
"""
sql_pk = """
SELECT kcu.column_name
FROM information_schema.table_constraints tc
JOIN information_schema.key_column_usage kcu
ON tc.constraint_name = kcu.constraint_name
AND tc.table_schema = kcu.table_schema
WHERE tc.constraint_type = 'PRIMARY KEY'
AND tc.table_schema = %s
AND tc.table_name = %s
ORDER BY kcu.ordinal_position;
"""
hook = self._get_pg_hook()
conn = hook.get_conn()
with conn.cursor() as cur:
cur.execute(sql_cols, (self.schema, self.source_table))
rows = cur.fetchall()
names = [getattr(d, "name", d[0]) for d in cur.description]
columns = [dict(zip(names, r)) for r in rows]
cur.execute(sql_pk, (self.schema, self.source_table))
pk_cols = [r[0] for r in cur.fetchall()]
return columns, pk_cols
################################################################################################
## Build Redshift CREATE TABLE DDL
################################################################################################
def _build_create_table(
self,
columns: List[dict],
pk_cols: List[str],
sortkey_norm: Union[str, List[str]],
diststyle_norm: str,
distkey: Optional[str],
) -> str:
"""
Build Redshift CREATE TABLE DDL statement based on columns and options.
:param columns: list of column dicts from _fetch_pg_columns_and_pk
:type columns: List[dict]
:param pk_cols: list of primary key column names
:type pk_cols: List[str]
:param sortkey_norm: normalized sortkey ("AUTO" or list of columns)
:type sortkey_norm: Union[str, List[str]]
:param diststyle_norm: normalized diststyle (AUTO | EVEN | ALL | KEY)
:type diststyle_norm: str
:param distkey: distribution key column name (if diststyle == KEY)
:type distkey: Optional[str]
:returns: Redshift CREATE TABLE DDL statement as string
"""
col_lines: List[str] = []
for col in columns:
rs_type = _pg_to_rs_type(
col, logger=self.log, identity_policy=self.identity_policy
)
is_pk = col["column_name"] in pk_cols
not_null = col["is_nullable"] == "NO" and (
not _as_bool(self.nullable_override) or is_pk
)
default = _normalize_default(col.get("column_default"))
# SUPER default fix
if rs_type == "SUPER" and default is not None:
super_default = self._json_default_for_super(default)
if super_default:
default = super_default
else:
# If we can't translate, drop the default to avoid syntax errors
self.log.warning(
"Dropping unsupported default for SUPER column %s.%s: %r",
col.get("table_schema"),
col.get("column_name"),
default,
)
default = None
line = f' "{col["column_name"]}" {rs_type}'
if default:
line += f" DEFAULT {default}"
if not_null:
line += " NOT NULL"
col_lines.append(line)
parts: List[str] = []
# TO-DO: Determine if we want to include this in the operator... this feels kind of like
# it should be a separate operator. (which does exist already)
# if self.create_schema_if_needed:
# parts.append(f'CREATE SCHEMA IF NOT EXISTS "{self.schema}";')
parts.append(
f'CREATE TABLE IF NOT EXISTS "{self.schema}"."{self.target_table}" ('
)
parts.append(",\n".join(col_lines))
if _as_bool(self.include_pk_metadata) and pk_cols:
parts.append(
",\n PRIMARY KEY (" + ", ".join([f'"{c}"' for c in pk_cols]) + ")"
)
parts.append(")")
# DISTSTYLE / DISTKEY: only emit if NOT AUTO
if diststyle_norm and diststyle_norm != "AUTO":
parts.append(f"DISTSTYLE {diststyle_norm}")
if diststyle_norm == "KEY":
if not distkey:
raise ValueError("DISTSTYLE KEY requires distkey=<column>.")
parts.append(f'DISTKEY("{distkey}")')
# SORTKEY: only emit if NOT AUTO
if isinstance(sortkey_norm, list) and sortkey_norm:
parts.append("SORTKEY(" + ", ".join([f'"{c}"' for c in sortkey_norm]) + ")")
return "\n".join(parts) + ";"
################################################################################################
## Execute or Preview DDL
################################################################################################
[docs]
def execute(self, context):
"""Main execution method for the operator.
Fetches Postgres table metadata, builds Redshift CREATE TABLE DDL, and either executes it or returns it in dry-run mode.
"""
cols, pk_cols = self._fetch_pg_columns_and_pk()
if not cols:
raise ValueError(
f"Source table {self.schema}.{self.source_table} not found or has no columns."
)
# Set statement name for redshift execution
# TODO: make this something you can pass in to the operator?
stmt_name = f"pg2rsddl_{self.target_table[:32]}"
# Normalize inputs
sortkey_norm = _norm_sortkey(self.sortkey) # "AUTO" or list[str]
diststyle_norm = (self.diststyle or "AUTO").strip().upper()
# Normalize distkey once (None / "", "null", "none" -> None)
raw_distkey = self.distkey
if isinstance(raw_distkey, str):
dk = raw_distkey.strip()
distkey = None if dk == "" or dk.lower() in {"null", "none"} else dk
else:
distkey = raw_distkey or None
# Case-insensitive column name map
colname_map = {c["column_name"].lower(): c["column_name"] for c in cols}
# Validate only when needed
if diststyle_norm == "KEY":
if not distkey:
raise ValueError("DISTSTYLE KEY requires distkey=<column>.")
if distkey.lower() not in colname_map:
raise ValueError(
f'DISTKEY column "{distkey}" not found. '
f"Available: {sorted(colname_map.values())}"
)
distkey = colname_map[distkey.lower()] # normalize to actual case
if isinstance(sortkey_norm, list) and sortkey_norm:
missing = [s for s in sortkey_norm if s.lower() not in colname_map]
if missing:
raise ValueError(
f"SORTKEY column(s) not found: {missing}. "
f"Available: {sorted(colname_map.values())}"
)
sortkey_norm = [colname_map[s.lower()] for s in sortkey_norm]
ddl = self._build_create_table(
cols, pk_cols, sortkey_norm, diststyle_norm, distkey
)
self.log.info("Generated Redshift DDL:\n%s", ddl)
# Build the full batch (optional drop + create)
statements: List[str] = []
if _as_bool(self.drop_if_exists):
statements.append(
f'DROP TABLE IF EXISTS "{self.schema}"."{self.target_table}";'
)
statements.append(ddl)
full_sql = _case_prefix(self.case_sensitive_identifiers) + "\n".join(statements)
if _as_bool(self.dry_run_mode):
return full_sql
rs_hook = self._get_rs_hook()
rs_hook.execute_query(
sql=full_sql,
workgroup_name=self.workgroup_name,
wait_for_completion=True,
database=self.rs_database,
statement_name=stmt_name,
)
return ddl
####################################################################################################
## Redshift Load from Postgres via S3 Operator
####################################################################################################
[docs]
class PostgresToRedshiftUnloadLoadOperator(BaseOperator):
"""
Unload table data from Postgres to S3 via aws_s3 extension, then COPY into Redshift staging
and INSERT…SELECT into the final table with minimal, type-aware shims.
- Use `_pg_to_rs_type` to determine the *final* RS column type.
- Adds transport-time shims for conversion only where needed:
* json/jsonb/arrays -> export as JSON text; stage VARCHAR; JSON_PARSE on load
* geometry/geography -> export hex EWKB; stage GEOMETRY/GEOGRAPHY; direct COPY load
* primitive types -> stage as VARCHAR; cast on load with NULLIF('', NULL)
:param pg_conn_id: Airflow connection ID for the source Postgres database.
:type pg_conn_id: str
:param rs_workgroup: Redshift workgroup name.
:type rs_workgroup: str
:param rs_database: Redshift database name.
:type rs_database: str
:param target_schema: Target schema in Redshift.
:type target_schema: str
:param source_schema: Source schema in Postgres.
:type source_schema: str
:param table: Table name to transfer.
:type table: str
:param s3_bucket: S3 bucket for staging data.
:type s3_bucket: str
:param s3_prefix: S3 prefix for staging data.
:type s3_prefix: str
:param redshift_iam_role: IAM role for Redshift to access S3.
:type redshift_iam_role: str
:param region: AWS region for S3.
:type region: str
:param header: Include header in CSV export.
:type header: bool
:param gzip: Gzip compress CSV files.
:type gzip: bool
:param where_clause: Optional WHERE clause for data filtering.
:type where_clause: str
:param cleanup_after_load: Whether to clean up temporary S3 files after loading.
:type cleanup_after_load: Union[str, bool]
:param max_errors: Maximum number of rows that can fail during COPY before the
load is aborted. Useful for tolerating a small number of rows with oversized
SUPER string values (>65KB Redshift limit). Default 0 (fail on first error).
:type max_errors: int
:param empty_as_null: When True (default), the CSV COPY uses
``EMPTYASNULL BLANKSASNULL TRIMBLANKS`` so any empty / whitespace CSV
field becomes NULL in staging. Set to False when the source PG
distinguishes ``''`` (empty string, exported as ``""``) from NULL
(exported as unquoted empty) and the target schema has NOT NULL text
columns — leaving these options on conflates the two and breaks the
downstream INSERT when an empty-string source row hits a NOT NULL
target column. JSON mode is unaffected (these options are CSV-only).
:type empty_as_null: Union[str, bool]
:param load_mode: How the target is mutated. One of:
``APPEND`` (insert-only, no purge — default for backward compat),
``TRUNCATE_INSERT`` (atomic ``DELETE FROM target`` + ``INSERT FROM stage``
in one Redshift transaction — full refresh),
``DELETE_INSERT`` (atomic ``DELETE FROM target WHERE delete_where_clause``
+ ``INSERT FROM stage`` in one Redshift transaction — predicate-based
replacement).
For refresh modes the operator owns the purge inside the transaction;
the DAG's ``optional_purge`` task must be check-only (no DELETE/TRUNCATE).
TRUNCATE is intentionally avoided because Redshift TRUNCATE auto-commits
and cannot participate in a transaction.
:type load_mode: str
:param delete_where_clause: SQL predicate (without leading ``WHERE``) used
as the DELETE filter in ``DELETE_INSERT`` mode. Required for
``DELETE_INSERT``; ignored otherwise. Note: timestamp-based predicates
(e.g. ``"updatedAt" >= cutoff``) cannot evict RS rows whose stored
timestamp is below the cutoff — use ``TRUNCATE_INSERT`` if PG rows are
updatable in place and you need parity.
:type delete_where_clause: Optional[str]
.. note::
When ``max_errors > 0`` and ``load_mode`` is a refresh mode
(``TRUNCATE_INSERT`` / ``DELETE_INSERT``), the operator **raises** if
COPY skipped any rows — refusing to mutate the target from an
incomplete stage. For ``APPEND``, the prior warn-only behavior is
preserved.
"""
template_fields = (
"source_schema",
"target_schema",
"table",
"s3_bucket",
"s3_prefix",
"pg_conn_id",
"rs_database",
"rs_workgroup",
"where_clause",
"aws_conn_id",
"region",
"redshift_iam_role",
"header",
"gzip",
"cleanup_after_load",
"case_sensitive_identifiers",
"max_errors",
"empty_as_null",
"load_mode",
"delete_where_clause",
)
ui_color = "#FFB03A"
# Modes where the operator owns the target purge inside the COPY+DELETE+INSERT
# transaction (atomic). When in one of these modes, the DAG's optional_purge
# task MUST NOT also purge — see dag.py optional_purge.
_REFRESH_MODES = frozenset({"TRUNCATE_INSERT", "DELETE_INSERT"})
def __init__(
self,
*,
pg_conn_id: str,
rs_workgroup: str,
rs_database: str,
target_schema: str,
source_schema: str,
table: str,
s3_bucket: str,
s3_prefix: str,
rs_stage_schema: Optional[str] = None,
redshift_iam_role: Optional[str],
region: str = "us-east-1",
header: bool = True,
gzip: bool = True,
where_clause: Optional[str] = None,
cleanup_after_load: Union[str, bool] = True,
aws_conn_id: str = "aws_default",
case_sensitive_identifiers: Union[str, bool] = True,
max_errors: int = 0,
empty_as_null: Union[str, bool] = True,
load_mode: str = "APPEND",
delete_where_clause: Optional[str] = None,
**kwargs,
):
super().__init__(**kwargs)
role_input = redshift_iam_role
if not role_input or not str(role_input).strip():
role_input = getattr(Config, "REDSHIFT_S3_ROLE_ARN", None)
self.pg_conn_id = pg_conn_id
self.rs_workgroup = rs_workgroup
self.rs_database = rs_database
self.target_schema = target_schema
self.source_schema = source_schema
self.table = table
self.s3_bucket = s3_bucket
self.s3_prefix = s3_prefix.rstrip("/")
self.rs_stage_schema = (rs_stage_schema or target_schema).strip()
self.redshift_iam_role = str(role_input or "").strip()
self.region = region
self.header = header
self.gzip = gzip
self.where_clause = (where_clause or "").strip() or None
self.cleanup_after_load = cleanup_after_load
self.aws_conn_id = aws_conn_id
self.case_sensitive_identifiers = case_sensitive_identifiers
self.max_errors = int(max_errors)
self.empty_as_null = empty_as_null
self.load_mode = (load_mode or "APPEND").strip().upper()
self.delete_where_clause = (delete_where_clause or "").strip() or None
if not _looks_templated(
self.redshift_iam_role
) and not self.redshift_iam_role.startswith("arn:aws:iam::"):
raise AirflowException(
"redshift_iam_role is required and must be a full IAM Role ARN "
"(e.g., arn:aws:iam::<acct>:role/<name>). "
"Tip: set DAG param 'redshift_iam_role' or Config.REDSHIFT_S3_ROLE_ARN."
)
####################################################################################################
## Get Postgres Column Info
####################################################################################################
def _get_pg_columns(self) -> List[Dict]:
sql = """
SELECT
c.column_name,
c.data_type,
c.udt_name,
c.is_nullable,
c.character_maximum_length,
c.numeric_precision,
c.numeric_scale,
(position('nextval(' in coalesce(c.column_default, '')) = 1) AS is_serial,
c.is_identity
FROM information_schema.columns c
WHERE c.table_schema = %s AND c.table_name = %s
ORDER BY c.ordinal_position;
"""
hook = PostgresHook(postgres_conn_id=self.pg_conn_id)
conn = hook.get_conn()
with conn.cursor() as cur:
cur.execute(sql, (self.source_schema, self.table))
rows = cur.fetchall()
names = [getattr(d, "name", d[0]) for d in cur.description]
return [dict(zip(names, r)) for r in rows]
####################################################################################################
## Helper Functions
####################################################################################################
@staticmethod
def _is_superish(col: Dict) -> bool:
t = (col.get("data_type") or "").lower()
u = (col.get("udt_name") or "").lower()
return t in {"json", "jsonb", "array"} or u.startswith("_")
@staticmethod
def _is_geom(col: Dict) -> bool:
u = (col.get("udt_name") or "").lower()
return u in {"geometry", "geography"}
def _sts_creds(self):
import boto3
session = boto3.Session()
creds = session.get_credentials()
if not creds:
raise AirflowException(
"No AWS credentials available for aws_s3 export. "
"Set an instance profile or Airflow AWS Connection."
)
f = creds.get_frozen_credentials()
return f.access_key, f.secret_key, f.token
def _run_rs_and_fetch(self, sql: str) -> list[dict]:
"""
Execute SQL via Redshift Data API (Serverless) and return rows as list[dict].
Avoids Airflow hook's silent result handling.
"""
import boto3
client = boto3.client("redshift-data", region_name=self.region)
resp = client.execute_statement(
WorkgroupName=self.rs_workgroup,
Database=self.rs_database,
Sql=sql,
WithEvent=False,
)
sid = resp["Id"]
# Poll until finished
while True:
d = client.describe_statement(Id=sid)
status = d["Status"]
if status in ("FINISHED", "FAILED", "ABORTED"):
break
time.sleep(1.0)
if status != "FINISHED":
raise AirflowException(d.get("Error", status))
# Ensure a result set is present (multi-statement batches won't have one)
if not d.get("HasResultSet", False):
raise AirflowException("Statement finished but returned no result set.")
res = client.get_statement_result(Id=sid)
cols = [c["name"] for c in res["ColumnMetadata"]]
rows: list[dict] = []
for r in res.get("Records", []):
row = {}
for name, cell in zip(cols, r):
((k, v),) = cell.items()
row[name] = v
rows.append(row)
return rows
####################################################################################################
## Transport Planner - Shim Layer
####################################################################################################
def _fetch_rs_target_columns(self) -> Dict[str, Dict]:
"""
Read target column types (from information_schema.columns) and identity flags (from svv_columns).
Works on Redshift Serverless where is_identity may not exist in information_schema.
"""
sql = f"""
WITH id_cols AS (
SELECT
LOWER(namespace_name) AS schemaname,
LOWER(relation_name) AS tablename,
LOWER(column_name) AS column_name,
MAX(CASE WHEN identity_id IS NOT NULL THEN 1 ELSE 0 END) AS is_identity
FROM svv_column_privileges
GROUP BY 1,2,3
)
SELECT
LOWER(ic.column_name) AS column_name,
ic.data_type AS data_type,
ic.character_maximum_length AS character_maximum_length,
ic.numeric_precision AS numeric_precision,
ic.numeric_scale AS numeric_scale,
COALESCE(id.is_identity, 0) AS is_identity
FROM information_schema.columns ic
LEFT JOIN id_cols id
ON id.schemaname = LOWER(ic.table_schema)
AND id.tablename = LOWER(ic.table_name)
AND id.column_name= LOWER(ic.column_name)
WHERE ic.table_schema = '{self.target_schema}'
AND ic.table_name = '{self.table}'
ORDER BY ic.ordinal_position;
"""
rows = self._run_rs_and_fetch(sql)
def _canon_type(dt: str, char_len, p, s) -> str:
dt_u = (dt or "").upper()
# --- floats: normalize here ---
if dt_u in {"FLOAT4"}:
return "REAL"
if dt_u in {"FLOAT8", "FLOAT"}:
return "DOUBLE PRECISION"
if dt_u in {"CHARACTER VARYING", "VARCHAR"}:
return f"VARCHAR({char_len if char_len and char_len > 0 else _MAX_VARCHAR})"
if dt_u in {"CHARACTER", "CHAR"}:
return f"CHAR({char_len if char_len and char_len > 0 else 1})"
if dt_u in {"DECIMAL", "NUMERIC"}:
return (
f"DECIMAL({p},{s})"
if (p is not None and s is not None)
else "DECIMAL(38,18)"
)
if dt_u == "TIMESTAMP WITH TIME ZONE":
return "TIMESTAMPTZ"
if dt_u in {"TIMESTAMP WITHOUT TIME ZONE", "TIMESTAMP"}:
return "TIMESTAMP"
# Pass-through for types we already like
passthru = {
"DOUBLE PRECISION",
"REAL",
"INTEGER",
"BIGINT",
"SMALLINT",
"BOOLEAN",
"DATE",
"TIME",
"SUPER",
"VARBYTE",
"GEOMETRY",
"GEOGRAPHY",
}
return dt_u if dt_u in passthru else (dt_u or "VARCHAR(255)")
out: Dict[str, Dict] = {}
for r in rows:
name = (r.get("column_name") or "").lower()
dt = r.get("data_type")
cl = r.get("character_maximum_length")
p = r.get("numeric_precision")
s = r.get("numeric_scale")
ident = bool(r.get("is_identity"))
out[name] = {
"type_sql": _canon_type(dt, cl, p, s),
"is_identity": ident,
"char_len": cl,
"numeric_precision": p,
"numeric_scale": s,
}
return out
def _plan_column(self, col: Dict, target_type_map: Dict[str, Dict]) -> Dict:
"""
Returns a dict with:
- name
- rs_final_type: from _pg_to_rs_type
- export_expr: PG-side SELECT expression
- stage_type: Redshift staging column type (usually VARCHAR(65535))
- load_expr: Redshift-side SELECT expression used in INSERT…SELECT
"""
import re
def _strip_identity(t: str) -> str:
# remove: <type> IDENTITY(…)
return re.sub(r"\s+IDENTITY\s*\([^)]*\)", "", t, flags=re.IGNORECASE)
name = col["column_name"]
name_l = name.lower()
qn = f'"{name}"'
rs_final = _pg_to_rs_type(col, logger=self.log)
rs_final_up = rs_final.upper()
# Default: export as text, stage as big VARCHAR, then CAST on load
export_expr = qn
stage_type = "VARCHAR(65535)"
# Figure out the actual target RS type (via introspection), fallback to mapping if missing.
tgt_info = target_type_map.get(name_l) or {}
rs_final = tgt_info.get("type_sql") or _pg_to_rs_type(col, logger=self.log)
rs_final_up = (rs_final or "").upper()
# If target column is identity -> you cannot set it explicitly. Omit it from INSERT.
if tgt_info.get("is_identity"):
self.log.warning(
"Target column %s is IDENTITY in Redshift; it will be omitted from INSERT (values will be auto-generated). "
"If you need to preserve PG values, drop IDENTITY from the target table DDL for this column.",
name,
)
# We still stage/export it (harmless), but mark omit=True so INSERT skips it.
return {
"name": name,
"rs_final_type": rs_final,
"export_expr": export_expr,
"stage_type": stage_type,
"load_expr": qn, # unused when omit=True
"omit": True,
}
# SUPER-like (json/jsonb/arrays)
if self._is_superish(col):
# Use CASE to export true NULL (empty cell) instead of literal 'null' string
export_expr = f"CASE WHEN {qn} IS NULL THEN NULL ELSE to_json({qn})::text END"
# SUPER columns in CSV must contain valid JSON. Redshift COPY parses JSON strings into SUPER directly.
# Using SUPER type allows up to 16MB per value (vs VARCHAR 65535 limit).
stage_type = "SUPER"
load_expr = f"{qn}" # No JSON_PARSE needed; COPY handles parsing into SUPER from CSV
# BYTEA (VARBYTE) handling
elif (col.get("udt_name") or "").lower() == "bytea":
# encode(NULL, 'hex') returns NULL in PG, but be explicit for safety
export_expr = f"CASE WHEN {qn} IS NULL THEN NULL ELSE encode({qn}, 'hex') END"
stage_type = "VARCHAR(65535)"
load_expr = f"CASE WHEN {qn} IS NULL THEN NULL ELSE FROM_HEX({qn}) END AS {qn}"
return {
"name": name,
"rs_final_type": rs_final,
"export_expr": export_expr,
"stage_type": stage_type,
"load_expr": load_expr,
}
# Geometry/Geography as hex EWKB
# Redshift COPY into GEOMETRY/GEOGRAPHY columns requires hex EWKB format
# (not WKT). ST_AsEWKB returns the binary, encode(...,'hex') makes it
# a hex string that COPY can parse directly into the native spatial type.
elif self._is_geom(col):
export_expr = (
f"CASE WHEN {qn} IS NULL THEN NULL "
f"WHEN OCTET_LENGTH(ST_AsEWKB({qn})) > {_MAX_GEOMETRY_BYTES} THEN NULL "
f"ELSE encode(ST_AsEWKB({qn}), 'hex') END"
)
if (col.get("udt_name") or "").lower() == "geography":
stage_type = "GEOGRAPHY"
else:
stage_type = "GEOMETRY"
load_expr = f"{qn}"
else:
# Primitives & strings — cast based on final RS type.
# We keep staging as VARCHAR and cast with NULLIF('', NULL) to be safe.
def cast_expr(sql_type: str) -> str:
# numeric / ints
if any(
sql_type.startswith(k) for k in ["SMALLINT", "INTEGER", "BIGINT"]
):
base = _strip_identity(sql_type)
return f"CAST(NULLIF({qn}, '') AS {base}) AS {qn}"
if sql_type.startswith("DECIMAL") or sql_type.startswith("NUMERIC"):
return f"CAST(NULLIF({qn}, '') AS {sql_type}) AS {qn}"
# floats
if (
sql_type.startswith("DOUBLE PRECISION")
or sql_type.startswith("REAL")
or sql_type.startswith("FLOAT")
):
return f"CAST(NULLIF({qn}, '') AS {sql_type}) AS {qn}"
# boolean
if sql_type == "BOOLEAN":
# Redshift is strict about varchar->boolean.
# Normalize common truthy/falsey strings and NULL/blank safely.
return (
f"CASE "
f"WHEN NULLIF(TRIM({qn}), '') IS NULL THEN NULL "
f"WHEN UPPER(TRIM({qn})) IN ('TRUE','T','YES','Y','1') THEN TRUE "
f"WHEN UPPER(TRIM({qn})) IN ('FALSE','F','NO','N','0') THEN FALSE "
f"ELSE NULL "
f"END AS {qn}"
)
# dates/times
if sql_type in {"DATE", "TIME", "TIMESTAMP", "TIMESTAMPTZ"}:
return f"CAST(NULLIF({qn}, '') AS {sql_type}) AS {qn}"
# varchars/chars/varbyte/etc -> no cast needed
return f"{qn}"
load_expr = cast_expr(rs_final_up)
return {
"name": name,
"rs_final_type": rs_final,
"export_expr": export_expr,
"stage_type": stage_type,
"load_expr": load_expr,
}
def _plan_columns(
self, columns: List[Dict], target_type_map: Dict[str, Dict]
) -> List[Dict]:
return [self._plan_column(c, target_type_map) for c in columns]
####################################################################################################
## PG to S3 Export via aws_s3
####################################################################################################
def _export_pg_to_s3(self, select_sql: str, job_prefix: str, json_mode: bool = False) -> Tuple[str, str]:
"""
Export SELECT to S3 via aws_s3. Uses 6-TEXT credentials signature on this build.
Returns the S3 prefix used (without trailing slash).
When json_mode=True, the select_sql is expected to output a single text
column (row_to_json). We use CSV format with control-character
delimiter/quote so the raw JSON passes through without escaping.
"""
ext = ".jsonl" if json_mode else ".csv"
key_prefix = f"{self.s3_prefix}/{self.source_schema}.{self.table}/{job_prefix}"
object_key = f"{key_prefix}/part-{uuid.uuid4().hex[:8]}{ext}"
if json_mode:
# Output raw NDJSON: use CSV format with control-char delimiter/quote
# so PG doesn't wrap or escape the JSON text. Valid JSON never
# contains literal 0x01/0x02 (they'd be \u0001/\u0002).
opts_parts = [
"format csv",
r"delimiter E'\x01'",
r"quote E'\x02'",
]
else:
opts_parts = [
"format csv",
f"header {'true' if self.header else 'false'}",
]
# Added logic to include compression in the opts text here if we later have this capability
# for now the dag will dictate if this is true or false
if self.gzip:
opts_parts.append("compression GZIP")
opts_text = ", ".join(opts_parts)
# Signatures we will use
sql_alltext_creds = """
SELECT aws_s3.query_export_to_s3(
%s, -- query
%s, -- bucket
%s, -- prefix
%s, -- region
%s, -- options (TEXT)
%s, -- aws_access_key_id
%s, -- aws_secret_access_key
%s -- session_token
);
"""
# We pass NULL there to keep AES256 (no KMS). We won't use this path unless you *want* KMS.
# We do not currently use KMS in the target bucket
sql_s3uri_kms = """
SELECT aws_s3.query_export_to_s3(
%s, -- query
aws_commons.create_s3_uri(%s, %s, %s), -- bucket, prefix, region
%s, -- options (TEXT)
NULL -- kms_key_id (NULL => SSE-S3/AES256)
);
"""
hook = PostgresHook(postgres_conn_id=self.pg_conn_id)
self.log.info(
"Exporting via aws_s3 to s3://%s/%s/ ...", self.s3_bucket, key_prefix
)
# Probe available signatures (for future-proofing)
sigs = {
r[0]
for r in hook.get_records(
"""
SELECT p.oid::regprocedure::text
FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
WHERE n.nspname = 'aws_s3' AND p.proname = 'query_export_to_s3';
"""
)
}
# Prefer the credentials variant on this build
sig_texts = {s.replace(" ", "") for s in sigs}
cred8 = "(text,text,text,text,text,text,text,text)"
if any(cred8 in s for s in sig_texts):
# Assume _sts_creds() returns (ak, sk, st) from a role with PutObject/List permission on the target bucket/prefix
ak, sk, st = self._sts_creds()
params = (
select_sql,
self.s3_bucket,
object_key,
self.region,
opts_text,
ak,
sk,
st,
)
hook.run(sql_alltext_creds, parameters=params)
else:
# Fallback to the KMS form but keep KMS NULL so bucket default AES256 is used
params = (
select_sql,
self.s3_bucket,
object_key,
self.region,
opts_text,
)
hook.run(sql_s3uri_kms, parameters=params)
return key_prefix, object_key
################################################################################################
## Format SQL for Staging and Copying
################################################################################################
def _stg_name(self) -> str:
# short, stable name per table
return f"__stg_{self.table}".lower()
def _stg_qual(self) -> str:
return f'"{self.rs_stage_schema}"."{self._stg_name()}"'
def _sql_prepare_stage(self, plan: list[dict]) -> str:
cols = ", ".join(f'"{p["name"]}" {p["stage_type"]}' for p in plan)
stg = self._stg_qual()
return f"""
CREATE SCHEMA IF NOT EXISTS "{self.rs_stage_schema}";
DROP TABLE IF EXISTS {stg};
CREATE TABLE {stg} ({cols});
"""
def _sql_copy_into(self, s3_object_key: str, json_mode: bool = False) -> str:
self.log.info(f"DEBUG _sql_copy_into: redshift_iam_role={self.redshift_iam_role!r}")
if not (
self.redshift_iam_role
and self.redshift_iam_role.startswith("arn:aws:iam::")
):
raise AirflowException(f"redshift_iam_role must be a valid role ARN. Got: {self.redshift_iam_role!r}")
iam_role = f"IAM_ROLE '{self.redshift_iam_role}'"
maxerror_opt = f"MAXERROR {self.max_errors}" if self.max_errors > 0 else ""
if json_mode:
# JSON COPY: each line is a JSON object with keys matching staging
# column names. Supports up to 4MB per value (vs CSV's 65KB limit).
opts = " ".join(
x
for x in [
"FORMAT AS JSON 'auto'",
"GZIP" if self.gzip else "",
"COMPUPDATE OFF",
maxerror_opt,
]
if x
)
else:
opts = " ".join(
x
for x in [
"FORMAT AS CSV",
"IGNOREHEADER 1" if self.header else "",
"GZIP" if self.gzip else "",
# EMPTYASNULL/BLANKSASNULL conflate PG's ``''`` (exported
# as ``""``) with PG NULL (exported as unquoted empty).
# Default keeps the legacy behavior; opt out via
# ``empty_as_null=False`` when the target has NOT NULL
# text columns and the source legitimately stores ``''``.
(
"EMPTYASNULL BLANKSASNULL TRIMBLANKS ACCEPTINVCHARS"
if _as_bool(self.empty_as_null)
else "ACCEPTINVCHARS"
),
"TIMEFORMAT 'auto' DATEFORMAT 'auto'",
"COMPUPDATE OFF",
maxerror_opt,
]
if x
)
stg = self._stg_qual()
return f"""
COPY {stg}
FROM 's3://{self.s3_bucket}/{s3_object_key}'
REGION '{self.region}'
{iam_role}
{opts};
"""
def _sql_insert_into(self, plan: list[dict]) -> str:
stg = self._stg_qual()
tgt = f'"{self.target_schema}"."{self.table}"'
# Skip columns marked omit=True (e.g., target identity columns)
effective_plan = [p for p in plan if not p.get("omit")]
if not effective_plan:
# No columns to insert (e.g., all were identity) – nothing to do.
return f"/* No loadable columns for {tgt}; skipping INSERT */ SELECT 1;"
cols = ", ".join(f'"{p["name"]}"' for p in effective_plan)
selects = ", ".join(p["load_expr"] for p in effective_plan)
return f"""
INSERT INTO {tgt} ({cols})
SELECT {selects}
FROM {stg};
"""
def _count_target_rows(self) -> int:
rows = self._run_rs_and_fetch(
f'SELECT COUNT(*) AS n FROM "{self.target_schema}"."{self.table}";'
)
return int(rows[0]["n"]) if rows else 0
def _latest_copy_metrics_for_file(self, s3_uri: str) -> dict:
rows = self._run_rs_and_fetch(
f"""
SELECT
query_id,
status,
COALESCE(loaded_rows, 0) AS loaded_rows,
error_count
FROM sys_load_history
WHERE data_source LIKE '%{s3_uri}%'
ORDER BY start_time DESC
LIMIT 1;
"""
)
return rows[0] if rows else {}
def _latest_insert_metrics(self, timeout_s: int = 60) -> dict:
# Find the most recent INSERT into this exact table from this user/db
q = f"""
SELECT
query_id,
start_time,
status,
trim(COALESCE(error_message,'')) AS error,
returned_rows -- row count affected for INSERT
FROM sys_query_history
WHERE 1 = 1
AND database_name = '{self.rs_database}'
AND query_text ILIKE '%%INSERT INTO "{self.target_schema}"."{self.table}"%%'
ORDER BY start_time DESC
LIMIT 1;
"""
t0 = time.time()
last = {}
while time.time() - t0 < timeout_s:
rows = self._run_rs_and_fetch(q)
if rows:
last = rows[0]
status = (last.get("status") or "").strip().lower()
if status in {"finished", "failed", "aborted"}:
break
time.sleep(1.0)
return last
################################################################################################
## Execute the Copy and Insert as a single transaction
################################################################################################
def _build_refresh_batch(self, plan: list[dict]) -> list[str]:
"""Build the ordered list of SQL statements that batch_execute_statement
will run as a single Redshift transaction.
Statement order:
(optional) SET enable_case_sensitive_identifier TO true
(refresh) DELETE FROM target [WHERE delete_where_clause]
INSERT INTO target (...) SELECT ... FROM stage
Notes:
* Do NOT include BEGIN/COMMIT — batch_execute_statement owns the
transaction and rolls back the entire batch on any failure.
* TRUNCATE is intentionally avoided: TRUNCATE auto-commits in
Redshift and cannot participate in a transaction. Bare
DELETE FROM target is the AWS-blessed transactional pattern for
atomic full-refresh.
* APPEND mode emits only the INSERT (no purge).
"""
mode = self.load_mode
if mode not in {"APPEND", "TRUNCATE_INSERT", "DELETE_INSERT"}:
raise AirflowException(
f"Invalid load_mode {mode!r}; must be APPEND, TRUNCATE_INSERT, "
f"or DELETE_INSERT."
)
# Guard: in refresh modes, refuse to DELETE the target if there are no
# loadable columns to re-insert. _sql_insert_into returns a harmless
# `SELECT 1` sentinel when every column is omitted (e.g., all identity);
# combined with a preceding DELETE that would silently empty the table.
if mode in self._REFRESH_MODES and not [p for p in plan if not p.get("omit")]:
raise AirflowException(
f"Refresh mode {mode} has no loadable columns for "
f"{self.target_schema}.{self.table}; refusing to purge target. "
f"Check identity-column detection or source/target schema drift."
)
tgt = f'"{self.target_schema}"."{self.table}"'
sqls: list[str] = []
# batch_execute_statement runs all statements in the same session, so
# SET enable_case_sensitive_identifier must be the first statement when
# any subsequent identifier needs case-sensitive resolution.
if _as_bool(self.case_sensitive_identifiers):
sqls.append("SET enable_case_sensitive_identifier TO true")
if mode == "TRUNCATE_INSERT":
sqls.append(f"DELETE FROM {tgt}")
elif mode == "DELETE_INSERT":
dwc = (self.delete_where_clause or "").strip()
if dwc.upper().startswith("WHERE "):
dwc = dwc[6:].strip()
if not dwc:
raise AirflowException(
"DELETE_INSERT requires a non-empty 'delete_where_clause'."
)
sqls.append(f"DELETE FROM {tgt} WHERE {dwc}")
# APPEND: no purge.
insert_sql = self._sql_insert_into(plan).strip().rstrip(";")
sqls.append(insert_sql)
return sqls
def _run_batch_in_transaction(self, sqls: list[str]) -> None:
"""Run a list of SQL statements as a single Redshift transaction.
Uses Redshift Data API's ``execute_statement`` with the statements
joined by semicolons and wrapped in ``BEGIN; ... COMMIT;``. Redshift
treats the wrapped block as one transaction: if any inner statement
fails, the explicit COMMIT is never reached and Redshift rolls back
the transaction.
We use ``execute_statement`` rather than ``batch_execute_statement``
because the latter requires the IAM permission
``redshift-data:BatchExecuteStatement``, which is not granted on our
default EC2 role today. ``execute_statement`` only requires
``redshift-data:ExecuteStatement`` (already granted) and supports
multi-statement SQL as a single submission. The transactional
semantics are equivalent — the API shape differs but the underlying
Redshift transaction behavior is the same.
Also enforces a soft total-payload size guard so we fail fast with a
clear message instead of an opaque Data API error if the combined SQL
approaches the 100KB per-statement limit.
"""
import boto3
# Soft guard well under the documented 100KB Data API per-statement limit.
_MAX_STMT_BYTES = 95_000
# Build single transactional payload: BEGIN; <stmt>; <stmt>; ... COMMIT;
body = "\n".join(s.rstrip(";").strip() + ";" for s in sqls)
payload = f"BEGIN;\n{body}\nCOMMIT;"
size = len(payload.encode("utf-8"))
self.log.info(
"Transactional payload size=%s bytes, %d inner statement(s)",
size,
len(sqls),
)
if size > _MAX_STMT_BYTES:
raise AirflowException(
f"Generated Redshift transactional payload is {size} bytes, "
f"exceeds soft limit of {_MAX_STMT_BYTES} bytes (Data API hard "
f"limit ~100KB). Reduce column-list width or split the operation."
)
client = boto3.client("redshift-data", region_name=self.region)
resp = client.execute_statement(
WorkgroupName=self.rs_workgroup,
Database=self.rs_database,
Sql=payload,
WithEvent=False,
)
sid = resp["Id"]
while True:
d = client.describe_statement(Id=sid)
status = d["Status"]
if status in ("FINISHED", "FAILED", "ABORTED"):
break
time.sleep(1.0)
if status != "FINISHED":
err = d.get("Error") or status
raise AirflowException(
f"Redshift transactional payload {status}: {err}"
)
def _copy_and_insert(self, s3_object_key: str, plan, json_mode: bool = False):
stg = self._stg_qual()
try:
# 1) Prepare durable stage (DROP + CREATE, idempotent; safe across
# cleanup_after_load=False debug sessions)
prep_sql = "\n".join(
[
_case_prefix(self.case_sensitive_identifiers),
self._sql_prepare_stage(plan),
]
)
self.rs.execute_query(
sql=prep_sql,
workgroup_name=self.rs_workgroup,
database=self.rs_database,
wait_for_completion=True,
)
# 2) COPY into stage
copy_sql = _case_prefix(
self.case_sensitive_identifiers
) + self._sql_copy_into(s3_object_key, json_mode=json_mode)
self.rs.execute_query(
sql=copy_sql,
workgroup_name=self.rs_workgroup,
database=self.rs_database,
wait_for_completion=True,
)
# 3) Sanity check rows in stage + enforce COPY error policy
stg_n = self._run_rs_and_fetch(f"SELECT COUNT(*) AS n FROM {stg};")
self.log.info("Staging rowcount just after COPY: %s", stg_n)
if self.max_errors > 0:
metrics = self._latest_copy_metrics_for_file(s3_object_key)
err_count = int(metrics.get("error_count", 0) or 0)
if err_count:
if self.load_mode in self._REFRESH_MODES:
# CRITICAL: refusing to mutate target from an incomplete
# stage. With atomic full-refresh, partial stage would
# commit incomplete data; better to fail loudly.
raise AirflowException(
f"COPY skipped {err_count} row(s) due to errors "
f"(MAXERROR={self.max_errors}); refusing to refresh "
f"target {self.target_schema}.{self.table} from an "
f"incomplete stage. Check sys_load_error_detail."
)
self.log.warning(
"COPY skipped %s row(s) due to errors (MAXERROR=%s). "
"Check sys_load_error_detail for details.",
err_count,
self.max_errors,
)
# 4) Atomic DELETE+INSERT (or INSERT-only for APPEND) in one txn.
# Owns the purge for TRUNCATE_INSERT/DELETE_INSERT — DAG's
# optional_purge must be a no-op in those modes.
self.log.info(
"Refresh mode=%s; submitting batch transaction (DELETE? + INSERT)",
self.load_mode,
)
self._run_batch_in_transaction(self._build_refresh_batch(plan))
finally:
# 5) (optional) cleanup stage. Wrap in try/except so a stage-drop
# failure does not mask the original transaction error.
if _as_bool(self.cleanup_after_load):
self.log.info("cleanup_after_load=True; dropping stage table %s", stg)
try:
drop_sql = (
_case_prefix(self.case_sensitive_identifiers)
+ f"DROP TABLE IF EXISTS {self._stg_qual()};"
)
self.rs.execute_query(
sql=drop_sql,
workgroup_name=self.rs_workgroup,
database=self.rs_database,
wait_for_completion=True,
)
except Exception:
self.log.exception(
"Failed to drop stage table %s during cleanup; manual "
"drop may be required.",
stg,
)
else:
self.log.info(
"cleanup_after_load=False; leaving stage table %s for inspection",
stg,
)
####################################################################################################
## S3 Cleanup
####################################################################################################
def _cleanup_s3_prefix(self, key_prefix: str) -> None:
"""
Delete all S3 objects under s3://{self.s3_bucket}/{key_prefix}/
Uses batched DeleteObjects (1000 at a time). No-op if nothing found.
NOTE: For versioned buckets this creates delete markers (does not purge versions).
"""
import boto3
def _norm(p: Optional[str]) -> str:
# normalize for comparison — strip leading/trailing slashes, coerce to str
return (str(p or "")).strip().lstrip("/").rstrip("/")
# Panic guards to prevent gratuitous deletions
norm_prefix = _norm(key_prefix)
# Guard: bucket root
if not norm_prefix:
raise AirflowException(
"Refusing to clean up S3 bucket root: prefix is empty. "
"This is a panic exit to prevent accidental mass deletes."
)
# Guard: exact match against configured panic roots (literal prefixes)
panic_values = {
_norm(x)
for x in getattr(Config, "PANIC_PREFIXES", [])
if isinstance(x, str)
}
if norm_prefix in panic_values:
raise AirflowException(
f"Refusing to clean up unsafe root prefix '{key_prefix}' "
f"(matches Config.PANIC_PREFIXES). Panic exit."
)
# Execute delete
s3 = boto3.client("s3")
paginator = s3.get_paginator("list_objects_v2")
safe_prefix = norm_prefix + "/"
total_deleted = 0
for page in paginator.paginate(Bucket=self.s3_bucket, Prefix=safe_prefix):
contents = page.get("Contents", [])
if not contents:
continue
batch = []
for obj in contents:
batch.append({"Key": obj["Key"]})
if len(batch) == 1000:
s3.delete_objects(Bucket=self.s3_bucket, Delete={"Objects": batch})
total_deleted += len(batch)
batch = []
if batch:
s3.delete_objects(Bucket=self.s3_bucket, Delete={"Objects": batch})
total_deleted += len(batch)
self.log.info(
"Cleaned up %s objects under s3://%s/%s",
total_deleted,
self.s3_bucket,
safe_prefix,
)
####################################################################################################
## Main Execute
####################################################################################################
[docs]
def execute(self, context):
# Log cleanup status
cleanup = _as_bool(self.cleanup_after_load)
self.log.info(
"cleanup_after_load (parsed)=%s raw=%r", cleanup, self.cleanup_after_load
)
# 0) inspect PG
cols = self._get_pg_columns()
if not cols:
raise ValueError(
f"Source table not found: {self.source_schema}.{self.table}"
)
# 1) inspect Redshift target types/identity
target_type_map = self._fetch_rs_target_columns()
# 2) build transport plan using the *actual target* RS types
plan = self._plan_columns(cols, target_type_map)
# 2b) Auto-detect JSON mode: if any column stages as SUPER, use JSON
# export/COPY to avoid CSV's 65KB per-field limit (JSON supports 4MB).
# Exception: if the table also has GEOMETRY/GEOGRAPHY columns, fall
# back to CSV — JSON COPY cannot ingest spatial types, and hex EWKB
# can exceed VARCHAR(65535) staging. CSV COPY handles GEOMETRY natively
# (up to ~1MB) while SUPER columns still work (with 65KB per-field cap).
has_super = any(
(p.get("stage_type") or "").upper() == "SUPER" for p in plan
)
has_spatial = any(
(p.get("stage_type") or "").upper() in {"GEOMETRY", "GEOGRAPHY"}
for p in plan
)
json_mode = has_super and not has_spatial
if has_super and has_spatial:
self.log.warning(
"Table has both SUPER and GEOMETRY/GEOGRAPHY columns. "
"Using CSV mode (GEOMETRY requires it). SUPER column values "
">65KB will be truncated by COPY."
)
if json_mode:
self.log.info(
"JSON mode enabled (table has SUPER columns); "
"export will use row_to_json + JSON COPY"
)
for p in plan:
st = (p.get("stage_type") or "").upper()
if st == "SUPER":
# SUPER columns: use raw column reference so row_to_json
# serializes jsonb/arrays as nested JSON (not double-encoded text).
p["export_expr"] = f'"{p["name"]}"'
else:
# Non-SUPER columns: cast to ::text so row_to_json emits
# JSON strings (not typed scalars like numbers/booleans).
# This ensures VARCHAR staging columns receive strings,
# matching CSV mode behavior for downstream CAST/NULLIF.
p["export_expr"] = f'({p["export_expr"]})::text'
# 3) PG SELECT with export_exprs
select_list = ", ".join(f'{p["export_expr"]} AS "{p["name"]}"' for p in plan)
base_from = f'FROM "{self.source_schema}"."{self.table}"'
where_sql = ""
if self.where_clause:
wc = self.where_clause
# avoid accidental double "WHERE"
where_sql = (
f" WHERE {wc[5:].strip()}"
if wc.strip().upper().startswith("WHERE ")
else f" WHERE {wc}"
)
select_sql = f"SELECT {select_list} {base_from}{where_sql};"
if json_mode:
# Wrap as NDJSON: one JSON object per row. row_to_json serializes
# all column types natively (jsonb as nested JSON, timestamps as
# ISO strings, NULLs as JSON null).
select_sql = (
f"SELECT row_to_json(t)::text "
f"FROM ({select_sql.rstrip(';')}) t;"
)
# 4) export to S3
job_prefix = uuid.uuid4().hex[:12]
s3_object_prefix, s3_object_key = self._export_pg_to_s3(
select_sql, job_prefix, json_mode=json_mode
)
# 5) COPY to staging and INSERT to final in Redshift
self.rs = RedshiftDataHook(
aws_conn_id=self.aws_conn_id, region_name=self.region
)
self.log.info("COPY FROM s3://%s/%s", self.s3_bucket, s3_object_key)
self._copy_and_insert(s3_object_key, plan, json_mode=json_mode)
# 6) optional S3 cleanup – only after successful INSERT
if cleanup:
try:
self._cleanup_s3_prefix(s3_object_prefix)
except Exception as e:
self.log.warning(
"S3 cleanup failed for s3://%s/%s: %s",
self.s3_bucket,
s3_object_prefix,
e,
)
self.log.info("Load complete for %s.%s", self.target_schema, self.table)
return {"s3_object_prefix": s3_object_prefix}