Source code for operators.pg_get_tenants

from config import Config
from airflow.models import BaseOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowException

DEFAULT_DB_POSTGRES = Config.BOSS_PLATFORM_DB
STAGE_DB_POSTGRES = Config.BOSS_PLATFORM_STAGE_DB




class PostgresGetTenantsOperator(BaseOperator):
    """
    An operator to retrieve a list of tenants from a Postgres database.

    :param use_stage_db: Whether to use the stage database. Defaults to True.
    :type use_stage_db: bool
    :param include_public: Whether to include the public schema in the list of tenants. Defaults to False.
    :type include_public: bool

    """

    ui_color = "#A80085"

    def __init__(
        self, use_stage_db: bool = True, include_public: bool = False, *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.use_stage_db = use_stage_db
        self.include_public = include_public



    def execute(self, context: dict) -> list[str]:
        """
        Execute the operator.

        :return: list of tenants
        :rtype: list[str]
        """

        def tenant_generator(tenants):
            for t in tenants:
                yield t[0]

        # Determine which database to use
        # Airflow dev runs to default (dev) always
        # Airflow prod can run to stage or default (prod)
        if not Config.PRODUCTION:
            dest_db = DEFAULT_DB_POSTGRES
        else:
            dest_db = STAGE_DB_POSTGRES if self.use_stage_db else DEFAULT_DB_POSTGRES

        query = f"""SELECT nspname
        FROM pg_catalog.pg_namespace
        where 1=1
            and nspname not ilike 'pg_%'
            and nspname not in (
                'data',
                'information_schema'{"" if self.include_public else ", 'public'"}
        )"""

        hook = PostgresHook(postgres_conn_id=dest_db)
        tenants = hook.get_records(query)
        if len(tenants) == 0:
            raise AirflowException(
                f"No tenants found in the specified database: {dest_db}"
            )

        return list(tenant_generator(tenants))