"""
Redshift Drop Table Operator
Drops a single Redshift table with safety guardrails. Optionally
authenticates as the cluster superuser (rsadmin) by fetching credentials
from AWS Secrets Manager — the same secret backs the standard Airflow
connection (``airflow/connections/{conn_id}``).
Why this operator exists
------------------------
Redshift does not support ``GRANT DROP``. DROP is owner-only. ctasifier
creates tables under the airflow IAM role, so federated SSO users cannot
drop those tables themselves. This operator (and its companion DAG
``redshift_drop_table``) fills that gap.
Case sensitivity
----------------
All identifier comparisons and the DROP statement are case-sensitive.
The ``enable_case_sensitive_identifier`` GUC is set on the SAME physical
DB session that performs the metadata lookup and the DROP, so quoted
mixed-case identifiers are honoured end-to-end (the GUC is a session
variable — running it through ``hook.run`` would land on a different
connection).
data@ready.net | 2026-05-14
"""
from contextlib import closing
from copy import deepcopy
from typing import List, Optional, Sequence, Union
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
# System / catalog schemas we always refuse to touch.
_SYSTEM_SCHEMAS_EXACT = {
"information_schema",
"pg_catalog",
"pg_toast",
"pg_internal",
"catalog_history",
}
_SYSTEM_SCHEMA_PREFIXES = (
"pg_",
"pg_temp_",
"aws_",
"svv_",
"stv_",
"stl_",
)
def _q_ident(parts: Union[str, List[str]]) -> str:
"""Quote a SQL identifier (or qualified identifier).
Mirrors the helper in ``redshift_ctas.py`` — escapes embedded
double-quotes by doubling them. Preserves the caller-provided case
so the DROP statement honours case-sensitive identifiers.
"""
def q1(tok: str) -> str:
s = str(tok).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)
def _parse_bool(name: str, v) -> bool:
"""Strict boolean parser for destructive params.
Accepts only native ``True``/``False`` or the exact strings
``"true"``/``"false"`` (case-insensitive, whitespace-stripped).
Anything else raises — the permissive coercion that turned typos
like ``"flase"`` or ``"None"`` into ``False`` (silently disabling
dry-run) is intentionally absent.
"""
from airflow.exceptions import AirflowException
if isinstance(v, bool):
return v
if isinstance(v, str):
s = v.strip().lower()
if s in ("true", "false"):
return s == "true"
raise AirflowException(
f"{name} must be a boolean (True/False or 'true'/'false'); got {v!r}"
)
[docs]
class RedshiftDropTableOperator(BaseOperator):
"""Drop a single Redshift table with safety guardrails.
By default authenticates as the airflow IAM role (via the standard
Airflow Postgres connection). When ``use_rsadmin_override=True``,
fetches the JSON secret named ``airflow/connections/{redshift_conn_id}``
from AWS Secrets Manager and authenticates as ``rsadmin``
(the cluster superuser) for that single task only — never persists
the password and never logs it.
Identifier comparisons and the DROP statement are case-sensitive
(``enable_case_sensitive_identifier`` is set on the same physical
session that performs the metadata lookup and the DROP).
:param redshift_conn_id: Airflow Postgres connection id for Redshift
(e.g. ``Config.REDSHIFT_PG``). Also used to derive the rsadmin
secret name when ``use_rsadmin_override`` is true.
:type redshift_conn_id: str
:param database: Redshift database within the cluster (e.g.
``"az_datamart"``, ``"dataops"``).
:type database: str
:param schema: Target schema (case-sensitive).
:type schema: str
:param table: Target table (case-sensitive).
:type table: str
:param dry_run: If true (default), log the DROP statement and table
metadata without executing. Operator default is true so the
reusable surface is non-destructive — set false explicitly.
:type dry_run: bool
:param use_rsadmin_override: If true, authenticate as rsadmin via
Secrets Manager (used when the table is not owned by the airflow
IAM role).
:type use_rsadmin_override: bool
:param confirm_target: REQUIRED for real runs (``dry_run=False``).
Must equal ``"{database}.{schema}.{table}"`` exactly. Operator
refuses to execute any DROP unless the confirmation matches.
Cheap defence against fat-finger mistakes.
:type confirm_target: Optional[str]
:param require_owner_check: If true (default), look up the table
owner and refuse to issue DDL when the connecting user is not
the owner (unless ``use_rsadmin_override`` is set, in which
case mismatch is expected and only logged).
:type require_owner_check: bool
:param forbid_uppercase_table: If true, refuse to drop CamelCase
/ mixed-case table names. Default false on the operator (general
purpose). The companion ``redshift_drop_table`` DAG sets this
to true to enforce its policy.
:type forbid_uppercase_table: bool
:param exclude_schemas: Optional list of additional schema names to
refuse. The system-schema list (``information_schema``,
``pg_catalog``, anything starting with ``pg_``, ``aws_``,
``svv_``, ``stv_``, ``stl_``) is always applied on top of this.
:type exclude_schemas: Optional[List[str]]
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"database",
"schema",
"table",
"dry_run",
"use_rsadmin_override",
"confirm_target",
)
ui_color = "#E87C3E"
def __init__(
self,
*,
redshift_conn_id: str,
database: str,
schema: str,
table: str,
dry_run: bool = True,
use_rsadmin_override: bool = False,
confirm_target: Optional[str] = None,
require_owner_check: bool = True,
forbid_uppercase_table: bool = False,
exclude_schemas: Optional[List[str]] = None,
**kwargs,
):
super().__init__(**kwargs)
self.redshift_conn_id = redshift_conn_id
self.database = database
self.schema = schema
self.table = table
self.dry_run = dry_run
self.use_rsadmin_override = use_rsadmin_override
self.confirm_target = confirm_target
self.require_owner_check = require_owner_check
self.forbid_uppercase_table = forbid_uppercase_table
self.exclude_schemas = exclude_schemas or []
# ---- validation ----------------------------------------------------
def _validate(self) -> None:
for name, val in (
("redshift_conn_id", self.redshift_conn_id),
("database", self.database),
("schema", self.schema),
("table", self.table),
):
if val is None or not str(val).strip():
raise AirflowException(f"{name} is required and cannot be empty")
# Case-sensitive: keep as-provided. Only strip whitespace.
self.database = str(self.database).strip()
self.schema = str(self.schema).strip()
self.table = str(self.table).strip()
if "." in self.table:
raise AirflowException(
f"`table` must be the bare table name (no '.'); got {self.table!r}. "
"Provide schema and table separately."
)
if self.forbid_uppercase_table and any(c.isupper() for c in self.table):
raise AirflowException(
f"Refusing to drop CamelCase / mixed-case table {self.table!r} "
"(forbid_uppercase_table=True)."
)
s_lower = self.schema.lower()
if s_lower in _SYSTEM_SCHEMAS_EXACT or any(
s_lower.startswith(p) for p in _SYSTEM_SCHEMA_PREFIXES
):
raise AirflowException(
f"Refusing to drop from system schema {self.schema!r}."
)
excluded = {str(s).strip() for s in self.exclude_schemas if str(s).strip()}
if self.schema in excluded:
raise AirflowException(
f"Schema {self.schema!r} is in exclude_schemas={sorted(excluded)}; refusing."
)
# ---- connection plumbing ------------------------------------------
def _build_hook(self):
"""Build a PostgresHook bound to either the airflow role or rsadmin.
For the override path: deep-copy the base Airflow Connection
(preserves SSL / keepalive extras), swap host/login/password/port,
and seed the hook via the supported ``connection=`` kwarg with a
fallback to the backing private attribute. The password lives
only on the in-memory Connection object — it is never persisted
to the Airflow metadata DB.
"""
from airflow.hooks.base import BaseHook
from airflow.providers.postgres.hooks.postgres import PostgresHook
if not self.use_rsadmin_override:
return PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database,
)
import json
import boto3
secret_name = f"airflow/connections/{self.redshift_conn_id}"
self.log.info(
"Override mode: fetching rsadmin credentials from Secrets Manager (%s)",
secret_name,
)
sm = boto3.client("secretsmanager")
secret_string = sm.get_secret_value(SecretId=secret_name)["SecretString"]
cfg = json.loads(secret_string)
host = cfg.get("host")
login = cfg.get("login")
password = cfg.get("password")
port = int(cfg.get("port", 5439))
if not (host and login and password):
raise AirflowException(
f"Secret {secret_name} missing one of host/login/password."
)
# Preserve extras (sslmode, keepalives, etc.) by deep-copying the
# base Connection and overriding only the auth/host fields.
base_conn = deepcopy(BaseHook.get_connection(self.redshift_conn_id))
base_conn.host = host
base_conn.login = login
base_conn.password = password
base_conn.port = port
base_conn.schema = self.database
try:
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database,
connection=base_conn,
)
except TypeError:
# Older provider versions do not accept ``connection=``.
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database,
)
hook._connection = base_conn # type: ignore[attr-defined]
return hook
# ---- metadata helpers (used inside a single live cursor) ----------
@staticmethod
def _fetch_owner(cur, schema: str, table: str) -> Optional[str]:
cur.execute(
"SELECT tableowner "
"FROM pg_catalog.pg_tables "
"WHERE schemaname = %(schema)s AND tablename = %(table)s "
"LIMIT 1;",
{"schema": schema, "table": table},
)
row = cur.fetchone()
return row[0] if row else None
def _fetch_size_stats(self, cur, db: str, schema: str, table: str) -> dict:
"""Best-effort size / rowcount enrichment via svv_table_info.
``svv_table_info`` only reports tables with at least one populated
block, may be permission-restricted, and only sees the currently
connected database. Treat any failure as "no stats available".
"""
try:
cur.execute(
'SELECT size, tbl_rows '
'FROM svv_table_info '
'WHERE database = %(db)s AND "schema" = %(schema)s '
' AND "table" = %(table)s '
'LIMIT 1;',
{"db": db, "schema": schema, "table": table},
)
row = cur.fetchone()
if not row:
return {"size_mb": None, "tbl_rows": None}
return {"size_mb": row[0], "tbl_rows": row[1]}
except Exception as e: # noqa: BLE001 — best-effort enrichment
self.log.warning("svv_table_info enrichment skipped: %s", e)
return {"size_mb": None, "tbl_rows": None}
# ---- execute -------------------------------------------------------
[docs]
def execute(self, context) -> dict:
# Strict boolean parse — fail-closed on garbage like "flase"/"None".
self.dry_run = _parse_bool("dry_run", self.dry_run)
self.use_rsadmin_override = _parse_bool(
"use_rsadmin_override", self.use_rsadmin_override
)
self._validate()
# Confirmation gate for real runs (cheap fat-finger defence).
# Required only when actually executing the DROP.
expected_confirm = f"{self.database}.{self.schema}.{self.table}"
if not self.dry_run:
actual = (self.confirm_target or "").strip()
if actual != expected_confirm:
raise AirflowException(
"Real run requires confirm_target to exactly equal "
f"{expected_confirm!r}; got {actual!r}. "
"Re-trigger the DAG with confirm_target set, or run "
"with dry_run=true first to inspect."
)
hook = self._build_hook()
qualified = _q_ident([self.schema, self.table])
dag_run = context.get("dag_run")
run_id = getattr(dag_run, "run_id", None)
# ---- single physical session: SET + lookup + DROP all here ----
with closing(hook.get_conn()) as conn:
conn.autocommit = True
with closing(conn.cursor()) as cur:
cur.execute("SET enable_case_sensitive_identifier TO true;")
cur.execute("SELECT current_user;")
current_user_row = cur.fetchone()
current_user = current_user_row[0] if current_user_row else None
owner = self._fetch_owner(cur, self.schema, self.table)
exists = owner is not None
stats = (
self._fetch_size_stats(
cur, self.database, self.schema, self.table
)
if exists
else {"size_mb": None, "tbl_rows": None}
)
self.log.info("=" * 60)
self.log.info("RedshiftDropTableOperator")
self.log.info(" database : %s", self.database)
self.log.info(" qualified target : %s", qualified)
self.log.info(" connecting user : %s", current_user)
self.log.info(" table exists : %s", exists)
self.log.info(" table owner : %s", owner)
self.log.info(" size (MB) : %s", stats["size_mb"])
self.log.info(" est. row count : %s", stats["tbl_rows"])
self.log.info(" override mode : %s", self.use_rsadmin_override)
self.log.info(" dry_run : %s", self.dry_run)
self.log.info(" dag_run.run_id : %s", run_id)
self.log.info("=" * 60)
# ---- DRY RUN: emit preview and exit ------------------
if self.dry_run:
preview_stmt = f"DROP TABLE {qualified};"
self.log.info("DRY RUN — would execute: %s", preview_stmt)
if not exists:
self.log.warning(
"Target %s does not exist (case-sensitive lookup). "
"A real run will FAIL fast unless the target is created first "
"or the case is corrected.",
qualified,
)
return {
"dropped": False,
"dry_run": True,
"qualified": qualified,
"exists": exists,
"owner": owner,
"size_mb": stats["size_mb"],
"tbl_rows": stats["tbl_rows"],
"current_user": current_user,
}
# ---- REAL RUN: fail-loud preflight checks ------------
if not exists:
raise AirflowException(
f"Refusing to drop: target {qualified} does not exist "
f"in database {self.database!r} (case-sensitive lookup). "
"Re-run with dry_run=true to inspect; check schema/table case."
)
if self.require_owner_check and owner and str(owner) != str(current_user):
if self.use_rsadmin_override:
self.log.info(
"Owner %r differs from current user %r; "
"proceeding under rsadmin override.",
owner,
current_user,
)
else:
raise AirflowException(
f"Owner mismatch for {qualified}: owner={owner!r}, "
f"current_user={current_user!r}. "
"Re-trigger with use_rsadmin_override=true if intended."
)
# IF EXISTS intentionally OMITTED: we already verified
# existence in the same session, and we want a true
# error if the target disappeared between lookup and DROP.
stmt = f"DROP TABLE {qualified};"
self.log.info("Executing: %s", stmt)
cur.execute(stmt)
self.log.info("DROP succeeded for %s", qualified)
return {
"dropped": True,
"dry_run": False,
"qualified": qualified,
"exists": exists,
"owner": owner,
"size_mb": stats["size_mb"],
"tbl_rows": stats["tbl_rows"],
"current_user": current_user,
}