"""
Redshift CTAS Operator
"""
from typing import Dict, List, Optional, Sequence, Union
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
ALLOWED_DISTSTYLE = {"EVEN", "KEY", "ALL", "AUTO"}
def _q_ident(parts: Union[str, List[str]]) -> str:
"""
Quote a SQL identifier or qualified identifier.
Escapes embedded double-quotes by doubling them.
:param parts: Single identifier string or list of [schema, table] parts.
:raises ValueError: If any identifier part is empty.
"""
def q1(tok: str) -> str:
s = str(tok)
s = s.replace('"', '""')
return f'"{s}"'
if isinstance(parts, (list, tuple)):
for p in parts:
if p is None or str(p).strip() == "":
raise ValueError(f"Empty identifier in sequence: {parts!r}")
return ".".join(q1(p) for p in parts)
if parts is None or str(parts).strip() == "":
raise ValueError(f"Empty identifier: {parts!r}")
return q1(parts)
[docs]
class RedshiftCTASOperator(BaseOperator):
"""
Execute CTAS (Create Table As Select) operations in Redshift with transaction support.
Wraps DROP TABLE + CREATE TABLE AS in a single transaction for atomic table replacement.
Supports Redshift-specific distribution and sort key options.
Note: The query parameter must be a single SELECT statement without a trailing semicolon.
If clean=False and the table exists, the operation will fail.
:param query: The SELECT query to use for creating the new table (no trailing semicolon).
:type query: str
:param target_table: Name of the target table to create.
:type target_table: str
:param target_schema: Target schema name. Can also be resolved from params for dynamic mapping.
:type target_schema: Optional[str]
:param redshift_conn_id: Airflow connection ID for Redshift.
:type redshift_conn_id: str
:param database: Redshift database to connect to. If None, uses the connection default.
:type database: Optional[str]
:param diststyle: Distribution style (EVEN, KEY, ALL, AUTO). If distkey is provided
without diststyle, KEY is used automatically.
:type diststyle: Optional[str]
:param distkey: Column name for DISTKEY. Only valid with diststyle KEY or unset.
:type distkey: Optional[str]
:param sortkey: List of column names for compound SORTKEY.
:type sortkey: Optional[List[str]]
:param clean: Whether to DROP existing table before creation. Defaults to True.
:type clean: bool
:param column_comments: Optional dict mapping column names to their descriptions.
If provided, runs COMMENT ON COLUMN statements after CTAS completes.
:type column_comments: Optional[Dict[str, str]]
:param table_description: Optional description for the table.
If provided, runs COMMENT ON TABLE statement after CTAS completes.
:type table_description: Optional[str]
:param grant_select_to: Optional list of Redshift users or groups to GRANT SELECT
on the newly created table. Runs after CTAS completes so that DROP+CREATE
cycles don't revoke access from BI service accounts.
Entries are plain usernames by default. To grant to a Redshift GROUP,
prefix the entry with ``GROUP:`` (case-insensitive)::
["mn_datamart", "GROUP:readydata"]
This generates::
GRANT SELECT ON TABLE "bi"."table" TO "mn_datamart";
GRANT SELECT ON TABLE "bi"."table" TO GROUP "readydata";
:type grant_select_to: Optional[List[str]]
"""
ui_color = "#BADA55"
template_fields: Sequence[str] = (
"query",
"target_table",
"target_schema",
"redshift_conn_id",
"database",
"diststyle",
"distkey",
"sortkey",
"clean",
"grant_select_to",
)
template_fields_renderers = {"query": "sql"}
def __init__(
self,
query: str,
target_table: str,
redshift_conn_id: str,
database: Optional[str] = None,
target_schema: Optional[str] = None,
diststyle: Optional[str] = None,
distkey: Optional[str] = None,
sortkey: Optional[List[str]] = None,
clean: bool = True,
column_comments: Optional[Dict[str, str]] = None,
table_description: Optional[str] = None,
grant_select_to: Optional[List[str]] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.query = query
self.target_table = target_table
self.target_schema = target_schema
self.redshift_conn_id = redshift_conn_id
self.database = database
self.diststyle = diststyle
self.distkey = distkey
self.sortkey = sortkey
self.clean = clean
self.column_comments = column_comments
self.table_description = table_description
self.grant_select_to = grant_select_to or []
def _apply_comments(self, hook: "PostgresHook") -> None:
"""Apply table and column comments via COMMENT ON statements."""
qualified_table = _q_ident([self._resolved_schema, self.target_table])
statements = []
if self.table_description:
escaped_desc = self.table_description.replace("'", "''")
statements.append(
f"COMMENT ON TABLE {qualified_table} IS '{escaped_desc}';"
)
if self.column_comments:
for col_name, description in self.column_comments.items():
escaped_desc = description.replace("'", "''")
statements.append(
f"COMMENT ON COLUMN {qualified_table}.{_q_ident(col_name)} IS '{escaped_desc}';"
)
if statements:
self.log.info("Applying %d comment(s) to %s", len(statements), qualified_table)
hook.run(statements, autocommit=True)
def _validate_inputs(self) -> None:
"""Validate operator inputs after template rendering."""
resolved_schema = self.target_schema or self.params.get("target_schema")
if not resolved_schema or not resolved_schema.strip():
raise AirflowException(
"target_schema must be provided via parameter or params for dynamic mapping"
)
self._resolved_schema = resolved_schema
if not self.target_table or not self.target_table.strip():
raise AirflowException("target_table cannot be empty")
if not self.query or not self.query.strip():
raise AirflowException("query cannot be empty")
clean_query = self.query.rstrip().rstrip(";")
if ";" in clean_query:
raise AirflowException(
"query must be a single SELECT statement (no semicolons except trailing)"
)
if self.sortkey is not None:
if isinstance(self.sortkey, str):
raise AirflowException(
"sortkey must be a list of column names, not a string"
)
if not all(isinstance(c, str) for c in self.sortkey):
raise AirflowException("sortkey must contain only strings")
if self.distkey:
if self.diststyle and self.diststyle.upper() in {"EVEN", "ALL", "AUTO"}:
raise AirflowException(
f"distkey is not compatible with diststyle '{self.diststyle}'. "
"Use diststyle='KEY' or omit diststyle."
)
def _build_ctas_sql(self) -> str:
"""Build the CREATE TABLE AS statement with distribution/sort options."""
qualified_table = _q_ident([self._resolved_schema, self.target_table])
parts = [f"CREATE TABLE {qualified_table}"]
effective_diststyle = self.diststyle
if self.distkey and not self.diststyle:
effective_diststyle = "KEY"
if effective_diststyle:
ds = effective_diststyle.upper()
if ds not in ALLOWED_DISTSTYLE:
raise AirflowException(
f"Invalid diststyle '{self.diststyle}'. Must be one of: {ALLOWED_DISTSTYLE}"
)
parts.append(f"DISTSTYLE {ds}")
if self.distkey:
parts.append(f"DISTKEY ({_q_ident(self.distkey)})")
if self.sortkey:
sortkey_cols = ", ".join(_q_ident(c) for c in self.sortkey)
parts.append(f"SORTKEY ({sortkey_cols})")
parts.append("AS")
query = self.query.rstrip().rstrip(";")
parts.append(query)
return "\n".join(parts)
[docs]
def execute(self, context: dict) -> dict:
"""
Execute the CTAS operation within a transaction.
Returns:
dict with keys: table, schema, created (pushed to XCom)
"""
try:
self._validate_inputs()
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None,
)
qualified_table = _q_ident([self._resolved_schema, self.target_table])
statements = []
# Enable case-sensitive identifiers for mixed-case column names
statements.append("SET enable_case_sensitive_identifier TO true;")
if self.clean:
self.log.info("Will DROP existing table %s if exists", qualified_table)
statements.append(f"DROP TABLE IF EXISTS {qualified_table};")
ctas_sql = self._build_ctas_sql()
self.log.info("Creating table %s via CTAS", qualified_table)
self.log.debug("CTAS SQL:\n%s", ctas_sql)
statements.append(f"{ctas_sql};")
hook.run(statements, autocommit=False)
self.log.info("Successfully created table %s", qualified_table)
if self.grant_select_to:
grant_stmts = []
for raw_grantee in self.grant_select_to:
grantee = str(raw_grantee).strip()
if not grantee:
raise AirflowException(
"grant_select_to contains an empty grantee"
)
if grantee.upper().startswith("GROUP:"):
group_name = grantee[6:].strip()
if not group_name:
raise AirflowException(
f"Empty group name in grant_select_to: {raw_grantee!r}"
)
grant_stmts.append(
f"GRANT SELECT ON TABLE {qualified_table} TO GROUP {_q_ident(group_name)};"
)
else:
grant_stmts.append(
f"GRANT SELECT ON TABLE {qualified_table} TO {_q_ident(grantee)};"
)
hook.run(grant_stmts, autocommit=True)
self.log.info(
"Granted SELECT on %s to: %s",
qualified_table,
", ".join(self.grant_select_to),
)
if self.column_comments or self.table_description:
self._apply_comments(hook)
self.log.info("Successfully applied comments to %s", qualified_table)
return {
"table": self.target_table,
"schema": self._resolved_schema,
"created": True,
}
except AirflowException:
raise
except Exception as e:
self.log.exception(
"CTAS failed for %s.%s", self._resolved_schema, self.target_table
)
raise AirflowException(
f"Failed to execute CTAS for {self._resolved_schema}.{self.target_table}: {e}"
) from e