From 773340a44cbfa1018a35abbfcc8a384c0d3e9d44 Mon Sep 17 00:00:00 2001 From: Julian LaNeve Date: Mon, 27 Jul 2026 11:26:55 -0400 Subject: [PATCH] Add Airflow UI plugin: Blueprint tab showing YAML and blueprint code Airflow's Code tab shows the loader file that builds YAML DAGs, not the YAML itself. This adds a plugin (registered via the airflow.plugins entry point) that puts a Blueprint tab on dag, dag run, task, and task instance pages in Airflow 3.1+: - dag and dag run: the source YAML, with each blueprint's Python underneath, syntax highlighted with line numbers - task and task instance: that step's config and blueprint source To make the lookups exact, build_all_airflow_dags now tags each DAG with blueprint: (opt out with source_tags=False) and records the path in each task's blueprint_step_config. Both survive serialization, so the plugin reads per-run rendered fields first (config as the task ran, param overrides included), then the serialized DAG, then the file on disk, and labels which one it shows. Values Airflow truncated under [core] max_templated_field_length fall through to the next source. The pages are served by a FastAPI app mounted at /blueprint. They render in Airflow's sandboxed iframes, so links drive the parent router through its history API, and the pages follow Airflow's theme toggle. On task group pages Airflow leaves {TASK_ID} unsubstituted; the group id is recovered from the Referer. On Airflow 2 the plugin registers nothing. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01TQnbVv8WAkqMzwKbQsbBER --- README.md | 13 + blueprint/builder.py | 67 ++- blueprint/plugin/__init__.py | 98 ++++ blueprint/plugin/app.py | 626 ++++++++++++++++++++++ blueprint/plugin/templates/base.html | 202 +++++++ blueprint/plugin/templates/dag_yaml.html | 29 + blueprint/plugin/templates/task_step.html | 41 ++ examples/advanced/README.md | 15 + pyproject.toml | 4 + tests/integration/test_dag_structure.py | 21 +- tests/integration/test_plugin.py | 96 ++++ tests/test_builder.py | 72 +++ tests/test_plugin.py | 324 +++++++++++ uv.lock | 4 +- 14 files changed, 1595 insertions(+), 17 deletions(-) create mode 100644 blueprint/plugin/__init__.py create mode 100644 blueprint/plugin/app.py create mode 100644 blueprint/plugin/templates/base.html create mode 100644 blueprint/plugin/templates/dag_yaml.html create mode 100644 blueprint/plugin/templates/task_step.html create mode 100644 tests/integration/test_plugin.py create mode 100644 tests/test_plugin.py diff --git a/README.md b/README.md index 8135774..297e71d 100644 --- a/README.md +++ b/README.md @@ -250,6 +250,19 @@ Every task instance gets two extra fields visible in Airflow's "Rendered Templat This makes it easy to understand what generated each task instance without leaving the Airflow UI. +## Airflow UI Plugin (Airflow 3) + +Installing `airflow-blueprint` registers an Airflow plugin that puts the YAML front and center in the UI. Airflow's built-in Code tab shows the loader file that built the DAG — the plugin adds a **Blueprint** tab (Airflow 3.1+) that shows the YAML and blueprint Python instead: + +- On a DAG or DAG run: the source YAML the DAG was built from, with the code of each blueprint it uses underneath. +- On a task or task instance: that step's config and its blueprint source. Where Airflow stored the run's rendered fields, the config is the exact one the task ran with — param overrides included — falling back to the serialized DAG and then the file on disk, labeled with which you're seeing. + +The pages are served by a FastAPI app mounted at `/blueprint` under the API server (Airflow 3.0+), so you can also open them directly. + +To map a DAG back to its YAML, `build_all_airflow_dags()` tags each DAG with `blueprint:` and records the path in each task's `blueprint_step_config`. Both survive DAG serialization, so the plugin resolves the file with one lookup instead of re-scanning the dags folder (it still falls back to a scan for DAGs built by older versions). Pass `source_tags=False` to skip the tag. + +On Airflow 2 the plugin loads but registers nothing. + ## Runtime Parameter Overrides Blueprints that set `supports_params = True` have their config fields registered as Airflow [DAG params](https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/params.html), namespaced as `{step}__{field}`. When you trigger a DAG from the Airflow UI, the trigger form shows those fields pre-filled with YAML defaults — users can override any value before running. diff --git a/blueprint/builder.py b/blueprint/builder.py index 4037fb6..0ef3585 100644 --- a/blueprint/builder.py +++ b/blueprint/builder.py @@ -30,6 +30,9 @@ DEFAULT_START_DATE = datetime(2024, 1, 1, tzinfo=timezone.utc) +SOURCE_TAG_PREFIX = "blueprint:" +_TAG_MAX_LEN = 100 + _PARAM_SCHEMA_KEYS = frozenset( { "type", @@ -186,11 +189,14 @@ def _ensure_start_date(dag_kwargs: dict[str, Any]) -> None: parsed = parsed.replace(tzinfo=timezone.utc) dag_kwargs["start_date"] = parsed - def build(self, config: DAGConfig) -> "DAG": + def build(self, config: DAGConfig, source: str | None = None) -> "DAG": """Build a DAG from a DAGConfig. Args: config: The parsed and validated DAG configuration + source: Path of the YAML file this DAG was built from, recorded + in each task's ``blueprint_step_config`` so the source + survives DAG serialization. Returns: A fully wired Airflow DAG @@ -235,7 +241,7 @@ def build(self, config: DAGConfig) -> "DAG": with dag: for step_name, step_config in config.steps.items(): - rendered[step_name] = self._render_step(step_name, step_config) + rendered[step_name] = self._render_step(step_name, step_config, source) for step_name, step_config in config.steps.items(): for dep_name in step_config.depends_on: @@ -274,7 +280,7 @@ def build_from_yaml( raw_config = yaml.safe_load(raw_content) dag_config = DAGConfig.model_validate(raw_config) - dag = self.build(dag_config) + dag = self.build(dag_config, source=str(yaml_path)) if self._on_dag_built: self._on_dag_built(dag, yaml_path) @@ -359,6 +365,7 @@ def _render_step( self, step_name: str, step_config: StepConfig, + source: str | None = None, ) -> TaskOrGroup: """Render a single step by instantiating its blueprint.""" from pydantic import ValidationError @@ -411,15 +418,14 @@ def _render_step( result = instance.render(validated_config) - step_yaml = yaml.dump( - { - "blueprint": step_config.blueprint, - "version": resolved_version, - **blueprint_config, - }, - default_flow_style=False, - sort_keys=False, - ) + step_context: dict[str, Any] = { + "blueprint": step_config.blueprint, + "version": resolved_version, + } + if source is not None: + step_context["source"] = source + step_context.update(blueprint_config) + step_yaml = yaml.dump(step_context, default_flow_style=False, sort_keys=False) source_code = bp_class.get_source_code() @@ -495,6 +501,33 @@ def _check_duplicate_dag_id(dag_id: str, yaml_path: Path, dag_id_to_file: dict[s raise DuplicateDAGIdError(dag_id, [dag_id_to_file[dag_id], yaml_path]) +def _relative_source(yaml_path: Path, search_path: Path) -> str: + """Return yaml_path relative to search_path, or as-is if not beneath it.""" + try: + return str(yaml_path.relative_to(search_path)) + except ValueError: + return str(yaml_path) + + +def _apply_source_tag(dag: "DAG", source: str) -> None: + """Tag a DAG with its source YAML path (``blueprint:``). + + The tag survives DAG serialization, so the Airflow UI plugin can map a + dag_id back to its YAML file without re-scanning the dags folder. Skipped + with a warning if the tag would exceed Airflow's tag length limit. + """ + tag = f"{SOURCE_TAG_PREFIX}{source}" + if len(tag) > _TAG_MAX_LEN: + logger.warning( + "Skipping source tag for DAG '%s': '%s' exceeds %d characters", + dag.dag_id, + tag, + _TAG_MAX_LEN, + ) + return + dag.tags = {*(dag.tags or ()), tag} + + def build_all_airflow_dags( search_path: str | Path | None = None, register_globals: dict | None = None, @@ -503,6 +536,7 @@ def build_all_airflow_dags( template_context: dict[str, Any] | None = None, bp_registry: BlueprintRegistry | None = None, on_dag_built: OnDagBuilt | None = None, + source_tags: bool = True, ) -> list["DAG"]: """Discover and build all DAGs from YAML files. @@ -530,6 +564,9 @@ def build_all_airflow_dags( on_dag_built: Optional callback invoked after each DAG is built. Receives the DAG and the Path to the source YAML file. Use this to apply post-processing such as access controls or tags. + source_tags: Whether to tag each DAG with its source YAML path + (``blueprint:``). The tag lets the + Airflow UI plugin resolve a DAG back to its YAML file. Returns: List of built DAGs @@ -583,7 +620,11 @@ def build_all_airflow_dags( dag_config = DAGConfig.model_validate(raw_config) _check_duplicate_dag_id(dag_config.dag_id, yaml_path, dag_id_to_file) - dag = builder.build(dag_config) + source = _relative_source(yaml_path, resolved_path) + dag = builder.build(dag_config, source=source) + + if source_tags: + _apply_source_tag(dag, source) if on_dag_built: on_dag_built(dag, yaml_path) diff --git a/blueprint/plugin/__init__.py b/blueprint/plugin/__init__.py new file mode 100644 index 0000000..42d11af --- /dev/null +++ b/blueprint/plugin/__init__.py @@ -0,0 +1,98 @@ +"""Airflow UI plugin: the YAML and blueprint code behind each DAG. + +Registered automatically via the ``airflow.plugins`` entry point when +airflow-blueprint is installed. On Airflow 3 it mounts a FastAPI app at +``/blueprint`` under the API server; on Airflow 3.1+ it adds a +"Blueprint" tab at every level of the UI: + +- DAG and DAG-run pages show the source YAML the DAG was built from, + with the blueprint Python underneath +- task and task-instance pages show that step's config and blueprint + source, resolved to what actually ran where possible + +On Airflow 2.x the plugin loads but registers nothing. +""" + +import logging +from typing import Any + +from airflow.plugins_manager import AirflowPlugin + +logger = logging.getLogger(__name__) + +URL_PREFIX = "/blueprint" + + +def _airflow_version() -> tuple[int, int]: + """Return the running Airflow (major, minor), or (0, 0) if unknown.""" + try: + from airflow import __version__ as airflow_version + from packaging.version import Version + + v = Version(airflow_version) + except Exception: + return (0, 0) + return (v.major, v.minor) + + +def _build_surfaces() -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Build the fastapi_apps and external_views lists for this Airflow version.""" + version = _airflow_version() + if version < (3, 0): + return [], [] + + try: + from blueprint.plugin.app import create_app + + app = create_app() + except ImportError as e: + logger.warning("Blueprint UI plugin disabled: %s", e) + return [], [] + + fastapi_apps = [{"app": app, "url_prefix": URL_PREFIX, "name": "Blueprint"}] + + if version < (3, 1): + return fastapi_apps, [] + + external_views = [ + { + "name": "Blueprint", + "href": f"{URL_PREFIX}/dags/{{DAG_ID}}/yaml", + "destination": "dag", + "url_route": "blueprint_dag", + }, + { + "name": "Blueprint", + "href": f"{URL_PREFIX}/dags/{{DAG_ID}}/yaml", + "destination": "dag_run", + "url_route": "blueprint_dag_run", + }, + { + "name": "Blueprint", + "href": f"{URL_PREFIX}/dags/{{DAG_ID}}/tasks/{{TASK_ID}}", + "destination": "task", + "url_route": "blueprint_task", + }, + { + "name": "Blueprint", + "href": ( + f"{URL_PREFIX}/dags/{{DAG_ID}}/tasks/{{TASK_ID}}" + "?run_id={RUN_ID}&map_index={MAP_INDEX}" + ), + "destination": "task_instance", + "url_route": "blueprint_task_instance", + }, + ] + return fastapi_apps, external_views + + +_fastapi_apps, _external_views = _build_surfaces() + + +class BlueprintPlugin(AirflowPlugin): + """Airflow plugin exposing Blueprint's YAML source and catalog views.""" + + name = "blueprint" + + fastapi_apps = _fastapi_apps + external_views = _external_views diff --git a/blueprint/plugin/app.py b/blueprint/plugin/app.py new file mode 100644 index 0000000..b8111d2 --- /dev/null +++ b/blueprint/plugin/app.py @@ -0,0 +1,626 @@ +"""FastAPI app serving the Blueprint UI pages inside Airflow. + +The app is mounted under the Airflow API server (see ``blueprint.plugin``) +and renders the source YAML and blueprint Python behind a DAG, and the +per-step config and code behind a task. + +DAG-to-YAML resolution prefers the ``blueprint:`` tag stamped by +``build_all_airflow_dags`` (one metadata-DB lookup), falling back to a +scan of the dags folder for older builds. +""" + +import logging +import os +import re +import time +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any, TypeVar +from urllib.parse import unquote + +import yaml + +from blueprint.builder import SOURCE_TAG_PREFIX, _relative_source +from blueprint.registry import BlueprintRegistry + +if TYPE_CHECKING: + from fastapi import FastAPI + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +_CACHE_TTL_SECONDS = 30.0 +_cache: dict[str, tuple[float, Any]] = {} + + +def _cached(key: str, factory: Callable[[], T]) -> T: + """Return a cached value, refreshing it after a short TTL.""" + now = time.monotonic() + hit = _cache.get(key) + if hit is not None and hit[0] > now: + return hit[1] + value = factory() + _cache[key] = (now + _CACHE_TTL_SECONDS, value) + return value + + +def clear_cache() -> None: + """Drop all cached scan results (used by tests).""" + _cache.clear() + + +def code_html(code: str, language: str, line_numbers: bool = False) -> str: + """Render code as highlighted HTML, mirroring Airflow's Code tab. + + Falls back to an escaped ``
`` if Pygments is unavailable.
+    """
+    try:
+        from pygments import highlight
+        from pygments.formatters.html import HtmlFormatter
+        from pygments.lexers import get_lexer_by_name
+
+        formatter = HtmlFormatter(
+            cssclass="highlight",
+            linenos="table" if line_numbers else False,
+        )
+        rendered = highlight(code, get_lexer_by_name(language, stripnl=False), formatter)
+    except Exception:
+        import html
+
+        logger.debug("Pygments highlighting failed", exc_info=True)
+        return f'
{html.escape(code)}
' + return f'
{rendered}
' + + +def pygments_stylesheet() -> str: + """Build highlight CSS: light by default, dark under the synced theme attr.""" + + def scoped(fmt: Any, scope: str) -> str: + # Pygments emits line-number rules unprefixed; drop them — the base + # stylesheet styles the gutter to match the surrounding theme. + return "\n".join( + line + for line in fmt.get_style_defs(scope).splitlines() + if "linenos" not in line and not line.startswith("pre ") + ) + + try: + from pygments.formatters.html import HtmlFormatter + + light = scoped(HtmlFormatter(style="xcode"), ".highlight") + dark_fmt = HtmlFormatter(style="github-dark") + dark = scoped(dark_fmt, ':root[data-theme="dark"] .highlight') + dark_media = scoped(dark_fmt, ':root:not([data-theme="light"]) .highlight') + except Exception: + logger.debug("Pygments stylesheet generation failed", exc_info=True) + return "" + return f"{light}\n{dark}\n@media (prefers-color-scheme: dark) {{\n{dark_media}\n}}" + + +def _default_dags_folder() -> Path: + """Resolve the dags folder from Airflow config, with an env fallback.""" + try: + from airflow.configuration import conf + + return Path(conf.get("core", "dags_folder")) + except Exception: + return Path(os.environ.get("AIRFLOW_HOME", ".")) / "dags" + + +def _parse_dag_yaml(path: Path) -> dict[str, Any] | None: + """Parse a DAG YAML file leniently. + + Jinja2 expressions are rendered with undefined variables collapsing to + empty strings, so files that need build-time context still yield their + ``dag_id`` and ``steps`` structure. Returns None if the file cannot be + parsed into a mapping. + """ + try: + raw = path.read_text(encoding="utf-8") + except OSError: + return None + + if "{{" in raw or "{%" in raw: + try: + import jinja2 + + from blueprint.loaders import _ContextProxy, _StubVarAccessor + + env = jinja2.Environment(undefined=jinja2.ChainableUndefined) + raw = env.from_string(raw).render( + env=os.environ, + var=_StubVarAccessor(), + context=_ContextProxy(), + ) + except Exception: + logger.debug("Tolerant Jinja2 render failed for %s", path, exc_info=True) + + try: + config = yaml.safe_load(raw) + except yaml.YAMLError: + return None + return config if isinstance(config, dict) else None + + +def _dag_sources(dags_folder: Path) -> dict[str, Path]: + """Map dag_id to its YAML path by scanning the dags folder.""" + + def scan() -> dict[str, Path]: + from blueprint.loaders import discover_yaml_files + + sources: dict[str, Path] = {} + if not dags_folder.exists(): + return sources + for path in discover_yaml_files(dags_folder, "*.dag.yaml"): + config = _parse_dag_yaml(path) + dag_id = (config or {}).get("dag_id") + if isinstance(dag_id, str): + sources[dag_id] = path + return sources + + return _cached(f"sources:{dags_folder}", scan) + + +def _source_from_tags(dag_id: str) -> str | None: + """Read the ``blueprint:`` tag for a DAG from the metadata DB.""" + try: + from airflow.models.dag import DagModel + from airflow.utils.session import create_session + + with create_session() as session: + dag_model = session.get(DagModel, dag_id) + if dag_model is None: + return None + for tag in dag_model.tags or []: + name = getattr(tag, "name", None) or str(tag) + if name.startswith(SOURCE_TAG_PREFIX): + return name[len(SOURCE_TAG_PREFIX) :] + except Exception: + logger.debug("Source tag lookup failed for %s", dag_id, exc_info=True) + return None + + +def resolve_dag_source(dag_id: str, dags_folder: Path) -> Path | None: + """Find the YAML file a DAG was built from. + + Tries the source tag first, then falls back to scanning the dags folder. + Tag values must resolve to a file inside the dags folder. + """ + rel = _source_from_tags(dag_id) + if rel: + candidate = Path(rel) + if not candidate.is_absolute(): + candidate = dags_folder / candidate + resolved = candidate.resolve() + if resolved.is_file() and resolved.is_relative_to(dags_folder.resolve()): + return resolved + return _dag_sources(dags_folder).get(dag_id) + + +def resolve_task_step( + dag_id: str, task_id: str, dags_folder: Path +) -> tuple[str, dict[str, Any]] | None: + """Find the YAML step that renders a task. + + A step rendering a TaskGroup produces task ids like ``step.child``, so + the step name is the task id itself or its prefix before a dot. + + Returns: + (step_name, step_dict) or None if the DAG or step cannot be found. + """ + path = resolve_dag_source(dag_id, dags_folder) + if path is None: + return None + config = _parse_dag_yaml(path) + steps = (config or {}).get("steps") + if not isinstance(steps, dict): + return None + candidates = [ + name + for name, step in steps.items() + if isinstance(step, dict) and (task_id == name or task_id.startswith(f"{name}.")) + ] + if not candidates: + return None + step_name = max(candidates, key=len) + return step_name, steps[step_name] + + +def _step_context_from_rtif( + dag_id: str, run_id: str, task_id: str, map_index: int +) -> tuple[str | None, str | None]: + """Fetch stamped step fields as rendered for a specific task instance. + + Airflow stores rendered template fields per run, so this is the exact + config (and blueprint source) the task ran with — including runtime + param overrides. Rows are pruned to the most recent runs per task, so + misses are expected. + + Returns: + (blueprint_step_config, blueprint_step_code), each None if absent. + """ + try: + from airflow.models.renderedtifields import RenderedTaskInstanceFields + from airflow.utils.session import create_session + from sqlalchemy import select + + with create_session() as session: + stmt = select(RenderedTaskInstanceFields).filter_by( + dag_id=dag_id, task_id=task_id, run_id=run_id, map_index=map_index + ) + row = session.scalars(stmt).first() + if row is None: + return None, None + fields = row.rendered_fields or {} + return ( + _clean_stamped_value(fields.get("blueprint_step_config")), + _clean_stamped_value(fields.get("blueprint_step_code")), + ) + except Exception: + logger.debug("RTIF lookup failed for %s/%s/%s", dag_id, run_id, task_id, exc_info=True) + return None, None + + +def _clean_stamped_value(value: Any) -> str | None: + """Drop non-strings and values Airflow truncated when storing them. + + Both rendered task fields and serialized DAGs replace template field + values longer than ``[core] max_templated_field_length`` with a + "Truncated. ..." placeholder — blueprint source files often exceed it, + so fall through to the next source instead of showing the placeholder. + """ + if not isinstance(value, str) or value.startswith("Truncated."): + return None + return value + + +def _step_context_from_serialized_dag(dag_id: str, task_id: str) -> tuple[str | None, str | None]: + """Fetch stamped step fields from the serialized DAG in the metadata DB. + + This is the config and blueprint source as built (resolved version + included) for the current DAG version — independent of what the files + on disk say now. + + Returns: + (blueprint_step_config, blueprint_step_code), each None if absent. + """ + try: + from airflow.models.serialized_dag import SerializedDagModel + from airflow.utils.session import create_session + + with create_session() as session: + sdm = SerializedDagModel.get(dag_id, session=session) + if sdm is None: + return None, None + task = sdm.dag.get_task(task_id) + return ( + _clean_stamped_value(getattr(task, "blueprint_step_config", None)), + _clean_stamped_value(getattr(task, "blueprint_step_code", None)), + ) + except Exception: + logger.debug("Serialized DAG lookup failed for %s/%s", dag_id, task_id, exc_info=True) + return None, None + + +def _parse_step_context(step_yaml: str) -> dict[str, Any] | None: + """Parse a stamped ``blueprint_step_config`` YAML string.""" + try: + data = yaml.safe_load(step_yaml) + except yaml.YAMLError: + return None + if isinstance(data, dict) and isinstance(data.get("blueprint"), str): + return data + return None + + +def _load_registry(dags_folder: Path) -> BlueprintRegistry: + """Build (and cache) a registry discovered from the dags folder.""" + + def make() -> BlueprintRegistry: + return BlueprintRegistry(template_dirs=[dags_folder]) + + return _cached(f"registry:{dags_folder}", make) + + +_GROUP_REFERER_RE = re.compile( + r"/dags/[^/]+(?:/runs/(?P[^/]+))?/tasks/group/(?P[^/?#]+)" +) + + +def group_from_referer(referer: str | None) -> tuple[str | None, str | None]: + """Recover (group_id, run_id) from the embedding page's URL. + + Task group pages render task-destination external views but leave the + ``{TASK_ID}`` token unsubstituted (there is no task in context). The + iframe request's Referer still names the group, and a group id equals + its step name. + """ + if not referer: + return None, None + m = _GROUP_REFERER_RE.search(referer) + if m is None: + return None, None + run = m.group("run") + return unquote(m.group("group")), unquote(run) if run else None + + +def dag_code_sections(config: dict[str, Any] | None, dags_folder: Path) -> list[dict[str, Any]]: + """Collect source code for each distinct blueprint a DAG's YAML uses.""" + import inspect + + steps = (config or {}).get("steps") + if not isinstance(steps, dict): + return [] + + reg = _load_registry(dags_folder) + sections: dict[tuple[str, int], dict[str, Any]] = {} + for step in steps.values(): + if not isinstance(step, dict): + continue + name = step.get("blueprint") + version = step.get("version") + if not isinstance(name, str): + continue + try: + cls = reg.get(name, version if isinstance(version, int) else None) + resolved = version if isinstance(version, int) else reg.get_latest_version(name) + key = (name, resolved) + if key in sections: + continue + sections[key] = { + "name": name, + "version": resolved, + "class": cls.__name__, + "location": _relative_source(Path(inspect.getfile(cls)), dags_folder), + "source_code": cls.get_source_code(), + } + except Exception: + logger.debug("Could not resolve blueprint %s for code section", name, exc_info=True) + return [sections[key] for key in sorted(sections)] + + +def _normalize_task_ref( + task_id: str, + run_id: str | None, + map_index: str | None, + referer: str | None, +) -> tuple[str, str | None, int]: + """Resolve unsubstituted UI tokens into a usable (task_id, run_id, map_index).""" + if "{" in task_id: + group_id, ref_run = group_from_referer(referer) + if group_id: + task_id = group_id + if not run_id or "{" in run_id: + run_id = ref_run + if run_id and "{" in run_id: + run_id = None + try: + mi = int(map_index) if map_index else -1 + except ValueError: + mi = -1 + return task_id, run_id, mi + + +def _stamped_step_fields( + dag_id: str, + task_id: str, + run_id: str | None, + map_index: int, +) -> tuple[str | None, str | None, str | None, str | None, str | None]: + """Fetch stamped step config and code, each with its provenance. + + Config and code fall back independently from the run's rendered fields + to the serialized DAG — a run's code copy is often truncated (see + ``_clean_rtif_value``) while its config is not. + + Returns: + (config, code, config_provenance, code_provenance, effective_run_id). + """ + stamped, code = None, None + provenance = None + code_provenance = None + if run_id: + stamped, code = _step_context_from_rtif(dag_id, run_id, task_id, map_index) + if stamped is None and " " in run_id: + # The UI substitutes {RUN_ID} without URL-encoding, so a "+" + # in timestamped run ids arrives as a space. + run_id = run_id.replace(" ", "+") + stamped, code = _step_context_from_rtif(dag_id, run_id, task_id, map_index) + if stamped: + provenance = "run" + if code: + code_provenance = "run" + if stamped is None or code is None: + ser_config, ser_code = _step_context_from_serialized_dag(dag_id, task_id) + if stamped is None and ser_config: + stamped = ser_config + provenance = "serialized" + if code is None and ser_code: + code = ser_code + code_provenance = "serialized" + return stamped, code, provenance, code_provenance, run_id + + +def task_step_context( + dag_id: str, + task_id: str, + run_id: str | None, + map_index: str | None, + dags_folder: Path, + referer: str | None = None, +) -> tuple[dict[str, Any], bool]: + """Build the task step page context via the metadata lookup ladder. + + Prefers the rendered fields of the given run (exact as-run config), then + the serialized DAG (as-built config), then the source YAML on disk. + An unsubstituted ``{TASK_ID}`` (task group pages) is recovered from the + Referer header, since a group id equals its step name. + + Returns: + (template context, whether a step was found). + """ + task_id, run_id, mi = _normalize_task_ref(task_id, run_id, map_index, referer) + + context: dict[str, Any] = { + "dag_id": dag_id, + "task_id": task_id, + "step_name": None, + "run_id": None, + "source_code": None, + "dags_folder": str(dags_folder), + } + + stamped, code, provenance, code_provenance, run_id = _stamped_step_fields( + dag_id, task_id, run_id, mi + ) + if provenance == "run": + context["run_id"] = run_id + + match = resolve_task_step(dag_id, task_id, dags_folder) + + data = _parse_step_context(stamped) if stamped else None + if data is not None: + if code is None: + code = _registry_source_code(data["blueprint"], data.get("version"), dags_folder) + code_provenance = "yaml" if code else None + context.update( + { + "step_name": match[0] if match else task_id.split(".", 1)[0], + "blueprint": data["blueprint"], + "version": data.get("version"), + "version_label": None, + "provenance": provenance, + "code_provenance": code_provenance, + "step_yaml": stamped, + "source_code": code, + } + ) + return context, True + + if match is not None: + step_name, step = match + blueprint_name = step.get("blueprint") + version = step.get("version") + version_label = "pinned" if version else "latest" + if isinstance(blueprint_name, str): + context["source_code"] = _registry_source_code(blueprint_name, version, dags_folder) + context["code_provenance"] = "yaml" if context["source_code"] else None + if version is None: + try: + version = _load_registry(dags_folder).get_latest_version(blueprint_name) + except Exception: + logger.debug("Could not resolve latest version", exc_info=True) + context.update( + { + "step_name": step_name, + "blueprint": blueprint_name, + "version": version, + "version_label": version_label, + "provenance": "yaml", + "step_yaml": yaml.dump( + {step_name: step}, default_flow_style=False, sort_keys=False + ), + } + ) + return context, True + + return context, False + + +def _registry_source_code(name: str, version: int | None, dags_folder: Path) -> str | None: + """Read a blueprint's current source file via the registry, fail-soft.""" + try: + return _load_registry(dags_folder).get(name, version).get_source_code() or None + except Exception: + logger.debug("Could not resolve blueprint %s from registry", name, exc_info=True) + return None + + +def create_app(dags_folder: Path | None = None) -> "FastAPI": + """Create the Blueprint UI FastAPI app. + + Args: + dags_folder: Directory holding DAG YAML files and blueprint modules. + Defaults to Airflow's configured dags folder, resolved per request. + + Returns: + The FastAPI app, ready to mount under the Airflow API server. + """ + import jinja2 + from fastapi import FastAPI, Request, Response + from fastapi.responses import HTMLResponse + + templates = jinja2.Environment( + loader=jinja2.FileSystemLoader(Path(__file__).parent / "templates"), + autoescape=True, + ) + + def _folder() -> Path: + return dags_folder if dags_folder is not None else _default_dags_folder() + + app = FastAPI(title="Blueprint", docs_url=None, redoc_url=None, openapi_url=None) + + @app.get("/static/pygments.css") + def pygments_css() -> Response: + return Response( + pygments_stylesheet(), + media_type="text/css", + headers={"Cache-Control": "max-age=3600"}, + ) + + @app.get("/dags/{dag_id}/yaml", response_class=HTMLResponse) + def dag_yaml_page(dag_id: str) -> HTMLResponse: + folder = _folder() + path = resolve_dag_source(dag_id, folder) + content_html = None + source = None + code_sections: list[dict[str, Any]] = [] + if path is not None: + source = _relative_source(path, folder) + try: + content_html = code_html( + path.read_text(encoding="utf-8"), "yaml", line_numbers=True + ) + except OSError: + logger.warning("Could not read %s", path, exc_info=True) + code_sections = dag_code_sections(_parse_dag_yaml(path), folder) + for section in code_sections: + section["source_html"] = code_html( + section["source_code"], "python", line_numbers=True + ) + html = templates.get_template("dag_yaml.html").render( + dag_id=dag_id, + source=source, + content_html=content_html, + code_sections=code_sections, + dags_folder=str(folder), + ) + return HTMLResponse(html, status_code=200 if content_html is not None else 404) + + @app.get("/dags/{dag_id}/tasks/{task_id}", response_class=HTMLResponse) + def task_step_page( + request: Request, + dag_id: str, + task_id: str, + run_id: str | None = None, + map_index: str | None = None, + ) -> HTMLResponse: + context, found = task_step_context( + dag_id, + task_id, + run_id, + map_index, + _folder(), + referer=request.headers.get("referer"), + ) + if context.get("step_yaml"): + context["step_yaml_html"] = code_html(context["step_yaml"], "yaml") + if context.get("source_code"): + context["source_code_html"] = code_html( + context["source_code"], "python", line_numbers=True + ) + html = templates.get_template("task_step.html").render(**context) + return HTMLResponse(html, status_code=200 if found else 404) + + return app diff --git a/blueprint/plugin/templates/base.html b/blueprint/plugin/templates/base.html new file mode 100644 index 0000000..14d5d37 --- /dev/null +++ b/blueprint/plugin/templates/base.html @@ -0,0 +1,202 @@ + + + + + + {% block title %}Blueprint{% endblock %} + + + + + +
+ {% block content %}{% endblock %} +
+ + + diff --git a/blueprint/plugin/templates/dag_yaml.html b/blueprint/plugin/templates/dag_yaml.html new file mode 100644 index 0000000..de50e86 --- /dev/null +++ b/blueprint/plugin/templates/dag_yaml.html @@ -0,0 +1,29 @@ +{% extends "base.html" %} +{% block title %}{{ dag_id }} — YAML{% endblock %} +{% block content %} +

{{ dag_id }}

+ +{% if content_html is not none %} +
+
{{ source }}
+ {{ content_html | safe }} +
+ +{% if code_sections %} +

Blueprint code

+

The Python that renders each step of this DAG.

+{% for bp in code_sections %} +
+ {{ bp.location }} + · {{ bp.name }} v{{ bp.version }} ({{ bp.class }}) + {{ bp.source_html | safe }} +
+{% endfor %} +{% endif %} +{% else %} +
No source YAML found — this DAG wasn't built by Blueprint.
+

Looked for a blueprint:<path> tag and scanned +{{ dags_folder }} for a matching +*.dag.yaml file.

+{% endif %} +{% endblock %} diff --git a/blueprint/plugin/templates/task_step.html b/blueprint/plugin/templates/task_step.html new file mode 100644 index 0000000..af64295 --- /dev/null +++ b/blueprint/plugin/templates/task_step.html @@ -0,0 +1,41 @@ +{% extends "base.html" %} +{% block title %}{{ task_id }} — Blueprint Step{% endblock %} +{% block content %} +

{{ task_id }}

+ +{% if step_name %} +

+ Rendered by blueprint {{ blueprint }} + {% if version %}v{{ version }}{% if version_label %} · {{ version_label }}{% endif %} + {% else %}latest{% endif %} + — step {{ step_name }} in + {{ dag_id }} +

+{{ step_yaml_html | safe }} +{% if provenance == "run" %} +

Config as rendered for run {{ run_id }} — +what this task actually ran with, param overrides included.

+{% elif provenance == "serialized" %} +

Config as built into the current serialized DAG. Open from a task +instance to see per-run values.

+{% else %} +

Step definition from the source YAML on disk; the DAG may have been +built from an earlier version of this file.

+{% endif %} + +{% if source_code %} +
+ {{ blueprint }} source + + {%- if code_provenance == "run" %}· as run + {%- elif code_provenance == "serialized" %}· as built + {%- else %}· current file{% endif %} + {{ source_code_html | safe }} +
+{% endif %} +{% else %} +
No blueprint step found for task {{ task_id }} — +it isn't rendered by Blueprint.
+{% endif %} +{% endblock %} diff --git a/examples/advanced/README.md b/examples/advanced/README.md index 5645df1..c745139 100644 --- a/examples/advanced/README.md +++ b/examples/advanced/README.md @@ -52,6 +52,21 @@ Custom `BlueprintDagArgs` subclass that converts a `priority` field into a DAG t semantics as Airflow's DAG processor. Here `dags/.airflowignore` lists the `drafts` directory, so `drafts/lunar_relay.dag.yaml` never builds. +### Airflow UI Plugin (Airflow 3 only) + +The `airflow-blueprint` package registers an Airflow plugin automatically. In +the Airflow 3 example you get a **Blueprint** tab (Airflow 3.1+) on every DAG, +DAG run, task, and task instance page: + +- DAG and DAG run: the source YAML the DAG was built from (resolved via the + `blueprint:` tag that `build_all_airflow_dags()` stamps on every DAG), + with the blueprint Python underneath. +- Task and task instance: that step's config and blueprint source, showing the + exact as-run values where Airflow stored them. + +The pages are served by a FastAPI app mounted at `/blueprint` under the API +server (Airflow 3.0+), e.g. http://localhost:8080/blueprint/dags/satellite_telemetry/yaml. + ### Programmatic Building (`dags/programmatic_dags.py`) Builds DAGs in a loop with the `Builder` API instead of YAML. One telemetry DAG diff --git a/pyproject.toml b/pyproject.toml index 558e2be..5c4c367 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "click>=8.0.0", "pydantic>=2.10.6", "rich>=13.9.4", + "pygments>=2.15.0", ] keywords = ["airflow", "dags", "templates", "blueprints", "automation"] classifiers = [ @@ -36,6 +37,9 @@ classifiers = [ [project.scripts] blueprint = "blueprint.cli:cli" +[project.entry-points."airflow.plugins"] +blueprint = "blueprint.plugin:BlueprintPlugin" + [project.urls] "Homepage" = "https://github.com/astronomer/blueprint" "Repository" = "https://github.com/astronomer/blueprint" diff --git a/tests/integration/test_dag_structure.py b/tests/integration/test_dag_structure.py index 2b3a261..a578a36 100644 --- a/tests/integration/test_dag_structure.py +++ b/tests/integration/test_dag_structure.py @@ -141,7 +141,12 @@ def test_dag_has_tags(self, api_client: AirflowAPI): resp = api_client.get("/dags/versioned_etl") assert resp.status_code == 200 tags = api_client.get_tags(resp.json()) - assert tags == {"team:data-eng", "critical", "callback-verified"} + assert tags == { + "team:data-eng", + "critical", + "callback-verified", + "blueprint:versioned.dag.yaml", + } def test_critical_tier_retries(self, api_client: AirflowAPI): tasks = self._get_tasks(api_client) @@ -223,7 +228,12 @@ def test_dag_has_default_team_tag(self, api_client: AirflowAPI): resp = api_client.get("/dags/explicit_naming") assert resp.status_code == 200 tags = api_client.get_tags(resp.json()) - assert tags == {"team:platform", "standard", "callback-verified"} + assert tags == { + "team:platform", + "standard", + "callback-verified", + "blueprint:explicit_naming.dag.yaml", + } class TestDagArgsRendering: @@ -260,7 +270,12 @@ def test_tier_tag_generated(self, api_client: AirflowAPI): def test_only_derived_tags_present(self, api_client: AirflowAPI): dag = self._get_dag(api_client) tags = api_client.get_tags(dag) - assert tags == {"team:analytics", "critical", "callback-verified"} + assert tags == { + "team:analytics", + "critical", + "callback-verified", + "blueprint:dag_args_test.dag.yaml", + } def test_owner_derived_from_team(self, api_client: AirflowAPI): tasks = self._get_tasks(api_client) diff --git a/tests/integration/test_plugin.py b/tests/integration/test_plugin.py new file mode 100644 index 0000000..48e35c3 --- /dev/null +++ b/tests/integration/test_plugin.py @@ -0,0 +1,96 @@ +"""Integration tests for the Blueprint UI plugin against a live Airflow. + +Verifies the plugin registers with the API server, the mounted FastAPI app +serves the YAML and blueprint-code pages, and DAGs carry the +``blueprint:`` source tag. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from .conftest import AirflowAPI + +pytestmark = pytest.mark.integration + + +class TestPluginRegistration: + def test_plugin_listed_by_api(self, api_client: AirflowAPI): + resp = api_client.get("/plugins") + assert resp.status_code == 200, resp.text + plugins = {p["name"]: p for p in resp.json()["plugins"]} + assert "blueprint" in plugins + + def test_external_views_registered(self, api_client: AirflowAPI): + resp = api_client.get("/plugins") + assert resp.status_code == 200, resp.text + (plugin,) = [p for p in resp.json()["plugins"] if p["name"] == "blueprint"] + views = plugin.get("external_views") or [] + destinations = {v.get("destination") for v in views} + assert destinations == {"dag", "dag_run", "task", "task_instance"} + assert {v.get("name") for v in views} == {"Blueprint"} + + +class TestDagYamlPage: + def test_shows_source_yaml(self, api_client: AirflowAPI): + resp = api_client.client().get("/blueprint/dags/simple_pipeline/yaml") + assert resp.status_code == 200, resp.text + assert "simple_pipeline" in resp.text + assert "simple.dag.yaml" in resp.text + assert "Blueprint code" in resp.text + assert "TransformConfig" in resp.text + + def test_unknown_dag_returns_404(self, api_client: AirflowAPI): + resp = api_client.client().get("/blueprint/dags/no_such_dag/yaml") + assert resp.status_code == 404 + assert "No source YAML found" in resp.text + + +class TestTaskStepPage: + def test_task_group_child_resolves_step(self, api_client: AirflowAPI): + resp = api_client.client().get("/blueprint/dags/simple_pipeline/tasks/process.clean") + assert resp.status_code == 200, resp.text + assert "transform" in resp.text + assert "process" in resp.text + + def test_serialized_dag_provenance_without_run(self, api_client: AirflowAPI): + resp = api_client.client().get("/blueprint/dags/simple_pipeline/tasks/process.clean") + assert resp.status_code == 200, resp.text + assert "current serialized DAG" in resp.text + assert "TransformConfig" in resp.text + # The blueprint file exceeds [core] max_templated_field_length, so the + # serialized code copy is truncated and the code falls back to disk. + assert "· current file" in resp.text + + def test_run_provenance_with_run_id(self, api_client: AirflowAPI): + from .test_dag_execution import _trigger_dag, _unpause_dag, _wait_for_dag_run + + dag_id = "simple_pipeline" + _unpause_dag(api_client, dag_id) + run_id = _trigger_dag(api_client, dag_id) + result = _wait_for_dag_run(api_client, dag_id, run_id) + assert result["state"] == "success", result + + resp = api_client.client().get( + f"/blueprint/dags/{dag_id}/tasks/process.clean", + params={"run_id": run_id, "map_index": "-1"}, + ) + assert resp.status_code == 200, resp.text + assert "as rendered for run" in resp.text + assert run_id in resp.text + + def test_unknown_task_returns_404(self, api_client: AirflowAPI): + resp = api_client.client().get("/blueprint/dags/simple_pipeline/tasks/nope") + assert resp.status_code == 404 + assert "No blueprint step found" in resp.text + + +class TestSourceTags: + def test_dag_carries_source_tag(self, api_client: AirflowAPI): + resp = api_client.get("/dags/simple_pipeline") + assert resp.status_code == 200, resp.text + tags = api_client.get_tags(resp.json()) + assert "blueprint:simple.dag.yaml" in tags diff --git a/tests/test_builder.py b/tests/test_builder.py index 835149a..c6970b3 100644 --- a/tests/test_builder.py +++ b/tests/test_builder.py @@ -1081,6 +1081,78 @@ def render(self, config): ) +class TestSourceMetadata: + STUB_BLUEPRINT = """ +from pydantic import BaseModel +from blueprint.core import Blueprint + +class StubConfig(BaseModel): + x: int = 1 + +class Stub(Blueprint[StubConfig]): + def render(self, config): + from airflow.operators.bash import BashOperator + return BashOperator(task_id=self.step_id, bash_command="echo ok") +""" + + def _build(self, tmp_path, **kwargs): + from blueprint.builder import build_all_airflow_dags + + (tmp_path / "blueprints.py").write_text(self.STUB_BLUEPRINT) + nested = tmp_path / "team" + nested.mkdir() + (nested / "meta.dag.yaml").write_text( + "dag_id: meta_test\nsteps:\n s1:\n blueprint: stub\n" + ) + return build_all_airflow_dags( + search_path=tmp_path, + register_globals={}, + render_templates=False, + **kwargs, + ) + + def test_source_tag_added(self, tmp_path): + (dag,) = self._build(tmp_path) + assert "blueprint:team/meta.dag.yaml" in dag.tags + + def test_source_tag_opt_out(self, tmp_path): + (dag,) = self._build(tmp_path, source_tags=False) + assert not [t for t in (dag.tags or []) if t.startswith("blueprint:")] + + def test_source_in_step_config(self, tmp_path): + (dag,) = self._build(tmp_path) + task = next(iter(dag.task_dict.values())) + parsed = yaml.safe_load(task.blueprint_step_config) + assert parsed["source"] == "team/meta.dag.yaml" + assert parsed["blueprint"] == "stub" + + def test_source_tag_skipped_when_too_long(self, tmp_path): + from blueprint.builder import build_all_airflow_dags + + (tmp_path / "blueprints.py").write_text(self.STUB_BLUEPRINT) + deep = tmp_path / ("d" * 120) + deep.mkdir() + (deep / "long.dag.yaml").write_text( + "dag_id: long_tag_test\nsteps:\n s1:\n blueprint: stub\n" + ) + (dag,) = build_all_airflow_dags( + search_path=tmp_path, + register_globals={}, + render_templates=False, + ) + assert not [t for t in (dag.tags or []) if t.startswith("blueprint:")] + + def test_build_without_source_omits_key(self, builder): + config = DAGConfig( + dag_id="no_source", + steps={"s1": StepConfig(blueprint="load", target_table="out")}, + ) + dag = builder.build(config) + task = next(iter(dag.task_dict.values())) + parsed = yaml.safe_load(task.blueprint_step_config) + assert "source" not in parsed + + class TestOnDagBuilt: def test_build_all_on_dag_built_called(self, tmp_path): from blueprint.builder import build_all_airflow_dags diff --git a/tests/test_plugin.py b/tests/test_plugin.py new file mode 100644 index 0000000..166e407 --- /dev/null +++ b/tests/test_plugin.py @@ -0,0 +1,324 @@ +"""Tests for the Airflow UI plugin: data helpers, FastAPI views, registration.""" + +import pytest + +pytest.importorskip("fastapi") + +from fastapi.testclient import TestClient + +from blueprint.plugin.app import ( + clear_cache, + create_app, + resolve_dag_source, + resolve_task_step, +) + +BLUEPRINTS_PY = ''' +from pydantic import BaseModel +from blueprint.core import Blueprint + +class ExtractConfig(BaseModel): + source_table: str + batch_size: int = 1000 + +class Extract(Blueprint[ExtractConfig]): + """Extract data from a source.""" + + def render(self, config): + from airflow.operators.bash import BashOperator + return BashOperator(task_id=self.step_id, bash_command="echo extract") + +class ExtractV2Config(BaseModel): + sources: list[str] + +class ExtractV2(Blueprint[ExtractV2Config]): + """Extract v2 with multi-source.""" + + def render(self, config): + from airflow.operators.bash import BashOperator + return BashOperator(task_id=self.step_id, bash_command="echo extract2") + +class LoadConfig(BaseModel): + target_table: str + +class Load(Blueprint[LoadConfig]): + """Load data to a target.""" + + def render(self, config): + from airflow.operators.bash import BashOperator + return BashOperator(task_id=self.step_id, bash_command="echo load") +''' + +PIPELINE_YAML = """\ +dag_id: plugin_pipeline +steps: + pull: + blueprint: extract + version: 1 + source_table: raw.events + push: + blueprint: load + depends_on: [pull] + target_table: warehouse.events +""" + +JINJA_YAML = """\ +dag_id: plugin_jinja +description: "{{ team_name }}" +steps: + pull: + blueprint: extract + version: 2 + sources: ["a"] +""" + + +@pytest.fixture +def dags_folder(tmp_path): + clear_cache() + (tmp_path / "blueprints.py").write_text(BLUEPRINTS_PY) + (tmp_path / "pipeline.dag.yaml").write_text(PIPELINE_YAML) + (tmp_path / "templated.dag.yaml").write_text(JINJA_YAML) + drafts = tmp_path / "drafts" + drafts.mkdir() + (drafts / "wip.dag.yaml").write_text("dag_id: plugin_wip\nsteps:\n s:\n blueprint: load\n") + (tmp_path / ".airflowignore").write_text("drafts\n") + yield tmp_path + clear_cache() + + +@pytest.fixture +def client(dags_folder): + return TestClient(create_app(dags_folder=dags_folder)) + + +class TestDagSourceResolution: + def test_resolves_by_scanning(self, dags_folder): + path = resolve_dag_source("plugin_pipeline", dags_folder) + assert path == dags_folder / "pipeline.dag.yaml" + + def test_resolves_templated_yaml(self, dags_folder): + path = resolve_dag_source("plugin_jinja", dags_folder) + assert path == dags_folder / "templated.dag.yaml" + + def test_unknown_dag_returns_none(self, dags_folder): + assert resolve_dag_source("nope", dags_folder) is None + + def test_airflowignore_respected(self, dags_folder): + assert resolve_dag_source("plugin_wip", dags_folder) is None + + +class TestStampedValueCleaning: + def test_truncated_value_dropped(self): + from blueprint.plugin.app import _clean_stamped_value + + truncated = ( + "Truncated. You can change this behaviour in " + '[core]max_templated_field_length. \'"""Blueprint definitions...\'' + ) + assert _clean_stamped_value(truncated) is None + + def test_normal_value_kept(self): + from blueprint.plugin.app import _clean_stamped_value + + assert _clean_stamped_value("blueprint: extract\n") == "blueprint: extract\n" + + def test_non_string_dropped(self): + from blueprint.plugin.app import _clean_stamped_value + + assert _clean_stamped_value(None) is None + assert _clean_stamped_value(123) is None + + +class TestTaskStepResolution: + def test_exact_task_id(self, dags_folder): + match = resolve_task_step("plugin_pipeline", "pull", dags_folder) + assert match is not None + step_name, step = match + assert step_name == "pull" + assert step["blueprint"] == "extract" + + def test_task_group_child(self, dags_folder): + match = resolve_task_step("plugin_pipeline", "pull.download", dags_folder) + assert match is not None + assert match[0] == "pull" + + def test_unknown_task(self, dags_folder): + assert resolve_task_step("plugin_pipeline", "nope", dags_folder) is None + + def test_unknown_dag(self, dags_folder): + assert resolve_task_step("nope", "pull", dags_folder) is None + + +class TestTaskStepPage: + def test_pinned_step(self, client): + resp = client.get("/dags/plugin_pipeline/tasks/pull") + assert resp.status_code == 200 + assert "extract" in resp.text + assert "v1 · pinned" in resp.text + assert "raw.events" in resp.text + assert "source YAML on disk" in resp.text + assert "ExtractConfig" in resp.text + assert "· current file" in resp.text + + def test_unpinned_step_resolves_latest(self, client): + resp = client.get("/dags/plugin_pipeline/tasks/push") + assert resp.status_code == 200 + assert "load" in resp.text + assert "v1 · latest" in resp.text + + def test_unknown_task_404(self, client): + resp = client.get("/dags/plugin_pipeline/tasks/nope") + assert resp.status_code == 404 + assert "No blueprint step found" in resp.text + + def test_group_page_recovers_step_from_referer(self, client): + resp = client.get( + "/dags/plugin_pipeline/tasks/{TASK_ID}", + headers={"referer": "http://localhost:8080/dags/plugin_pipeline/tasks/group/pull"}, + ) + assert resp.status_code == 200 + assert "extract" in resp.text + assert "v1 · pinned" in resp.text + + def test_group_page_referer_with_run(self, client, monkeypatch): + import blueprint.plugin.app as app_module + + seen = {} + + def fake_rtif(_dag_id, run_id, task_id, _map_index): + seen.update(run_id=run_id, task_id=task_id) + return None, None + + monkeypatch.setattr(app_module, "_step_context_from_rtif", fake_rtif) + resp = client.get( + "/dags/plugin_pipeline/tasks/{TASK_ID}", + params={"run_id": "{RUN_ID}"}, + headers={ + "referer": ( + "http://localhost:8080/dags/plugin_pipeline" + "/runs/manual__2026-07-24T00%3A00%3A00%2B00%3A00/tasks/group/pull" + ) + }, + ) + assert resp.status_code == 200 + assert seen == {"run_id": "manual__2026-07-24T00:00:00+00:00", "task_id": "pull"} + + def test_unsubstituted_token_without_referer_404(self, client): + resp = client.get("/dags/plugin_pipeline/tasks/{TASK_ID}") + assert resp.status_code == 404 + + def test_unsubstituted_tokens_ignored(self, client): + resp = client.get( + "/dags/plugin_pipeline/tasks/pull", + params={"run_id": "{RUN_ID}", "map_index": "{MAP_INDEX}"}, + ) + assert resp.status_code == 200 + assert "source YAML on disk" in resp.text + + def test_serialized_context_preferred_over_yaml(self, client, monkeypatch): + import blueprint.plugin.app as app_module + + stamped = "blueprint: extract\nversion: 2\nsource: pipeline.dag.yaml\nsources:\n- a\n" + monkeypatch.setattr( + app_module, + "_step_context_from_serialized_dag", + lambda *_args: (stamped, "# serialized copy\nclass ExtractV2: ..."), + ) + resp = client.get("/dags/plugin_pipeline/tasks/pull") + assert resp.status_code == 200 + assert "v2" in resp.text + assert "current serialized DAG" in resp.text + assert "# serialized copy" in resp.text + assert "· as built" in resp.text + + def test_run_context_preferred_when_run_id_given(self, client, monkeypatch): + import blueprint.plugin.app as app_module + + calls = {} + + def fake_rtif(_dag_id, run_id, _task_id, map_index): + calls.update(run_id=run_id, map_index=map_index) + return "blueprint: extract\nversion: 1\nsource_table: overridden.events\n", None + + monkeypatch.setattr(app_module, "_step_context_from_rtif", fake_rtif) + resp = client.get( + "/dags/plugin_pipeline/tasks/pull", + params={"run_id": "manual__2026-07-24", "map_index": "-1"}, + ) + assert resp.status_code == 200 + assert calls == {"run_id": "manual__2026-07-24", "map_index": -1} + assert "overridden.events" in resp.text + assert "as rendered for run" in resp.text + + def test_run_id_plus_recovered_from_space(self, client, monkeypatch): + import blueprint.plugin.app as app_module + + seen = [] + + def fake_rtif(_dag_id, run_id, _task_id, _map_index): + seen.append(run_id) + return ("blueprint: extract\nversion: 1\n", None) if "+" in run_id else (None, None) + + monkeypatch.setattr(app_module, "_step_context_from_rtif", fake_rtif) + resp = client.get( + "/dags/plugin_pipeline/tasks/pull", + params={"run_id": "scheduled__2026-07-24T00:00:00 00:00"}, + ) + assert resp.status_code == 200 + assert seen == [ + "scheduled__2026-07-24T00:00:00 00:00", + "scheduled__2026-07-24T00:00:00+00:00", + ] + assert "as rendered for run" in resp.text + + +class TestDagYamlPage: + def test_shows_yaml(self, client): + resp = client.get("/dags/plugin_pipeline/yaml") + assert resp.status_code == 200 + assert "plugin_pipeline" in resp.text + assert "warehouse.events" in resp.text + assert "pipeline.dag.yaml" in resp.text + + def test_shows_blueprint_code_sections(self, client): + resp = client.get("/dags/plugin_pipeline/yaml") + assert resp.status_code == 200 + assert "Blueprint code" in resp.text + assert "ExtractConfig" in resp.text + assert "LoadConfig" in resp.text + + def test_unknown_dag_404(self, client): + resp = client.get("/dags/nope/yaml") + assert resp.status_code == 404 + assert "No source YAML found" in resp.text + + +class TestPluginRegistration: + def test_surfaces_on_airflow_31(self, monkeypatch): + import blueprint.plugin as plugin_module + + monkeypatch.setattr(plugin_module, "_airflow_version", lambda: (3, 1)) + fastapi_apps, external_views = plugin_module._build_surfaces() + assert fastapi_apps[0]["url_prefix"] == "/blueprint" + destinations = {v["destination"] for v in external_views} + assert destinations == {"dag", "dag_run", "task", "task_instance"} + (dag_view,) = [v for v in external_views if v["destination"] == "dag"] + assert "{DAG_ID}" in dag_view["href"] + assert {v["name"] for v in external_views} == {"Blueprint"} + (task_view,) = [v for v in external_views if v["destination"] == "task"] + assert "{TASK_ID}" in task_view["href"] + + def test_no_external_views_on_airflow_30(self, monkeypatch): + import blueprint.plugin as plugin_module + + monkeypatch.setattr(plugin_module, "_airflow_version", lambda: (3, 0)) + fastapi_apps, external_views = plugin_module._build_surfaces() + assert fastapi_apps + assert external_views == [] + + def test_disabled_on_airflow_2(self, monkeypatch): + import blueprint.plugin as plugin_module + + monkeypatch.setattr(plugin_module, "_airflow_version", lambda: (2, 10)) + assert plugin_module._build_surfaces() == ([], []) diff --git a/uv.lock b/uv.lock index bcba841..29030e5 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.13'", @@ -38,6 +38,7 @@ dependencies = [ { name = "apache-airflow" }, { name = "click" }, { name = "pydantic" }, + { name = "pygments" }, { name = "pyyaml" }, { name = "rich" }, ] @@ -60,6 +61,7 @@ requires-dist = [ { name = "apache-airflow", specifier = ">=2.5.0" }, { name = "click", specifier = ">=8.0.0" }, { name = "pydantic", specifier = ">=2.10.6" }, + { name = "pygments", specifier = ">=2.15.0" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "rich", specifier = ">=13.9.4" }, ]