Source code for operators.redshift.get_database

from airflow.exceptions import AirflowException
from typing import Any, Sequence

from operators.redshift.base_redshift import BaseRedshiftOperator




class RedshiftGetDatabaseOperator(BaseRedshiftOperator):
    """
    Operator to retrieve a list of Redshift database names matching a LIKE pattern.

    Queries svv_redshift_databases via the Redshift Data API. Useful for
    fanning out a CTAS across all databases matching a naming convention
    (e.g. all '*_datamart' databases for ctasifier all_tenants mode with
    tenant_discovery='databases_matching').

    Unlike RedshiftGetTenantsOperator (which discovers schemas inside a
    single database), this operator discovers DATABASES across the cluster.

    :param redshift_workgroup: Redshift workgroup name (e.g. Config.REDSHIFT_WORKGROUP).
    :type redshift_workgroup: str
    :param database: Redshift database to issue the discovery query against
        (just a routing target for the Data API; the result set covers all
        databases visible to svv_redshift_databases).
    :type database: str
    :param database_pattern: SQL LIKE pattern for the database_name column,
        e.g. '%_datamart'. Required.
    :type database_pattern: str
    :param region: AWS region for the Redshift Data API client.
    :type region: str
    :param additional_exclude_databases: Optional collection or string of
        database names to exclude from the result set. Accepts list/tuple/set,
        a JSON string (e.g. '["dev","staging"]'), a Python-literal string
        (e.g. "['dev']"), or a single plain string (e.g. "dev"). This is
        CaSe SeNsItIvE.
    :type additional_exclude_databases: Any
    """

    template_fields: Sequence[str] = (
        *BaseRedshiftOperator.template_fields,
        "database_pattern",
        "additional_exclude_databases",
    )

    def __init__(
        self,
        redshift_workgroup: str,
        database: str,
        database_pattern: str,
        region: str = "us-east-1",
        additional_exclude_databases: Any = None,
        *args,
        **kwargs,
    ):
        super().__init__(
            redshift_workgroup=redshift_workgroup,
            database=database,
            region=region,
            *args,
            **kwargs,
        )
        self.database_pattern = database_pattern
        # Store raw; coerce at execution time after templating has rendered.
        self.additional_exclude_databases = additional_exclude_databases



    def execute(self, context):
        """Retrieve the list of database names from Redshift."""
        if not self.database_pattern or not str(self.database_pattern).strip():
            raise AirflowException(
                "RedshiftGetDatabaseOperator requires database_pattern "
                "(SQL LIKE pattern, e.g. '%_datamart')."
            )

        user_excludes = self._coerce_excludes(self.additional_exclude_databases)
        # Redshift databases are user-named, so there's no baseline exclude
        # set (unlike schemas, which include information_schema/aws_*/etc).
        exclude_databases = sorted(set(user_excludes))

        self.log.info(
            "RedshiftGetDatabaseOperator pattern=%r excludes=%s (raw=%r)",
            self.database_pattern,
            exclude_databases,
            self.additional_exclude_databases,
        )

        where_excl = ""
        if exclude_databases:
            excl_clause = ", ".join(
                self.sql_literal(d) for d in exclude_databases
            )
            where_excl = f"    AND database_name NOT IN ({excl_clause})\n"

        sql = (
            "SELECT database_name\n"
            "FROM svv_redshift_databases\n"
            f"WHERE database_name LIKE {self.sql_literal(self.database_pattern)}\n"
            f"{where_excl}"
            "ORDER BY database_name ASC;"
        )

        try:
            rows = self._rs_fetch(sql)
            databases = [r["database_name"] for r in rows]
        except Exception as e:
            self.log.exception(
                "Error retrieving database names from Redshift "
                "(workgroup=%s database=%s pattern=%r): %s",
                self.redshift_workgroup,
                self.database,
                self.database_pattern,
                e,
            )
            raise AirflowException(
                "Failed to retrieve Redshift database names "
                f"(workgroup='{self.redshift_workgroup}', "
                f"database='{self.database}', "
                f"pattern={self.database_pattern!r})."
            ) from e

        if not databases:
            raise AirflowException(
                "No Redshift databases matched pattern "
                f"{self.database_pattern!r} "
                f"(workgroup='{self.redshift_workgroup}')."
            )

        self.log.info("Found %d database(s): %s", len(databases), databases)
        return databases