From 380e41e899a31722d760a7873240a49e24e29025 Mon Sep 17 00:00:00 2001 From: Maxim Hammer Date: Fri, 23 Oct 2026 12:41:28 +0200 Subject: [PATCH 1/2] contrib/dbt-factory: port v0.4.0, add split by tag Re-vendor databricks-dbt-factory v0.4.0 (c830e20) into the example and the template, and point both NOTICE files at it. This brings the runner fix for manifests with dates (upstream #42) and split by tag (#43). SPLIT_BY_TAGS in resources/__init__.py exposes the split: one job per dbt tag, run in dependency order by the main job. Empty keeps one job. Tested: vendor-sync check, ruff 0.9.1, make test in the example and in generated projects, and a deploy and run on a serverless workspace with and without the split (dbt-core 1.12.5, dbt-databricks 1.12.7). Co-authored-by: Isaac --- contrib/dbt_factory/NOTICE | 2 +- contrib/dbt_factory/README.md | 7 + contrib/dbt_factory/resources/__init__.py | 120 ++++++++--- .../src/databricks_dbt_factory/__version__.py | 2 +- .../src/databricks_dbt_factory/dbt_factory.py | 40 +++- .../databricks_dbt_factory/job_partition.py | 132 ++++++++++++ .../notebook/run_dbt_command.py | 3 +- .../src/databricks_dbt_factory/utils.py | 3 + .../dbt_factory/tests/test_load_resources.py | 90 +++++++- contrib/templates/dbt-factory/README.md | 3 + .../template/{{.project_name}}/NOTICE | 2 +- .../template/{{.project_name}}/README.md.tmpl | 23 +- .../resources/__init__.py.tmpl | 198 +++++++++++------- .../src/databricks_dbt_factory/__version__.py | 2 +- .../src/databricks_dbt_factory/dbt_factory.py | 40 +++- .../databricks_dbt_factory/job_partition.py | 132 ++++++++++++ .../notebook/run_dbt_command.py | 3 +- .../src/databricks_dbt_factory/utils.py | 3 + .../tests/test_load_resources.py.tmpl | 106 +++++++++- 19 files changed, 766 insertions(+), 145 deletions(-) create mode 100644 contrib/dbt_factory/src/databricks_dbt_factory/job_partition.py create mode 100644 contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/job_partition.py diff --git a/contrib/dbt_factory/NOTICE b/contrib/dbt_factory/NOTICE index 702735ef..dda67852 100644 --- a/contrib/dbt_factory/NOTICE +++ b/contrib/dbt_factory/NOTICE @@ -2,7 +2,7 @@ This project includes code adapted from the "databricks-dbt-factory" library. Source: https://github.com/mwojtyczka/databricks-dbt-factory Full history: https://github.com/mwojtyczka/databricks-dbt-factory/commits/main - Adapted from: commit 29866edb82de28c33fe888be67d451068d4d7f94 (v0.3.5) + Adapted from: commit c830e20e9fb5196d7a10c42cf054cfd4d91eba46 (v0.4.0) Location: src/databricks_dbt_factory/ The code under src/databricks_dbt_factory/ originates from that repository (trimmed to the diff --git a/contrib/dbt_factory/README.md b/contrib/dbt_factory/README.md index 37eaae09..aebabfb0 100644 --- a/contrib/dbt_factory/README.md +++ b/contrib/dbt_factory/README.md @@ -153,6 +153,13 @@ A few knobs are exposed as constants at the top of `resources/__init__.py`: test node, for maximum per-test visibility). * `ENVIRONMENT_KEY` — the serverless environment key (default `Default`). * `EXTRA_DBT_COMMAND_OPTIONS` — extra options appended to every generated dbt command. +* `SPLIT_BY_TAGS` — split the job into one job per dbt tag listed here (exact tags or globs such as + `"domain_*"`). Each resource runs in the `dbt_factory_job_` job of its tag, or in + `dbt_factory_job_default` if it carries none, and tests run in the job of what they test. + `dbt_factory_job` then runs those jobs in dependency order and keeps the schedule. Default `[]` + (one job). Use it when a project outgrows the 1,000-task limit of a job. It works like the CLI's + `--split-by-tag`; see [Splitting into several jobs](https://github.com/mwojtyczka/databricks-dbt-factory#splitting-into-several-jobs) + for the full rules, including which splits are refused. The dbt target, warehouse, catalog, and schema are configured in `dbt_profiles/profiles.yml` and selected per bundle target via `--target ${bundle.target}`. diff --git a/contrib/dbt_factory/resources/__init__.py b/contrib/dbt_factory/resources/__init__.py index 4c6a7e5c..9d2e3a79 100644 --- a/contrib/dbt_factory/resources/__init__.py +++ b/contrib/dbt_factory/resources/__init__.py @@ -3,8 +3,9 @@ At ``databricks bundle deploy`` time the Databricks CLI calls :func:`load_resources`, which reads the dbt manifest and expands it into a Databricks job with one task per dbt object -(model / seed / snapshot / test) using the vendored ``databricks_dbt_factory`` core. No -per-model YAML is generated or checked in — the job graph is built on the fly from the manifest. +(model / seed / snapshot / test), optionally split into one job per dbt tag, using the vendored +``databricks_dbt_factory`` core. No per-model YAML is generated or checked in — the job graph is +built on the fly from the manifest. This module is the only integration glue: it configures the vendored ``databricks_dbt_factory`` core (under ``src/databricks_dbt_factory``) for serverless notebook tasks and returns a @@ -20,6 +21,7 @@ from databricks_dbt_factory.dbt_factory import DbtFactory from databricks_dbt_factory.dbt_task import DbtTaskOptions, TaskType +from databricks_dbt_factory.job_partition import JobPartition from databricks_dbt_factory.task_factory import ( DbtDependencyResolver, ModelTaskFactory, @@ -52,9 +54,15 @@ # context must match the supplied manifest. EXTRA_DBT_COMMAND_OPTIONS = "" +# Split the job into one job per dbt tag matching these patterns, each an exact tag or a glob such as +# "domain_*"; resources without a matching tag run in a `default` job. The JOB_NAME job then runs those +# jobs in dependency order. Empty keeps one job. Works like the databricks-dbt-factory CLI's +# --split-by-tag; see "Splitting into several jobs" in its README for how tests are placed and what is refused. +SPLIT_BY_TAGS = [] + # Serverless base-environment file, generated at deploy time from the pinned # dbt-databricks version and synced with the bundle. The generated job points its -# serverless environment at this file (see build_job) so Databricks pre-builds the +# serverless environment at this file (see _job) so Databricks pre-builds the # environment once instead of installing dbt on every task. `${workspace.file_path}` # is substituted by the CLI at deploy time to the bundle's synced files root. SERVERLESS_ENV_FILE = "dbt_serverless_env.yaml" @@ -69,8 +77,8 @@ PROFILES_DIRECTORY = "dbt_profiles" -def _build_tasks(target: str) -> list[dict]: - """Reads the dbt manifest and returns the list of Databricks task dicts (one per dbt node).""" +def _build_factory(target: str) -> DbtFactory: + """Configures the factory to generate serverless notebook tasks that run the given dbt target.""" resolver = DbtDependencyResolver() task_options = DbtTaskOptions( environment_key=ENVIRONMENT_KEY, @@ -88,9 +96,17 @@ def _build_tasks(target: str) -> list[dict]: "test": TestTaskFactory(resolver, task_options, dbt_options), } - factory = DbtFactory(task_factories, bundle_tests=BUNDLE_TESTS) - manifest = read_dbt_manifest(MANIFEST_PATH) - return factory.create_tasks(manifest) + return DbtFactory(task_factories, bundle_tests=BUNDLE_TESTS) + + +def _build_tasks(target: str) -> list[dict]: + """Reads the dbt manifest and returns the list of Databricks task dicts (one per dbt node).""" + return _build_factory(target).create_tasks(read_dbt_manifest(MANIFEST_PATH)) + + +def _build_partitions(target: str) -> list[JobPartition]: + """Reads the dbt manifest and returns its tasks split into one job per tag matching SPLIT_BY_TAGS.""" + return _build_factory(target).create_partitions(read_dbt_manifest(MANIFEST_PATH), SPLIT_BY_TAGS) def _dbt_databricks_dependency() -> str: @@ -157,28 +173,67 @@ def _write_serverless_environment_file() -> None: ) from exc -def build_job(target: str) -> Job: - """Builds the Databricks job that runs the dbt project on serverless via notebook tasks.""" - return Job.from_dict( - { - "name": JOB_NAME, - "queue": {"enabled": True}, - "trigger": { - # Run this job every day, exactly one day from the last run; - # see https://docs.databricks.com/api/workspace/jobs/create#trigger - "periodic": {"interval": 1, "unit": "DAYS"}, - }, - "tasks": _build_tasks(target), - "environments": [ - { - "environment_key": ENVIRONMENT_KEY, - # Pre-built base env (synced with the bundle), not inline deps — - # the two are mutually exclusive; the file pins dbt-databricks. - "spec": {"base_environment": SERVERLESS_ENV_PATH}, - } - ], - } - ) +# How SPLIT_BY_TAGS shares out the job settings, as the databricks-dbt-factory CLI's --split-by-tag does: +# the JOB_NAME job keeps the schedule, triggers, notifications and health rules, the per-tag jobs keep +# the compute, and all of them keep everything else (such as the queue). +_ORCHESTRATOR_ONLY = frozenset( + { + "schedule", + "trigger", + "continuous", + "email_notifications", + "webhook_notifications", + "notification_settings", + "health", + } +) +_SPLIT_JOB_ONLY = frozenset({"environments", "job_clusters"}) + + +def _job(name: str, tasks: list[dict], without: frozenset[str] = frozenset()) -> Job: + """Builds a job that runs `tasks` with the settings below, leaving out those named in `without`.""" + settings = { + "name": name, + "queue": {"enabled": True}, + "trigger": { + # Run this job every day, exactly one day from the last run; + # see https://docs.databricks.com/api/workspace/jobs/create#trigger + "periodic": {"interval": 1, "unit": "DAYS"}, + }, + "tasks": tasks, + "environments": [ + { + "environment_key": ENVIRONMENT_KEY, + # Pre-built base env (synced with the bundle), not inline deps — + # the two are mutually exclusive; the file pins dbt-databricks. + "spec": {"base_environment": SERVERLESS_ENV_PATH}, + } + ], + } + return Job.from_dict({key: value for key, value in settings.items() if key not in without}) + + +def build_jobs(target: str) -> dict[str, Job]: + """ + Builds the Databricks jobs that run the dbt project on serverless via notebook tasks, by resource key. + + That is the JOB_NAME job, or with SPLIT_BY_TAGS one `_` job per tag plus the JOB_NAME + job, which runs each of them once the jobs it depends on have succeeded. + """ + if not SPLIT_BY_TAGS: + return {JOB_NAME: _job(JOB_NAME, _build_tasks(target))} + jobs, run_jobs = {}, [] + for partition in _build_partitions(target): + key = f"{JOB_NAME}_{partition.name}" + jobs[key] = _job(key, partition.tasks, without=_ORCHESTRATOR_ONLY) + run_jobs.append( + { + "task_key": partition.name, + "depends_on": [{"task_key": upstream} for upstream in partition.depends_on], + "run_job_task": {"job_id": f"${{resources.jobs.{key}.id}}"}, + } + ) + return {JOB_NAME: _job(JOB_NAME, run_jobs, without=_SPLIT_JOB_ONLY), **jobs} def load_resources(bundle: Bundle) -> Resources: @@ -186,9 +241,10 @@ def load_resources(bundle: Bundle) -> Resources: Called by the Databricks CLI during `bundle deploy` to load resources defined in Python. Reads the dbt manifest, writes the serverless base-environment file (synced with the bundle), - and registers the generated job. After deployment this function is not used. + and registers the generated jobs. After deployment this function is not used. """ _write_serverless_environment_file() resources = Resources() - resources.add_job(JOB_NAME, build_job(bundle.target)) + for key, job in build_jobs(bundle.target).items(): + resources.add_job(key, job) return resources diff --git a/contrib/dbt_factory/src/databricks_dbt_factory/__version__.py b/contrib/dbt_factory/src/databricks_dbt_factory/__version__.py index a8d4557d..6a9beea8 100644 --- a/contrib/dbt_factory/src/databricks_dbt_factory/__version__.py +++ b/contrib/dbt_factory/src/databricks_dbt_factory/__version__.py @@ -1 +1 @@ -__version__ = "0.3.5" +__version__ = "0.4.0" diff --git a/contrib/dbt_factory/src/databricks_dbt_factory/dbt_factory.py b/contrib/dbt_factory/src/databricks_dbt_factory/dbt_factory.py index 39cc0788..a26a6bc7 100644 --- a/contrib/dbt_factory/src/databricks_dbt_factory/dbt_factory.py +++ b/contrib/dbt_factory/src/databricks_dbt_factory/dbt_factory.py @@ -1,12 +1,13 @@ import heapq -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from dataclasses import dataclass, field, replace from pathlib import PurePosixPath, PureWindowsPath from typing import cast from databricks_dbt_factory.dbt_task import DbtTask +from databricks_dbt_factory.job_partition import JobPartition, TaskPlacement, partition_tasks from databricks_dbt_factory.task_factory import TaskFactory, TestTaskFactory -from databricks_dbt_factory.utils import DYNAMIC_VALUE_REFERENCE, build_task_key_maps +from databricks_dbt_factory.utils import DYNAMIC_VALUE_REFERENCE, MAX_JOB_TASKS, build_task_key_maps # The `unique_id` prefixes of resources a test can be attached to. _DBT_TEST_TARGET_PREFIXES = ("model.", "seed.", "snapshot.", "source.") @@ -108,14 +109,18 @@ def create_tasks(self, dbt_manifest: dict) -> list[dict]: Raises: ValueError: If the generated job would exceed Databricks' 1,000-task limit. """ - tasks = self._create_tasks(dbt_manifest) - if len(tasks) > 1_000: + tasks, _ = self._create_tasks(dbt_manifest) + if len(tasks) > MAX_JOB_TASKS: raise ValueError( - f"Databricks jobs support at most 1,000 tasks; this manifest generates {len(tasks):,}. " - f"Reduce the generated resources or enable test bundling." + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; this manifest generates {len(tasks):,}. " + f"Reduce the generated resources, enable test bundling, or split the job by tag." ) return [task.to_dict() for task in tasks] + def create_partitions(self, dbt_manifest: dict, split_by_tags: Sequence[str]) -> list[JobPartition]: + """Generates the tasks split into one job per dbt tag matching `split_by_tags`; see `partition_tasks`.""" + return partition_tasks(*self._create_tasks(dbt_manifest), split_by_tags) + _GATEABLE_TYPES = frozenset({"model", "seed", "snapshot"}) _DBT_TEST_TARGET_PREFIXES = _DBT_TEST_TARGET_PREFIXES @@ -678,7 +683,7 @@ def _unaddressable(node_info: dict) -> ValueError: f"resources, so the task could run another task's resource." ) - def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: + def _create_tasks(self, dbt_manifest: dict) -> tuple[list[DbtTask], dict[str, TaskPlacement]]: """ Builds `DbtTask` instances from the manifest, applying the bundling and gating policies. @@ -686,7 +691,8 @@ def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: dbt_manifest (dict): Parsed dbt manifest content. Returns: - list[DbtTask]: `DbtTask` instances (not yet rendered to dicts). + tuple[list[DbtTask], dict[str, TaskPlacement]]: `DbtTask` instances (not yet rendered to dicts), and + the placement of each by task key. """ dbt_nodes = self._enabled_only(dbt_manifest.get("nodes", {})) dbt_sources = self._enabled_only(dbt_manifest.get("sources", {})) @@ -752,7 +758,23 @@ def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: elif "test" in self.task_factories: tasks.extend(self._build_unit_test_tasks(dbt_unit_tests, task_keys, peers)) - return tasks + return tasks, self._placements(peers, task_keys, bundled_test_keys) + + @staticmethod + def _placements( + peers: dict, task_keys: dict[str, str], bundled_test_keys: dict[str, str] + ) -> dict[str, TaskPlacement]: + """Places each task by its resource's tags; tests, and the bundles running them, follow their dependencies.""" + bundles = set(bundled_test_keys.values()) + placements = {} + for full_name, task_key in (*task_keys.items(), *bundled_test_keys.items()): + info = peers[full_name] + placements[task_key] = TaskPlacement( + f"{info.get('resource_type')} {info.get('name')!r} ({info.get('original_file_path')})", + frozenset(info.get("tags") or ()), + task_key in bundles or info["resource_type"] in {"test", "unit_test"}, + ) + return placements @staticmethod def _enabled_only(entries: dict) -> dict: diff --git a/contrib/dbt_factory/src/databricks_dbt_factory/job_partition.py b/contrib/dbt_factory/src/databricks_dbt_factory/job_partition.py new file mode 100644 index 00000000..d96a693c --- /dev/null +++ b/contrib/dbt_factory/src/databricks_dbt_factory/job_partition.py @@ -0,0 +1,132 @@ +import re +from collections.abc import Sequence +from dataclasses import dataclass, replace +from fnmatch import fnmatchcase +from graphlib import CycleError, TopologicalSorter + +from databricks_dbt_factory.dbt_task import DbtTask +from databricks_dbt_factory.utils import MAX_JOB_TASKS, MAX_TASK_KEY_LENGTH + +DEFAULT_JOB = "default" + + +@dataclass(frozen=True) +class JobPartition: + """One job of a split project: its tasks, and the jobs that must succeed before it runs.""" + + name: str + tasks: list[dict] + depends_on: list[str] + + +@dataclass(frozen=True) +class TaskPlacement: + """The resource a task runs, its tags, and whether the task follows its dependencies instead of its tags.""" + + resource: str + tags: frozenset[str] + follows_dependencies: bool + + +def partition_tasks( + tasks: list[DbtTask], placements: dict[str, TaskPlacement], patterns: Sequence[str] +) -> list[JobPartition]: + """ + Splits tasks into one job per tag matching `patterns`, ordered so each job runs after those it depends on. + + A task joins the job of the one matching tag it carries, or the default job. A task that follows its + dependencies joins the last of their jobs, so tests run after every resource they test. + + Raises: + ValueError: when a resource carries several matching tags, a pattern matches no tag, two tags name the + same job, jobs would depend on each other, or a job exceeds Databricks' task limit. + """ + dependencies = {task.task_key: list(task.depends_on or []) for task in tasks} + tags = { + task.task_key: _split_tag(placements[task.task_key], patterns) + for task in tasks + if not (placements[task.task_key].follows_dependencies and dependencies[task.task_key]) + } + owner = _job_names(tags, patterns) + anchors = _anchors(dependencies, owner) + order = _job_order(owner, dependencies, anchors) + rank = {job: index for index, job in enumerate(order)} + for key in dependencies.keys() - owner.keys(): + owner[key] = max((owner[anchor] for anchor in anchors[key]), key=rank.__getitem__) + + jobs: dict[str, list[DbtTask]] = {job: [] for job in order} + upstream: dict[str, set[str]] = {job: set() for job in order} + for task in tasks: + job = owner[task.task_key] + jobs[job].append(replace(task, depends_on=[d for d in dependencies[task.task_key] if owner[d] == job])) + upstream[job].update(owner[d] for d in dependencies[task.task_key] if owner[d] != job) + for job, job_tasks in jobs.items(): + if len(job_tasks) > MAX_JOB_TASKS: + raise ValueError( + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; the {job!r} job generates " + f"{len(job_tasks):,}. Split it further by tag." + ) + if len(order) > MAX_JOB_TASKS: + raise ValueError( + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; the tags split into {len(order):,} jobs. " + f"Split by fewer tags." + ) + return [JobPartition(job, [task.to_dict() for task in jobs[job]], sorted(upstream[job])) for job in order] + + +def _split_tag(placement: TaskPlacement, patterns: Sequence[str]) -> str | None: + matched = sorted(tag for tag in placement.tags if any(fnmatchcase(tag, pattern) for pattern in patterns)) + if len(matched) > 1: + raise ValueError( + f"Cannot split by tag: {placement.resource} carries the split tags {', '.join(matched)}. " + f"Keep one, so it belongs to one job." + ) + return matched[0] if matched else None + + +def _job_names(tags: dict[str, str | None], patterns: Sequence[str]) -> dict[str, str]: + used = set(tags.values()) + for pattern in patterns: + if not any(tag is not None and fnmatchcase(tag, pattern) for tag in used): + raise ValueError( + f"--split-by-tag {pattern!r} matches no tag of a resource to split; tests run in the job of the " + f"resources they test, whatever their own tags." + ) + names = {tag: DEFAULT_JOB if tag is None else re.sub(r"[^A-Za-z0-9_]", "_", tag) for tag in used} + owners: dict[str, str | None] = {} + for tag in sorted(used, key=str): + if len(names[tag]) > MAX_TASK_KEY_LENGTH: + raise ValueError( + f"Cannot split by tag: tag {tag!r} is longer than the {MAX_TASK_KEY_LENGTH} characters Databricks " + f"allows in a task key. Use a shorter tag." + ) + if owners.setdefault(names[tag], tag) != tag: + labels = [f"tag {t!r}" if t is not None else "untagged resources" for t in (owners[names[tag]], tag)] + raise ValueError(f"Cannot split by tag: {' and '.join(labels)} would both run as job {names[tag]!r}.") + return {key: names[tag] for key, tag in tags.items()} + + +def _anchors(dependencies: dict[str, list[str]], placed: dict[str, str]) -> dict[str, set[str]]: + anchors: dict[str, set[str]] = {} + for key in TopologicalSorter(dependencies).static_order(): + anchors[key] = {key} if key in placed else set().union(*(anchors[d] for d in dependencies[key])) + return anchors + + +def _job_order(owner: dict[str, str], dependencies: dict[str, list[str]], anchors: dict[str, set[str]]) -> list[str]: + upstream: dict[str, set[str]] = {job: set() for job in owner.values()} + witness: dict[tuple[str, str], tuple[str, str]] = {} + for key, job in owner.items(): + for source in (anchor for d in dependencies[key] for anchor in anchors[d]): + if owner[source] != job: + upstream[job].add(owner[source]) + witness.setdefault((owner[source], job), (key, source)) + try: + return list(TopologicalSorter({job: sorted(upstream[job]) for job in sorted(upstream)}).static_order()) + except CycleError as error: + cycle = error.args[1] + needs = "; ".join(f"{key} needs {source}" for key, source in (witness[p] for p in zip(cycle, cycle[1:]))) + raise ValueError( + f"Cannot split by tag: jobs {' -> '.join(cycle)} would depend on each other ({needs}). " + f"Retag the resources so dependencies between jobs flow one way." + ) from error diff --git a/contrib/dbt_factory/src/databricks_dbt_factory/notebook/run_dbt_command.py b/contrib/dbt_factory/src/databricks_dbt_factory/notebook/run_dbt_command.py index accb8f6a..9742b649 100644 --- a/contrib/dbt_factory/src/databricks_dbt_factory/notebook/run_dbt_command.py +++ b/contrib/dbt_factory/src/databricks_dbt_factory/notebook/run_dbt_command.py @@ -68,9 +68,10 @@ if os.path.exists(prebuilt_manifest_path): try: from dbt.contracts.graph.manifest import Manifest + from dbt.parser.manifest import extended_mashumuro_decoder with open(prebuilt_manifest_path, "rb") as f: - manifest = Manifest.from_msgpack(f.read()) + manifest = Manifest.from_msgpack(f.read(), decoder=extended_mashumuro_decoder) manifest.build_flat_graph() print(f"[dbt-factory] injecting pre-built manifest from {prebuilt_manifest_path} (skipping dbt parse)") except Exception as e: diff --git a/contrib/dbt_factory/src/databricks_dbt_factory/utils.py b/contrib/dbt_factory/src/databricks_dbt_factory/utils.py index 02ff51a2..fd675205 100644 --- a/contrib/dbt_factory/src/databricks_dbt_factory/utils.py +++ b/contrib/dbt_factory/src/databricks_dbt_factory/utils.py @@ -9,6 +9,9 @@ # Databricks caps task keys at 100 characters (letters, numbers, underscores, hyphens). MAX_TASK_KEY_LENGTH = 100 +# Databricks caps jobs at 1,000 tasks. +MAX_JOB_TASKS = 1_000 + # dbt resource types whose task key is `_` (the type doubles as the suffix). _SUFFIXED_TYPES = frozenset({"model", "seed", "snapshot"}) diff --git a/contrib/dbt_factory/tests/test_load_resources.py b/contrib/dbt_factory/tests/test_load_resources.py index 648a421b..ee739059 100644 --- a/contrib/dbt_factory/tests/test_load_resources.py +++ b/contrib/dbt_factory/tests/test_load_resources.py @@ -3,7 +3,8 @@ Exercises the whole glue path against the committed dbt manifest (no Databricks workspace required): reading the manifest, wiring the factory, generating notebook tasks, and registering -the job on a Resources object (which runs Job.from_dict, validating the task dict schema). +the job on a Resources object (which runs Job.from_dict, validating the task dict schema). The +split by tag runs against a small tagged project that dbt parses offline. """ import json @@ -13,6 +14,7 @@ import pytest import yaml from databricks.bundles.core import Bundle +from dbt.cli.main import dbtRunner import resources from databricks_dbt_factory.utils import generate_task_key @@ -33,6 +35,13 @@ def _chdir_to_project_root(monkeypatch): monkeypatch.chdir(PROJECT_ROOT) +def _jobs_running_dbt(jobs: dict) -> list: + """The generated jobs whose tasks run dbt: the job itself, or with SPLIT_BY_TAGS the per-tag jobs it runs.""" + dbt_jobs = [job for job in jobs.values() if any(task.notebook_task for task in job.tasks)] + assert dbt_jobs, "no generated job runs dbt" + return dbt_jobs + + def test_build_tasks_covers_every_buildable_node(): # Every model/seed/snapshot node gets its own run task regardless of the BUNDLE_TESTS mode # (tests may be collapsed into per-resource `_test` tasks when bundling is on). @@ -71,16 +80,16 @@ def test_environment_pins_installed_dbt_databricks(): assert spec["dependencies"] == [f"dbt-databricks=={version('dbt-databricks')}"] # The generated job points its serverless environment at the synced env file. - job = resources.build_job("dev") - assert job.environments[0].spec.base_environment == resources.SERVERLESS_ENV_PATH + for job in _jobs_running_dbt(resources.build_jobs("dev")): + assert job.environments[0].spec.base_environment == resources.SERVERLESS_ENV_PATH def test_load_resources_registers_the_job(): result = resources.load_resources(Bundle(target="dev")) try: assert resources.JOB_NAME in result.jobs - job = result.jobs[resources.JOB_NAME] - assert job.environments[0].environment_key == resources.ENVIRONMENT_KEY + for job in _jobs_running_dbt(result.jobs): + assert job.environments[0].environment_key == resources.ENVIRONMENT_KEY # load_resources writes the base-environment file to the bundle root for # the bundle to sync; it pins dbt-databricks to the installed version. @@ -93,6 +102,77 @@ def test_load_resources_registers_the_job(): (PROJECT_ROOT / resources.SERVERLESS_ENV_FILE).unlink(missing_ok=True) +TAGGED_PROJECT = { + "dbt_project.yml": 'name: probe\nprofile: probe\nversion: "1.0"\nconfig-version: 2\n', + "profiles.yml": ( + "probe:\n target: dev\n outputs:\n dev:\n type: databricks\n host: example.databricks.com\n" + " http_path: /sql/1.0/warehouses/x\n token: dummy\n schema: default\n" + ), + "models/stg_orders.sql": "{{ config(tags=['staging']) }}\nselect 1 as order_id", + "models/orders.sql": "{{ config(tags=['marts']) }}\nselect * from {{ ref('stg_orders') }}", + "models/legacy.sql": "select 1 as id", + "models/schema.yml": ( + "version: 2\nmodels:\n - name: stg_orders\n columns:\n - name: order_id\n data_tests: [not_null]\n" + ), +} + + +@pytest.fixture +def tagged_manifest(tmp_path, monkeypatch): + """Points the hook at the manifest dbt parses from a project tagging its models `staging`, `marts` or nothing.""" + for path, content in TAGGED_PROJECT.items(): + (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / path).write_text(content) + result = dbtRunner().invoke(["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)]) + assert result.success, result.exception + monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + + +def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["staging", "marts"]) + monkeypatch.setattr(resources, "BUNDLE_TESTS", False) # the layout below names each test's own task + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + name = resources.JOB_NAME + assert set(jobs) == {name, f"{name}_staging", f"{name}_marts", f"{name}_default"} + run_jobs = {task["task_key"]: task for task in jobs[name]["tasks"]} + assert {key: task["run_job_task"]["job_id"] for key, task in run_jobs.items()} == { + tag: f"${{resources.jobs.{name}_{tag}.id}}" for tag in ("staging", "marts", "default") + } + assert {key: [d["task_key"] for d in task.get("depends_on", [])] for key, task in run_jobs.items()} == { + "staging": [], + "marts": ["staging"], + "default": [], + } + + # The job keeps the schedule and runs no dbt itself; the per-tag jobs get the compute instead. + assert jobs[name]["trigger"] and not {"environments", "job_clusters"} & set(jobs[name]) + for tag in ("staging", "marts", "default"): + job = jobs[f"{name}_{tag}"] + assert "trigger" not in job + assert job["queue"] == {"enabled": True} + assert job["environments"][0]["spec"]["base_environment"] == resources.SERVERLESS_ENV_PATH + + # Each model runs in the job of its tag and tests follow what they test; no task depends on another job's task. + placed = { + tag: {task["task_key"] for task in jobs[f"{name}_{tag}"]["tasks"]} for tag in ("staging", "marts", "default") + } + assert placed == { + "staging": {"stg_orders_model", "not_null_stg_orders_order_id_test"}, + "marts": {"orders_model"}, + "default": {"legacy_model"}, + } + for tag, keys in placed.items(): + for task in jobs[f"{name}_{tag}"]["tasks"]: + assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + + +def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["no_such_tag"]) + with pytest.raises(ValueError, match="matches no tag"): + resources.build_jobs("dev") + + def test_dependency_pin_rejects_non_pypi_versions(monkeypatch): # A local or dev build cannot be pip-installed from PyPI, so the deploy must fail early # with a clear message. diff --git a/contrib/templates/dbt-factory/README.md b/contrib/templates/dbt-factory/README.md index 639df99b..3807a395 100644 --- a/contrib/templates/dbt-factory/README.md +++ b/contrib/templates/dbt-factory/README.md @@ -89,6 +89,9 @@ $ databricks bundle run _job | `environment_key` | Key of the serverless environment used by the generated job (serverless only). | | `extra_dbt_command_options` | Extra options appended to every generated dbt command. | +To split the job into one job per dbt tag once your models carry tags, set `SPLIT_BY_TAGS` in the +generated `resources/__init__.py` (see the generated project's README, "Splitting into several jobs"). + ### Compute The generated job supports two compute types, chosen by the `use_serverless` prompt: diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/NOTICE b/contrib/templates/dbt-factory/template/{{.project_name}}/NOTICE index 702735ef..dda67852 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/NOTICE +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/NOTICE @@ -2,7 +2,7 @@ This project includes code adapted from the "databricks-dbt-factory" library. Source: https://github.com/mwojtyczka/databricks-dbt-factory Full history: https://github.com/mwojtyczka/databricks-dbt-factory/commits/main - Adapted from: commit 29866edb82de28c33fe888be67d451068d4d7f94 (v0.3.5) + Adapted from: commit c830e20e9fb5196d7a10c42cf054cfd4d91eba46 (v0.4.0) Location: src/databricks_dbt_factory/ The code under src/databricks_dbt_factory/ originates from that repository (trimmed to the diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl index 077baaad..90c4c8f6 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl @@ -86,7 +86,7 @@ are selected per bundle target via `--target ${bundle.target}`. **You must size the cluster yourself by editing a file — there is no `bundle init` prompt for it.** The job cluster ships with default settings; to change them, open [`resources/__init__.py`](resources/__init__.py) and edit the `job_clusters` block inside -`build_job`: +`_job`: ```python "new_cluster": { @@ -103,7 +103,8 @@ models as SQL to the **SQL warehouse** set by `http_path` in `dbt_profiles/profi warehouse — not this cluster — executes the queries. Because the tasks are I/O-bound (submit SQL, wait), the default is a **single-node** cluster: adding workers does **not** speed dbt up. Bump `node_type_id` only if orchestration itself needs more headroom, and size the **SQL warehouse** for -your query load. All tasks share this one cluster for the whole run. +your query load. All tasks of a job share this one cluster for the whole run; with `SPLIT_BY_TAGS` +(below), each per-tag job gets its own. **When to do this:** the cluster is created from this block at `databricks bundle deploy` time (not at `bundle init`). So edit it **before your first deploy** (`make deploy` / `databricks bundle @@ -129,6 +130,24 @@ the version there, re-run `make setup`, and the next deploy uses it — so what matches what you tested with. {{- end}} +### Splitting into several jobs + +To split the job into one job per dbt tag, list the tags in `SPLIT_BY_TAGS` in +`resources/__init__.py`, each an exact tag or a glob such as `"domain_*"`: + +```python +SPLIT_BY_TAGS = ["staging", "marts"] +``` + +Each resource then runs in the `{{.project_name}}_job_` job of its tag, or in +`{{.project_name}}_job_default` if it carries none, and tests run in the job of what they test. +`{{.project_name}}_job` becomes the job that runs those jobs in dependency order: it keeps the +schedule, and each per-tag job gets the {{if eq .use_serverless "no"}}job cluster{{else}}serverless environment{{end}}. +Use it when a project outgrows the 1,000-task limit of a job. It works like the `--split-by-tag` +option of the databricks-dbt-factory CLI; see +[Splitting into several jobs](https://github.com/mwojtyczka/databricks-dbt-factory#splitting-into-several-jobs) +for the full rules, including which splits are refused. + ## Tests ``` diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl index 90a9c717..627e0704 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl @@ -3,8 +3,9 @@ PyDABs integration for databricks-dbt-factory. At ``databricks bundle deploy`` time the Databricks CLI calls :func:`load_resources`, which reads the dbt manifest and expands it into a Databricks job with one task per dbt object -(model / seed / snapshot / test) using the vendored ``databricks_dbt_factory`` core. No -per-model YAML is generated or checked in — the job graph is built on the fly from the manifest. +(model / seed / snapshot / test), optionally split into one job per dbt tag, using the vendored +``databricks_dbt_factory`` core. No per-model YAML is generated or checked in — the job graph is +built on the fly from the manifest. This module is the only integration glue: it configures the vendored ``databricks_dbt_factory`` core (under ``src/databricks_dbt_factory``) for {{if eq .use_serverless "no"}}notebook tasks on a job cluster{{else}}serverless notebook tasks{{end}} and returns a @@ -22,6 +23,7 @@ from databricks.bundles.jobs import Job from databricks_dbt_factory.dbt_factory import DbtFactory from databricks_dbt_factory.dbt_task import DbtTaskOptions, TaskType +from databricks_dbt_factory.job_partition import JobPartition from databricks_dbt_factory.task_factory import ( DbtDependencyResolver, ModelTaskFactory, @@ -60,10 +62,16 @@ BUNDLE_TESTS = {{if eq .bundle_tests "yes"}}True{{else}}False{{end}} # `printf "%q"` renders a safely-quoted string literal so values containing quotes stay valid Python. EXTRA_DBT_COMMAND_OPTIONS = {{ printf "%q" .extra_dbt_command_options }} +# Split the job into one job per dbt tag matching these patterns, each an exact tag or a glob such as +# "domain_*"; resources without a matching tag run in a `default` job. The JOB_NAME job then runs those +# jobs in dependency order. Empty keeps one job. Works like the databricks-dbt-factory CLI's +# --split-by-tag; see "Splitting into several jobs" in its README for how tests are placed and what is refused. +SPLIT_BY_TAGS = [] + {{ if eq .use_serverless "yes" -}} # Serverless base-environment file, generated at deploy time from the pinned # dbt-databricks version and synced with the bundle. The generated job points its -# serverless environment at this file (see build_job) so Databricks pre-builds the +# serverless environment at this file (see _job) so Databricks pre-builds the # environment once instead of installing dbt on every task. `${workspace.file_path}` # is substituted by the CLI at deploy time to the bundle's synced files root. SERVERLESS_ENV_FILE = "dbt_serverless_env.yaml" @@ -79,8 +87,8 @@ PROJECT_DIRECTORY = os.path.relpath(".", os.path.dirname(RUNNER_NOTEBOOK_PATH)) PROFILES_DIRECTORY = "dbt_profiles" -def _build_tasks(target: str) -> list[dict]: - """Reads the dbt manifest and returns the list of Databricks task dicts (one per dbt node).""" +def _build_factory(target: str) -> DbtFactory: + """Configures the factory to generate {{if eq .use_serverless "no"}}notebook tasks on a job cluster{{else}}serverless notebook tasks{{end}} that run the given dbt target.""" resolver = DbtDependencyResolver() task_options = DbtTaskOptions( {{- if eq .use_serverless "no"}} @@ -102,10 +110,29 @@ def _build_tasks(target: str) -> list[dict]: "test": TestTaskFactory(resolver, task_options, dbt_options), } - factory = DbtFactory(task_factories, bundle_tests=BUNDLE_TESTS) - manifest = read_dbt_manifest(MANIFEST_PATH) + return DbtFactory(task_factories, bundle_tests=BUNDLE_TESTS) + + +def _build_tasks(target: str) -> list[dict]: + """Reads the dbt manifest and returns the list of Databricks task dicts (one per dbt node).""" +{{- if eq .use_serverless "no"}} + return _add_dbt_library(_build_factory(target).create_tasks(read_dbt_manifest(MANIFEST_PATH))) +{{- else}} + return _build_factory(target).create_tasks(read_dbt_manifest(MANIFEST_PATH)) +{{- end}} + + +def _build_partitions(target: str) -> list[JobPartition]: + """Reads the dbt manifest and returns its tasks split into one job per tag matching SPLIT_BY_TAGS.""" {{- if eq .use_serverless "no"}} - tasks = factory.create_tasks(manifest) + partitions = _build_factory(target).create_partitions(read_dbt_manifest(MANIFEST_PATH), SPLIT_BY_TAGS) + for partition in partitions: + _add_dbt_library(partition.tasks) + return partitions + + +def _add_dbt_library(tasks: list[dict]) -> list[dict]: + """Adds dbt-databricks to the libraries of every task, and returns the tasks.""" # A job cluster does not get the serverless base environment, so dbt-databricks must be # installed on the cluster. Listing it as a library on every task installs it once on the # shared job cluster (idempotent) and guarantees it is present no matter which task the @@ -115,7 +142,7 @@ def _build_tasks(target: str) -> list[dict]: task["libraries"] = [library] return tasks {{- else}} - return factory.create_tasks(manifest) + return _build_factory(target).create_partitions(read_dbt_manifest(MANIFEST_PATH), SPLIT_BY_TAGS) {{- end}} @@ -185,82 +212,109 @@ def _write_serverless_environment_file() -> None: {{ end -}} -def build_job(target: str) -> Job: +# How SPLIT_BY_TAGS shares out the job settings, as the databricks-dbt-factory CLI's --split-by-tag does: +# the JOB_NAME job keeps the schedule, triggers, notifications and health rules, the per-tag jobs keep +# the compute, and all of them keep everything else (such as the queue). +_ORCHESTRATOR_ONLY = frozenset( + { + "schedule", + "trigger", + "continuous", + "email_notifications", + "webhook_notifications", + "notification_settings", + "health", + } +) +_SPLIT_JOB_ONLY = frozenset({"environments", "job_clusters"}) + + +def _job(name: str, tasks: list[dict], without: frozenset[str] = frozenset()) -> Job: + """Builds a job that runs `tasks` with the settings below, leaving out those named in `without`.""" + settings = { + "name": name, + "queue": {"enabled": True}, + "trigger": { + # Run this job every day, exactly one day from the last run; + # see https://docs.databricks.com/api/workspace/jobs/create#trigger + "periodic": {"interval": 1, "unit": "DAYS"}, + }, + "tasks": tasks, {{- if eq .use_serverless "no"}} - """Builds the Databricks job that runs the dbt project on a job cluster via notebook tasks.""" - return Job.from_dict( - { - "name": JOB_NAME, - "queue": {"enabled": True}, - "trigger": { - # Run this job every day, exactly one day from the last run; - # see https://docs.databricks.com/api/workspace/jobs/create#trigger - "periodic": {"interval": 1, "unit": "DAYS"}, - }, - "tasks": _build_tasks(target), - # The job cluster every task runs on. It only orchestrates — each task runs the runner - # notebook, which submits SQL to the warehouse in dbt_profiles/profiles.yml and waits — - # so it is a single-node cluster by default: extra workers do not speed dbt up (the - # warehouse runs the queries). Bump node_type_id if orchestration itself needs more - # headroom, and scale the SQL warehouse for query load. Changes take effect on the next - # `databricks bundle deploy`. - "job_clusters": [ - { - "job_cluster_key": JOB_CLUSTER_KEY, - "new_cluster": { - # `spark_version` holds the Databricks Runtime version. "18.x-scala2.13" is - # DBR 18 LTS; its ".x" auto-updates to the latest maintenance release. Use a - # specific release (e.g. "18.2.x-scala2.13") only if you want to pin it. - "spark_version": "18.x-scala2.13", - "node_type_id": "{{smallest_node_type}}", - # Single node (driver only): no workers, since the tasks are I/O-bound - # orchestration, not Spark compute. - "num_workers": 0, - "spark_conf": { - "spark.databricks.cluster.profile": "singleNode", - "spark.master": "local[*]", - }, - "custom_tags": {"ResourceClass": "SingleNode"}, + # The job cluster every task runs on. It only orchestrates — each task runs the runner + # notebook, which submits SQL to the warehouse in dbt_profiles/profiles.yml and waits — + # so it is a single-node cluster by default: extra workers do not speed dbt up (the + # warehouse runs the queries). Bump node_type_id if orchestration itself needs more + # headroom, and scale the SQL warehouse for query load. Changes take effect on the next + # `databricks bundle deploy`. With SPLIT_BY_TAGS, each per-tag job gets its own cluster. + "job_clusters": [ + { + "job_cluster_key": JOB_CLUSTER_KEY, + "new_cluster": { + # `spark_version` holds the Databricks Runtime version. "18.x-scala2.13" is + # DBR 18 LTS; its ".x" auto-updates to the latest maintenance release. Use a + # specific release (e.g. "18.2.x-scala2.13") only if you want to pin it. + "spark_version": "18.x-scala2.13", + "node_type_id": "{{smallest_node_type}}", + # Single node (driver only): no workers, since the tasks are I/O-bound + # orchestration, not Spark compute. + "num_workers": 0, + "spark_conf": { + "spark.databricks.cluster.profile": "singleNode", + "spark.master": "local[*]", }, - } - ], - } - ) + "custom_tags": {"ResourceClass": "SingleNode"}, + }, + } + ], {{- else}} - """Builds the Databricks job that runs the dbt project on serverless via notebook tasks.""" - return Job.from_dict( - { - "name": JOB_NAME, - "queue": {"enabled": True}, - "trigger": { - # Run this job every day, exactly one day from the last run; - # see https://docs.databricks.com/api/workspace/jobs/create#trigger - "periodic": {"interval": 1, "unit": "DAYS"}, - }, - "tasks": _build_tasks(target), - "environments": [ - { - "environment_key": ENVIRONMENT_KEY, - # Pre-built base env (synced with the bundle), not inline deps — - # the two are mutually exclusive; the file pins dbt-databricks. - "spec": {"base_environment": SERVERLESS_ENV_PATH}, - } - ], - } - ) + "environments": [ + { + "environment_key": ENVIRONMENT_KEY, + # Pre-built base env (synced with the bundle), not inline deps — + # the two are mutually exclusive; the file pins dbt-databricks. + "spec": {"base_environment": SERVERLESS_ENV_PATH}, + } + ], {{- end}} + } + return Job.from_dict({key: value for key, value in settings.items() if key not in without}) + + +def build_jobs(target: str) -> dict[str, Job]: + """ + Builds the Databricks jobs that run the dbt project on {{if eq .use_serverless "no"}}a job cluster{{else}}serverless{{end}} via notebook tasks, by resource key. + + That is the JOB_NAME job, or with SPLIT_BY_TAGS one `_` job per tag plus the JOB_NAME + job, which runs each of them once the jobs it depends on have succeeded. + """ + if not SPLIT_BY_TAGS: + return {JOB_NAME: _job(JOB_NAME, _build_tasks(target))} + jobs, run_jobs = {}, [] + for partition in _build_partitions(target): + key = f"{JOB_NAME}_{partition.name}" + jobs[key] = _job(key, partition.tasks, without=_ORCHESTRATOR_ONLY) + run_jobs.append( + { + "task_key": partition.name, + "depends_on": [{"task_key": upstream} for upstream in partition.depends_on], + "run_job_task": {"job_id": f"${{`{{`}}resources.jobs.{key}.id{{`}}`}}"}, + } + ) + return {JOB_NAME: _job(JOB_NAME, run_jobs, without=_SPLIT_JOB_ONLY), **jobs} def load_resources(bundle: Bundle) -> Resources: """ Called by the Databricks CLI during `bundle deploy` to load resources defined in Python. - Reads the dbt manifest{{if eq .use_serverless "no"}} and registers the generated job{{else}}, writes the serverless base-environment file (synced with the bundle), - and registers the generated job{{end}}. After deployment this function is not used. + Reads the dbt manifest{{if eq .use_serverless "no"}} and registers the generated jobs{{else}}, writes the serverless base-environment file (synced with the bundle), + and registers the generated jobs{{end}}. After deployment this function is not used. """ {{- if eq .use_serverless "yes"}} _write_serverless_environment_file() {{- end}} resources = Resources() - resources.add_job(JOB_NAME, build_job(bundle.target)) + for key, job in build_jobs(bundle.target).items(): + resources.add_job(key, job) return resources diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/__version__.py b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/__version__.py index a8d4557d..6a9beea8 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/__version__.py +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/__version__.py @@ -1 +1 @@ -__version__ = "0.3.5" +__version__ = "0.4.0" diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/dbt_factory.py b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/dbt_factory.py index 39cc0788..a26a6bc7 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/dbt_factory.py +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/dbt_factory.py @@ -1,12 +1,13 @@ import heapq -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from dataclasses import dataclass, field, replace from pathlib import PurePosixPath, PureWindowsPath from typing import cast from databricks_dbt_factory.dbt_task import DbtTask +from databricks_dbt_factory.job_partition import JobPartition, TaskPlacement, partition_tasks from databricks_dbt_factory.task_factory import TaskFactory, TestTaskFactory -from databricks_dbt_factory.utils import DYNAMIC_VALUE_REFERENCE, build_task_key_maps +from databricks_dbt_factory.utils import DYNAMIC_VALUE_REFERENCE, MAX_JOB_TASKS, build_task_key_maps # The `unique_id` prefixes of resources a test can be attached to. _DBT_TEST_TARGET_PREFIXES = ("model.", "seed.", "snapshot.", "source.") @@ -108,14 +109,18 @@ def create_tasks(self, dbt_manifest: dict) -> list[dict]: Raises: ValueError: If the generated job would exceed Databricks' 1,000-task limit. """ - tasks = self._create_tasks(dbt_manifest) - if len(tasks) > 1_000: + tasks, _ = self._create_tasks(dbt_manifest) + if len(tasks) > MAX_JOB_TASKS: raise ValueError( - f"Databricks jobs support at most 1,000 tasks; this manifest generates {len(tasks):,}. " - f"Reduce the generated resources or enable test bundling." + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; this manifest generates {len(tasks):,}. " + f"Reduce the generated resources, enable test bundling, or split the job by tag." ) return [task.to_dict() for task in tasks] + def create_partitions(self, dbt_manifest: dict, split_by_tags: Sequence[str]) -> list[JobPartition]: + """Generates the tasks split into one job per dbt tag matching `split_by_tags`; see `partition_tasks`.""" + return partition_tasks(*self._create_tasks(dbt_manifest), split_by_tags) + _GATEABLE_TYPES = frozenset({"model", "seed", "snapshot"}) _DBT_TEST_TARGET_PREFIXES = _DBT_TEST_TARGET_PREFIXES @@ -678,7 +683,7 @@ def _unaddressable(node_info: dict) -> ValueError: f"resources, so the task could run another task's resource." ) - def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: + def _create_tasks(self, dbt_manifest: dict) -> tuple[list[DbtTask], dict[str, TaskPlacement]]: """ Builds `DbtTask` instances from the manifest, applying the bundling and gating policies. @@ -686,7 +691,8 @@ def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: dbt_manifest (dict): Parsed dbt manifest content. Returns: - list[DbtTask]: `DbtTask` instances (not yet rendered to dicts). + tuple[list[DbtTask], dict[str, TaskPlacement]]: `DbtTask` instances (not yet rendered to dicts), and + the placement of each by task key. """ dbt_nodes = self._enabled_only(dbt_manifest.get("nodes", {})) dbt_sources = self._enabled_only(dbt_manifest.get("sources", {})) @@ -752,7 +758,23 @@ def _create_tasks(self, dbt_manifest: dict) -> list[DbtTask]: elif "test" in self.task_factories: tasks.extend(self._build_unit_test_tasks(dbt_unit_tests, task_keys, peers)) - return tasks + return tasks, self._placements(peers, task_keys, bundled_test_keys) + + @staticmethod + def _placements( + peers: dict, task_keys: dict[str, str], bundled_test_keys: dict[str, str] + ) -> dict[str, TaskPlacement]: + """Places each task by its resource's tags; tests, and the bundles running them, follow their dependencies.""" + bundles = set(bundled_test_keys.values()) + placements = {} + for full_name, task_key in (*task_keys.items(), *bundled_test_keys.items()): + info = peers[full_name] + placements[task_key] = TaskPlacement( + f"{info.get('resource_type')} {info.get('name')!r} ({info.get('original_file_path')})", + frozenset(info.get("tags") or ()), + task_key in bundles or info["resource_type"] in {"test", "unit_test"}, + ) + return placements @staticmethod def _enabled_only(entries: dict) -> dict: diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/job_partition.py b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/job_partition.py new file mode 100644 index 00000000..d96a693c --- /dev/null +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/job_partition.py @@ -0,0 +1,132 @@ +import re +from collections.abc import Sequence +from dataclasses import dataclass, replace +from fnmatch import fnmatchcase +from graphlib import CycleError, TopologicalSorter + +from databricks_dbt_factory.dbt_task import DbtTask +from databricks_dbt_factory.utils import MAX_JOB_TASKS, MAX_TASK_KEY_LENGTH + +DEFAULT_JOB = "default" + + +@dataclass(frozen=True) +class JobPartition: + """One job of a split project: its tasks, and the jobs that must succeed before it runs.""" + + name: str + tasks: list[dict] + depends_on: list[str] + + +@dataclass(frozen=True) +class TaskPlacement: + """The resource a task runs, its tags, and whether the task follows its dependencies instead of its tags.""" + + resource: str + tags: frozenset[str] + follows_dependencies: bool + + +def partition_tasks( + tasks: list[DbtTask], placements: dict[str, TaskPlacement], patterns: Sequence[str] +) -> list[JobPartition]: + """ + Splits tasks into one job per tag matching `patterns`, ordered so each job runs after those it depends on. + + A task joins the job of the one matching tag it carries, or the default job. A task that follows its + dependencies joins the last of their jobs, so tests run after every resource they test. + + Raises: + ValueError: when a resource carries several matching tags, a pattern matches no tag, two tags name the + same job, jobs would depend on each other, or a job exceeds Databricks' task limit. + """ + dependencies = {task.task_key: list(task.depends_on or []) for task in tasks} + tags = { + task.task_key: _split_tag(placements[task.task_key], patterns) + for task in tasks + if not (placements[task.task_key].follows_dependencies and dependencies[task.task_key]) + } + owner = _job_names(tags, patterns) + anchors = _anchors(dependencies, owner) + order = _job_order(owner, dependencies, anchors) + rank = {job: index for index, job in enumerate(order)} + for key in dependencies.keys() - owner.keys(): + owner[key] = max((owner[anchor] for anchor in anchors[key]), key=rank.__getitem__) + + jobs: dict[str, list[DbtTask]] = {job: [] for job in order} + upstream: dict[str, set[str]] = {job: set() for job in order} + for task in tasks: + job = owner[task.task_key] + jobs[job].append(replace(task, depends_on=[d for d in dependencies[task.task_key] if owner[d] == job])) + upstream[job].update(owner[d] for d in dependencies[task.task_key] if owner[d] != job) + for job, job_tasks in jobs.items(): + if len(job_tasks) > MAX_JOB_TASKS: + raise ValueError( + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; the {job!r} job generates " + f"{len(job_tasks):,}. Split it further by tag." + ) + if len(order) > MAX_JOB_TASKS: + raise ValueError( + f"Databricks jobs support at most {MAX_JOB_TASKS:,} tasks; the tags split into {len(order):,} jobs. " + f"Split by fewer tags." + ) + return [JobPartition(job, [task.to_dict() for task in jobs[job]], sorted(upstream[job])) for job in order] + + +def _split_tag(placement: TaskPlacement, patterns: Sequence[str]) -> str | None: + matched = sorted(tag for tag in placement.tags if any(fnmatchcase(tag, pattern) for pattern in patterns)) + if len(matched) > 1: + raise ValueError( + f"Cannot split by tag: {placement.resource} carries the split tags {', '.join(matched)}. " + f"Keep one, so it belongs to one job." + ) + return matched[0] if matched else None + + +def _job_names(tags: dict[str, str | None], patterns: Sequence[str]) -> dict[str, str]: + used = set(tags.values()) + for pattern in patterns: + if not any(tag is not None and fnmatchcase(tag, pattern) for tag in used): + raise ValueError( + f"--split-by-tag {pattern!r} matches no tag of a resource to split; tests run in the job of the " + f"resources they test, whatever their own tags." + ) + names = {tag: DEFAULT_JOB if tag is None else re.sub(r"[^A-Za-z0-9_]", "_", tag) for tag in used} + owners: dict[str, str | None] = {} + for tag in sorted(used, key=str): + if len(names[tag]) > MAX_TASK_KEY_LENGTH: + raise ValueError( + f"Cannot split by tag: tag {tag!r} is longer than the {MAX_TASK_KEY_LENGTH} characters Databricks " + f"allows in a task key. Use a shorter tag." + ) + if owners.setdefault(names[tag], tag) != tag: + labels = [f"tag {t!r}" if t is not None else "untagged resources" for t in (owners[names[tag]], tag)] + raise ValueError(f"Cannot split by tag: {' and '.join(labels)} would both run as job {names[tag]!r}.") + return {key: names[tag] for key, tag in tags.items()} + + +def _anchors(dependencies: dict[str, list[str]], placed: dict[str, str]) -> dict[str, set[str]]: + anchors: dict[str, set[str]] = {} + for key in TopologicalSorter(dependencies).static_order(): + anchors[key] = {key} if key in placed else set().union(*(anchors[d] for d in dependencies[key])) + return anchors + + +def _job_order(owner: dict[str, str], dependencies: dict[str, list[str]], anchors: dict[str, set[str]]) -> list[str]: + upstream: dict[str, set[str]] = {job: set() for job in owner.values()} + witness: dict[tuple[str, str], tuple[str, str]] = {} + for key, job in owner.items(): + for source in (anchor for d in dependencies[key] for anchor in anchors[d]): + if owner[source] != job: + upstream[job].add(owner[source]) + witness.setdefault((owner[source], job), (key, source)) + try: + return list(TopologicalSorter({job: sorted(upstream[job]) for job in sorted(upstream)}).static_order()) + except CycleError as error: + cycle = error.args[1] + needs = "; ".join(f"{key} needs {source}" for key, source in (witness[p] for p in zip(cycle, cycle[1:]))) + raise ValueError( + f"Cannot split by tag: jobs {' -> '.join(cycle)} would depend on each other ({needs}). " + f"Retag the resources so dependencies between jobs flow one way." + ) from error diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/notebook/run_dbt_command.py b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/notebook/run_dbt_command.py index accb8f6a..9742b649 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/notebook/run_dbt_command.py +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/notebook/run_dbt_command.py @@ -68,9 +68,10 @@ if os.path.exists(prebuilt_manifest_path): try: from dbt.contracts.graph.manifest import Manifest + from dbt.parser.manifest import extended_mashumuro_decoder with open(prebuilt_manifest_path, "rb") as f: - manifest = Manifest.from_msgpack(f.read()) + manifest = Manifest.from_msgpack(f.read(), decoder=extended_mashumuro_decoder) manifest.build_flat_graph() print(f"[dbt-factory] injecting pre-built manifest from {prebuilt_manifest_path} (skipping dbt parse)") except Exception as e: diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/utils.py b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/utils.py index 02ff51a2..fd675205 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/utils.py +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/src/databricks_dbt_factory/utils.py @@ -9,6 +9,9 @@ # Databricks caps task keys at 100 characters (letters, numbers, underscores, hyphens). MAX_TASK_KEY_LENGTH = 100 +# Databricks caps jobs at 1,000 tasks. +MAX_JOB_TASKS = 1_000 + # dbt resource types whose task key is `_` (the type doubles as the suffix). _SUFFIXED_TYPES = frozenset({"model", "seed", "snapshot"}) diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl index ad4a7267..8382335d 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl @@ -3,7 +3,8 @@ Offline test for the PyDABs integration in resources/__init__.py. Exercises the whole glue path against the committed dbt manifest (no Databricks workspace required): reading the manifest, wiring the factory, generating notebook tasks, and registering -the job on a Resources object (which runs Job.from_dict, validating the task dict schema). +the job on a Resources object (which runs Job.from_dict, validating the task dict schema). The +split by tag runs against a small tagged project that dbt parses offline. """ import json @@ -15,6 +16,7 @@ import pytest import yaml {{- end}} from databricks.bundles.core import Bundle +from dbt.cli.main import dbtRunner import resources from databricks_dbt_factory.utils import generate_task_key @@ -35,6 +37,13 @@ def _chdir_to_project_root(monkeypatch): monkeypatch.chdir(PROJECT_ROOT) +def _jobs_running_dbt(jobs: dict) -> list: + """The generated jobs whose tasks run dbt: the job itself, or with SPLIT_BY_TAGS the per-tag jobs it runs.""" + dbt_jobs = [job for job in jobs.values() if any(task.notebook_task for task in job.tasks)] + assert dbt_jobs, "no generated job runs dbt" + return dbt_jobs + + def test_build_tasks_covers_every_buildable_node(): # Every model/seed/snapshot node gets its own run task regardless of the BUNDLE_TESTS mode # (tests may be collapsed into per-resource `_test` tasks when bundling is on). @@ -78,16 +87,16 @@ def test_tasks_pin_dbt_databricks_as_a_library(): def test_job_cluster_is_defined_on_the_job(): - job = resources.build_job("dev") - assert len(job.job_clusters) == 1 - assert job.job_clusters[0].job_cluster_key == resources.JOB_CLUSTER_KEY + for job in _jobs_running_dbt(resources.build_jobs("dev")): + assert len(job.job_clusters) == 1 + assert job.job_clusters[0].job_cluster_key == resources.JOB_CLUSTER_KEY def test_load_resources_registers_the_job(): result = resources.load_resources(Bundle(target="dev")) assert resources.JOB_NAME in result.jobs - job = result.jobs[resources.JOB_NAME] - assert job.job_clusters[0].job_cluster_key == resources.JOB_CLUSTER_KEY + for job in _jobs_running_dbt(result.jobs): + assert job.job_clusters[0].job_cluster_key == resources.JOB_CLUSTER_KEY {{- else -}} def test_generated_tasks_are_serverless_notebook_tasks(): tasks = resources._build_tasks("dev") @@ -109,16 +118,16 @@ def test_environment_pins_installed_dbt_databricks(): assert spec["dependencies"] == [f"dbt-databricks=={version('dbt-databricks')}"] # The generated job points its serverless environment at the synced env file. - job = resources.build_job("dev") - assert job.environments[0].spec.base_environment == resources.SERVERLESS_ENV_PATH + for job in _jobs_running_dbt(resources.build_jobs("dev")): + assert job.environments[0].spec.base_environment == resources.SERVERLESS_ENV_PATH def test_load_resources_registers_the_job(): result = resources.load_resources(Bundle(target="dev")) try: assert resources.JOB_NAME in result.jobs - job = result.jobs[resources.JOB_NAME] - assert job.environments[0].environment_key == resources.ENVIRONMENT_KEY + for job in _jobs_running_dbt(result.jobs): + assert job.environments[0].environment_key == resources.ENVIRONMENT_KEY # load_resources writes the base-environment file to the bundle root for # the bundle to sync; it pins dbt-databricks to the installed version. @@ -132,6 +141,83 @@ def test_load_resources_registers_the_job(): {{- end}} +TAGGED_PROJECT = { + "dbt_project.yml": 'name: probe\nprofile: probe\nversion: "1.0"\nconfig-version: 2\n', + "profiles.yml": ( + "probe:\n target: dev\n outputs:\n dev:\n type: databricks\n host: example.databricks.com\n" + " http_path: /sql/1.0/warehouses/x\n token: dummy\n schema: default\n" + ), + "models/stg_orders.sql": "{{`{{`}} config(tags=['staging']) {{`}}`}}\nselect 1 as order_id", + "models/orders.sql": "{{`{{`}} config(tags=['marts']) {{`}}`}}\nselect * from {{`{{`}} ref('stg_orders') {{`}}`}}", + "models/legacy.sql": "select 1 as id", + "models/schema.yml": ( + "version: 2\nmodels:\n - name: stg_orders\n columns:\n - name: order_id\n data_tests: [not_null]\n" + ), +} + + +@pytest.fixture +def tagged_manifest(tmp_path, monkeypatch): + """Points the hook at the manifest dbt parses from a project tagging its models `staging`, `marts` or nothing.""" + for path, content in TAGGED_PROJECT.items(): + (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / path).write_text(content) + result = dbtRunner().invoke(["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)]) + assert result.success, result.exception + monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + + +def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["staging", "marts"]) + monkeypatch.setattr(resources, "BUNDLE_TESTS", False) # the layout below names each test's own task + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + name = resources.JOB_NAME + assert set(jobs) == {name, f"{name}_staging", f"{name}_marts", f"{name}_default"} + run_jobs = {task["task_key"]: task for task in jobs[name]["tasks"]} + assert {key: task["run_job_task"]["job_id"] for key, task in run_jobs.items()} == { + tag: f"${{`{{`}}resources.jobs.{name}_{tag}.id{{`}}`}}" for tag in ("staging", "marts", "default") + } + assert {key: [d["task_key"] for d in task.get("depends_on", [])] for key, task in run_jobs.items()} == { + "staging": [], + "marts": ["staging"], + "default": [], + } + + # The job keeps the schedule and runs no dbt itself; the per-tag jobs get the compute instead. + assert jobs[name]["trigger"] and not {"environments", "job_clusters"} & set(jobs[name]) + for tag in ("staging", "marts", "default"): + job = jobs[f"{name}_{tag}"] + assert "trigger" not in job + assert job["queue"] == {"enabled": True} +{{- if eq .use_serverless "no"}} + assert job["job_clusters"][0]["job_cluster_key"] == resources.JOB_CLUSTER_KEY + for task in job["tasks"]: + assert task["libraries"] == [{"pypi": {"package": f"dbt-databricks=={version('dbt-databricks')}"}}] +{{- else}} + assert job["environments"][0]["spec"]["base_environment"] == resources.SERVERLESS_ENV_PATH +{{- end}} + + # Each model runs in the job of its tag and tests follow what they test; no task depends on another job's task. + placed = { + tag: {task["task_key"] for task in jobs[f"{name}_{tag}"]["tasks"]} for tag in ("staging", "marts", "default") + } + assert placed == { + "staging": {"stg_orders_model", "not_null_stg_orders_order_id_test"}, + "marts": {"orders_model"}, + "default": {"legacy_model"}, + } + for tag, keys in placed.items(): + for task in jobs[f"{name}_{tag}"]["tasks"]: + assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + + +def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["no_such_tag"]) + with pytest.raises(ValueError, match="matches no tag"): + resources.build_jobs("dev") + + def test_dependency_pin_rejects_non_pypi_versions(monkeypatch): # A local or dev build cannot be pip-installed from PyPI, so the deploy must fail early # with a clear message. From 9e277addc19f164369179b6fab5765bb1c546517 Mon Sep 17 00:00:00 2001 From: peterpark-db Date: Fri, 23 Oct 2026 12:20:18 -0700 Subject: [PATCH 2/2] contrib/dbt-factory: fix split-by-tag docs Replace stale `build_job` references with `_job`, document sanitized per-tag job names and that a failed job stops every downstream job, and test split placement with bundled tests, globs, sanitized tag names, and the refused splits (two split tags, job cycle, colliding job names). --- contrib/dbt_factory/README.md | 7 +- .../dbt_factory/tests/test_load_resources.py | 113 +++++++++++++++--- .../template/{{.project_name}}/README.md.tmpl | 12 +- .../resources/__init__.py.tmpl | 2 +- .../tests/test_load_resources.py.tmpl | 113 +++++++++++++++--- 5 files changed, 208 insertions(+), 39 deletions(-) diff --git a/contrib/dbt_factory/README.md b/contrib/dbt_factory/README.md index aebabfb0..701404e2 100644 --- a/contrib/dbt_factory/README.md +++ b/contrib/dbt_factory/README.md @@ -155,8 +155,11 @@ A few knobs are exposed as constants at the top of `resources/__init__.py`: * `EXTRA_DBT_COMMAND_OPTIONS` — extra options appended to every generated dbt command. * `SPLIT_BY_TAGS` — split the job into one job per dbt tag listed here (exact tags or globs such as `"domain_*"`). Each resource runs in the `dbt_factory_job_` job of its tag, or in - `dbt_factory_job_default` if it carries none, and tests run in the job of what they test. - `dbt_factory_job` then runs those jobs in dependency order and keeps the schedule. Default `[]` + `dbt_factory_job_default` if it carries none, and tests run in the job of what they test. In job + names, characters other than letters, digits and `_` in a tag become `_` (`finance-core` runs in + `dbt_factory_job_finance_core`). `dbt_factory_job` then runs those jobs in dependency order and keeps + the schedule. A job waits for whole upstream jobs, not just for the models it needs, so when any task + of a job fails, none of the jobs downstream of it run. Default `[]` (one job). Use it when a project outgrows the 1,000-task limit of a job. It works like the CLI's `--split-by-tag`; see [Splitting into several jobs](https://github.com/mwojtyczka/databricks-dbt-factory#splitting-into-several-jobs) for the full rules, including which splits are refused. diff --git a/contrib/dbt_factory/tests/test_load_resources.py b/contrib/dbt_factory/tests/test_load_resources.py index ee739059..b4716b8a 100644 --- a/contrib/dbt_factory/tests/test_load_resources.py +++ b/contrib/dbt_factory/tests/test_load_resources.py @@ -113,19 +113,42 @@ def test_load_resources_registers_the_job(): "models/legacy.sql": "select 1 as id", "models/schema.yml": ( "version: 2\nmodels:\n - name: stg_orders\n columns:\n - name: order_id\n data_tests: [not_null]\n" + " - name: orders\n columns:\n - name: order_id\n data_tests:\n - relationships:\n" + " arguments:\n to: ref('stg_orders')\n field: order_id\n" ), } +RELATIONSHIPS_TEST = "relationships_orders_order_id__order_id__ref_stg_orders__test" @pytest.fixture -def tagged_manifest(tmp_path, monkeypatch): +def parse_tagged_project(tmp_path, monkeypatch): + """Returns a function that points the hook at the manifest dbt parses from TAGGED_PROJECT with `changes` applied.""" + + def parse(changes: dict[str, str] | None = None) -> None: + for path, content in {**TAGGED_PROJECT, **(changes or {})}.items(): + (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / path).write_text(content) + args = ["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)] + result = dbtRunner().invoke(args) + assert result.success, result.exception + monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + + return parse + + +@pytest.fixture +def tagged_manifest(parse_tagged_project): """Points the hook at the manifest dbt parses from a project tagging its models `staging`, `marts` or nothing.""" - for path, content in TAGGED_PROJECT.items(): - (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) - (tmp_path / path).write_text(content) - result = dbtRunner().invoke(["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)]) - assert result.success, result.exception - monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + parse_tagged_project() + + +def _placed_tasks(jobs: dict, tags: tuple[str, ...]) -> dict[str, set[str]]: + """The task keys of each per-tag job, checking that no task depends on a task of another job.""" + placed = {tag: {task["task_key"] for task in jobs[f"{resources.JOB_NAME}_{tag}"]["tasks"]} for tag in tags} + for tag, keys in placed.items(): + for task in jobs[f"{resources.JOB_NAME}_{tag}"]["tasks"]: + assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + return placed def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkeypatch): @@ -153,18 +176,41 @@ def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkey assert job["queue"] == {"enabled": True} assert job["environments"][0]["spec"]["base_environment"] == resources.SERVERLESS_ENV_PATH - # Each model runs in the job of its tag and tests follow what they test; no task depends on another job's task. - placed = { - tag: {task["task_key"] for task in jobs[f"{name}_{tag}"]["tasks"]} for tag in ("staging", "marts", "default") - } - assert placed == { + # Each model runs in the job of its tag, and a test runs in the last job of what it tests. + assert _placed_tasks(jobs, ("staging", "marts", "default")) == { "staging": {"stg_orders_model", "not_null_stg_orders_order_id_test"}, - "marts": {"orders_model"}, + "marts": {"orders_model", RELATIONSHIPS_TEST}, "default": {"legacy_model"}, } - for tag, keys in placed.items(): - for task in jobs[f"{name}_{tag}"]["tasks"]: - assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + + +def test_split_by_tags_runs_bundled_tests_in_the_job_of_their_resource(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["staging", "marts"]) + monkeypatch.setattr(resources, "BUNDLE_TESTS", True) + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + assert _placed_tasks(jobs, ("staging", "marts", "default")) == { + "staging": {"stg_orders_model", "stg_orders_test"}, + "marts": {"orders_model", RELATIONSHIPS_TEST}, + "default": {"legacy_model"}, + } + + +def test_split_by_a_glob_runs_one_job_per_matching_tag(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["sta*", "mar*"]) + name = resources.JOB_NAME + assert set(resources.build_jobs("dev")) == {name, f"{name}_staging", f"{name}_marts", f"{name}_default"} + + +def test_split_job_names_replace_characters_a_task_key_cannot_hold(parse_tagged_project, monkeypatch): + parse_tagged_project({"models/legacy.sql": "{{ config(tags=['finance-core']) }}\nselect 1 as id"}) + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["finance-core"]) + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + name = resources.JOB_NAME + assert set(jobs) == {name, f"{name}_finance_core", f"{name}_default"} + assert {task["task_key"] for task in jobs[name]["tasks"]} == {"finance_core", "default"} + assert _placed_tasks(jobs, ("finance_core",)) == {"finance_core": {"legacy_model"}} def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypatch): @@ -173,6 +219,41 @@ def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypat resources.build_jobs("dev") +@pytest.mark.parametrize( + "changes, split_by_tags, error", + [ + pytest.param( + {"models/legacy.sql": "{{ config(tags=['staging', 'marts']) }}\nselect 1 as id"}, + ["staging", "marts"], + "carries the split tags marts, staging", + id="resource-with-two-split-tags", + ), + pytest.param( + {"models/legacy.sql": "{{ config(tags=['staging']) }}\nselect * from {{ ref('orders') }}"}, + ["staging", "marts"], + "would depend on each other", + id="jobs-in-a-cycle", + ), + pytest.param( + { + "models/legacy.sql": "{{ config(tags=['finance_core']) }}\nselect 1 as id", + "models/ledger.sql": "{{ config(tags=['finance-core']) }}\nselect 1 as id", + }, + ["finance*"], + "would both run as job 'finance_core'", + id="tags-naming-the-same-job", + ), + ], +) +def test_split_that_cannot_be_placed_fails_instead_of_deploying( + parse_tagged_project, monkeypatch, changes, split_by_tags, error +): + parse_tagged_project(changes) + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", split_by_tags) + with pytest.raises(ValueError, match=error): + resources.build_jobs("dev") + + def test_dependency_pin_rejects_non_pypi_versions(monkeypatch): # A local or dev build cannot be pip-installed from PyPI, so the deploy must fail early # with a clear message. diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl index 90c4c8f6..b57ba265 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/README.md.tmpl @@ -16,9 +16,9 @@ It combines two pieces: YAML is checked in** — the task graph is generated on the fly from the manifest. {{ if eq .use_serverless "no" -}} -This project runs on a **job cluster** (classic compute) using the **notebook task type**. A single -job cluster is created once per run and **shared by every task**, then torn down when the run -finishes. It is defined in `resources/__init__.py` (the `job_clusters` block in `build_job`) — see +This project runs on a **job cluster** (classic compute) using the **notebook task type**. A job +cluster is created once per job run and **shared by every task of the job**, then torn down when the +run finishes. It is defined in `resources/__init__.py` (the `job_clusters` block in `_job`) — see [Sizing the job cluster](#sizing-the-job-cluster) below. {{- else -}} This project uses **serverless compute** and the **notebook task type** for the fastest task @@ -140,9 +140,13 @@ SPLIT_BY_TAGS = ["staging", "marts"] ``` Each resource then runs in the `{{.project_name}}_job_` job of its tag, or in -`{{.project_name}}_job_default` if it carries none, and tests run in the job of what they test. +`{{.project_name}}_job_default` if it carries none, and tests run in the job of what they test. In job +names, characters other than letters, digits and `_` in a tag become `_`, so tag `finance-core` runs +in `{{.project_name}}_job_finance_core`. `{{.project_name}}_job` becomes the job that runs those jobs in dependency order: it keeps the schedule, and each per-tag job gets the {{if eq .use_serverless "no"}}job cluster{{else}}serverless environment{{end}}. +A job waits for whole upstream jobs, not just for the models it needs: when any task of a job fails, +none of the jobs downstream of it run. Use it when a project outgrows the 1,000-task limit of a job. It works like the `--split-by-tag` option of the databricks-dbt-factory CLI; see [Splitting into several jobs](https://github.com/mwojtyczka/databricks-dbt-factory#splitting-into-several-jobs) diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl index 627e0704..a05a18ec 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/resources/__init__.py.tmpl @@ -43,7 +43,7 @@ MANIFEST_PATH = os.environ.get("DBT_MANIFEST_PATH", "target/manifest.json") JOB_NAME = "{{.project_name}}_job" {{ if eq .use_serverless "no" -}} -# Key of the job cluster (defined in build_job below) that every task runs on. +# Key of the job cluster (defined in _job below) that every task runs on. JOB_CLUSTER_KEY = "dbt_cluster" {{- else -}} # Key of the serverless environment (defined on the job below). diff --git a/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl b/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl index 8382335d..85fe3c09 100644 --- a/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl +++ b/contrib/templates/dbt-factory/template/{{.project_name}}/tests/test_load_resources.py.tmpl @@ -152,19 +152,42 @@ TAGGED_PROJECT = { "models/legacy.sql": "select 1 as id", "models/schema.yml": ( "version: 2\nmodels:\n - name: stg_orders\n columns:\n - name: order_id\n data_tests: [not_null]\n" + " - name: orders\n columns:\n - name: order_id\n data_tests:\n - relationships:\n" + " arguments:\n to: ref('stg_orders')\n field: order_id\n" ), } +RELATIONSHIPS_TEST = "relationships_orders_order_id__order_id__ref_stg_orders__test" @pytest.fixture -def tagged_manifest(tmp_path, monkeypatch): +def parse_tagged_project(tmp_path, monkeypatch): + """Returns a function that points the hook at the manifest dbt parses from TAGGED_PROJECT with `changes` applied.""" + + def parse(changes: dict[str, str] | None = None) -> None: + for path, content in {**TAGGED_PROJECT, **(changes or {})}.items(): + (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / path).write_text(content) + args = ["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)] + result = dbtRunner().invoke(args) + assert result.success, result.exception + monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + + return parse + + +@pytest.fixture +def tagged_manifest(parse_tagged_project): """Points the hook at the manifest dbt parses from a project tagging its models `staging`, `marts` or nothing.""" - for path, content in TAGGED_PROJECT.items(): - (tmp_path / path).parent.mkdir(parents=True, exist_ok=True) - (tmp_path / path).write_text(content) - result = dbtRunner().invoke(["parse", "--quiet", "--project-dir", str(tmp_path), "--profiles-dir", str(tmp_path)]) - assert result.success, result.exception - monkeypatch.setattr(resources, "MANIFEST_PATH", str(tmp_path / "target" / "manifest.json")) + parse_tagged_project() + + +def _placed_tasks(jobs: dict, tags: tuple[str, ...]) -> dict[str, set[str]]: + """The task keys of each per-tag job, checking that no task depends on a task of another job.""" + placed = {tag: {task["task_key"] for task in jobs[f"{resources.JOB_NAME}_{tag}"]["tasks"]} for tag in tags} + for tag, keys in placed.items(): + for task in jobs[f"{resources.JOB_NAME}_{tag}"]["tasks"]: + assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + return placed def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkeypatch): @@ -198,18 +221,41 @@ def test_split_by_tags_runs_one_job_per_tag_from_the_job(tagged_manifest, monkey assert job["environments"][0]["spec"]["base_environment"] == resources.SERVERLESS_ENV_PATH {{- end}} - # Each model runs in the job of its tag and tests follow what they test; no task depends on another job's task. - placed = { - tag: {task["task_key"] for task in jobs[f"{name}_{tag}"]["tasks"]} for tag in ("staging", "marts", "default") - } - assert placed == { + # Each model runs in the job of its tag, and a test runs in the last job of what it tests. + assert _placed_tasks(jobs, ("staging", "marts", "default")) == { "staging": {"stg_orders_model", "not_null_stg_orders_order_id_test"}, - "marts": {"orders_model"}, + "marts": {"orders_model", RELATIONSHIPS_TEST}, "default": {"legacy_model"}, } - for tag, keys in placed.items(): - for task in jobs[f"{name}_{tag}"]["tasks"]: - assert {d["task_key"] for d in task.get("depends_on", [])} <= keys + + +def test_split_by_tags_runs_bundled_tests_in_the_job_of_their_resource(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["staging", "marts"]) + monkeypatch.setattr(resources, "BUNDLE_TESTS", True) + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + assert _placed_tasks(jobs, ("staging", "marts", "default")) == { + "staging": {"stg_orders_model", "stg_orders_test"}, + "marts": {"orders_model", RELATIONSHIPS_TEST}, + "default": {"legacy_model"}, + } + + +def test_split_by_a_glob_runs_one_job_per_matching_tag(tagged_manifest, monkeypatch): + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["sta*", "mar*"]) + name = resources.JOB_NAME + assert set(resources.build_jobs("dev")) == {name, f"{name}_staging", f"{name}_marts", f"{name}_default"} + + +def test_split_job_names_replace_characters_a_task_key_cannot_hold(parse_tagged_project, monkeypatch): + parse_tagged_project({"models/legacy.sql": "{{`{{`}} config(tags=['finance-core']) {{`}}`}}\nselect 1 as id"}) + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", ["finance-core"]) + jobs = {key: job.as_dict() for key, job in resources.build_jobs("dev").items()} + + name = resources.JOB_NAME + assert set(jobs) == {name, f"{name}_finance_core", f"{name}_default"} + assert {task["task_key"] for task in jobs[name]["tasks"]} == {"finance_core", "default"} + assert _placed_tasks(jobs, ("finance_core",)) == {"finance_core": {"legacy_model"}} def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypatch): @@ -218,6 +264,41 @@ def test_split_by_a_tag_no_resource_carries_fails_instead_of_deploying(monkeypat resources.build_jobs("dev") +@pytest.mark.parametrize( + "changes, split_by_tags, error", + [ + pytest.param( + {"models/legacy.sql": "{{`{{`}} config(tags=['staging', 'marts']) {{`}}`}}\nselect 1 as id"}, + ["staging", "marts"], + "carries the split tags marts, staging", + id="resource-with-two-split-tags", + ), + pytest.param( + {"models/legacy.sql": "{{`{{`}} config(tags=['staging']) {{`}}`}}\nselect * from {{`{{`}} ref('orders') {{`}}`}}"}, + ["staging", "marts"], + "would depend on each other", + id="jobs-in-a-cycle", + ), + pytest.param( + { + "models/legacy.sql": "{{`{{`}} config(tags=['finance_core']) {{`}}`}}\nselect 1 as id", + "models/ledger.sql": "{{`{{`}} config(tags=['finance-core']) {{`}}`}}\nselect 1 as id", + }, + ["finance*"], + "would both run as job 'finance_core'", + id="tags-naming-the-same-job", + ), + ], +) +def test_split_that_cannot_be_placed_fails_instead_of_deploying( + parse_tagged_project, monkeypatch, changes, split_by_tags, error +): + parse_tagged_project(changes) + monkeypatch.setattr(resources, "SPLIT_BY_TAGS", split_by_tags) + with pytest.raises(ValueError, match=error): + resources.build_jobs("dev") + + def test_dependency_pin_rejects_non_pypi_versions(monkeypatch): # A local or dev build cannot be pip-installed from PyPI, so the deploy must fail early # with a clear message.