Source code for operators.postgres.utils.get_tenants

from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowException
from typing import Any
import json, ast




class GetTenantsOperator(BaseOperator):
    """
    Operator to retrieve a list of tenant schema names from a Postgres database.

    :param postgres_conn_id: Airflow connection id for the Postgres database.
    :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. '["public","demo"]'), a Python-literal string (e.g. "['public']"),
           or a single plain string (e.g. "public"). This is CaSe SeNsItIvE.
    """

    template_fields = ("postgres_conn_id", "additional_exclude_schemas")
    ui_color = "#4D6A42"

    def __init__(
        self,
        postgres_conn_id: str,
        additional_exclude_schemas: Any = None,  # <-- store raw, coerce later
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.postgres_conn_id = postgres_conn_id
        # Store raw; coerce at execution time after templating has rendered.
        self.additional_exclude_schemas = additional_exclude_schemas

    @staticmethod
    def _coerce_excludes(raw: Any) -> list[str]:
        """Turn any templated input into a list[str]."""
        if raw is None:
            return []
        if isinstance(raw, (list, tuple, set)):
            return [str(x) for x in raw if str(x).strip()]
        if isinstance(raw, str):
            s = raw.strip()
            if not s:
                return []
            # Try JSON then Python literal; otherwise treat as single name
            try:
                v = json.loads(s)
            except Exception:
                try:
                    v = ast.literal_eval(s)
                except Exception:
                    return [s]
            if isinstance(v, (list, tuple, set)):
                return [str(x) for x in v if str(x).strip()]
            return [str(v)]
        return [str(raw)]



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

        # Always exclude information_schema; add any user-provided schemas
        base_excludes = {
            "information_schema",
            "public",
            "aws_commons",
            "aws_s3",
            "test_base",
        }
        exclude_schemas = tuple(sorted(base_excludes | set(user_excludes)))

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

        hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
        sql = """
            SELECT schema_name
            FROM information_schema.schemata
            WHERE schema_name NOT IN %(exclude)s
              AND schema_name NOT LIKE 'pg_%%'
            ORDER BY schema_name ASC;
        """
        try:
            records = hook.get_records(sql, {"exclude": exclude_schemas})
            tenants = [row[0] for row in records]
        except Exception as e:
            self.log.exception(
                "Error retrieving tenant schemas with conn_id '%s': %s",
                self.postgres_conn_id,
                e,
            )
            raise AirflowException(
                f"Failed to retrieve tenant schemas from Postgres (conn_id='{self.postgres_conn_id}')."
            ) from e

        if not tenants:
            raise AirflowException(
                f"No tenant schemas found in Postgres (conn_id='{self.postgres_conn_id}')."
            )

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