From e52941661468904f6e65e530fcb5042741fe5d3e Mon Sep 17 00:00:00 2001 From: Eugene Date: Wed, 22 Jul 2026 17:45:00 +0200 Subject: [PATCH] feat: support multiple Git DAG bundles --- README.md | 46 ++++++++++ bin/etl-workbench | 188 +++++++++++++++++++++++++++++++---------- tests/test_launcher.py | 88 +++++++++++++++++++ 3 files changed, 276 insertions(+), 46 deletions(-) create mode 100644 tests/test_launcher.py diff --git a/README.md b/README.md index d2c1a0b..ffbf6d1 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ Useful options: --ref VERSION branch, tag, or commit; default: main --subdir PATH DAG directory; default: dags --image IMAGE use a prebuilt pipeline image +--bundle-manifest FILE load several Git DAG sources; requires --image --env FILE pipeline-owned runtime environment --external-db do not start local PostgreSQL --external-objects do not start local object storage @@ -64,6 +65,51 @@ ignored `.workbench/runtime.env` with mode `0600`. The private key is used by Docker BuildKit and the local Airflow container; it is not copied into the image. Host-key checking uses `~/.ssh/known_hosts` by default. +## Several product sources in one Airflow + +One Airflow can load DAG entrypoints from several independent Git repositories. +Use a versioned JSON manifest when a shared factory serves several trusted +products: + +```json +{ + "version": 1, + "sources": [ + { + "name": "learning-platform", + "repository": "git@github.com:example/learning-platform.git", + "ref": "main", + "subdir": "airflow/dags" + }, + { + "name": "beavers-data", + "repository": "git@github.com:example/beavers-data-pipelines.git", + "ref": "main", + "subdir": "dags" + } + ] +} +``` + +Then start the factory with an image which contains the compatible Python +packages of **every** listed product: + +```bash +./bin/etl-workbench \ + --bundle-manifest trusted-products.json \ + --image trusted-airflow-pipelines:2026-07-22 \ + --ssh-key ~/.ssh/id_ed25519 +``` + +The factory creates one Git Connection per source and configures Airflow's +native `GitDagBundle` list. A Git bundle provides DAG files only; it must never +install arbitrary dependencies at parse time. The shared image is therefore an +explicit release artifact, built and tested from pinned product revisions. + +Keep source-specific Connections, object prefixes and Pools named by product. +That separates operational ownership inside one trusted Airflow, but does not +turn this local workbench into an isolation boundary for untrusted code. + ## Pipeline repository contract The smallest repository contains one or more DAG files: diff --git a/bin/etl-workbench b/bin/etl-workbench index e1a8040..9dd9fbb 100755 --- a/bin/etl-workbench +++ b/bin/etl-workbench @@ -2,6 +2,7 @@ from __future__ import annotations import argparse +from dataclasses import dataclass import hashlib import json import os @@ -15,6 +16,17 @@ STATE = ROOT / ".workbench" BASE_IMAGE = "etl-workbench:local" +@dataclass(frozen=True) +class DagBundle: + """One Git-backed DAG source loaded by the shared Airflow runtime.""" + + name: str + repository: str + ref: str + subdir: str + git_connection: str + + def run(command: list[str], *, env: dict[str, str] | None = None) -> None: print("+", " ".join(command), flush=True) subprocess.run(command, cwd=ROOT, env=env, check=True) @@ -31,10 +43,14 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Build and start a Git-backed Airflow pipeline locally." ) - parser.add_argument("repository", help="Git repository URL") + parser.add_argument("repository", nargs="?", help="single pipeline Git repository URL") parser.add_argument("--ref", default="main", help="branch, tag, or commit") parser.add_argument("--subdir", default="dags", help="DAG directory in the repository") - parser.add_argument("--image", help="use an existing pipeline image instead of building") + parser.add_argument("--image", help="use an existing Airflow image instead of building one pipeline image") + parser.add_argument( + "--bundle-manifest", + help="JSON manifest for multiple Git DAG sources; requires --image with all source packages installed", + ) parser.add_argument("--env", help="pipeline-owned environment file") parser.add_argument("--ssh-key", help="private SSH key for Git and the image build") parser.add_argument("--known-hosts", default="~/.ssh/known_hosts") @@ -46,80 +62,159 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def bundle_config(args: argparse.Namespace, connection_id: str) -> str: - kwargs: dict[str, object] = { - "tracking_ref": args.ref, - "subdir": args.subdir, - } - kwargs["git_conn_id"] = connection_id +def bundle_config(bundles: list[DagBundle]) -> str: return json.dumps( [ { - "name": "pipeline", + "name": bundle.name, "classpath": "airflow.providers.git.bundles.git.GitDagBundle", - "kwargs": kwargs, + "kwargs": { + "tracking_ref": bundle.ref, + "subdir": bundle.subdir, + "git_conn_id": bundle.git_connection, + }, } + for bundle in bundles ], separators=(",", ":"), ) +def bundle_manifest(path: Path) -> list[DagBundle]: + try: + payload = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise SystemExit(f"bundle manifest cannot be read: {path}") from error + if not isinstance(payload, dict) or payload.get("version") != 1: + raise SystemExit("bundle manifest must be an object with version: 1") + raw_sources = payload.get("sources") + if not isinstance(raw_sources, list) or not raw_sources: + raise SystemExit("bundle manifest must contain a non-empty sources array") + + bundles: list[DagBundle] = [] + names: set[str] = set() + connection_ids: set[str] = set() + for raw in raw_sources: + if not isinstance(raw, dict): + raise SystemExit("each bundle source must be an object") + name = raw.get("name") + repository = raw.get("repository") + ref = raw.get("ref", "main") + subdir = raw.get("subdir", "dags") + if ( + not isinstance(name, str) + or not name + or not all(character.isalnum() or character in "_-" for character in name) + or not isinstance(repository, str) + or not repository.strip() + or not isinstance(ref, str) + or not ref.strip() + or not isinstance(subdir, str) + or not subdir.strip() + ): + raise SystemExit("each source needs non-empty name, repository, ref and subdir values") + if name in names: + raise SystemExit(f"bundle source name is repeated: {name}") + names.add(name) + connection_id = f"workbench_git_{name.replace('-', '_')}" + if connection_id in connection_ids: + raise SystemExit( + "bundle source names must produce distinct Git connection IDs: " + f"{name}" + ) + connection_ids.add(connection_id) + bundles.append( + DagBundle( + name=name, + repository=repository.strip(), + ref=ref.strip(), + subdir=subdir.strip(), + git_connection=connection_id, + ) + ) + return bundles + + +def runtime_connection_line(bundle: DagBundle, ssh_key: Path | None) -> str: + connection: dict[str, object] = {"conn_type": "git", "host": bundle.repository} + if ssh_key: + connection["extra"] = { + "private_key": ssh_key.read_text(), + "strict_host_key_checking": "yes", + "known_hosts_file": "/var/lib/airflow/.ssh/known_hosts", + } + environment_name = "AIRFLOW_CONN_" + bundle.git_connection.upper() + return environment_name + "=" + json.dumps(connection, separators=(",", ":")) + + def main() -> None: args = parse_args() if args.ssh_key and args.git_connection: raise SystemExit("use either --ssh-key or --git-connection, not both") + if args.repository and args.bundle_manifest: + raise SystemExit("use either one repository or --bundle-manifest, not both") + if not args.repository and not args.bundle_manifest: + raise SystemExit("provide one repository or --bundle-manifest") + if args.bundle_manifest and args.git_connection: + raise SystemExit("--git-connection is only available for the single-repository mode") pipeline_env = existing_file(args.env, "pipeline environment file") if args.env else None ssh_key = existing_file(args.ssh_key, "SSH key") if args.ssh_key else None known_hosts = existing_file(args.known_hosts, "known_hosts file") if ssh_key else None - image = args.image - if not image: - run(["docker", "build", "-t", BASE_IMAGE, "."]) - digest = hashlib.sha256(f"{args.repository}\0{args.ref}".encode()).hexdigest()[:12] - image = f"etl-pipeline:{digest}" - command = [ - "docker", - "build", - "-t", - image, - "--build-arg", - f"ETL_WORKBENCH_IMAGE={BASE_IMAGE}", - "-f", - "Dockerfile.airflow", + runtime_lines: list[str] + if args.bundle_manifest: + manifest = existing_file(args.bundle_manifest, "bundle manifest") + if not args.image: + raise SystemExit("--bundle-manifest requires --image with every source package installed") + image = args.image + bundles = bundle_manifest(manifest) + runtime_lines = [runtime_connection_line(bundle, ssh_key) for bundle in bundles] + else: + assert args.repository is not None + connection_id = args.git_connection or "workbench_git" + bundles = [ + DagBundle( + name="pipeline", + repository=args.repository, + ref=args.ref, + subdir=args.subdir, + git_connection=connection_id, + ) ] - if ssh_key: - command.extend(["--ssh", f"default={ssh_key}"]) - command.append(f"{args.repository}#{args.ref}") - run(command) + image = args.image + if not image: + run(["docker", "build", "-t", BASE_IMAGE, "."]) + digest = hashlib.sha256(f"{args.repository}\0{args.ref}".encode()).hexdigest()[:12] + image = f"etl-pipeline:{digest}" + command = [ + "docker", + "build", + "-t", + image, + "--build-arg", + f"ETL_WORKBENCH_IMAGE={BASE_IMAGE}", + "-f", + "Dockerfile.airflow", + ] + if ssh_key: + command.extend(["--ssh", f"default={ssh_key}"]) + command.append(f"{args.repository}#{args.ref}") + run(command) + runtime_lines = [] + if not args.git_connection: + runtime_lines.append(runtime_connection_line(bundles[0], ssh_key)) - connection_id = args.git_connection or "workbench_git" - runtime_lines: list[str] = [] compose_files = ["-f", "compose.yaml"] compose_env = os.environ.copy() - if not args.git_connection: - connection = { - "conn_type": "git", - "host": args.repository, - } - if ssh_key: - connection["extra"] = { - "private_key": ssh_key.read_text(), - "strict_host_key_checking": "yes", - "known_hosts_file": "/var/lib/airflow/.ssh/known_hosts", - } - runtime_lines.append( - "AIRFLOW_CONN_WORKBENCH_GIT=" + json.dumps(connection, separators=(",", ":")) - ) - if ssh_key: compose_files.extend(["-f", "compose.git-ssh.yaml"]) compose_env["GIT_KNOWN_HOSTS_PATH"] = str(known_hosts) runtime_lines.append( "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST=" - + bundle_config(args, connection_id) + + bundle_config(bundles) ) STATE.mkdir(exist_ok=True) runtime_env = STATE / "runtime.env" @@ -139,6 +234,7 @@ def main() -> None: command.extend(["up", "-d", "--wait"]) run(command, env=compose_env) print("Airflow: http://127.0.0.1:18080", flush=True) + print("DAG bundles: " + ", ".join(bundle.name for bundle in bundles), flush=True) if __name__ == "__main__": diff --git a/tests/test_launcher.py b/tests/test_launcher.py new file mode 100644 index 0000000..1da5e96 --- /dev/null +++ b/tests/test_launcher.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import importlib.util +from importlib.machinery import SourceFileLoader +import json +from pathlib import Path +import sys +import tempfile +import unittest + + +ROOT = Path(__file__).resolve().parents[1] +LOADER = SourceFileLoader("etl_workbench_launcher", str(ROOT / "bin" / "etl-workbench")) +SPEC = importlib.util.spec_from_loader("etl_workbench_launcher", LOADER) +assert SPEC is not None and SPEC.loader is not None +launcher = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = launcher +SPEC.loader.exec_module(launcher) + + +class BundleManifestTests(unittest.TestCase): + def test_multiple_sources_get_stable_independent_git_connections(self) -> None: + payload = { + "version": 1, + "sources": [ + { + "name": "learning-platform", + "repository": "git@github.com:example/learning-platform.git", + "ref": "main", + "subdir": "airflow/dags", + }, + { + "name": "beavers-data", + "repository": "git@github.com:example/beavers-data-pipelines.git", + "ref": "release-2026-07", + "subdir": "dags", + }, + ], + } + with tempfile.TemporaryDirectory() as directory: + manifest = Path(directory) / "bundles.json" + manifest.write_text(json.dumps(payload)) + bundles = launcher.bundle_manifest(manifest) + + self.assertEqual([bundle.git_connection for bundle in bundles], ["workbench_git_learning_platform", "workbench_git_beavers_data"]) + configuration = json.loads(launcher.bundle_config(bundles)) + self.assertEqual([item["name"] for item in configuration], ["learning-platform", "beavers-data"]) + self.assertEqual(configuration[0]["kwargs"], { + "tracking_ref": "main", + "subdir": "airflow/dags", + "git_conn_id": "workbench_git_learning_platform", + }) + self.assertEqual( + launcher.runtime_connection_line(bundles[1], None), + 'AIRFLOW_CONN_WORKBENCH_GIT_BEAVERS_DATA={"conn_type":"git","host":"git@github.com:example/beavers-data-pipelines.git"}', + ) + + def test_manifest_rejects_repeated_source_names(self) -> None: + payload = { + "version": 1, + "sources": [ + {"name": "product", "repository": "https://example.test/a.git"}, + {"name": "product", "repository": "https://example.test/b.git"}, + ], + } + with tempfile.TemporaryDirectory() as directory: + manifest = Path(directory) / "bundles.json" + manifest.write_text(json.dumps(payload)) + with self.assertRaisesRegex(SystemExit, "repeated"): + launcher.bundle_manifest(manifest) + + def test_manifest_rejects_connection_id_collisions(self) -> None: + payload = { + "version": 1, + "sources": [ + {"name": "learning-platform", "repository": "https://example.test/a.git"}, + {"name": "learning_platform", "repository": "https://example.test/b.git"}, + ], + } + with tempfile.TemporaryDirectory() as directory: + manifest = Path(directory) / "bundles.json" + manifest.write_text(json.dumps(payload)) + with self.assertRaisesRegex(SystemExit, "distinct Git connection IDs"): + launcher.bundle_manifest(manifest) + + +if __name__ == "__main__": + unittest.main()