from typing import Any, Dict, List, Optional
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
[docs]
class AddSupersetDashboardOperator(BaseOperator):
"""Create a dataset and one or more dashboards in Superset.
Order of operations:
1. Create the dataset described in ``dataset``.
2. For each entry in ``dashboards``, create a dashboard.
Returns a dict ``{"dataset_id": <int>, "dashboard_ids": [<int>, ...]}``.
:param dataset: Dict describing the dataset to create. Required keys:
``database_id`` and ``table_name``. Optional: ``schema``, ``owners``,
``extra`` (merged into the Superset payload).
:type dataset: dict
:param dashboards: List of dicts describing dashboards to create. Each
dict requires ``dashboard_title``. Optional: ``slug``, ``owners``,
``published`` (default False), ``extra`` (merged into the payload).
:type dashboards: list[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
"""
template_fields = ("dataset", "dashboards", "conn_id", "base_url")
ui_color = "#fce4ec"
@apply_defaults
def __init__(
self,
dataset: Dict[str, Any],
dashboards: List[Dict[str, Any]],
conn_id: str = "superset_default",
base_url: Optional[str] = None,
timeout: int = 30,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.dataset = dataset
self.dashboards = dashboards
self.conn_id = conn_id
self.base_url = base_url
self.timeout = timeout
[docs]
def execute(self, context) -> Dict[str, Any]:
from hooks.superset_hook import SupersetHook
if not self.dataset:
raise AirflowException("Missing required field: dataset")
if not isinstance(self.dataset, dict):
raise AirflowException("dataset must be a dict")
if not self.dashboards or not isinstance(self.dashboards, list):
raise AirflowException("dashboards must be a non-empty list of dicts")
database_id = self.dataset.get("database_id")
table_name = self.dataset.get("table_name")
if database_id is None or not table_name:
raise AirflowException(
"dataset requires 'database_id' and 'table_name' keys"
)
hook = SupersetHook(
conn_id=self.conn_id, base_url=self.base_url, timeout=self.timeout
)
self.log.info(
"Creating Superset dataset '%s' on database_id=%s",
table_name,
database_id,
)
dataset_resp = hook.add_dataset(
database_id=database_id,
table_name=table_name,
schema=self.dataset.get("schema"),
owners=self.dataset.get("owners"),
extra=self.dataset.get("extra"),
)
dataset_id = (dataset_resp or {}).get("id")
if dataset_id is None:
raise AirflowException(
f"Dataset create did not return an id: {dataset_resp}"
)
self.log.info("Created dataset id=%s", dataset_id)
dashboard_ids: List[int] = []
for idx, dash in enumerate(self.dashboards):
if not isinstance(dash, dict):
raise AirflowException(
f"dashboards[{idx}] must be a dict, got {type(dash).__name__}"
)
title = dash.get("dashboard_title")
if not title:
raise AirflowException(
f"dashboards[{idx}] missing required key 'dashboard_title'"
)
self.log.info("Creating Superset dashboard '%s'", title)
dash_resp = hook.add_dashboard(
dashboard_title=title,
slug=dash.get("slug"),
owners=dash.get("owners"),
published=dash.get("published", False),
extra=dash.get("extra"),
)
dash_id = (dash_resp or {}).get("id")
if dash_id is None:
raise AirflowException(
f"Dashboard create did not return an id: {dash_resp}"
)
self.log.info("Created dashboard id=%s", dash_id)
dashboard_ids.append(dash_id)
return {"dataset_id": dataset_id, "dashboard_ids": dashboard_ids}