Source code for operators.superset.add_superset_db_connection

from typing import Any, Dict, Optional

from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class AddSupersetDbConnectionOperator(BaseOperator):
    """Create a new database connection in Superset via the REST API.

    :param database_name: Display label shown in the Superset UI. This is
        NOT the actual database name connected to — the real DB is determined
        entirely by the ``sqlalchemy_uri``. Free-form; spaces allowed.
    :type database_name: str
    :param sqlalchemy_uri: SQLAlchemy URI for the database connection.
    :type sqlalchemy_uri: str
    :param expose_in_sqllab: Whether to expose the database in SQL Lab.
    :type expose_in_sqllab: bool
    :param allow_ctas: Whether to allow CTAS statements.
    :type allow_ctas: bool
    :param allow_cvas: Whether to allow CVAS statements.
    :type allow_cvas: bool
    :param allow_dml: Whether to allow DML statements.
    :type allow_dml: bool
    :param extra: Optional dict merged into the Superset payload (e.g.
        ``{"masked_encrypted_extra": "..."}`` or ``{"extra": "..."}``).
    :type extra: dict
    :param conn_id: Airflow connection id for the Superset instance.
    :type conn_id: str
    :param base_url: Optional full URL override for the Superset API (e.g.
        ``"https://dap-tx.bead.example.com"``). When provided, overrides the
        connection's host/schema.
    :type base_url: str
    :param timeout: HTTP request timeout in seconds.
    :type timeout: int
    :param owner_email: Optional canonical Ready email (e.g.
        ``majid@ready.net``). When set, resolves to a single Superset
        user and adds them to the DB's owner list alongside the API
        user. DAP stores emails with a per-tenant ``+<uuid>`` suffix;
        this resolver matches both shapes. Important when the API user
        is going to be disabled later in the run — the resolved
        end-user is what keeps DAP's owner-based visibility filter
        from hiding the DB from listings post-disable.
    :type owner_email: str
    """

    template_fields = (
        "database_name",
        "sqlalchemy_uri",
        "extra",
        "conn_id",
        "base_url",
        "owner_email",
    )
    ui_color = "#fce4ec"

    @apply_defaults
    def __init__(
        self,
        database_name: str,
        sqlalchemy_uri: str,
        expose_in_sqllab: bool = True,
        allow_ctas: bool = False,
        allow_cvas: bool = False,
        allow_dml: bool = False,
        extra: Optional[Dict[str, Any]] = None,
        conn_id: str = "superset_default",
        base_url: Optional[str] = None,
        timeout: int = 30,
        owner_email: Optional[str] = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.database_name = database_name
        self.sqlalchemy_uri = sqlalchemy_uri
        self.expose_in_sqllab = expose_in_sqllab
        self.allow_ctas = allow_ctas
        self.allow_cvas = allow_cvas
        self.allow_dml = allow_dml
        self.extra = extra
        self.conn_id = conn_id
        self.base_url = base_url
        self.timeout = timeout
        self.owner_email = owner_email



    def execute(self, context) -> Dict[str, Any]:
        from hooks.superset_hook import SupersetHook

        if not self.database_name:
            raise AirflowException("Missing required field: database_name")
        if not self.sqlalchemy_uri:
            raise AirflowException("Missing required field: sqlalchemy_uri")

        hook = SupersetHook(
            conn_id=self.conn_id, base_url=self.base_url, timeout=self.timeout
        )

        extra_owners = []
        if self.owner_email:
            user = hook.find_user_by_email(self.owner_email)
            extra_owners.append(int(user["id"]))
            self.log.info(
                "Resolved owner_email '%s' to Superset user id=%s email=%s",
                self.owner_email, extra_owners[0], user.get("email"),
            )

        self.log.info("Adding Superset database connection '%s'", self.database_name)
        result = hook.add_db_connection(
            database_name=self.database_name,
            sqlalchemy_uri=self.sqlalchemy_uri,
            expose_in_sqllab=self.expose_in_sqllab,
            allow_ctas=self.allow_ctas,
            allow_cvas=self.allow_cvas,
            allow_dml=self.allow_dml,
            extra=self.extra,
            owners=extra_owners or None,
        )
        self.log.info("Add database result: %s", result)
        return result