contrib/dbt-factory: port v0.4.0, add split by tag

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

contrib/dbt_factory/README.md

@@ -153,6 +153,16 @@ 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_<tag>` job of its tag, or in
+  `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.
 
 The dbt target, warehouse, catalog, and schema are configured in `dbt_profiles/profiles.yml`
 and selected per bundle target via `--target ${bundle.target}`.

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,38 +173,78 @@ 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_NAME>_<tag>` 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, 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

contrib/dbt_factory/src/databricks_dbt_factory/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.4.0"

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,15 +683,16 @@ 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.
 
         Args:
             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: