"""
S3 to Redshift Parquet Operator
Ingests Parquet files from S3 into Redshift with optional column mapping
and geometry field creation.
data@ready.net | 2026-01-15
"""
from typing import Optional, Dict, List, Literal, Union, Sequence
import logging
import time
from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
from airflow.exceptions import AirflowException
from config import Config
from util import Util
WriteMode = Literal["append", "overwrite"]
GeometryType = Literal["latlon", "wkt"]
def _as_bool(v: Union[str, bool, None]) -> bool:
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 _case_prefix(flag: Union[str, bool]) -> str:
return "SET enable_case_sensitive_identifier TO true;\n" if _as_bool(flag) else ""
[docs]
class S3ToRedshiftOperator(BaseOperator):
"""
Ingest Parquet files from S3 into Redshift.
Supports:
- Column mapping (rename columns during ingestion)
- Geometry creation from lat/lon or WKT columns
- Append or overwrite write modes
:param s3_path: S3 path to parquet files (e.g., s3://bucket/prefix/)
:param rs_workgroup: Redshift Serverless workgroup name
:param rs_database: Redshift database name
:param target_schema: Target schema in Redshift
:param target_table: Target table name in Redshift
:param redshift_iam_role: IAM role ARN for Redshift to access S3
:param column_mapping: Optional dict mapping source columns to target columns
Example: {"src_col": "tgt_col", "old_name": "new_name"}
:param geometry_type: Type of geometry to create: "latlon" or "wkt"
:param geometry_config: Config for geometry creation:
- For latlon: {"lat_column": "lat", "lon_column": "lon", "geometry_column": "geom", "srid": 4326}
- For wkt: {"wkt_column": "wkt", "geometry_column": "geom", "srid": 4326}
:param write_mode: "append" to add rows, "overwrite" to replace table
:param region: AWS region
:param aws_conn_id: Airflow AWS connection ID
:param case_sensitive_identifiers: Enable case-sensitive identifiers in Redshift
"""
template_fields: Sequence[str] = (
"s3_path",
"rs_workgroup",
"rs_database",
"target_schema",
"target_table",
"redshift_iam_role",
"column_mapping",
"geometry_type",
"geometry_config",
"write_mode",
"region",
"aws_conn_id",
)
ui_color = "#E87C3E"
def __init__(
self,
*,
s3_path: str,
rs_workgroup: str,
rs_database: str,
target_schema: str,
target_table: str,
redshift_iam_role: Optional[str] = None,
column_mapping: Optional[Dict[str, str]] = None,
geometry_type: Optional[GeometryType] = None,
geometry_config: Optional[Dict[str, str]] = None,
reproject_srid: Optional[int] = 4326,
h3_index_resolution: Optional[int] = 9,
write_mode: WriteMode = "append",
region: str = "us-east-1",
aws_conn_id: str = "aws_default",
case_sensitive_identifiers: Union[str, bool] = True,
**kwargs,
):
super().__init__(**kwargs)
role_input = redshift_iam_role
if not role_input or not str(role_input).strip():
role_input = Config.REDSHIFT_S3_ROLE_ARN
self.s3_path = s3_path.rstrip("/")
self.rs_workgroup = rs_workgroup
self.rs_database = rs_database
self.target_schema = target_schema
self.target_table = target_table
self.redshift_iam_role = str(role_input or "").strip()
self.column_mapping = column_mapping or {}
self.geometry_type = geometry_type
self.geometry_config = geometry_config or {}
self.write_mode = write_mode
self.region = region
self.aws_conn_id = aws_conn_id
self.case_sensitive_identifiers = case_sensitive_identifiers
if self.geometry_config:
self.reproject_srid = reproject_srid or self.geometry_config.get(
"srid", 4326
)
self.h3_index_resolution = h3_index_resolution
self._validate_params()
def _validate_params(self):
if not self.redshift_iam_role.startswith("arn:aws:iam::"):
if "{{" not in self.redshift_iam_role:
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 Config.REDSHIFT_S3_ROLE_ARN or pass redshift_iam_role param."
)
if self.geometry_type == "latlon":
required = {"lat_column", "lon_column"}
missing = required - set(self.geometry_config.keys())
if missing:
raise AirflowException(
f"geometry_config missing required keys for latlon: {missing}"
)
if self.write_mode not in ("append", "overwrite"):
raise AirflowException(
f"write_mode must be 'append' or 'overwrite', got: {self.write_mode}"
)
def _run_rs_and_fetch(self, sql: str) -> List[Dict]:
import boto3
self.log.info(f"Running SQL: {sql}")
client = boto3.client("redshift-data", region_name=self.region)
resp = client.execute_statement(
WorkgroupName=self.rs_workgroup,
Database=self.rs_database,
Sql=sql,
)
stmt_id = resp["Id"]
while True:
desc = client.describe_statement(Id=stmt_id)
status = desc["Status"]
if status in ("FINISHED", "FAILED", "ABORTED"):
break
time.sleep(0.5)
if status == "FAILED":
raise AirflowException(
f"Redshift query failed: {desc.get('Error', 'Unknown error')}"
)
if not desc.get("HasResultSet"):
return []
result = client.get_statement_result(Id=stmt_id)
columns = [c["name"] for c in result["ColumnMetadata"]]
rows = []
for rec in result["Records"]:
row = {}
for i, field in enumerate(rec):
val = None
for k, v in field.items():
if k != "isNull":
val = v
break
row[columns[i]] = val
rows.append(row)
return rows
def _run_rs(self, sql: str) -> None:
self.log.info(f"Running SQL: {sql}")
self.rs.execute_query(
sql=sql,
workgroup_name=self.rs_workgroup,
database=self.rs_database,
wait_for_completion=True,
poll_interval=30,
)
def _qual(self, schema: str, table: str) -> str:
return f'"{schema}"."{table}"'.lower()
def _target_qual(self) -> str:
return self._qual(self.target_schema, self.target_table)
def _stage_qual(self) -> str:
return self._qual(self.get_spectrum_schema(), f"{self.target_table}_temp")
[docs]
def get_spectrum_schema(self) -> str:
return "spectrum_readydatalake" if Config.PRODUCTION else "spectrum_shoreline"
def _get_stage_columns(self) -> List[str]:
sql = f"""
SELECT columnname as column_name
FROM svv_external_columns
WHERE schemaname = '{self.get_spectrum_schema()}'
AND tablename = lower('{self.target_table}')
"""
rows = self._run_rs_and_fetch(sql)
return [r["column_name"] for r in rows]
def _build_select_columns(self, stage_columns: List[str]) -> str:
select_parts = []
geom_col = self.geometry_config.get("target_column", "geom")
srid = self.geometry_config.get("srid", 4326)
self.log.info(f"Geometry config: {self.geometry_config}")
self.log.info(f"geom_col : {geom_col}")
self.log.info(f"srid : {srid}")
self.log.info(f"reproject_srid : {self.reproject_srid}")
skip_cols = set()
if self.geometry_type == "latlon":
skip_cols = {
self.geometry_config.get("lat_column"),
self.geometry_config.get("lon_column"),
}
elif self.geometry_type == "wkt":
skip_cols = {self.geometry_config.get("input_column")}
elif self.geometry_type == "wkb":
skip_cols = {self.geometry_config.get("input_column")}
elif self.geometry_type == "ewkb":
skip_cols = {self.geometry_config.get("input_column")}
skip_cols.add("geometry_bbox") # in case it is a geoparquet
skip_cols.add("geometry") # in case it is a geoparquet
for col in stage_columns:
if col in skip_cols:
continue
target_col = self.column_mapping.get(col, col)
select_parts.append(f'"{col}" AS "{target_col}"')
if self.geometry_type == "latlon":
lat_col = self.geometry_config["lat_column"]
lon_col = self.geometry_config["lon_column"]
geom_expr = f'ST_Transform(ST_SetSRID(ST_Point("{lon_col}", "{lat_col}"), {srid}), {self.reproject_srid}) AS "{geom_col}"'
select_parts.append(geom_expr)
elif self.geometry_type == "wkt":
wkt_col = self.geometry_config["input_column"]
geom_expr = f'ST_Transform(ST_GeomFromText("{wkt_col}", {srid}), {self.reproject_srid}) AS "{geom_col}"'
select_parts.append(geom_expr)
elif self.geometry_type == "wkb":
wkb_col = self.geometry_config["input_column"]
geom_expr = f'ST_Transform(ST_GeomFromWKB("{wkb_col}", {srid}), {self.reproject_srid}) AS "{geom_col}"'
select_parts.append(geom_expr)
elif self.geometry_type == "ewkb":
ewkb_col = self.geometry_config["input_column"]
geom_expr = f'ST_Transform(ST_SetSRID(ST_GeomFromEWKB("{ewkb_col}"),{srid} ), {self.reproject_srid}) AS "{geom_col}"'
select_parts.append(geom_expr)
self.log.info(f"Select columns: {select_parts}")
return ",\n ".join(select_parts)
def _sql_insert_from_stage(self, stage_columns: List[str]) -> str:
select_cols = self._build_select_columns(stage_columns)
return f"""
{self._sql_create_table(select_cols)};
INSERT INTO {self._target_qual()}
SELECT
{select_cols}
FROM {self._stage_qual()};
"""
def _sql_add_h3_index(self):
geometry_column = self.geometry_config.get("target_column", "geom")
return f"""
ALTER TABLE {self._target_qual()}
ADD h3_index BIGINT;
UPDATE {self._target_qual()}
SET h3_index = H3_FromPoint(ST_Centroid({geometry_column}), {self.h3_index_resolution});
"""
def _sql_overwrite_from_stage(self, stage_columns: List[str]) -> str:
select_cols = self._build_select_columns(stage_columns)
return f"""
BEGIN TRANSACTION;
{self._sql_create_table(select_cols)};
DELETE FROM {self._target_qual()};
INSERT INTO {self._target_qual()}
SELECT
{select_cols}
FROM {self._stage_qual()};
END TRANSACTION;
"""
def _sql_create_table(self, select_cols: str) -> str:
return f"""
DROP TABLE IF EXISTS {self._target_qual()};
CREATE TABLE {self._target_qual()}
AS
SELECT {select_cols} from {self._stage_qual()}
"""
[docs]
def ingest_to_athena(self):
import awswrangler as wr
Util.remove_datalake_table(tgt_table=self.target_table)
Util.remove_datalake_table(tgt_table=f"{self.target_table}_temp")
# Read GeoParquet (geometry becomes WKB binary)
df = wr.s3.read_parquet(self.s3_path)
# Write to Parquet + register table in Glue Catalog (Athena)
wr.s3.to_parquet(
df=df,
path=f"s3://{Config.DATA_LAKE_BUCKET}/raw/{self.target_schema}/{self.target_table}/",
database=Config.DATA_LAKE_DB,
dataset=True,
index=False,
mode=self.write_mode,
table=self.target_table,
)
wr.athena.create_ctas_table(
sql=f"SELECT to_hex(ST_AsBinary(st_geomfrombinary(geometry))) AS geometry_hex, * FROM {self.target_table}",
database=Config.DATA_LAKE_DB,
workgroup=Config.DEFAULT_WORKGROUP,
ctas_table=f"{self.target_table}_temp",
wait=True,
storage_format="PARQUET",
write_compression="snappy",
s3_output=f"s3://{Config.DATA_LAKE_BUCKET}/raw/{self.target_schema}/{self.target_table}_temp/",
)
# Now queryable in Athena as: SELECT * FROM your_database.your_table_name
[docs]
def execute(self, context):
context["params"]["map_name"] = self._target_qual()
context["map_index_template"] = "{{ params.map_name }}"
self.log.info(
"Starting S3 to Redshift ingestion: %s -> %s.%s",
self.s3_path,
self.target_schema,
self.target_table,
)
self.log.info("Write mode: %s", self.write_mode)
self.ingest_to_athena()
if self.column_mapping:
self.log.info("Column mapping: %s", self.column_mapping)
if self.geometry_type:
self.log.info(
"Geometry type: %s, config: %s",
self.geometry_type,
self.geometry_config,
)
self.rs = RedshiftDataHook(
aws_conn_id=self.aws_conn_id,
region_name=self.region,
)
case_prefix = _case_prefix(self.case_sensitive_identifiers)
try:
stage_columns = self._get_stage_columns()
self.log.info("Stage columns: %s", stage_columns)
if self.write_mode == "overwrite":
self.log.info("Overwriting target table with data from staging...")
self._run_rs(
case_prefix + self._sql_overwrite_from_stage(stage_columns)
)
Util.remove_datalake_table(tgt_table=f"{self.target_table}_temp")
else:
self.log.info("Appending data from staging to target table...")
raise NotImplementedError("Append mode is not implemented yet.")
# self._run_rs(case_prefix + self._sql_insert_from_stage(stage_columns))
self._run_rs(case_prefix + self._sql_add_h3_index())
row_count = self._run_rs_and_fetch(
f"SELECT COUNT(*) AS cnt FROM {self._target_qual()};"
)
final_count = row_count[0]["cnt"] if row_count else 0
self.log.info("Final row count in target table: %s", final_count)
row_count = self._run_rs_and_fetch(
f"SELECT COUNT(*) AS cnt FROM {self._target_qual()} where geom is null;"
)
null_geom_count = row_count[0]["cnt"] if row_count else 0
if null_geom_count == row_count:
raise Exception(
"All the geometires seem to be null, They need to be checked"
)
except Exception as e:
self.log.error("Error during S3 to Redshift ingestion: %s", e)
raise
self.log.info("S3 to Redshift ingestion complete.")
return {
"s3_path": self.s3_path,
"target_schema": self.target_schema,
"target_table": self.target_table,
"write_mode": self.write_mode,
"row_count": final_count,
}