From 75b38b8803b99291aa58f89dc599a25f4598d632 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 27 Jul 2026 12:51:45 +0200 Subject: [PATCH 1/4] Support installation via entry-points --- README.md | 71 +++++ blueprint/builder.py | 14 +- blueprint/cli.py | 85 +++++- blueprint/loaders.py | 49 ++- blueprint/registry.py | 171 +++++++++-- blueprint/utils.py | 6 +- examples/README.md | 2 +- examples/advanced/airflow2/Dockerfile | 8 +- examples/advanced/airflow2/Tiltfile | 6 + examples/advanced/airflow3/Dockerfile | 8 +- examples/advanced/airflow3/Tiltfile | 6 + examples/advanced/dags/shared_blueprints.yaml | 8 + examples/shared-blueprints/README.md | 30 ++ examples/shared-blueprints/pyproject.toml | 20 ++ .../shared_blueprints/__init__.py | 6 + .../shared_blueprints/example.py | 28 ++ pyproject.toml | 6 + .../entry_point_test_blueprints/__init__.py | 6 + .../entrypoint_bp_test.py | 19 ++ tests/entry_point_package/pyproject.toml | 21 ++ tests/integration/conftest.py | 4 +- .../project/dags/entry_point_test.dag.yaml | 5 + tests/integration/test_cli.py | 23 +- .../integration/test_entry_point_discovery.py | 72 +++++ tests/test_cli.py | 2 +- tests/test_errors.py | 13 + tests/test_loaders.py | 2 +- tests/test_registry.py | 287 ++++++++++++++++-- tests/test_utils.py | 18 ++ uv.lock | 15 +- 30 files changed, 935 insertions(+), 76 deletions(-) create mode 100644 examples/advanced/dags/shared_blueprints.yaml create mode 100644 examples/shared-blueprints/README.md create mode 100644 examples/shared-blueprints/pyproject.toml create mode 100644 examples/shared-blueprints/shared_blueprints/__init__.py create mode 100644 examples/shared-blueprints/shared_blueprints/example.py create mode 100644 tests/entry_point_package/entry_point_test_blueprints/__init__.py create mode 100644 tests/entry_point_package/entry_point_test_blueprints/entrypoint_bp_test.py create mode 100644 tests/entry_point_package/pyproject.toml create mode 100644 tests/integration/project/dags/entry_point_test.dag.yaml create mode 100644 tests/integration/test_entry_point_discovery.py diff --git a/README.md b/README.md index 8135774..b2f8a0e 100644 --- a/README.md +++ b/README.md @@ -241,6 +241,77 @@ steps: table: orders ``` +## Sharing Blueprints Across Teams + +In larger organizations, it's common that one "data engineering" team implements Blueprint +templates and multiple other "dag authoring" teams leverage those templates. The data +engineering team can publish those templates using a pip-installable package. It's a bad practice +to copy-paste template code across repositories since that quickly goes out of sync. + +Publishing Blueprint templates in a shared package requires declaring an entry point under the +`airflow_blueprint.blueprints` group in your package's `pyproject.toml`: + +```toml +# pyproject.toml +[project.entry-points."airflow_blueprint.blueprints"] +company_blueprints = "company_blueprints" +``` + +To leverage the Blueprint templates, install the package: + +```bash +pip install company-blueprints # or add it to requirements.txt +``` + +The Blueprint templates from the shared package become discoverable with the Blueprint CLI when the package is installed in your Python environment: + +```bash +blueprint list +``` + +A few things worth knowing: + +**Collisions give an error** + +Two templates with the same name and version will raise a `DuplicateBlueprintError`. This also happens when templates are stored in different locations (such as locally and in package). + +**Keep the `entry-point` target scoped to code defining Blueprint templates** + +Every submodule under the `entry-point` gets scanned for Blueprint templates on every DAG parsing cycle. This could include unnecessary code. If your project contains other folders with non-Blueprint code, for example: + +``` +my_project/ +├── pyproject.toml +└── my_project/ + ├── __init__.py + ├── utils/ # Utility code, not Blueprint templates + │ └── ... + ├── operators/ # Custom Airflow operators, not Blueprint templates + │ └── ... + └── blueprints/ # <-- Only this contains Blueprint templates + ├── __init__.py + ├── extract.py + └── load.py +``` + +A top-level `entry-point` will look like so: + +```toml +# pyproject.toml +[project.entry-points."airflow_blueprint.blueprints"] +my_project = "my_project" +``` + +And limiting the `entry-point` to a subfolder is done like so: + +```toml +# pyproject.toml +[project.entry-points."airflow_blueprint.blueprints"] +my_project = "my_project.blueprints" +``` + +Note that this applies to the discoverability of Blueprint templates. A template can still import from another module that's not included in the `entry-point`. + ## Airflow Rendered Templates Every task instance gets two extra fields visible in Airflow's "Rendered Template" tab: diff --git a/blueprint/builder.py b/blueprint/builder.py index 4037fb6..58790bd 100644 --- a/blueprint/builder.py +++ b/blueprint/builder.py @@ -503,6 +503,7 @@ def build_all_airflow_dags( template_context: dict[str, Any] | None = None, bp_registry: BlueprintRegistry | None = None, on_dag_built: OnDagBuilt | None = None, + discover_entry_points: bool = True, ) -> list["DAG"]: """Discover and build all DAGs from YAML files. @@ -530,6 +531,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. + discover_entry_points: Whether to also discover blueprints from installed packages + advertising themselves via the ``airflow_blueprint.blueprints`` entry-point + group. Ignored when ``bp_registry`` is supplied directly. Returns: List of built DAGs @@ -553,7 +557,11 @@ def build_all_airflow_dags( if bp_registry is None: caller_file = _get_caller_file() exclude = {Path(caller_file)} if caller_file else set() - bp_registry = BlueprintRegistry(template_dirs=[resolved_path], exclude_files=exclude) + bp_registry = BlueprintRegistry( + template_dirs=[resolved_path], + exclude_files=exclude, + discover_entry_points=discover_entry_points, + ) bp_registry.discover(force=True) builder = Builder(bp_registry=bp_registry) @@ -611,6 +619,7 @@ def build_all_dags( template_context: dict[str, Any] | None = None, bp_registry: BlueprintRegistry | None = None, on_dag_built: OnDagBuilt | None = None, + discover_entry_points: bool = True, ) -> list["DAG"]: """Deprecated alias for ``build_all_airflow_dags``. @@ -641,6 +650,7 @@ def build_all_dags( template_context=template_context, bp_registry=bp_registry, on_dag_built=on_dag_built, + discover_entry_points=discover_entry_points, ) @@ -652,6 +662,7 @@ def build_all( template_context: dict[str, Any] | None = None, bp_registry: BlueprintRegistry | None = None, on_dag_built: OnDagBuilt | None = None, + discover_entry_points: bool = True, ) -> list["DAG"]: """Deprecated alias for ``build_all_airflow_dags``.""" warnings.warn( @@ -674,6 +685,7 @@ def build_all( template_context=template_context, bp_registry=bp_registry, on_dag_built=on_dag_built, + discover_entry_points=discover_entry_points, ) diff --git a/blueprint/cli.py b/blueprint/cli.py index 01bd7dd..544d24c 100644 --- a/blueprint/cli.py +++ b/blueprint/cli.py @@ -47,14 +47,24 @@ def _get_configs_to_check(path: str | None) -> list[Path]: return discover_yaml_files(Path(), "*.dag.yaml") -def _validate_config(config_path: Path, template_dir: str | None) -> tuple[bool, str | None]: +def _validate_config( + config_path: Path, template_dir: str | None, discover_entry_points: bool = True +) -> tuple[bool, str | None]: """Validate a single configuration file. + Args: + config_path: Path to the .dag.yaml file to validate. + template_dir: Directory containing blueprint files. + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points. + Returns: tuple of (success, dag_id) """ try: - result = validate_yaml(str(config_path), template_dir=template_dir) + result = validate_yaml( + str(config_path), template_dir=template_dir, discover_entry_points=discover_entry_points + ) except Exception as e: console.print(f"[red]FAIL[/red] {config_path}") if hasattr(e, "_format_message") and callable(e._format_message): @@ -89,7 +99,12 @@ def _check_duplicate_dag_ids(dag_ids_to_files: dict[str, list[Path]]) -> bool: @cli.command() @click.argument("path", required=False, type=click.Path(exists=True)) @click.option("--template-dir", default=None, help="Directory containing blueprint files") -def lint(path: str | None, template_dir: str | None): +@click.option( + "--entry-points/--no-entry-points", + default=True, + help="Discover blueprints from installed packages via entry points.", +) +def lint(path: str | None, template_dir: str | None, entry_points: bool): """Validate DAG YAML definitions. If PATH is provided, validate a specific file. @@ -107,7 +122,9 @@ def lint(path: str | None, template_dir: str | None): valid_count = 0 for config_path in configs_to_check: - success, dag_id = _validate_config(config_path, template_dir) + success, dag_id = _validate_config( + config_path, template_dir, discover_entry_points=entry_points + ) if success and dag_id: if dag_id in dag_ids_to_files: @@ -127,9 +144,14 @@ def lint(path: str | None, template_dir: str | None): @cli.command("list") @click.option("--template-dir", default=None, help="Directory containing blueprint files") -def list_blueprints(template_dir: str | None): +@click.option( + "--entry-points/--no-entry-points", + default=True, + help="Discover blueprints from installed packages via entry points.", +) +def list_blueprints(template_dir: str | None, entry_points: bool): """List available blueprints.""" - blueprints = discover_blueprints(template_dir) + blueprints = discover_blueprints(template_dir, discover_entry_points=entry_points) if not blueprints: console.print("[yellow]No blueprints found.[/yellow]") @@ -157,10 +179,19 @@ def list_blueprints(template_dir: str | None): @click.argument("blueprint_name") @click.option("--version", "-v", type=int, default=None, help="Specific version (default: latest)") @click.option("--template-dir", default=None, help="Directory containing blueprint files") -def describe(blueprint_name: str, version: int | None, template_dir: str | None): +@click.option( + "--entry-points/--no-entry-points", + default=True, + help="Discover blueprints from installed packages via entry points.", +) +def describe( + blueprint_name: str, version: int | None, template_dir: str | None, entry_points: bool +): """Show blueprint parameters and documentation.""" try: - info = get_blueprint_info(blueprint_name, template_dir, version=version) + info = get_blueprint_info( + blueprint_name, template_dir, version=version, discover_entry_points=entry_points + ) except Exception as e: console.print(f"[red]Error:[/red] {e}") sys.exit(1) @@ -213,11 +244,22 @@ def describe(blueprint_name: str, version: int | None, template_dir: str | None) console.print(yaml_syntax) -def _get_registry(template_dir: str | None) -> BlueprintRegistry: - """Get a BlueprintRegistry for the given template directory.""" +def _get_registry( + template_dir: str | None, discover_entry_points: bool = True +) -> BlueprintRegistry: + """Get a BlueprintRegistry for the given template directory. + + Args: + template_dir: Directory containing blueprint files. + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points. + + Returns: + A BlueprintRegistry with discovery already run. + """ from blueprint.loaders import get_registry - return get_registry(template_dir) + return get_registry(template_dir, discover_entry_points=discover_entry_points) def _get_trigger_rule_values() -> list[str]: @@ -304,11 +346,17 @@ def _build_dag_yaml_schema(dag_args_schema: dict) -> dict: @click.option("--dag-args", "dag_args", is_flag=True, help="Emit schema for DAG-level YAML fields") @click.option("--output", "-o", type=click.Path(), help="Output file (default: stdout)") @click.option("--template-dir", default=None, help="Directory containing blueprint files") +@click.option( + "--entry-points/--no-entry-points", + default=True, + help="Discover blueprints from installed packages via entry points.", +) def schema( blueprint_name: str | None, dag_args: bool, output: str | None, template_dir: str | None, + entry_points: bool, ): """Generate JSON Schema for a blueprint's configuration. @@ -327,7 +375,7 @@ def schema( sys.exit(1) try: - reg = _get_registry(template_dir) + reg = _get_registry(template_dir, discover_entry_points=entry_points) except Exception as e: console.print(f"[red]Error:[/red] {e}") sys.exit(1) @@ -446,9 +494,14 @@ def _collect_parameters(info: dict[str, Any]) -> dict[str, object]: @cli.command() @click.option("--template-dir", default=None, help="Directory containing blueprint files") @click.option("--output-dir", default=".", help="Output directory for YAML config") -def new(template_dir: str | None, output_dir: str): +@click.option( + "--entry-points/--no-entry-points", + default=True, + help="Discover blueprints from installed packages via entry points.", +) +def new(template_dir: str | None, output_dir: str, entry_points: bool): """Interactively create a new DAG YAML definition.""" - blueprints = discover_blueprints(template_dir) + blueprints = discover_blueprints(template_dir, discover_entry_points=entry_points) if not blueprints: console.print("[red]No blueprints found.[/red]") @@ -457,14 +510,14 @@ def new(template_dir: str | None, output_dir: str): selected = _select_blueprint(blueprints) console.print(f"\n[green]Selected:[/green] {selected['name']}") - info = get_blueprint_info(selected["name"], template_dir) + info = get_blueprint_info(selected["name"], template_dir, discover_entry_points=entry_points) dag_id = console.input("\nDAG ID: ") if not dag_id: console.print("[red]DAG ID is required[/red]") sys.exit(1) - reg = _get_registry(template_dir) + reg = _get_registry(template_dir, discover_entry_points=entry_points) dag_args_cls = reg.get_dag_args() dag_args_schema = dag_args_cls.get_schema() dag_args_params = dag_args_schema.get("properties", {}) diff --git a/blueprint/loaders.py b/blueprint/loaders.py index 19510b6..3e31a41 100644 --- a/blueprint/loaders.py +++ b/blueprint/loaders.py @@ -184,6 +184,7 @@ def load_blueprint( blueprint_name: str, template_dir: str | None = None, version: int | None = None, + discover_entry_points: bool = True, ) -> type[Blueprint]: """Load a blueprint class by name and optional version. @@ -191,24 +192,31 @@ def load_blueprint( blueprint_name: Name of the blueprint (e.g., 'extract') template_dir: Directory containing blueprint files version: Specific version (None for latest) + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points Returns: The Blueprint class """ - reg = get_registry(template_dir) + reg = get_registry(template_dir, discover_entry_points=discover_entry_points) return reg.get(blueprint_name, version) -def discover_blueprints(template_dir: str | None = None) -> list[dict[str, Any]]: +def discover_blueprints( + template_dir: str | None = None, + discover_entry_points: bool = True, +) -> list[dict[str, Any]]: """Discover all available blueprints. Args: template_dir: Directory containing blueprint files + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points Returns: List of blueprint information dictionaries """ - reg = get_registry(template_dir) + reg = get_registry(template_dir, discover_entry_points=discover_entry_points) return reg.list_blueprints() @@ -241,6 +249,7 @@ def get_blueprint_info( blueprint_name: str, template_dir: str | None = None, version: int | None = None, + discover_entry_points: bool = True, ) -> dict[str, Any]: """Get detailed information about a specific blueprint. @@ -248,23 +257,28 @@ def get_blueprint_info( blueprint_name: Name of the blueprint template_dir: Directory containing blueprint files version: Specific version (None for latest) + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points Returns: Dictionary with blueprint information including schema """ - reg = get_registry(template_dir) + reg = get_registry(template_dir, discover_entry_points=discover_entry_points) return reg.get_blueprint_info(blueprint_name, version) def validate_yaml( path: str, template_dir: str | None = None, + discover_entry_points: bool = True, ) -> dict[str, Any]: """Validate a DAG YAML file without building the DAG. Args: path: Path to the .dag.yaml file template_dir: Directory containing blueprint files + discover_entry_points: Whether to also discover blueprints from installed packages via + entry points Returns: The parsed and validated DAGConfig as a dict @@ -276,7 +290,7 @@ def validate_yaml( dag_config = DAGConfig.model_validate(config) - reg = get_registry(template_dir) + reg = get_registry(template_dir, discover_entry_points=discover_entry_points) builder = Builder(bp_registry=reg) builder.validate_dependencies(dag_config) @@ -294,10 +308,27 @@ def validate_yaml( return dag_config.model_dump() -def get_registry(template_dir: str | None = None) -> BlueprintRegistry: - """Get or create a BlueprintRegistry for the given template directory.""" - if template_dir: - temp_registry = BlueprintRegistry(template_dirs=[Path(template_dir)]) +def get_registry( + template_dir: str | None = None, + discover_entry_points: bool = True, +) -> BlueprintRegistry: + """Get or create a BlueprintRegistry for the given template directory. + + Args: + template_dir: Directory containing blueprint files. If given, a fresh registry + scoped to that directory is built and returned; otherwise the module-level + singleton registry is used. + discover_entry_points: Whether to also discover blueprints from installed packages + via entry points. + + Returns: + A BlueprintRegistry with discovery already run. + """ + if template_dir or not discover_entry_points: + temp_registry = BlueprintRegistry( + template_dirs=[Path(template_dir)] if template_dir else None, + discover_entry_points=discover_entry_points, + ) temp_registry.discover(force=True) return temp_registry diff --git a/blueprint/registry.py b/blueprint/registry.py index a4e65db..5c5f482 100644 --- a/blueprint/registry.py +++ b/blueprint/registry.py @@ -1,11 +1,16 @@ """Global registry for Blueprint discovery and management with version tracking.""" import ast +import importlib.metadata import importlib.util +import inspect import logging import os +import pkgutil import sys +from collections.abc import Iterator from pathlib import Path +from types import ModuleType from typing import Any from blueprint.core import Blueprint, BlueprintDagArgs, DefaultDagArgs @@ -21,6 +26,18 @@ _BLUEPRINT_BASE_NAMES = frozenset({"Blueprint", "BlueprintDagArgs"}) +# Shared-blueprint packages advertise themselves under this entry-point group so +# BlueprintRegistry can discover them once installed, with no per-repo config: +# +# [project.entry-points."airflow_blueprint.blueprints"] +# company_blueprints = "company_blueprints" # noqa: ERA001 +# +# The value must be a plain dotted module/package path (no "module:attr" syntax). +# The module (and, if it's a package, every submodule) is scanned the same way +# a locally discovered file is: any Blueprint/BlueprintDagArgs subclass defined +# directly in it gets registered. +_ENTRY_POINT_GROUP = "airflow_blueprint.blueprints" + def _defines_blueprint_subclass(py_file: Path) -> bool: """Return True if the file's source defines a Blueprint or BlueprintDagArgs subclass. @@ -64,6 +81,7 @@ def __init__( self, template_dirs: list[Path] | None = None, exclude_files: set[Path] | None = None, + discover_entry_points: bool = True, ) -> None: self._blueprints: dict[str, dict[int, type[Blueprint]]] = {} self._blueprint_locations: dict[str, dict[int, str]] = {} @@ -73,6 +91,7 @@ def __init__( self._discovery_in_progress = False self._template_dirs = template_dirs self._exclude_files = {p.resolve() for p in exclude_files} if exclude_files else set() + self._discover_entry_points = discover_entry_points def get_template_dirs(self) -> list[Path]: """Get all template directories to search.""" @@ -110,6 +129,7 @@ def discover(self, force: bool = False) -> None: self._discovery_in_progress = True try: + self._discover_from_entry_points() for template_dir in self.get_template_dirs(): if template_dir.exists(): self._discover_in_directory(template_dir) @@ -141,34 +161,131 @@ def _discover_in_directory(self, directory: Path) -> None: module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) - - for name in dir(module): - obj = getattr(module, name) - if ( - isinstance(obj, type) - and issubclass(obj, Blueprint) - and obj is not Blueprint - and obj.__module__ == module_name - ): - self._register_class(obj, py_file) - elif ( - isinstance(obj, type) - and issubclass(obj, BlueprintDagArgs) - and obj not in (BlueprintDagArgs, DefaultDagArgs) - and obj.__module__ == module_name - ): - self._register_dag_args(obj, py_file) + self._register_module_classes(module, str(py_file.resolve())) except (DuplicateBlueprintError, MultipleDagArgsError, ValueError): raise except (ImportError, SyntaxError) as e: logger.warning("Failed to load %s: %s", py_file, e) - def _register_class(self, cls: type[Blueprint], py_file: Path) -> None: - """Register a blueprint class with its parsed name and version.""" - bp_name, version = cls.parse_name_and_version() + def _discover_from_entry_points(self) -> None: + """Discover blueprints from installed packages via entry points. + + Packages advertise themselves under the ``airflow_blueprint.blueprints`` entry-point group; + this imports each advertised module (and, recursively, every submodule of a package) and + registers any Blueprint/BlueprintDagArgs subclasses defined directly in it. Load failures + for an individual entry point or submodule are logged and skipped rather than raised: a + broken shared package must not take down every DAG in the deployment, only the (normal, + scoped) BlueprintNotFoundError for DAGs that actually reference its blueprints. + """ + if not self._discover_entry_points: + return + + seen_modules: set[str] = set() + + for ep in importlib.metadata.entry_points(group=_ENTRY_POINT_GROUP): + dist_name = ep.dist.name if ep.dist else "unknown package" + try: + loaded = ep.load() + except Exception as e: + logger.warning( + "Failed to load entry point '%s' (%s) from %s: %s", + ep.name, + ep.value, + dist_name, + e, + ) + continue - location = str(py_file.resolve()) + if not inspect.ismodule(loaded): + logger.warning( + "Entry point '%s' (%s) from %s does not resolve to a module; skipping", + ep.name, + ep.value, + dist_name, + ) + continue + + for module in self._iter_entry_point_modules(loaded): + if module.__name__ in seen_modules: + continue + seen_modules.add(module.__name__) + self._register_module_classes(module, module.__name__) + + def _iter_entry_point_modules(self, module: ModuleType) -> Iterator[ModuleType]: + """Yield module and, if it is a package, every importable submodule. + + A submodule that fails to import is logged and skipped rather than raised, so one + broken submodule does not prevent its siblings (or a broken subpackage's siblings) + from being scanned. + + Args: + module: The top-level module or package resolved from an entry point. + + Yields: + The module itself, then each importable submodule in turn. + """ + yield module + + module_path = getattr(module, "__path__", None) + if module_path is None: + return + + def _on_error(name: str) -> None: + logger.warning("Failed to import submodule '%s' of '%s'", name, module.__name__) + + for _finder, name, _is_pkg in pkgutil.walk_packages( + module_path, prefix=f"{module.__name__}.", onerror=_on_error + ): + if name.rsplit(".", 1)[-1].startswith("_"): + continue + try: + submodule = importlib.import_module(name) + except Exception as e: + logger.warning("Failed to import submodule '%s': %s", name, e) + continue + yield submodule + + def _register_module_classes(self, module: ModuleType, location: str) -> None: + """Register any Blueprint/BlueprintDagArgs subclasses defined directly in module. + + Args: + module: The already-imported module to scan. + location: Human-readable source of this module (a resolved file path for + directory-scanned modules, a dotted module name for entry-point-discovered + ones), recorded for error messages and `blueprint list` output. + """ + module_name = module.__name__ + for name in dir(module): + obj = getattr(module, name) + if ( + isinstance(obj, type) + and issubclass(obj, Blueprint) + and obj is not Blueprint + and obj.__module__ == module_name + ): + self._register_class(obj, location) + elif ( + isinstance(obj, type) + and issubclass(obj, BlueprintDagArgs) + and obj not in (BlueprintDagArgs, DefaultDagArgs) + and obj.__module__ == module_name + ): + self._register_dag_args(obj, location) + + def _register_class(self, cls: type[Blueprint], location: str) -> None: + """Register a blueprint class with its parsed name and version. + + Args: + cls: The Blueprint subclass to register. + location: Human-readable source of this class, recorded for error messages + and `blueprint list` output. + + Raises: + DuplicateBlueprintError: If a class with the same name and version is + already registered. + """ + bp_name, version = cls.parse_name_and_version() if bp_name not in self._blueprints: self._blueprints[bp_name] = {} @@ -182,10 +299,16 @@ def _register_class(self, cls: type[Blueprint], py_file: Path) -> None: self._blueprints[bp_name][version] = cls self._blueprint_locations[bp_name][version] = location - def _register_dag_args(self, cls: type[BlueprintDagArgs], py_file: Path) -> None: - """Register the single BlueprintDagArgs template, tracking its location.""" - location = str(py_file.resolve()) + def _register_dag_args(self, cls: type[BlueprintDagArgs], location: str) -> None: + """Register the single BlueprintDagArgs template, tracking its location. + Args: + cls: The BlueprintDagArgs subclass to register. + location: Human-readable source of this class, recorded for error messages. + + Raises: + MultipleDagArgsError: If a BlueprintDagArgs template is already registered. + """ if self._dag_args is not None: raise MultipleDagArgsError([self._dag_args_location or "", location]) diff --git a/blueprint/utils.py b/blueprint/utils.py index 810adc0..c0b881f 100644 --- a/blueprint/utils.py +++ b/blueprint/utils.py @@ -16,10 +16,14 @@ def display_path(path: str | Path, base: Path | None = None) -> str: Returns: The path relative to ``base`` when it is below it, otherwise the - absolute path. A falsy ``path`` is returned unchanged. + absolute path. A falsy ``path`` is returned unchanged, as is a + non-absolute ``path`` (e.g. a dotted module name for a blueprint + discovered from an installed package rather than a file). """ if not path: return str(path) + if not Path(path).is_absolute(): + return str(path) base = (base or Path.cwd()).resolve() resolved = Path(path).resolve() try: diff --git a/examples/README.md b/examples/README.md index d9bb1bc..a3c1e80 100644 --- a/examples/README.md +++ b/examples/README.md @@ -11,7 +11,7 @@ tilt up ## [Advanced](advanced/) -Space-themed example demonstrating many Blueprint features: versioning, runtime params, custom DAG args, Jinja2 templating, `on_dag_built` callbacks, programmatic DAG building with the `Builder` API, and more. +Space-themed example demonstrating many Blueprint features: versioning, runtime params, custom DAG args, Jinja2 templating, `on_dag_built` callbacks, programmatic DAG building with the `Builder` API, blueprints shared via an installed package ([shared-blueprints/](shared-blueprints/)), and more. ```bash cd examples/advanced/airflow3 diff --git a/examples/advanced/airflow2/Dockerfile b/examples/advanced/airflow2/Dockerfile index 24038c1..f2b6209 100644 --- a/examples/advanced/airflow2/Dockerfile +++ b/examples/advanced/airflow2/Dockerfile @@ -8,11 +8,17 @@ COPY ./examples/advanced/requirements.txt ${AIRFLOW_HOME}/requirements.txt COPY ./pyproject.toml ${AIRFLOW_HOME}/blueprint/pyproject.toml COPY ./README.md ${AIRFLOW_HOME}/blueprint/README.md COPY ./blueprint ${AIRFLOW_HOME}/blueprint/blueprint +COPY ./examples/shared-blueprints ${AIRFLOW_HOME}/shared-blueprints RUN uv pip install --system -e "${AIRFLOW_HOME}/blueprint" +# A central team's blueprint package, installed like any other dependency. +# Its blueprints are discovered automatically via the entry point it +# declares in shared-blueprints/pyproject.toml -- no config needed here. +RUN uv pip install --system -e "${AIRFLOW_HOME}/shared-blueprints" + RUN uv pip install --system -r "${AIRFLOW_HOME}/requirements.txt" -RUN chown -R astro:astro ${AIRFLOW_HOME}/blueprint +RUN chown -R astro:astro ${AIRFLOW_HOME}/blueprint ${AIRFLOW_HOME}/shared-blueprints USER astro diff --git a/examples/advanced/airflow2/Tiltfile b/examples/advanced/airflow2/Tiltfile index f991573..71f5c93 100644 --- a/examples/advanced/airflow2/Tiltfile +++ b/examples/advanced/airflow2/Tiltfile @@ -9,6 +9,7 @@ dc_resource('triggerer', resource_deps=['db-init']) sync_pyproj_toml = sync('../../../pyproject.toml', '/usr/local/airflow/blueprint/pyproject.toml') sync_readme = sync('../../../README.md', '/usr/local/airflow/blueprint/README.md') sync_src = sync('../../../blueprint', '/usr/local/airflow/blueprint/blueprint') +sync_shared_blueprints = sync('../../shared-blueprints', '/usr/local/airflow/shared-blueprints') docker_build( 'blueprint-advanced-airflow2', @@ -19,9 +20,14 @@ docker_build( sync_pyproj_toml, sync_src, sync_readme, + sync_shared_blueprints, run( 'cd /usr/local/airflow/blueprint && uv pip install -e .', trigger=['pyproject.toml'] ), + run( + 'cd /usr/local/airflow/shared-blueprints && uv pip install -e .', + trigger=['../../shared-blueprints/pyproject.toml'] + ), ] ) diff --git a/examples/advanced/airflow3/Dockerfile b/examples/advanced/airflow3/Dockerfile index 91fe5b1..b92c029 100644 --- a/examples/advanced/airflow3/Dockerfile +++ b/examples/advanced/airflow3/Dockerfile @@ -8,11 +8,17 @@ COPY ./examples/advanced/requirements.txt ${AIRFLOW_HOME}/requirements.txt COPY ./pyproject.toml ${AIRFLOW_HOME}/blueprint/pyproject.toml COPY ./README.md ${AIRFLOW_HOME}/blueprint/README.md COPY ./blueprint ${AIRFLOW_HOME}/blueprint/blueprint +COPY ./examples/shared-blueprints ${AIRFLOW_HOME}/shared-blueprints RUN uv pip install --system -e "${AIRFLOW_HOME}/blueprint" +# A central team's blueprint package, installed like any other dependency. +# Its blueprints are discovered automatically via the entry point it +# declares in shared-blueprints/pyproject.toml -- no config needed here. +RUN uv pip install --system -e "${AIRFLOW_HOME}/shared-blueprints" + RUN uv pip install --system -r "${AIRFLOW_HOME}/requirements.txt" -RUN chown -R astro:astro ${AIRFLOW_HOME}/blueprint +RUN chown -R astro:astro ${AIRFLOW_HOME}/blueprint ${AIRFLOW_HOME}/shared-blueprints USER astro diff --git a/examples/advanced/airflow3/Tiltfile b/examples/advanced/airflow3/Tiltfile index f00ef72..050d96f 100644 --- a/examples/advanced/airflow3/Tiltfile +++ b/examples/advanced/airflow3/Tiltfile @@ -10,6 +10,7 @@ dc_resource('triggerer', resource_deps=['db-migration']) sync_pyproj_toml = sync('../../../pyproject.toml', '/usr/local/airflow/blueprint/pyproject.toml') sync_readme = sync('../../../README.md', '/usr/local/airflow/blueprint/README.md') sync_src = sync('../../../blueprint', '/usr/local/airflow/blueprint/blueprint') +sync_shared_blueprints = sync('../../shared-blueprints', '/usr/local/airflow/shared-blueprints') docker_build( 'blueprint-advanced-airflow3', @@ -20,9 +21,14 @@ docker_build( sync_pyproj_toml, sync_src, sync_readme, + sync_shared_blueprints, run( 'cd /usr/local/airflow/blueprint && uv pip install -e .', trigger=['pyproject.toml'] ), + run( + 'cd /usr/local/airflow/shared-blueprints && uv pip install -e .', + trigger=['../../shared-blueprints/pyproject.toml'] + ), ] ) diff --git a/examples/advanced/dags/shared_blueprints.yaml b/examples/advanced/dags/shared_blueprints.yaml new file mode 100644 index 0000000..8229c2a --- /dev/null +++ b/examples/advanced/dags/shared_blueprints.yaml @@ -0,0 +1,8 @@ +dag_id: shared_blueprints_example +schedule: "@daily" + +steps: + example: + blueprint: example + foo: "Hello" + bar: "World" diff --git a/examples/shared-blueprints/README.md b/examples/shared-blueprints/README.md new file mode 100644 index 0000000..e0e9042 --- /dev/null +++ b/examples/shared-blueprints/README.md @@ -0,0 +1,30 @@ +# Shared Blueprints Example Package + +This folder contains an installable package demonstrating how to distribute Blueprint templates via a Python package. This avoids having to copy-paste code to multiple repositories that want to leverage the same Blueprint templates. + +Each repository using this package must still define `.dag.yaml` templates and a `loader` file. + +## How it works + +`pyproject.toml` declares an entry point: + +```toml +[project.entry-points."airflow_blueprint.blueprints"] +shared_blueprints = "shared_blueprints" +``` + +Projects that install this `shared-blueprints` package will then be able to use the additional templates. The `BlueprintRegistry` auto-discovers templates from the `shared-blueprints` package. Listing the templates using `blueprint list` also displays the templates from the package: + +```bash +$ blueprint list + Available Blueprints +┏━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ +┃ Name ┃ Versions ┃ Description ┃ Class ┃ Location ┃ +┡━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ +│ example │ 1 │ Example Blueprint template from shared package. │ Example │ shared_blueprints.example │ +├─────────┼──────────┼─────────────────────────────────────────────────┼─────────┼───────────────────────────┤ +│ extract │ 1 │ Pull data from a source system. │ Extract │ dags/blueprints.py │ +├─────────┼──────────┼─────────────────────────────────────────────────┼─────────┼───────────────────────────┤ +│ load │ 1 │ Load data into a destination. │ Load │ dags/blueprints.py │ +└─────────┴──────────┴─────────────────────────────────────────────────┴─────────┴───────────────────────────┘ +``` diff --git a/examples/shared-blueprints/pyproject.toml b/examples/shared-blueprints/pyproject.toml new file mode 100644 index 0000000..ce00554 --- /dev/null +++ b/examples/shared-blueprints/pyproject.toml @@ -0,0 +1,20 @@ +[project] +name = "shared-blueprints" +version = "0.1.0" +description = "Example central-team blueprint package, discovered via entry points" +requires-python = ">=3.10" +dependencies = [ + "airflow-blueprint", +] + +# Advertising under this group ("airflow_blueprint.blueprints") is what lets BlueprintRegistry +# discover this package's blueprints automatically once installed. +[project.entry-points."airflow_blueprint.blueprints"] +shared_blueprints = "shared_blueprints" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["shared_blueprints"] diff --git a/examples/shared-blueprints/shared_blueprints/__init__.py b/examples/shared-blueprints/shared_blueprints/__init__.py new file mode 100644 index 0000000..b019783 --- /dev/null +++ b/examples/shared-blueprints/shared_blueprints/__init__.py @@ -0,0 +1,6 @@ +"""Example central-team blueprint package. + +Demonstrates a shared-blueprints package discovered by a downstream project +purely by being installed -- see examples/advanced, which uses the +``example`` blueprint defined here without any local .py file. +""" diff --git a/examples/shared-blueprints/shared_blueprints/example.py b/examples/shared-blueprints/shared_blueprints/example.py new file mode 100644 index 0000000..3ba3255 --- /dev/null +++ b/examples/shared-blueprints/shared_blueprints/example.py @@ -0,0 +1,28 @@ +"""Example Blueprint template.""" + +try: + # Airflow 3 + from airflow.providers.standard.operators.bash import BashOperator + from airflow.sdk import TaskGroup +except ImportError: + # Airflow 2 + from airflow.operators.bash import BashOperator + from airflow.utils.task_group import TaskGroup + +from blueprint import BaseModel, Blueprint, TaskOrGroup + + +class ExampleConfig(BaseModel): + foo: str + bar: str + + +class Example(Blueprint[ExampleConfig]): + """Example Blueprint template from shared package.""" + + def render(self, config: ExampleConfig) -> TaskOrGroup: + with TaskGroup(group_id=self.step_id) as group: + foo = BashOperator(task_id="foo", bash_command=f"echo {config.foo}") + bar = BashOperator(task_id="bar", bash_command=f"echo {config.bar}") + foo >> bar + return group diff --git a/pyproject.toml b/pyproject.toml index 558e2be..550cc64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,8 +73,14 @@ dev = [ "ty>=0.0.1a15", "playwright>=1.58.0", "httpx>=0.28.0", + "entry-point-test-blueprints", ] +# Local package installed into the dev venv so entry-point discovery is exercised +# against a real installed distribution instead of only mocks. +[tool.uv.sources] +entry-point-test-blueprints = { path = "tests/entry_point_package", editable = true } + [tool.ruff] target-version = "py310" line-length = 100 diff --git a/tests/entry_point_package/entry_point_test_blueprints/__init__.py b/tests/entry_point_package/entry_point_test_blueprints/__init__.py new file mode 100644 index 0000000..ec0258e --- /dev/null +++ b/tests/entry_point_package/entry_point_test_blueprints/__init__.py @@ -0,0 +1,6 @@ +"""Test-only blueprint package, installed and discovered via its entry point. + +Exists solely so the integration suite can prove that BlueprintRegistry discovers +blueprints from an installed package end to end, without depending on anything +under ``examples/``. +""" diff --git a/tests/entry_point_package/entry_point_test_blueprints/entrypoint_bp_test.py b/tests/entry_point_package/entry_point_test_blueprints/entrypoint_bp_test.py new file mode 100644 index 0000000..4cfab58 --- /dev/null +++ b/tests/entry_point_package/entry_point_test_blueprints/entrypoint_bp_test.py @@ -0,0 +1,19 @@ +"""A blueprint for testing the entry-point packaging mechanism.""" + +from blueprint import BaseModel, Blueprint, TaskOrGroup + + +class EntryPointBpTestConfig(BaseModel): + message: str + + +class EntryPointBpTest(Blueprint[EntryPointBpTestConfig]): + """Test-only blueprint shipped via an installed package entry point.""" + + def render(self, config: EntryPointBpTestConfig) -> TaskOrGroup: + try: + from airflow.providers.standard.operators.bash import BashOperator + except ImportError: + from airflow.operators.bash import BashOperator + + return BashOperator(task_id=self.step_id, bash_command=f"echo {config.message}") diff --git a/tests/entry_point_package/pyproject.toml b/tests/entry_point_package/pyproject.toml new file mode 100644 index 0000000..eb96623 --- /dev/null +++ b/tests/entry_point_package/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "entry-point-test-blueprints" +version = "0.0.0" +description = "Test-only package that exercises entry-point blueprint discovery." +requires-python = ">=3.10" +dependencies = [ + "airflow-blueprint", +] + +# This is what makes the package's blueprints discoverable once installed, and is +# the whole point of the fixture: BlueprintRegistry reads this group at discovery +# time. See tests/integration/test_entry_point_discovery.py. +[project.entry-points."airflow_blueprint.blueprints"] +entry_point_test_blueprints = "entry_point_test_blueprints" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["entry_point_test_blueprints"] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index da895dc..cfbe699 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -20,6 +20,7 @@ INTEGRATION_DIR = Path(__file__).parent PROJECT_DIR = INTEGRATION_DIR / "project" REPO_ROOT = Path(__file__).resolve().parents[2] +ENTRY_POINT_PACKAGE_DIR = REPO_ROOT / "tests" / "entry_point_package" HEALTH_CHECK_TIMEOUT = 120 HEALTH_CHECK_INTERVAL = 2 @@ -33,6 +34,7 @@ "explicit_naming", "params_test", "context_test", + "entry_point_test", } @@ -186,7 +188,7 @@ def airflow_env(): port = str(_find_free_port()) req_file = PROJECT_DIR / "requirements.txt" - req_file.write_text(f"-e {REPO_ROOT}\n") + req_file.write_text(f"-e {REPO_ROOT}\n-e {ENTRY_POINT_PACKAGE_DIR}\n") _run_astro("dev", "kill", "--standalone", check=False) result = _run_astro( diff --git a/tests/integration/project/dags/entry_point_test.dag.yaml b/tests/integration/project/dags/entry_point_test.dag.yaml new file mode 100644 index 0000000..43515d0 --- /dev/null +++ b/tests/integration/project/dags/entry_point_test.dag.yaml @@ -0,0 +1,5 @@ +dag_id: entry_point_test +steps: + example: + blueprint: entry_point_bp_test + message: "Hello from test package" diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py index 16e62c2..5785bae 100644 --- a/tests/integration/test_cli.py +++ b/tests/integration/test_cli.py @@ -6,6 +6,7 @@ from __future__ import annotations +import os import subprocess import pytest @@ -17,13 +18,20 @@ DAGS_DIR = str(INTEGRATION_DIR / "project" / "dags") -def _run_blueprint(*args: str) -> subprocess.CompletedProcess: +def _run_blueprint(*args: str, columns: int | None = None) -> subprocess.CompletedProcess: """Run a blueprint CLI command against the test project's dags.""" + env = None + if columns is not None: + env = { + **os.environ, + "COLUMNS": str(columns), + } # Could avoid columns if we'd have JSON output return subprocess.run( ["uv", "run", "blueprint", *args], capture_output=True, text=True, check=False, + env=env, ) @@ -40,6 +48,19 @@ def test_shows_versions(self): assert "1" in result.stdout assert "2" in result.stdout + def test_lists_entry_point_sourced_blueprint_with_dotted_location(self): + """Verify if the entry-point blueprints are discovered.""" + result = _run_blueprint("list", "--template-dir", DAGS_DIR, columns=200) + assert result.returncode == 0, f"blueprint list failed:\n{result.stderr}" + assert "entry_point_bp_test" in result.stdout.lower() + assert "entry_point_test_blueprints.entrypoint_bp_test" in result.stdout + assert DAGS_DIR not in result.stdout + + def test_no_entry_points_flag_hides_installed_package_blueprint(self): + result = _run_blueprint("list", "--template-dir", DAGS_DIR, "--no-entry-points") + assert result.returncode == 0, f"blueprint list failed:\n{result.stderr}" + assert "entry_point_test_blueprints.entrypoint_bp_test" not in result.stdout + class TestDescribe: def test_describe_extract(self): diff --git a/tests/integration/test_entry_point_discovery.py b/tests/integration/test_entry_point_discovery.py new file mode 100644 index 0000000..97d2455 --- /dev/null +++ b/tests/integration/test_entry_point_discovery.py @@ -0,0 +1,72 @@ +"""Tier 1: blueprint discovery from an installed package via entry points. + +The project's requirements.txt (written by the `airflow_env` fixture) installs +`tests/entry_point_package` alongside this repo. That package declares its `EntryPointBpTest` +blueprint under the `airflow_blueprint.blueprints` entry-point group and ships no YAML of its own +-- `dags/entry_point_test.dag.yaml` references it by name with no corresponding local .py file, +so the DAG only parses if BlueprintRegistry's entry-point discovery is actually working end to end +against a real installed package (not just a mocked one). +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +import pytest + +from .conftest import DAG_PARSE_TIMEOUT, HEALTH_CHECK_INTERVAL, PROJECT_DIR + +if TYPE_CHECKING: + from .conftest import AirflowAPI + +pytestmark = pytest.mark.integration + +ENTRY_POINT_DAG_ID = "entry_point_test" + + +class TestEntryPointDiscovery: + """Verify a blueprint from an installed package (no local .py) is discovered.""" + + def test_entry_point_sourced_dag_parses(self, api_client: AirflowAPI): + deadline = time.monotonic() + DAG_PARSE_TIMEOUT + dag_ids: set[str] = set() + while time.monotonic() < deadline: + dag_ids = api_client.get_dag_ids() + if ENTRY_POINT_DAG_ID in dag_ids: + break + time.sleep(HEALTH_CHECK_INTERVAL) + + assert ENTRY_POINT_DAG_ID in dag_ids, ( + f"Airflow did not discover '{ENTRY_POINT_DAG_ID}'. Its blueprint " + "(entry_point_bp_test) exists only in the installed test package " + "-- entry-point discovery must be resolving it for this DAG to " + "parse at all." + ) + + def test_no_import_errors_for_entry_point_dag(self, api_client: AirflowAPI): + resp = api_client.get("/importErrors") + assert resp.status_code == 200, resp.text + offending = [ + e + for e in resp.json().get("import_errors", []) + if "entry_point_test" in (e.get("filename") or "") + ] + assert not offending, f"Import errors for the entry-point DAG: {offending}" + + def test_no_local_py_file_defines_the_probe_blueprint(self): + """Guard against a future 'fix' that quietly adds a local copy of the blueprint, + which would defeat the entire point of this test suite without any single assertion + above failing. + """ + dags_dir = PROJECT_DIR / "dags" + offenders = [ + py_file + for py_file in dags_dir.rglob("*.py") + if "class EntryPointBpTest(" in py_file.read_text() + ] + assert not offenders, ( + f"Found a local copy of EntryPointBpTest in {offenders} -- the " + "entry-point DAG must resolve its blueprint purely from the " + "installed test package." + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 56e73f6..4d62cce 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -24,7 +24,7 @@ def test_cli_version(self): def test_list_command_empty(self, tmp_path): runner = CliRunner() - result = runner.invoke(cli, ["list", "--template-dir", str(tmp_path)]) + result = runner.invoke(cli, ["list", "--template-dir", str(tmp_path), "--no-entry-points"]) assert result.exit_code == 0 assert "No blueprints found" in result.output diff --git a/tests/test_errors.py b/tests/test_errors.py index 92037dd..478cfee 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -153,6 +153,19 @@ def test_duplicate_locations(self): message = str(error) assert "Duplicate blueprint name 'my_blueprint'" in message + def test_duplicate_location_from_installed_package_shown_unchanged(self): + """A dotted-module location (from entry-point discovery) isn't a real file. + + display_path() must render it as-is rather than mangling it into a + fabricated absolute filesystem path (see blueprint/utils.py). + """ + error = DuplicateBlueprintError( + "my_blueprint", + locations=["templates/etl.py", "company_blueprints.etl"], + ) + message = str(error) + assert "company_blueprints.etl" in message + class TestDuplicateDAGIdError: """Test duplicate DAG ID error.""" diff --git a/tests/test_loaders.py b/tests/test_loaders.py index 9742998..9392952 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -169,7 +169,7 @@ def render(self, config): def test_discover_empty_dir(self, tmp_path): template_dir = tmp_path / "empty" template_dir.mkdir() - assert discover_blueprints(str(template_dir)) == [] + assert discover_blueprints(str(template_dir), discover_entry_points=False) == [] class TestGetBlueprintInfo: diff --git a/tests/test_registry.py b/tests/test_registry.py index 83e58dd..919668f 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -1,5 +1,8 @@ """Tests for the version-aware Blueprint registry.""" +import importlib +import importlib.metadata +import logging from pathlib import Path import pytest @@ -8,6 +11,7 @@ from blueprint.core import Blueprint, DefaultDagArgs from blueprint.errors import ( BlueprintNotFoundError, + DuplicateBlueprintError, InvalidVersionError, MultipleDagArgsError, NonContiguousVersionError, @@ -15,6 +19,44 @@ from blueprint.registry import BlueprintRegistry, _defines_blueprint_subclass +def _blueprint_source(class_name: str, config_name: str = "Config") -> str: + """Generate minimal Blueprint subclass source for entry-point discovery tests.""" + return f""" +from pydantic import BaseModel +from blueprint.core import Blueprint + +class {config_name}(BaseModel): + x: int = 1 + +class {class_name}(Blueprint[{config_name}]): + def render(self, config): + pass +""" + + +class _FakeEntryPoint: + """ + Simulate installed entry points without having to package and installing distributions. + """ + + def __init__(self, name, value, dist_name=None, load_fn=None): + self.name = name + self.value = value + self._dist_name = dist_name + self._load_fn = load_fn or (lambda: importlib.import_module(value)) + + def load(self): + return self._load_fn() + + @property + def dist(self): + if self._dist_name is None: + return None + from types import SimpleNamespace + + return SimpleNamespace(name=self._dist_name) + + class SimpleConfig(BaseModel): name: str @@ -42,7 +84,9 @@ class TestBlueprintRegistry: @pytest.fixture def reg(self): - return BlueprintRegistry() + # Hermetic: directory-scan tests shouldn't pick up whatever this dev + # venv happens to have installed under the entry-point group. + return BlueprintRegistry(discover_entry_points=False) @pytest.fixture def temp_blueprints(self, tmp_path): @@ -225,12 +269,12 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) with pytest.raises(DuplicateBlueprintError, match="dup"): reg.discover(force=True) def test_template_dirs_constructor(self, temp_blueprints): - reg = BlueprintRegistry(template_dirs=[temp_blueprints]) + reg = BlueprintRegistry(template_dirs=[temp_blueprints], discover_entry_points=False) reg.discover(force=True) blueprints = reg.list_blueprints() @@ -239,7 +283,7 @@ def test_template_dirs_constructor(self, temp_blueprints): assert "load" in names def test_template_dirs_constructor_overrides_defaults(self, temp_blueprints): - reg = BlueprintRegistry(template_dirs=[temp_blueprints]) + reg = BlueprintRegistry(template_dirs=[temp_blueprints], discover_entry_points=False) dirs = reg.get_template_dirs() assert dirs == [temp_blueprints] @@ -302,7 +346,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) cls = reg.get("extract") @@ -340,7 +384,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) with pytest.raises(DuplicateBlueprintError, match="extract"): reg.discover(force=True) @@ -367,7 +411,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) with pytest.raises(NonContiguousVersionError, match="Missing versions: 2"): reg.discover(force=True) @@ -389,7 +433,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) with pytest.raises(NonContiguousVersionError, match="extract"): reg.discover(force=True) @@ -420,7 +464,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) versions = reg.get_all_versions_info("extract") @@ -446,7 +490,7 @@ def render(self, config): pass """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) assert reg.get_dag_args() is DefaultDagArgs @@ -468,7 +512,7 @@ def render(self, config) -> dict[str, Any]: return {"schedule": config.schedule} if config.schedule else {} """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) dag_args_cls = reg.get_dag_args() @@ -505,7 +549,7 @@ def render(self, config) -> dict[str, Any]: return {} """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) with pytest.raises(MultipleDagArgsError): reg.discover(force=True) @@ -526,7 +570,7 @@ def render(self, config) -> dict[str, Any]: return {} """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) assert reg.get_dag_args() is not DefaultDagArgs @@ -557,7 +601,7 @@ def render(self, config) -> dict[str, Any]: return {} """) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) bp_names = [bp["name"] for bp in reg.list_blueprints()] @@ -646,7 +690,7 @@ def test_non_blueprint_file_is_not_executed(self, tmp_path): " pass\n" ) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) bp_names = [bp["name"] for bp in reg.list_blueprints()] @@ -670,7 +714,7 @@ def test_blueprint_file_is_still_executed(self, tmp_path): " pass\n" ) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) bp_names = [bp["name"] for bp in reg.list_blueprints()] @@ -693,9 +737,218 @@ def test_list_blueprints_reports_absolute_location(self, tmp_path): " pass\n" ) - reg = BlueprintRegistry(template_dirs=[template_dir]) + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=False) reg.discover(force=True) location = reg.list_blueprints()[0]["locations"][1] assert Path(location).is_absolute() assert Path(location) == (template_dir / "etl.py").resolve() + + +class TestEntryPointDiscovery: + """Test discovering blueprints from installed packages via entry points.""" + + def _patch_entry_points(self, monkeypatch, eps): + monkeypatch.setattr(importlib.metadata, "entry_points", lambda **_: eps) + + def test_entry_point_discovery_basic(self, tmp_path, monkeypatch): + """Checks that a single installed entry-point module is discovered as a blueprint.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_basic_mod.py").write_text(_blueprint_source("BasicBp")) + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("basic", "_ep_basic_mod")]) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + reg.discover(force=True) + + blueprints = reg.list_blueprints() + assert [bp["name"] for bp in blueprints] == ["basic_bp"] + assert blueprints[0]["locations"][1] == "_ep_basic_mod" + + def test_entry_point_discovery_recursive_package(self, tmp_path, monkeypatch): + """Checks that when an entry point targets a package, its blueprint submodules are found too.""" + monkeypatch.syspath_prepend(str(tmp_path)) + pkg_dir = tmp_path / "_ep_recursive_pkg" + pkg_dir.mkdir() + (pkg_dir / "__init__.py").write_text("from .primary import Primary # noqa: F401\n") + (pkg_dir / "primary.py").write_text(_blueprint_source("Primary", "PrimaryConfig")) + (pkg_dir / "secondary.py").write_text(_blueprint_source("Secondary", "SecondaryConfig")) + + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("recursive", "_ep_recursive_pkg")]) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + reg.discover(force=True) # would raise DuplicateBlueprintError if double-registered + + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"primary", "secondary"} + + def test_entry_point_discovery_multiple_entry_points(self, tmp_path, monkeypatch): + """Checks that blueprints from more than one installed entry point are all discovered.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_multi_a.py").write_text(_blueprint_source("MultiA")) + (tmp_path / "_ep_multi_b.py").write_text(_blueprint_source("MultiB")) + + self._patch_entry_points( + monkeypatch, + [_FakeEntryPoint("a", "_ep_multi_a"), _FakeEntryPoint("b", "_ep_multi_b")], + ) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + reg.discover(force=True) + + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"multi_a", "multi_b"} + + def test_entry_point_duplicate_with_local_directory_raises(self, tmp_path, monkeypatch): + """Checks that discovery fails when the same blueprint exists both locally and in an installed package.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_dup_local.py").write_text(_blueprint_source("DupLocal")) + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("local", "_ep_dup_local")]) + + template_dir = tmp_path / "dags" + template_dir.mkdir() + (template_dir / "blueprints.py").write_text(_blueprint_source("DupLocal")) + + reg = BlueprintRegistry(template_dirs=[template_dir], discover_entry_points=True) + with pytest.raises(DuplicateBlueprintError, match="dup_local"): + reg.discover(force=True) + + def test_entry_point_duplicate_across_two_entry_points_raises(self, tmp_path, monkeypatch): + """Checks that discovery fails when two installed packages export the same blueprint name and version.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_dup_a.py").write_text(_blueprint_source("DupSame", "DupSameConfigA")) + (tmp_path / "_ep_dup_b.py").write_text(_blueprint_source("DupSame", "DupSameConfigB")) + + self._patch_entry_points( + monkeypatch, + [_FakeEntryPoint("a", "_ep_dup_a"), _FakeEntryPoint("b", "_ep_dup_b")], + ) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + with pytest.raises(DuplicateBlueprintError, match="dup_same"): + reg.discover(force=True) + + def test_entry_point_broken_module_logs_and_continues(self, tmp_path, monkeypatch, caplog): + """Checks that one broken entry point is skipped without preventing good packages from loading.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_good.py").write_text(_blueprint_source("GoodBp")) + + def _raise(): + msg = "boom" + raise ImportError(msg) + + self._patch_entry_points( + monkeypatch, + [ + _FakeEntryPoint("broken", "_ep_missing", load_fn=_raise), + _FakeEntryPoint("good", "_ep_good"), + ], + ) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + with caplog.at_level(logging.WARNING): + reg.discover(force=True) + + assert any("broken" in rec.getMessage() for rec in caplog.records) + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"good_bp"} + + def test_entry_point_non_module_target_logs_and_skips(self, tmp_path, monkeypatch, caplog): + """Checks that an entry point pointing to the wrong kind of object is ignored instead of crashing discovery.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_good2.py").write_text(_blueprint_source("GoodBp2")) + + self._patch_entry_points( + monkeypatch, + [ + _FakeEntryPoint("bad", "_ep_bad:attr", load_fn=lambda: 42), + _FakeEntryPoint("good", "_ep_good2"), + ], + ) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + with caplog.at_level(logging.WARNING): + reg.discover(force=True) + + assert any("does not resolve to a module" in rec.getMessage() for rec in caplog.records) + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"good_bp2"} + + def test_entry_point_broken_leaf_submodule_continues_others( + self, tmp_path, monkeypatch, caplog + ): + """Checks that one bad module inside a package does not stop other modules in that package from being discovered.""" + monkeypatch.syspath_prepend(str(tmp_path)) + pkg_dir = tmp_path / "_ep_broken_leaf_pkg" + pkg_dir.mkdir() + (pkg_dir / "__init__.py").write_text("") + (pkg_dir / "broken.py").write_text("raise ImportError('leaf boom')\n") + (pkg_dir / "fine.py").write_text(_blueprint_source("FineBp")) + + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("leaf", "_ep_broken_leaf_pkg")]) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + with caplog.at_level(logging.WARNING): + reg.discover(force=True) + + assert any("broken" in rec.getMessage() for rec in caplog.records) + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"fine_bp"} + + def test_entry_point_broken_subpackage_continues_siblings(self, tmp_path, monkeypatch, caplog): + """Checks that a broken subpackage does not stop discovery from reaching its sibling subpackages.""" + monkeypatch.syspath_prepend(str(tmp_path)) + pkg_dir = tmp_path / "_ep_broken_subpkg_pkg" + pkg_dir.mkdir() + (pkg_dir / "__init__.py").write_text("") + + broken_sub = pkg_dir / "broken_sub" + broken_sub.mkdir() + (broken_sub / "__init__.py").write_text("raise RuntimeError('subpackage boom')\n") + + fine_sub = pkg_dir / "fine_sub" + fine_sub.mkdir() + (fine_sub / "__init__.py").write_text("") + (fine_sub / "bp.py").write_text(_blueprint_source("FineSub")) + + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("subpkg", "_ep_broken_subpkg_pkg")]) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + with caplog.at_level(logging.WARNING): + reg.discover(force=True) + + names = {bp["name"] for bp in reg.list_blueprints()} + assert names == {"fine_sub"} + + def test_discover_entry_points_false_disables_discovery(self, tmp_path, monkeypatch): + """Checks that entry-point discovery is completely skipped when the feature is turned off.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_disabled_mod.py").write_text(_blueprint_source("DisabledBp")) + + calls = [] + + def _fake_entry_points(**kwargs): + calls.append(kwargs) + return [_FakeEntryPoint("disabled", "_ep_disabled_mod")] + + monkeypatch.setattr(importlib.metadata, "entry_points", _fake_entry_points) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=False) + reg.discover(force=True) + + assert reg.list_blueprints() == [] + assert calls == [] + + def test_entry_point_module_reused_across_force_rediscovery(self, tmp_path, monkeypatch): + """Checks that rediscovering entry points reuses the already imported module instead of creating a new class object.""" + monkeypatch.syspath_prepend(str(tmp_path)) + (tmp_path / "_ep_cache_mod.py").write_text(_blueprint_source("CacheBp")) + self._patch_entry_points(monkeypatch, [_FakeEntryPoint("cache", "_ep_cache_mod")]) + + reg = BlueprintRegistry(template_dirs=[], discover_entry_points=True) + reg.discover(force=True) + cls_first = reg.get("cache_bp") + + reg.discover(force=True) + cls_second = reg.get("cache_bp") + + assert cls_first is cls_second diff --git a/tests/test_utils.py b/tests/test_utils.py index 95d92b7..70f67df 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -28,3 +28,21 @@ def test_explicit_base_overrides_cwd(self, tmp_path, monkeypatch): monkeypatch.chdir(other) nested = tmp_path / "dags" / "bp.py" assert display_path(nested, base=tmp_path) == str(Path("dags") / "bp.py") + + def test_non_absolute_dotted_module_name_returned_unchanged(self): + """A blueprint discovered from an installed package has a dotted module name as its + location, not a real file. It must pass through as-is rather than being resolved as + if it were a relative filesystem path. + """ + assert display_path("company_blueprints.extract") == "company_blueprints.extract" + + def test_non_absolute_dotted_module_name_unaffected_by_differing_base(self, tmp_path): + """ + Entry-point locations are Python module names, not file paths. Changing the base directory + should not rewrite or resolve them. + """ + other_dir = tmp_path / "dags" + other_dir.mkdir() + assert display_path("company_blueprints.extract", base=other_dir) == ( + "company_blueprints.extract" + ) diff --git a/uv.lock b/uv.lock index bcba841..df18dcd 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'", @@ -44,6 +44,7 @@ dependencies = [ [package.dev-dependencies] dev = [ + { name = "entry-point-test-blueprints" }, { name = "httpx" }, { name = "hypothesis" }, { name = "playwright" }, @@ -66,6 +67,7 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ + { name = "entry-point-test-blueprints", editable = "tests/entry_point_package" }, { name = "httpx", specifier = ">=0.28.0" }, { name = "hypothesis", specifier = ">=6.113.0" }, { name = "playwright", specifier = ">=1.58.0" }, @@ -711,6 +713,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/ee/bf0adb559ad3c786f12bcbc9296b3f5675f529199bef03e2df281fa1fadb/email_validator-2.2.0-py3-none-any.whl", hash = "sha256:561977c2d73ce3611850a06fa56b414621e0c8faa9d66f2611407d87465da631", size = 33521, upload-time = "2024-06-20T11:30:28.248Z" }, ] +[[package]] +name = "entry-point-test-blueprints" +version = "0.0.0" +source = { editable = "tests/entry_point_package" } +dependencies = [ + { name = "airflow-blueprint" }, +] + +[package.metadata] +requires-dist = [{ name = "airflow-blueprint" }] + [[package]] name = "exceptiongroup" version = "1.3.0" From 1dd03482d7d02ccbfb07802951ba8209485d182a Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 27 Jul 2026 13:18:19 +0200 Subject: [PATCH 2/4] Remove unnecessary text --- examples/shared-blueprints/shared_blueprints/__init__.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/examples/shared-blueprints/shared_blueprints/__init__.py b/examples/shared-blueprints/shared_blueprints/__init__.py index b019783..e69de29 100644 --- a/examples/shared-blueprints/shared_blueprints/__init__.py +++ b/examples/shared-blueprints/shared_blueprints/__init__.py @@ -1,6 +0,0 @@ -"""Example central-team blueprint package. - -Demonstrates a shared-blueprints package discovered by a downstream project -purely by being installed -- see examples/advanced, which uses the -``example`` blueprint defined here without any local .py file. -""" From 7b18494456d98da0e4ddf1be58c99ffa68b72912 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 27 Jul 2026 15:15:48 +0200 Subject: [PATCH 3/4] Remove more text --- .../entry_point_test_blueprints/__init__.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/entry_point_package/entry_point_test_blueprints/__init__.py b/tests/entry_point_package/entry_point_test_blueprints/__init__.py index ec0258e..e69de29 100644 --- a/tests/entry_point_package/entry_point_test_blueprints/__init__.py +++ b/tests/entry_point_package/entry_point_test_blueprints/__init__.py @@ -1,6 +0,0 @@ -"""Test-only blueprint package, installed and discovered via its entry point. - -Exists solely so the integration suite can prove that BlueprintRegistry discovers -blueprints from an installed package end to end, without depending on anything -under ``examples/``. -""" From f94cccaa34201c18ef3066e7ddb71d19211fe457 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 27 Jul 2026 15:26:09 +0200 Subject: [PATCH 4/4] Clean up docstring --- .../integration/test_entry_point_discovery.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/integration/test_entry_point_discovery.py b/tests/integration/test_entry_point_discovery.py index 97d2455..7709763 100644 --- a/tests/integration/test_entry_point_discovery.py +++ b/tests/integration/test_entry_point_discovery.py @@ -1,11 +1,8 @@ -"""Tier 1: blueprint discovery from an installed package via entry points. +""" +Test blueprint discovery from an installed package via entry points. -The project's requirements.txt (written by the `airflow_env` fixture) installs -`tests/entry_point_package` alongside this repo. That package declares its `EntryPointBpTest` -blueprint under the `airflow_blueprint.blueprints` entry-point group and ships no YAML of its own --- `dags/entry_point_test.dag.yaml` references it by name with no corresponding local .py file, -so the DAG only parses if BlueprintRegistry's entry-point discovery is actually working end to end -against a real installed package (not just a mocked one). +Validate if Blueprints from installed packages can be referenced by name with no corresponding +local .py file. """ from __future__ import annotations @@ -54,7 +51,7 @@ def test_no_import_errors_for_entry_point_dag(self, api_client: AirflowAPI): ] assert not offending, f"Import errors for the entry-point DAG: {offending}" - def test_no_local_py_file_defines_the_probe_blueprint(self): + def test_no_local_py_file_defines_the_test_blueprint(self): """Guard against a future 'fix' that quietly adds a local copy of the blueprint, which would defeat the entire point of this test suite without any single assertion above failing. @@ -66,7 +63,6 @@ def test_no_local_py_file_defines_the_probe_blueprint(self): if "class EntryPointBpTest(" in py_file.read_text() ] assert not offenders, ( - f"Found a local copy of EntryPointBpTest in {offenders} -- the " - "entry-point DAG must resolve its blueprint purely from the " - "installed test package." + f"Found a local copy of EntryPointBpTest in {offenders}. The entry-point DAG must " + "resolve its blueprint purely from the installed test package." )