Source code for operators.redshift.get_tenants

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

from operators.redshift.base_redshift import BaseRedshiftOperator




class RedshiftGetTenantsOperator(BaseRedshiftOperator):
    """
    Operator to retrieve a list of tenant schema names from a Redshift database.

    Queries svv_all_schemas via the Redshift Data API, filtered to a single
    source database (default 'grapi_sync'). Mirrors the API of
    GetTenantsOperator (Postgres version) so DAGs can swap discovery backends
    with minimal changes.

    :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
        (does not have to match tenant_source_database; this is just where
        the SELECT runs).
    :type database: str
    :param tenant_source_database: Value used in the
        ``svv_all_schemas.database_name = ...`` filter. Defaults to
        ``'grapi_sync'``.
    :type tenant_source_database: str
    :param region: AWS region for the Redshift Data API client.
    :type region: str
    :param additional_exclude_schemas: Optional collection or string of schema
        names to exclude in addition to defaults. Accepts list/tuple/set, a
        JSON string (e.g. '["casf","test_base"]'), a Python-literal string
        (e.g. "['casf']"), or a single plain string (e.g. "casf"). This is
        CaSe SeNsItIvE.
    :type additional_exclude_schemas: Any
    """

    template_fields: Sequence[str] = (
        *BaseRedshiftOperator.template_fields,
        "tenant_source_database",
        "additional_exclude_schemas",
    )

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



    def execute(self, context):
        """Retrieve the list of tenant schema names from Redshift."""
        user_excludes = self._coerce_excludes(self.additional_exclude_schemas)

        # Mirror the Postgres GetTenantsOperator default-exclusion set so the
        # two operators behave consistently.
        base_excludes = {"information_schema", "public", "aws_commons", "aws_s3"}
        exclude_schemas = sorted(base_excludes | set(user_excludes))

        self.log.info(
            "RedshiftGetTenantsOperator using excludes=%s (raw=%r) tenant_source_database=%s",
            exclude_schemas,
            self.additional_exclude_schemas,
            self.tenant_source_database,
        )

        exclude_clause = ", ".join(self.sql_literal(s) for s in exclude_schemas)
        sql = (
            "SELECT schema_name\n"
            "FROM svv_all_schemas\n"
            f"WHERE database_name = {self.sql_literal(self.tenant_source_database)}\n"
            f"    AND schema_name NOT IN ({exclude_clause})\n"
            "    AND schema_name NOT LIKE 'pg_%'\n"
            "ORDER BY schema_name ASC;"
        )

        try:
            rows = self._rs_fetch(sql)
            tenants = [r["schema_name"] for r in rows]
        except Exception as e:
            self.log.exception(
                "Error retrieving tenant schemas from Redshift "
                "(workgroup=%s database=%s tenant_source_database=%s): %s",
                self.redshift_workgroup,
                self.database,
                self.tenant_source_database,
                e,
            )
            raise AirflowException(
                "Failed to retrieve tenant schemas from Redshift "
                f"(workgroup='{self.redshift_workgroup}', "
                f"database='{self.database}', "
                f"tenant_source_database='{self.tenant_source_database}')."
            ) from e

        if not tenants:
            raise AirflowException(
                "No tenant schemas found in Redshift "
                f"(workgroup='{self.redshift_workgroup}', "
                f"database='{self.database}', "
                f"tenant_source_database='{self.tenant_source_database}')."
            )

        self.log.info("Found tenant schemas: %s", tenants)
        return tenants