Adding job runs support to PyDABs

.nextchanges/bundles/pydabs-job-runs.md

@@ -0,0 +1 @@
+* Added PyDABs (Python) support for `job_runs`: `Resources.add_job_run` and the `job_run_mutator` decorator.

acceptance/bundle/python/job-runs-support/databricks.yml

@@ -0,0 +1,10 @@
+bundle:
+  name: my_project
+
+sync: {paths: []} # don't need to copy files
+
+python:
+  resources:
+    - "resources:load_resources"
+  mutators:
+    - "mutators:update_job_run"

acceptance/bundle/python/job-runs-support/mutators.py

@@ -0,0 +1,14 @@
+from dataclasses import replace
+
+from databricks.bundles.core import job_run_mutator
+from databricks.bundles.job_runs import JobRun
+
+
+@job_run_mutator
+def update_job_run(job_run: JobRun) -> JobRun:
+    assert isinstance(job_run.job_parameters, dict)
+
+    return replace(
+        job_run,
+        job_parameters={**job_run.job_parameters, "mutated": "true"},
+    )

acceptance/bundle/python/job-runs-support/output.txt

@@ -0,0 +1,47 @@
+
+>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json
+{
+  "experimental": {
+    "python": {
+      "mutators": [
+        "mutators:update_job_run"
+      ],
+      "resources": [
+        "resources:load_resources"
+      ]
+    }
+  },
+  "resources": {
+    "job_runs": {
+      "my_run": {
+        "job_id": "${resources.jobs.my_job.id}",
+        "job_parameters": {
+          "environment": "test",
+          "mutated": "true"
+        },
+        "lifecycle": {
+          "triggers": [
+            {
+              "on_bundle_deploy": true
+            }
+          ]
+        }
+      }
+    },
+    "jobs": {
+      "my_job": {
+        "deployment": {
+          "kind": "BUNDLE",
+          "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/my_project/default/state/metadata.json"
+        },
+        "edit_mode": "UI_LOCKED",
+        "format": "MULTI_TASK",
+        "max_concurrent_runs": 1,
+        "name": "My Job",
+        "queue": {
+          "enabled": true
+        }
+      }
+    }
+  }
+}

acceptance/bundle/python/job-runs-support/resources.py

@@ -0,0 +1,22 @@
+from databricks.bundles.core import Resources, Variable
+from databricks.bundles.job_runs import JobRun, JobRunLifecycle, JobRunTrigger
+from databricks.bundles.jobs import Job
+
+
+def load_resources() -> Resources:
+    resources = Resources()
+
+    resources.add_job(
+        "my_job",
+        Job(name="My Job"),
+    )
+    resources.add_job_run(
+        "my_run",
+        JobRun(
+            job_id=Variable(path="resources.jobs.my_job.id", type=int),
+            job_parameters={"environment": "test"},
+            lifecycle=JobRunLifecycle(triggers=[JobRunTrigger(on_bundle_deploy=True)]),
+        ),
+    )
+
+    return resources

acceptance/bundle/python/job-runs-support/script

@@ -0,0 +1,5 @@
+
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq "pick(.experimental.python, .resources)"
+
+rm -fr .databricks __pycache__

acceptance/bundle/python/job-runs-support/test.toml

@@ -0,0 +1,7 @@
+Cloud = false # tests don't interact with APIs
+
+# job_runs are only supported in the current version of the wheel
+EnvMatrix.PYDAB_VERSION = ["current"]
+
+# job_runs are only supported on the direct deployment engine
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

python/codegen/codegen/generated_wiring.py

@@ -56,7 +56,7 @@ def _wired_resources() -> list[_WiredResource]:
         resources.append(
             _WiredResource(
                 class_name=class_name,
-                singular_name=class_name.lower(),
+                singular_name=packages.get_snake_case_name(ref),
                 plural_name=namespace,
                 model_module=packages.get_package(namespace, ref),
             )

python/codegen/codegen/packages.py

@@ -4,6 +4,7 @@
 # All supported resource types and their namespace
 RESOURCE_NAMESPACE = {
     "resources.Job": "jobs",
+    "resources.JobRun": "job_runs",
     "resources.Pipeline": "pipelines",
     "resources.Catalog": "catalogs",
     "resources.Schema": "schemas",
@@ -41,6 +42,11 @@ def get_class_name(ref: str) -> str:
     return RENAMES.get(name, name)
 
 
+def get_snake_case_name(ref: str) -> str:
+    class_name = get_class_name(ref)
+    return re.sub(r"(?<!^)(?=[A-Z])", "_", class_name).lower()
+
+
 def is_resource(ref: str) -> bool:
     return ref in RESOURCE_TYPES
 
@@ -60,7 +66,6 @@ def get_package(namespace: str, ref: str) -> Optional[str]:
     if full_name in PRIMITIVES:
         return None
 
-    [_, name] = full_name.split(".")
-    package_name = re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
+    package_name = get_snake_case_name(ref)
 
     return f"{get_root_package(namespace)}._models.{package_name}"

python/codegen/codegen_tests/test_generated_dataclass.py

@@ -81,3 +81,7 @@ def test_get_class_name_string():
     class_name = packages.get_class_name("#/$defs/string")
 
     assert class_name == "str"
+
+
+def test_get_snake_case_name():
+    assert packages.get_snake_case_name("resources.JobRun") == "job_run"