"""
================================================================================
Redshift Database Management Operators
================================================================================
This module contains custom Airflow operators for managing Redshift databases,
schemas, and permissions for both *internal* Redshift objects and the
auto-mounted AWS Glue Data Catalog (`awsdatacatalog`).
Operators
---------
1. ValidateOrCreateDatabaseOperator
--------------------------------
- Ensures a Redshift database exists.
- Optionally creates the database (with a specific owner) if missing.
- Typical use: bootstrap environment (e.g., make sure "dev" DB exists).
2. ValidateOrCreateSchemaOperator
------------------------------
- Ensures a schema exists inside a Redshift database.
- Optionally creates the schema with an owner if missing.
- Typical use: provision tenant or project schemas before loading data.
3. SetAwsRolePermissiveGrantsOperator
-----------------------------------
- Manages *internal* schema permissions in Redshift.
- Creates a DB group if missing, optionally adds IAMR users to it, and applies permissive grants (USAGE, CREATE, DML) on the schema.
- Can also set default privileges for future tables owned by specific users.
- Typical use: ensure federated IAM users mapped to DbGroups have the right privileges in application schemas.
4. SetAwsRoleAwsDataCatalogGrantsOperator
---------------------------------------
- Baseline grants for the auto-mounted Glue Data Catalog (`awsdatacatalog`).
- Grants USAGE on the external database and specified Glue DBs (schemas).
- Two modes:
* lake_formation_mode=True → USAGE-only; rely on LF policies for SELECT.
* lake_formation_mode=False → Also grant SELECT on existing/future tables.
- Typical use: initial setup so Redshift users/groups can see external data.
5. RefreshAwsDataCatalogGrantsOperator
------------------------------------
- Refresh / sweeper for external catalog permissions.
- Grants USAGE on catalog + schemas, and (optionally) SELECT on all tables.
- Intended to be re-run whenever new external tables are added outside Redshift (e.g., Athena CTAS, awswrangler), since ALTER DEFAULT PRIVILEGES does not apply to those.
- Typical use:
* Add as a final task in DAGs that create Glue tables.
* Or schedule a nightly sweep to keep catalog permissions fresh.
6. SetCaseSensitiveIdentifierOperator
------------------------------------
- Sets the case sensitivity for identifiers in Redshift.
- Typical use: ensure consistent casing for schema/table names.
7. CreateRedshiftUserOperator
---------------------------
- Creates a Redshift user with password if it doesn't exist.
- Optionally retrieves password from AWS Secrets Manager.
- Grants CONNECT on the target database.
- Typical use: provision service accounts for external tools (e.g., CARTO, Tableau)
or internal services (ETL pipelines, reporting).
8. SetSchemaReadGrantsOperator
----------------------------
- Grants read-only access on a schema (USAGE, SELECT, EXECUTE).
- Sets default privileges for future tables/functions.
- Typical use: least-privilege access for any user or service that only needs
to read data or execute functions (BI tools, analysts, external integrations).
Usage Patterns
--------------
- **Environment bootstrap:**
ValidateOrCreateDatabaseOperator → ValidateOrCreateSchemaOperator
- **Internal schema access:**
SetAwsRolePermissiveGrantsOperator
- **External catalog access:**
SetAwsRoleAwsDataCatalogGrantsOperator (baseline)
RefreshAwsDataCatalogGrantsOperator (ongoing refresh)
- **Service account setup (external or internal):**
CreateRedshiftUserOperator → ValidateOrCreateSchemaOperator
→ SetSchemaReadGrantsOperator (read-only schemas)
→ SetAwsRolePermissiveGrantsOperator (read-write schemas)
All operators are idempotent and safe to re-run.
Notes
-----
- For Lake Formation–governed setups, keep `SetAwsRoleAwsDataCatalogGrantsOperator`
in `lake_formation_mode=True` (USAGE-only) and manage SELECT in LF policies.
- For IAM-only / Redshift-grants setups, set `lake_formation_mode=False`
and pair with RefreshAwsDataCatalogGrantsOperator for future-proofing.
- Ensure the Redshift workgroup IAM role has the necessary S3 permissions
to access external tables.
================================================================================
"""
import ast
import json
import logging
from typing import Iterable, List, Optional, Sequence, Tuple, Union
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
from airflow.providers.postgres.hooks.postgres import PostgresHook
####################################################################################################
### NOTE NOTE NOTE #################################################################################
####################################################################################################
## PostgresHook works for Redshift clusters and serverless.
## It is used here when executing operations on the non-default schema because the redshift_sql hook
## does not support overriding the schema.
## The RedshiftSQLHook is used when connecting to the "admin" database to create a database.
####################################################################################################
#########
# Utility functions to quote identifiers and users
#########
def _q_ident(parts: Union[str, Sequence[str]]) -> str:
"""
Quote a SQL identifier or a qualified identifier.
- Always double-quotes tokens and escapes internal quotes.
- Accepts a string or a sequence of strings (schema, table, column).
Example: _q_ident(["public", "users"]) -> "public"."users"
"""
def q1(tok: str) -> str:
s = str(tok)
# Double any embedded double-quotes
s = s.replace('"', '""')
return f'"{s}"'
if isinstance(parts, (list, tuple)):
tokens = [q1(p) for p in parts if p is not None and str(p) != ""]
if not tokens:
raise ValueError("Empty identifier sequence.")
return ".".join(tokens)
if parts is None or str(parts) == "":
raise ValueError(f"Empty identifier: {parts!r}")
return q1(str(parts))
def _q_user(name: str) -> str:
"""Double-quote a username (works for IAMR:... names)."""
return '"' + name.replace('"', '""') + '"'
def _collect_principals(
hook: PostgresHook,
*,
group_name: Optional[str],
iamr_users: Optional[Iterable[str]],
include_current_user: bool,
require_nonempty: bool = True,
log: Optional[logging.Logger] = None,
) -> List[str]:
"""
Build a list of Redshift principals (CURRENT_USER, GROUP <name>, explicit users),
validating that explicit users exist before including them.
Returns a list like: ["CURRENT_USER", "GROUP \"readydata\"", "\"IAMR:...\""]
"""
logger = log or logging.getLogger(__name__)
principals: List[str] = []
if include_current_user:
row = hook.get_first("SELECT current_user;")
if row and row[0]:
principals.append(_q_user(row[0]))
else:
logger.warning("Could not resolve current_user; skipping.")
if group_name:
principals.append(f"GROUP {_q_ident(group_name)}")
for u in iamr_users or []:
exists = hook.get_first(
"SELECT 1 FROM pg_user WHERE usename=%s LIMIT 1", parameters=(u,)
)
if exists:
principals.append(_q_user(u))
else:
logger.warning("User %s not found in Redshift; skipping.", u)
if require_nonempty and not principals:
raise AirflowException(
"No valid principals to grant (provide group_name/iamr_users or set include_current_user=True)."
)
return principals
def _coerce_glue_db_list(val) -> Union[str, List[str]]:
"""
Accepts:
- "ALL"
- JSON array string (e.g., '["shoreline","readydatalake"]')
- comma-separated string (e.g., "shoreline, readydatalake")
- list/tuple
Returns "ALL" or a list[str].
"""
if val is None:
return []
if isinstance(val, str):
s = val.strip()
if s.upper() == "ALL":
return "ALL"
if s.startswith("["):
try:
parsed = json.loads(s)
except Exception:
parsed = ast.literal_eval(s)
return [x for x in map(str, parsed) if x.strip()]
if "," in s:
return [x.strip() for x in s.split(",") if x.strip()]
return [s] # single name as string
# list/tuple/etc.
return [str(x).strip() for x in val if str(x).strip()]
def _coerce_bool(val) -> bool:
"""
Coerce a value to bool, handling Jinja-rendered strings.
Jinja renders booleans as the strings "True" / "False", both of which
are truthy in Python. This function normalises them.
"""
if isinstance(val, bool):
return val
if isinstance(val, str):
return val.strip().lower() in ("true", "1")
return bool(val)
def build_redshift_hook(
redshift_conn_id: str,
database: Optional[str] = None,
use_rsadmin_override: bool = False,
log: Optional[logging.Logger] = None,
) -> PostgresHook:
"""
Build a PostgresHook for Redshift (works for clusters and serverless).
By default the hook uses the connection's own identity — on our Airflow
that is the EC2 IAM role (e.g. ``IAMR:airflow_prod_ec2_role``). When
``use_rsadmin_override=True``, rsadmin (cluster superuser) credentials are
fetched from AWS Secrets Manager at ``airflow/connections/<conn_id>`` and
used instead, so privileged DDL/DCL that only a superuser or object owner
may run — e.g. ``CREATE DATABASE ... WITH OWNER`` and
``GRANT ... ON DATABASE`` against a database owned by rsadmin — succeeds.
The fetched password lives only on the in-memory Connection object; it is
never persisted to the Airflow metadata DB. Mirrors the override pattern in
``operators/redshift/redshift_drop_table.py``.
:param redshift_conn_id: Airflow connection id for Redshift.
:param database: Database to bind the hook to (overrides the conn default).
:param use_rsadmin_override: If True, authenticate as rsadmin via Secrets Manager.
:param log: Optional logger for informational messages.
"""
logger = log or logging.getLogger(__name__)
if not use_rsadmin_override:
return PostgresHook(
postgres_conn_id=redshift_conn_id,
schema=database or None,
)
import boto3
from copy import deepcopy
from airflow.hooks.base import BaseHook
secret_name = f"airflow/connections/{redshift_conn_id}"
logger.info(
"rsadmin override: fetching superuser credentials from Secrets Manager (%s)",
secret_name,
)
sm = boto3.client("secretsmanager")
cfg = json.loads(sm.get_secret_value(SecretId=secret_name)["SecretString"])
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."
)
# Deep-copy the base Connection and override only the auth/host fields.
base_conn = deepcopy(BaseHook.get_connection(redshift_conn_id))
base_conn.host = host
base_conn.login = login
base_conn.password = password
base_conn.port = port
base_conn.schema = database or base_conn.schema
# CRITICAL: strip provider/auth extras (e.g. iam, serverless, cluster
# identifier) so the Postgres/Redshift provider does NOT swap our explicit
# rsadmin login/password for an IAM token derived from the EC2 role. We
# supply an explicit host/port/login/password, so only client-level psycopg2
# extras should survive.
_SAFE_EXTRA_KEYS = {
"sslmode",
"sslrootcert",
"sslcert",
"sslkey",
"connect_timeout",
"application_name",
"options",
"keepalives",
"keepalives_idle",
"keepalives_interval",
"keepalives_count",
}
try:
extra = base_conn.extra_dejson or {}
except Exception:
extra = {}
safe_extra = {k: v for k, v in extra.items() if k in _SAFE_EXTRA_KEYS}
dropped = sorted(set(extra) - set(safe_extra))
if dropped:
logger.info("rsadmin override: dropping non-client conn extras: %s", dropped)
base_conn.extra = json.dumps(safe_extra) if safe_extra else None
try:
hook = PostgresHook(
postgres_conn_id=redshift_conn_id,
schema=database or None,
connection=base_conn,
)
except TypeError:
# Older provider versions do not accept ``connection=``.
hook = PostgresHook(
postgres_conn_id=redshift_conn_id,
schema=database or None,
)
hook._connection = base_conn # type: ignore[attr-defined]
# Fail loudly if we did not actually obtain a superuser session. This catches
# a misconfigured secret OR provider extras that silently downgraded us to
# the IAM/EC2 role — in which case privileged DDL/DCL would not behave as
# rsadmin.
who = hook.get_first(
"SELECT current_user, usesuper FROM pg_user WHERE usename = current_user;"
)
if not who:
raise AirflowException("rsadmin override: could not resolve current_user.")
current_user, is_super = who[0], bool(who[1])
logger.info(
"rsadmin override connected as %s (superuser=%s).", current_user, is_super
)
if not is_super:
raise AirflowException(
f"rsadmin override expected a superuser session but connected as "
f"{current_user!r} (superuser=False). Verify Secrets Manager secret "
f"{secret_name} holds rsadmin credentials and that provider IAM/"
f"serverless extras are not overriding them."
)
return hook
def _discover_glue_dbs(hook: PostgresHook, database_name: str) -> List[str]:
"""
Return Glue DB (schema) names visible under the auto-mounted external DB.
Handles slight differences in SHOW output shape across engine builds.
"""
rows = hook.get_records(f"SHOW SCHEMAS FROM DATABASE {_q_ident(database_name)};")
# Prefer r[1] (common on recent serverless), fall back to r[0]
vals = [
r[1]
for r in rows
if isinstance(r, (list, tuple)) and len(r) > 1 and isinstance(r[1], str)
]
if not vals:
vals = [r[0] for r in rows if r and isinstance(r[0], str)]
return sorted(set(vals))
[docs]
class ValidateOrCreateDatabaseOperator(BaseOperator):
"""
Ensure a Redshift database exists. If missing and `create_if_not_exists=True`,
create it (optionally with a specific owner). Ownership is only set at creation time.
:param redshift_conn_id: Airflow connection id for Redshift.
:type redshift_conn_id: str
:param database: Name of the database to ensure.
:type database: str
:param db_owner: Optional DB owner to set at creation (not changed if DB already exists).
:type db_owner: str
:param create_if_not_exists: Create the DB if it doesn't exist.
:type create_if_not_exists: bool
:param admin_database: Database to connect to when issuing CREATE DATABASE. Defaults to 'dev'.
:type admin_database: str
:param use_rsadmin_override: If True, authenticate as rsadmin (cluster superuser)
via AWS Secrets Manager instead of the connection's own identity (the Airflow
EC2 IAM role). Required when the executing role is not a superuser and you need
to ``CREATE DATABASE ... WITH OWNER <other_user>`` (e.g. owner rsadmin), which
only a superuser may do. Defaults to False (backwards compatible).
:type use_rsadmin_override: bool
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"database",
"db_owner",
"admin_database",
"create_if_not_exists",
"use_rsadmin_override",
)
ui_color = "#57D1FA"
def __init__(
self,
redshift_conn_id: str,
database: str,
db_owner: Optional[str] = "rsadmin",
create_if_not_exists: bool = False,
admin_database: Optional[str] = "dev",
use_rsadmin_override: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.database = database
self.db_owner = db_owner
self.create_if_not_exists = create_if_not_exists
self.admin_database = admin_database
self.use_rsadmin_override = use_rsadmin_override
def _db_exists(self, hook) -> bool:
row = hook.get_first(
"SELECT 1 FROM pg_database WHERE datname = %s LIMIT 1",
parameters=(self.database,),
)
return bool(row)
[docs]
def execute(self, context):
self.create_if_not_exists = _coerce_bool(self.create_if_not_exists)
self.use_rsadmin_override = _coerce_bool(self.use_rsadmin_override)
self.log.info(
"Connecting to Redshift database: %s (rsadmin override: %s)",
self.admin_database or "<conn default>",
self.use_rsadmin_override,
)
# Connect to the admin database (default 'dev'). When use_rsadmin_override
# is set, authenticate as the rsadmin superuser so CREATE DATABASE with a
# specific owner is permitted; otherwise use the connection's own identity.
if self.use_rsadmin_override:
hook = build_redshift_hook(
redshift_conn_id=self.redshift_conn_id,
database=self.admin_database or None,
use_rsadmin_override=True,
log=self.log,
)
else:
hook = RedshiftSQLHook(
redshift_conn_id=self.redshift_conn_id,
database=self.admin_database if self.admin_database else None,
)
if self._db_exists(hook):
self.log.info("Database %s already exists; nothing to do.", self.database)
return {"database": self.database, "created": False}
if not self.create_if_not_exists:
raise AirflowException(
f"Database {self.database} does not exist and create_if_not_exists is False."
)
db_ident = _q_ident(self.database)
if self.db_owner:
sql = f"CREATE DATABASE {db_ident} WITH OWNER {_q_ident(self.db_owner)};"
else:
sql = f"CREATE DATABASE {db_ident};"
self.log.info(
"Creating database %s (admin connection DB: %s)",
self.database,
self.admin_database or "<connection default>",
)
try:
# CREATE DATABASE must run outside a transaction in many engines; use autocommit.
hook.run(sql, autocommit=True)
except Exception as e:
# Race-safe: if someone else created it, continue.
if "already exists" not in str(e).lower():
raise
self.log.info(
"Database %s appeared concurrently; continuing.", self.database
)
self.log.info("Database %s ensured.", self.database)
return {"database": self.database, "created": True}
[docs]
class ValidateOrCreateSchemaOperator(BaseOperator):
"""
Create a schema if it doesn't exist. Optionally set the owner at creation.
(Ownership is not changed if the schema already exists.)
:param redshift_conn_id: Airflow connection id for Redshift (cluster or serverless).
:type redshift_conn_id: str
:param schema: Name of the schema to create.
:type schema: str
:param schema_owner: Optional user or role to own the schema (default is the current user).
:type schema_owner: Optional[str]
:param create_if_not_exists: If True, create the schema if it doesn't exist.
:type create_if_not_exists: bool
:param database: Database to connect to (default 'dev').
:type database: str
"""
template_fields: Sequence[str] = (
"schema",
"schema_owner",
"redshift_conn_id",
"database",
)
ui_color = "#5AFE57"
def __init__(
self,
redshift_conn_id: str,
schema: str,
schema_owner: Optional[str] = None,
create_if_not_exists: bool = True,
database: str = "dev",
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.schema = schema
self.schema_owner = schema_owner
self.create_if_not_exists = create_if_not_exists
self.database = database
def _schema_exists(self, hook: PostgresHook) -> bool:
row = hook.get_first(
"SELECT 1 FROM pg_namespace WHERE nspname = %s LIMIT 1",
parameters=(self.schema,),
)
return bool(row)
[docs]
def execute(self, context):
self.log.info(
"Connecting to Redshift database: %s", self.database or "<conn default>"
)
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None, # override only if provided
)
if self._schema_exists(hook):
self.log.info(
"Schema %s already exists on %s; nothing to do.",
self.schema,
self.database,
)
return {"schema": self.schema, "database": self.database, "created": False}
elif not self.create_if_not_exists:
raise AirflowException(
f"Schema {self.schema} does not exist and create_if_not_exists is False"
)
else:
self.log.info(
"Schema %s does not exist in %s; creating it.",
self.schema,
self.database,
)
schema = _q_ident(self.schema)
if self.schema_owner:
sql = f"CREATE SCHEMA IF NOT EXISTS {schema} AUTHORIZATION {_q_ident(self.schema_owner)};"
else:
sql = f"CREATE SCHEMA IF NOT EXISTS {schema};"
try:
self.log.info("Creating schema %s in %s", self.schema, self.database)
hook.run(sql)
except Exception as e:
# Race-safe: if someone created it moments ago, continue
if "already exists" not in str(e).lower():
raise
self.log.info(
"Schema %s appeared concurrently in %s; continuing.",
self.schema,
self.database,
)
self.log.info("Schema %s ensured in %s.", self.schema, self.database)
return {"schema": self.schema, "database": self.database, "created": True}
[docs]
class SetAwsRolePermissiveGrantsOperator(BaseOperator):
"""
Ensure a Redshift DB group exists, (optionally) add specific users to it,
and grant permissive privileges on a schema to that group.
Typical use: map Identity Center / Google SSO sessions into the group via DbGroups.
If you also pass `iamr_users`, this operator will try to add those *existing* users
to the group (skips any that don't exist yet).
:param redshift_conn_id: Airflow connection id for Redshift (cluster or serverless).
:type redshift_conn_id: str
:param schema: Target schema to grant.
:type schema: str
:param group_name: DB group name to manage (e.g., 'data_team_rw').
:type group_name: str
:param iamr_users: Optional iterable of DB usernames (e.g., 'IAMR:AWSReservedSSO_...'). These must already exist in the database (user must have logged in once).
:type iamr_users: Optional[Iterable[str]]
:param grant_default_privileges: If True, set default privileges for future tables/sequences (applies to objects created by the user that runs this).
:type grant_default_privileges: bool
:param default_privilege_owners: If `grant_default_privileges` is True, the list of users whose
future objects will have default privileges granted to the group. If empty, no default
privileges are set.
:type default_privilege_owners: Optional[Iterable[str]]
:param database: Database to connect to (default 'dev').
:type database: str
:param grant_database_privileges: Whether to grant CREATE/TEMP on the database (default True).
:type grant_database_privileges: bool
"""
template_fields: Sequence[str] = (
"schema",
"group_name",
"iamr_users",
"redshift_conn_id",
"database",
"default_privilege_owners",
)
ui_color = "#DEFA75"
def __init__(
self,
redshift_conn_id: str,
schema: str,
group_name: str,
iamr_users: Optional[Iterable[str]] = None,
grant_default_privileges: bool = True,
default_privilege_owners: Optional[Iterable[str]] = None,
database: str = "dev",
grant_database_privileges: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.schema = schema
self.group_name = group_name
self.iamr_users = list(iamr_users) if iamr_users else []
self.grant_default_privileges = grant_default_privileges
self.default_privilege_owners = list(default_privilege_owners or [])
self.database = database
self.grant_database_privileges = grant_database_privileges
#########
# Helper methods to check existence of group and users
# These are used to avoid unnecessary errors when the group or users don't exist.
#########
def _group_exists(self, hook: PostgresHook) -> bool:
sql = "SELECT 1 FROM pg_group WHERE groname = %s LIMIT 1"
row = hook.get_first(sql, parameters=(self.group_name,))
return bool(row)
def _user_exists(self, hook: PostgresHook, username: str) -> bool:
sql = "SELECT 1 FROM pg_user WHERE usename = %s LIMIT 1"
row = hook.get_first(sql, parameters=(username,))
return bool(row)
def _user_in_group(self, hook: PostgresHook, username: str) -> bool:
row = hook.get_first(
"""SELECT 1
FROM pg_group g
JOIN pg_user u
ON u.usesysid = ANY (g.grolist)
WHERE g.groname = %s
AND u.usename = %s
LIMIT 1""",
parameters=(self.group_name, username),
)
return bool(row)
[docs]
def execute(self, context):
self.log.info(
"Connecting to Redshift database: %s", self.database or "<conn default>"
)
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None, # override only if provided
)
# Collect principals (group + iamr users + current_user if you like)
principals = _collect_principals(
hook,
group_name=self.group_name,
iamr_users=self.iamr_users,
include_current_user=True,
log=self.log,
)
#########
# 0) Database-level grants (CONNECT/TEMP)
#########
if self.grant_database_privileges:
db_ident = _q_ident(self.database)
for p in principals:
try:
hook.run(f"GRANT CREATE ON DATABASE {db_ident} TO {p};")
hook.run(f"GRANT TEMP ON DATABASE {db_ident} TO {p};")
except Exception as e:
self.log.warning(
"Granting DB privileges to %s failed (continuing): %s", p, e
)
#########
# 1) Ensure the group exists
#########
if not self._group_exists(hook):
try:
self.log.info("Creating DB group %s", self.group_name)
hook.run(f"CREATE GROUP {_q_ident(self.group_name)};")
except Exception as e:
if "already exists" not in str(e).lower():
raise
self.log.info(
"Group %s already exists in %s (race). Continuing.",
self.group_name,
self.database,
)
else:
self.log.info(
"DB group %s already exists in %s; skipping create.",
self.group_name,
self.database,
)
#########
# 2) Optionally add users to the group (only those that already exist)
#########
if self.iamr_users:
to_add, missing = [], []
for u in self.iamr_users:
if not self._user_exists(hook, u):
missing.append(u)
elif not self._user_in_group(hook, u):
to_add.append(u)
if to_add:
users_csv = ", ".join(_q_user(u) for u in to_add)
self.log.info(
"Adding %d user(s) to group %s in %s",
len(to_add),
self.group_name,
self.database,
)
hook.run(
f"ALTER GROUP {_q_ident(self.group_name)} ADD USER {users_csv};"
)
if missing:
self.log.warning(
"Skipping %d user(s) that don't exist yet: %s. Tip: prefer DbGroups mapping.",
len(missing),
", ".join(missing),
)
#########
# 3) Grants for schema + existing objects
#########
schema = _q_ident(self.schema)
group = _q_ident(self.group_name)
statements: List[str] = []
for p in principals:
statements.append(f"GRANT USAGE, CREATE ON SCHEMA {schema} TO {p};")
statements.append(
# TRUNCATE is a separate Redshift privilege from DELETE. It must be granted
# explicitly because table-owner-implicit TRUNCATE only works for the table
# creator. Required for any DAG that calls TRUNCATE (e.g. pg_to_redshift_data
# optional_purge) when running as a non-owner principal.
f"GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA {schema} TO {p};"
)
# The following block of ifs will do the following:
# 1. If grant_default_privileges is False, skip all default privilege changes.
# 2. If grant_default_privileges is True but default_privilege_owners is empty, log and skip.
# 3. If grant_default_privileges is True and default_privilege_owners is non-empty, proceed.
# For each owner, ensure the owner exists, and if so, add the default privileges.
# If an owner does not exist, log a warning and skip that owner.
if self.grant_default_privileges:
if not self.default_privilege_owners:
self.log.info(
"grant_default_privileges=True but default_privilege_owners is empty; "
"skipping default privilege changes."
)
else:
for owner in self.default_privilege_owners:
if not self._user_exists(hook, owner):
self.log.warning(
"Owner %r does not exist; skipping default privileges.",
owner,
)
continue
owner_q = _q_user(owner)
# Redshift supports default privileges for TABLES only (not sequences).
# TRUNCATE included so future tables created by the new airflow EC2 roles
# can be truncated by other readydata members (sync DAGs run as group
# principals, not as the table owner).
statements += [
f"ALTER DEFAULT PRIVILEGES FOR USER {owner_q} IN SCHEMA {schema} "
f"GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO GROUP {group};"
]
self.log.info(
"Applying grants on schema %s to group %s", self.schema, self.group_name
)
hook.run(statements, autocommit=True)
self.log.info("Done.")
return {
"schema": self.schema,
"database": self.database,
"group_name": self.group_name,
"statements_applied": len(statements),
}
[docs]
class SetAwsRoleAwsDataCatalogGrantsOperator(BaseOperator):
"""
Grant IAMR users/groups access to the auto-mounted Glue Data Catalog in Redshift
(database: awsdatacatalog), plus selected Glue databases (as schemas).
Modes:
- lake_formation_mode=True -> Grant USAGE on database and schemas; LF governs table access.
- lake_formation_mode=False -> Also grant SELECT on existing tables (and optional defaults).
:param redshift_conn_id: Airflow connection id for Redshift.
:type redshift_conn_id: str
:param iamr_users: Optional iterable of DB usernames (e.g., 'IAMR:...').
:type iamr_users: Optional[Iterable[str]]
:param group_name: Optional DB group to grant.
:type group_name: Optional[str]
:param glue_databases: Iterable of Glue DB names (schemas) OR the string 'ALL' to auto-discover.
:type glue_databases: Optional[Iterable[str] | str]
:param lake_formation_mode: If True, rely on LF for table-level access (USAGE-only in Redshift).
:type lake_formation_mode: bool
:param grant_default_select: If False, skip default privileges. (Default True when not LF mode.)
:type grant_default_select: bool
:param database_name: The auto-mounted external database name (default 'awsdatacatalog').
:type database_name: str
:param database: Optional Redshift database to connect to for issuing statements.
:type database: Optional[str]
:param include_current_user: If True, include the current user in the grants.
:type include_current_user: bool
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"iamr_users",
"group_name",
"glue_databases",
"lake_formation_mode",
"grant_default_select",
"database_name",
"database",
"include_current_user",
)
ui_color = "#CB80FD"
def __init__(
self,
redshift_conn_id: str,
iamr_users: Optional[Iterable[str]] = None,
group_name: Optional[str] = None,
glue_databases: Optional[Iterable[str] | str] = "ALL",
lake_formation_mode: bool = True,
grant_default_select: bool = True,
database_name: str = "awsdatacatalog",
database: Optional[str] = "dev",
include_current_user: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.iamr_users = list(iamr_users or [])
self.group_name = group_name
self.glue_databases = glue_databases
self.lake_formation_mode = lake_formation_mode
self.grant_default_select = grant_default_select
self.database_name = database_name
self.database = database
self.include_current_user = include_current_user
[docs]
def execute(self, context):
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None,
)
# 0) Sanity check
who = hook.get_first("SELECT current_user;")
self.log.info(
"Catalog grants will include CURRENT_USER=%s", who[0] if who else "?"
)
# 1) Collect principals to grant
principals = _collect_principals(
hook,
group_name=self.group_name,
iamr_users=self.iamr_users,
include_current_user=self.include_current_user,
log=self.log,
)
# 2) Grant USAGE on the external database (awsdatacatalog) — tolerant if unsupported
db_ident = _q_ident(self.database_name)
try:
for p in principals:
hook.run(f"GRANT USAGE ON DATABASE {db_ident} TO {p};")
except Exception as e:
# Some engine builds don't support database-level USAGE; it's non-fatal.
self.log.debug("Ignoring GRANT USAGE ON DATABASE error (non-fatal): %s", e)
# 3) Determine Glue DBs (schemas) to grant
raw = self.glue_databases
coerced = _coerce_glue_db_list(raw)
if coerced == "ALL":
glue_dbs = _discover_glue_dbs(hook, self.database_name)
if not glue_dbs:
self.log.warning(
"No Glue databases discovered under %s.", self.database_name
)
return {
"database_name": self.database_name,
"glue_dbs": [],
"statements_applied": 0,
}
else:
glue_dbs = coerced
if not glue_dbs:
self.log.warning("glue_databases is empty; nothing to grant.")
return {
"database_name": self.database_name,
"glue_dbs": [],
"statements_applied": 0,
}
# 4) Grant per-Glue DB (as schema inside the external database)
# Three-part naming: awsdatacatalog.<glue_db>.<table>
# Batch USAGE/SELECT statements to reduce round-trips; handle default privs separately.
batched_statements: List[str] = []
for glue_db in glue_dbs:
schema_3p = f"{db_ident}.{_q_ident(glue_db)}"
# Always grant USAGE on the schema so principals can see it
for p in principals:
batched_statements.append(f"GRANT USAGE ON SCHEMA {schema_3p} TO {p};")
# If not using LF permissions, also grant SELECT on existing tables
if not self.lake_formation_mode:
for p in principals:
batched_statements.append(
f"GRANT SELECT ON ALL TABLES IN SCHEMA {schema_3p} TO {p};"
)
if self.grant_default_select and self.group_name:
try:
hook.run(
f"ALTER DEFAULT PRIVILEGES IN SCHEMA {schema_3p} "
f"GRANT SELECT ON TABLES TO GROUP {_q_ident(self.group_name)};"
)
except Exception as e:
self.log.info(
"Default privileges not applied for %s (likely external schema): %s",
schema_3p,
e,
)
# Apply batched grants (idempotent)
if batched_statements:
self.log.info(
"Applying %d grant statements across %d Glue DB(s) in %s.",
len(batched_statements),
len(glue_dbs),
self.database_name,
)
hook.run(batched_statements, autocommit=True)
self.log.info(
"Catalog grants complete (db=%s, schemas=%s, LF_mode=%s).",
self.database_name,
glue_dbs or "(none)",
self.lake_formation_mode,
)
return {
"database_name": self.database_name,
"glue_dbs": glue_dbs,
"statements_applied": len(batched_statements),
}
[docs]
class RefreshAwsDataCatalogGrantsOperator(BaseOperator):
"""
Ensure Redshift principals can see and SELECT all current external tables in the
auto-mounted Glue Data Catalog database (default: 'awsdatacatalog'), and keep
things fresh after new tables are added by Athena/awswrangler.
Actions (idempotent):
- GRANT USAGE ON DATABASE <catalog>
- For each Glue DB (schema) -> GRANT USAGE ON SCHEMA
- For each Glue DB (schema) -> GRANT SELECT ON ALL TABLES IN SCHEMA
Targets:
- A DB GROUP (recommended) via group_name
- And/or explicit users via iamr_users (only if they already exist in Redshift)
:param redshift_conn_id: Airflow connection to Redshift (works with Serverless).
:type redshift_conn_id: str
:param group_name: Optional DB group to grant (e.g., 'data_team_rw').
:type group_name: Optional[str]
:param iamr_users: Optional iterable of Redshift usernames (e.g., 'IAMR:...').
:type iamr_users: Optional[Iterable[str]]
:param glue_databases: List of Glue DB names (schemas), or "ALL" to auto-discover.
:type glue_databases: Optional[Iterable[str] | str]
:param database_name: Catalog database name (default 'awsdatacatalog').
:type database_name: str
:param grant_select: If False, only USAGE grants are applied (default True).
:type grant_select: bool
:param database: Optional Redshift database to connect to for issuing statements.
:type database: Optional[str]
:param include_current_user: If True, include the current user in the grants.
:type include_current_user: bool
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"group_name",
"iamr_users",
"glue_databases",
"database_name",
"grant_select",
"database",
"include_current_user",
)
ui_color = "#F859FB"
def __init__(
self,
redshift_conn_id: str,
group_name: Optional[str] = None,
iamr_users: Optional[Iterable[str]] = None,
glue_databases: Optional[Iterable[str] | str] = "ALL",
database_name: str = "awsdatacatalog",
grant_select: bool = True,
database: Optional[str] = "dev",
include_current_user: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.group_name = group_name
self.iamr_users = list(iamr_users or [])
self.glue_databases = glue_databases
self.database_name = database_name
self.grant_select = grant_select
self.database = database
self.include_current_user = include_current_user
#################################################################################################
# Main execution
#################################################################################################
[docs]
def execute(self, context):
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None,
)
principals = _collect_principals(
hook,
group_name=self.group_name,
iamr_users=self.iamr_users,
include_current_user=self.include_current_user,
log=self.log,
)
db_ident = _q_ident(self.database_name)
# USAGE on the catalog DB — tolerant if unsupported
try:
for p in principals:
hook.run(f"GRANT USAGE ON DATABASE {db_ident} TO {p};")
except Exception as e:
self.log.debug("Ignoring GRANT USAGE ON DATABASE error (non-fatal): %s", e)
# Which Glue DBs (schemas)?
raw = self.glue_databases
coerced = _coerce_glue_db_list(raw)
if coerced == "ALL":
glue_dbs = _discover_glue_dbs(hook, self.database_name)
if not glue_dbs:
self.log.warning(
"No Glue databases discovered under %s.", self.database_name
)
return {
"database_name": self.database_name,
"glue_dbs": [],
"statements_applied": 0,
}
else:
glue_dbs = coerced
if not glue_dbs:
self.log.warning("glue_databases is empty; nothing to grant.")
return {
"database_name": self.database_name,
"glue_dbs": [],
"statements_applied": 0,
}
# Apply schema USAGE and table SELECT
statements: List[str] = []
for gdb in glue_dbs:
schema_3p = f"{db_ident}.{_q_ident(gdb)}"
for p in principals:
statements.append(f"GRANT USAGE ON SCHEMA {schema_3p} TO {p};")
if self.grant_select:
for p in principals:
statements.append(
f"GRANT SELECT ON ALL TABLES IN SCHEMA {schema_3p} TO {p};"
)
if statements:
self.log.info(
"Applying %d grant statements across %d Glue DB(s) in %s.",
len(statements),
len(glue_dbs),
self.database_name,
)
hook.run(statements, autocommit=True)
self.log.info("Catalog grants refreshed.")
return {
"database_name": self.database_name,
"glue_dbs": glue_dbs,
"statements_applied": len(statements),
}
[docs]
class SetCaseSensitiveIdentifierOperator(BaseOperator):
"""
Set (or clear) the Redshift user-level parameter `enable_case_sensitive_identifier`
for many users at once. Useful for "baseline" runs so every new session for those
users is case-sensitive by default without a per-session SET.
Actions (idempotent):
- ALTER USER "<user>" SET enable_case_sensitive_identifier TO true|false
Selection:
- users="ALL" (default) -> act on every user from pg_user (minus exclusions)
- users=[...] -> act only on the provided usernames
- include_iam_users_only=True -> limit to IAM-authenticated principals (IAMR:/IAMA:)
:param redshift_conn_id: Airflow connection id for Redshift (cluster or serverless).
:type redshift_conn_id: str
:param value: Whether to set the parameter to TRUE (default) or FALSE.
:type value: bool
:param users: "ALL" to target all users, or an iterable of specific usernames.
:type users: Union[str, Iterable[str]]
:param exclude_users: Optional iterable of usernames to skip (applied after selection).
:type exclude_users: Optional[Iterable[str]]
:param include_iam_users_only: If True, restrict to users whose names start with IAMR:/IAMA:.
:type include_iam_users_only: bool
:param include_superusers: If True, include superusers in the update (default False).
:type include_superusers: bool
:param additional_protected_users: Additional usernames to skip (beyond built-in service accounts).
:type additional_protected_users: Optional[Iterable[str]]
:param dry_run: If True, log planned changes without executing (default False).
:type dry_run: bool
:param database: Optional database to connect to when issuing statements.
:type database: Optional[str]
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"users",
"exclude_users",
"include_iam_users_only",
"database",
"value",
)
ui_color = "#ABA7E5"
def __init__(
self,
redshift_conn_id: str,
value: bool = True,
users: Union[str, Iterable[str]] = "ALL",
exclude_users: Optional[Iterable[str]] = None,
include_iam_users_only: bool = False,
include_superusers: bool = False,
additional_protected_users: Optional[Iterable[str]] = None,
dry_run: bool = False,
database: Optional[str] = "dev",
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.value = bool(value)
self.users = users
self.exclude_users = set(exclude_users or [])
self.include_iam_users_only = include_iam_users_only
self.include_superusers = include_superusers
self.additional_protected_users = set(
map(str.strip, additional_protected_users or [])
)
self.database = database
self.dry_run = dry_run
################################################################################################
# Helpers
################################################################################################
@staticmethod
def _coerce_user_list(val) -> List[str]:
"""
Accepts "ALL" or an iterable of usernames; returns a concrete list (or ["ALL"] marker).
"""
if isinstance(val, str):
s = val.strip()
if s.upper() == "ALL":
return ["ALL"]
return [s] if s else []
return [str(x).strip() for x in (val or []) if str(x).strip()]
def _fetch_all_users(self, hook: PostgresHook) -> List[str]:
"""
Pull every user from pg_user; optionally filter to IAM users only.
"""
rows = hook.get_records("SELECT usename FROM pg_user;")
names = [r[0] for r in rows if r and isinstance(r[0], str)]
if self.include_iam_users_only:
# Redshift maps IAM role/user sessions with IAMR:/IAMA: prefixes.
names = [u for u in names if u.startswith("IAMR:") or u.startswith("IAMA:")]
return sorted(set(names))
def _list_db_users(self, hook: PostgresHook) -> List[Tuple[str, bool]]:
# returns [(usename, usesuper), ...]
rows = hook.get_records("SELECT usename, usesuper FROM pg_user;")
return [(r[0], bool(r[1])) for r in rows if r and r[0]]
def _resolve_targets(self, hook: PostgresHook) -> Tuple[List[str], List[str]]:
"""Return (to_change, skipped) after applying filters & safety rules."""
all_users = self._list_db_users(hook) # [(name, is_super)]
service_accounts = {"rsadmin", "rdsdb", "rds_iam", "rds_superuser"}
# seed protected set
protected = set(u.lower() for u in service_accounts) | set(
u.lower() for u in self.additional_protected_users
)
# choose the universe
if isinstance(self.users, str) and self.users.strip().upper() == "ALL":
candidates = [u for (u, _) in all_users]
else:
wanted = {str(u).strip() for u in (self.users or []) if str(u).strip()}
candidates = [u for (u, _) in all_users if u in wanted]
# IAM-only filter
if self.include_iam_users_only:
candidates = [
u for u in candidates if u.startswith("IAMR:") or u.startswith("IAMA:")
]
# protect superusers + service accounts unless overridden
name_to_super = {u: s for (u, s) in all_users}
if not self.include_superusers:
candidates, skipped = (
[
u
for u in candidates
if not name_to_super.get(u, False) and u.lower() not in protected
],
[
u
for u in candidates
if name_to_super.get(u, False) or u.lower() in protected
],
)
else:
skipped = []
return candidates, skipped
################################################################################################
# Main execution
################################################################################################
[docs]
def execute(self, context):
"""
Apply `enable_case_sensitive_identifier` for the selected users.
Safety rails:
- Skips superusers and common service accounts by default
(unless include_superusers=True).
- Can restrict to IAMR users only (include_iam_users_only=True).
- Supports dry_run to preview changes.
Returns a small summary dict for logging/metrics.
"""
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None, # override DB only if provided
)
# Resolve target set and who we will skip
targets, skipped = self._resolve_targets(hook)
val = "true" if self.value else "false"
if not targets:
self.log.info(
"No eligible users to update (value=%s). Skipped: %s",
val,
", ".join(skipped) or "(none)",
)
return {
"updated": 0,
"planned": 0,
"skipped": skipped,
"applied_value": self.value,
"database": self.database,
}
# Build ALTER USER statements
stmts = [
f"ALTER USER {_q_user(u)} SET enable_case_sensitive_identifier TO {val};"
for u in targets
]
self.log.info(
"Planned %d ALTER USER statements (dry_run=%s). "
"Skipping %d protected users: %s",
len(stmts),
self.dry_run,
len(skipped),
", ".join(skipped) or "(none)",
)
if self.dry_run:
# Preview only; don't run anything
return {
"updated": 0,
"planned": len(stmts),
"targets": targets,
"skipped": skipped,
"applied_value": self.value,
"database": self.database,
}
# Execute as a batch (consistent with other operators)
hook.run(stmts, autocommit=True)
self.log.info(
"Applied enable_case_sensitive_identifier=%s to %d user(s) in %s.",
val,
len(targets),
self.database or "<connection default>",
)
if skipped:
self.log.info("Skipped (protected): %s", ", ".join(skipped))
return {
"updated": len(targets),
"planned": len(stmts),
"targets": targets,
"skipped": skipped,
"applied_value": self.value,
"database": self.database,
}
[docs]
class CreateRedshiftUserOperator(BaseOperator):
"""
Create a Redshift user with password if it doesn't exist.
Optionally retrieves password from AWS Secrets Manager.
Grants CONNECT (and optionally TEMP) on the target database.
Password can be provided in two ways (one is required):
- `secret_arn`: Retrieve from AWS Secrets Manager (recommended for production).
The secret must already exist and contain a JSON key named 'password'.
- `password`: Pass directly as a string (less secure, suitable for dev/testing).
:param redshift_conn_id: Airflow connection id for Redshift.
:type redshift_conn_id: str
:param username: Name of the user to create.
:type username: str
:param password: Password for the user. Ignored if secret_arn is provided.
:type password: Optional[str]
:param secret_arn: AWS Secrets Manager ARN containing password. The secret must
already exist and contain a JSON object with a 'password' key.
:type secret_arn: Optional[str]
:param grant_connect: Deprecated/ignored. Redshift does not support GRANT CONNECT.
Users can connect by default once created. Kept for backwards compatibility.
:type grant_connect: bool
:param grant_temp: Grant TEMPORARY (create temp tables) on the database (default False).
:type grant_temp: bool
:param database: Database to connect to and grant on.
:type database: str
:param createdb: Allow the user to create databases (default False).
:type createdb: bool
:param connection_limit: Max concurrent connections for this user (default unlimited).
:type connection_limit: Optional[int]
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"username",
"password",
"secret_arn",
"database",
"grant_temp",
"createdb",
)
ui_color = "#5EEDED"
def __init__(
self,
redshift_conn_id: str,
username: str,
password: Optional[str] = None,
secret_arn: Optional[str] = None,
grant_connect: bool = True,
grant_temp: bool = False,
database: str = "dev",
createdb: bool = False,
connection_limit: Optional[int] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.username = username
self.password = password
self.secret_arn = secret_arn
self.grant_connect = grant_connect
self.grant_temp = grant_temp
self.database = database
self.createdb = createdb
self.connection_limit = connection_limit
def _get_password(self) -> str:
"""Retrieve password from Secrets Manager or use provided password."""
if self.secret_arn:
import boto3 # Lazy load: heavy AWS SDK
client = boto3.client("secretsmanager")
response = client.get_secret_value(SecretId=self.secret_arn)
secret = json.loads(response["SecretString"])
pwd = secret.get("password")
if not pwd:
raise AirflowException(
f"Secret {self.secret_arn} does not contain 'password' key."
)
return pwd
elif self.password:
return self.password
else:
raise AirflowException(
"Either 'password' or 'secret_arn' must be provided."
)
def _user_exists(self, hook: PostgresHook) -> bool:
row = hook.get_first(
"SELECT 1 FROM pg_user WHERE usename = %s LIMIT 1",
parameters=(self.username,),
)
return bool(row)
[docs]
def execute(self, context):
self.grant_temp = _coerce_bool(self.grant_temp)
self.createdb = _coerce_bool(self.createdb)
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None,
)
created = False
if self._user_exists(hook):
self.log.info("User %s already exists; skipping creation.", self.username)
else:
pwd = self._get_password()
user_ident = _q_ident(self.username)
options = [f"PASSWORD '********'"] # Placeholder for logging
if self.createdb:
options.append("CREATEDB")
if self.connection_limit is not None:
options.append(f"CONNECTION LIMIT {int(self.connection_limit)}")
self.log.info(
"Creating user %s in database %s with options: %s",
self.username,
self.database,
" ".join(options),
)
real_options = [f"PASSWORD '{pwd}'"]
if self.createdb:
real_options.append("CREATEDB")
if self.connection_limit is not None:
real_options.append(f"CONNECTION LIMIT {int(self.connection_limit)}")
sql = f"CREATE USER {user_ident} {' '.join(real_options)};"
# Redshift CREATE USER does not accept a parameterized password, so
# the plaintext is in the SQL string. Log a redacted form for audit
# and suppress DbApiHook's unredacted auto-log of the same statement.
self.log.info(
"Running statement: CREATE USER %s %s;",
user_ident,
" ".join(options),
)
original_log_sql = hook.log_sql
hook.log_sql = False
try:
hook.run(sql, handler=lambda cur: None) # Suppress result logging
finally:
hook.log_sql = original_log_sql
created = True
db_ident = _q_ident(self.database)
user_ident = _q_ident(self.username)
if self.grant_temp:
hook.run(f"GRANT TEMPORARY ON DATABASE {db_ident} TO {user_ident};")
self.log.info("Granted TEMPORARY on %s to %s.", self.database, self.username)
self.log.info("User %s ensured with requested privileges.", self.username)
return {"user": self.username, "database": self.database, "created": created}
[docs]
class SetSchemaReadGrantsOperator(BaseOperator):
"""
Grant read-only access on a schema: USAGE, SELECT on tables, EXECUTE on functions.
Sets default privileges for future objects created by specified owners.
This operator follows least-privilege principles for service accounts that only
need to read data or execute functions (e.g., Analytics Toolbox).
:param redshift_conn_id: Airflow connection id for Redshift.
:type redshift_conn_id: str
:param schema: Target schema to grant read access on.
:type schema: str
:param grantee_users: List of usernames to grant access to.
:type grantee_users: Optional[Iterable[str]]
:param grantee_group: Group name to grant access to.
:type grantee_group: Optional[str]
:param grant_select: Grant SELECT on all tables (default True).
:type grant_select: bool
:param grant_execute: Grant EXECUTE on all functions (default True).
:type grant_execute: bool
:param set_default_privileges: Set default privileges for future objects (default True).
:type set_default_privileges: bool
:param default_privilege_owners: Users whose future objects get default grants.
:type default_privilege_owners: Optional[Iterable[str]]
:param database: Database to connect to.
:type database: str
"""
template_fields: Sequence[str] = (
"redshift_conn_id",
"schema",
"grantee_users",
"grantee_group",
"default_privilege_owners",
"database",
)
ui_color = "#A55A55"
def __init__(
self,
redshift_conn_id: str,
schema: str,
grantee_users: Optional[Iterable[str]] = None,
grantee_group: Optional[str] = None,
grant_select: bool = True,
grant_execute: bool = True,
set_default_privileges: bool = True,
default_privilege_owners: Optional[Iterable[str]] = None,
database: str = "dev",
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.redshift_conn_id = redshift_conn_id
self.schema = schema
self.grantee_users = list(grantee_users) if grantee_users else []
self.grantee_group = grantee_group
self.grant_select = grant_select
self.grant_execute = grant_execute
self.set_default_privileges = set_default_privileges
self.default_privilege_owners = list(default_privilege_owners or [])
self.database = database
def _collect_grantees(self, hook: PostgresHook) -> List[str]:
"""Build list of grantees, validating users exist."""
grantees: List[str] = []
if self.grantee_group:
grantees.append(f"GROUP {_q_ident(self.grantee_group)}")
for u in self.grantee_users:
exists = hook.get_first(
"SELECT 1 FROM pg_user WHERE usename = %s LIMIT 1",
parameters=(u,),
)
if exists:
grantees.append(_q_user(u))
else:
self.log.warning("User %s not found; skipping.", u)
if not grantees:
raise AirflowException(
"No valid grantees provided (grantee_users or grantee_group required)."
)
return grantees
[docs]
def execute(self, context):
hook = PostgresHook(
postgres_conn_id=self.redshift_conn_id,
schema=self.database or None,
)
grantees = self._collect_grantees(hook)
schema_ident = _q_ident(self.schema)
statements: List[str] = []
for g in grantees:
statements.append(f"GRANT USAGE ON SCHEMA {schema_ident} TO {g};")
if self.grant_select:
for g in grantees:
statements.append(
f"GRANT SELECT ON ALL TABLES IN SCHEMA {schema_ident} TO {g};"
)
if self.grant_execute:
for g in grantees:
statements.append(
f"GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA {schema_ident} TO {g};"
)
if statements:
self.log.info(
"Applying %d read grant statements on schema %s.",
len(statements),
self.schema,
)
hook.run(statements, autocommit=True)
if self.set_default_privileges and self.default_privilege_owners:
for owner in self.default_privilege_owners:
owner_exists = hook.get_first(
"SELECT 1 FROM pg_user WHERE usename = %s LIMIT 1",
parameters=(owner,),
)
if not owner_exists:
self.log.warning(
"Default privilege owner %s not found; skipping.", owner
)
continue
owner_ident = _q_user(owner)
for g in grantees:
if self.grant_select:
hook.run(
f"ALTER DEFAULT PRIVILEGES FOR USER {owner_ident} IN SCHEMA {schema_ident} "
f"GRANT SELECT ON TABLES TO {g};"
)
if self.grant_execute:
hook.run(
f"ALTER DEFAULT PRIVILEGES FOR USER {owner_ident} IN SCHEMA {schema_ident} "
f"GRANT EXECUTE ON FUNCTIONS TO {g};"
)
self.log.info(
"Default privileges set for owners: %s",
", ".join(self.default_privilege_owners),
)
self.log.info(
"Read grants complete on schema %s in database %s.",
self.schema,
self.database,
)
return {
"schema": self.schema,
"database": self.database,
"grantees": self.grantee_users
+ ([self.grantee_group] if self.grantee_group else []),
"statements_applied": len(statements),
}