From 954b51f221747d3e6ea240138875e51ba0e1097e Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Mon, 17 Mar 2025 16:33:29 +0100
Subject: [PATCH 1/2] [Python] Fix issues with multiple mutators

---
 .../python/databricks/bundles/build.py        |  1 +
 .../python/databricks_tests/test_build.py     | 66 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/experimental/python/databricks/bundles/build.py b/experimental/python/databricks/bundles/build.py
index d9ddc626b8..d8a6d947cc 100644
--- a/experimental/python/databricks/bundles/build.py
+++ b/experimental/python/databricks/bundles/build.py
@@ -128,6 +128,7 @@ def _apply_mutators(
                             ("resources", "jobs", resource_name), location
                         )
                     resources.jobs[resource_name] = new_job
+                    job = new_job
             except Exception as exc:
                 mutator_name = mutator.function.__name__
 
diff --git a/experimental/python/databricks_tests/test_build.py b/experimental/python/databricks_tests/test_build.py
index 980cbbe76e..2d311728a2 100644
--- a/experimental/python/databricks_tests/test_build.py
+++ b/experimental/python/databricks_tests/test_build.py
@@ -1,9 +1,11 @@
 import sys
+from dataclasses import replace
 from io import StringIO
 from pathlib import Path
 
 from databricks.bundles.build import (
     _append_resources,
+    _apply_mutators,
     _Args,
     _Conf,
     _load_object,
@@ -20,6 +22,7 @@
     Location,
     Resources,
     Severity,
+    job_mutator,
 )
 from databricks.bundles.jobs import Job
 
@@ -296,3 +299,66 @@ def test_conf_from_dict():
         ],
         venv_path="venv",
     )
+
+
+def test_mutators():
+    bundle = Bundle(target="default")
+    resources = Resources()
+    resources.add_job("job_0", Job(tags={"tag": "value"}))
+
+    @job_mutator
+    def add_first_tag(bundle: Bundle, job: Job) -> Job:
+        tags = bundle.resolve_variable(job.tags)
+
+        return replace(job, tags={"first": "tag", **tags})
+
+    @job_mutator
+    def add_second_tag(bundle: Bundle, job: Job) -> Job:
+        tags = bundle.resolve_variable(job.tags)
+
+        return replace(job, tags={"second": "tag", **tags})
+
+    new_resources, diagnostics = _apply_mutators(
+        bundle=bundle,
+        resources=resources,
+        mutator_functions=[add_first_tag, add_second_tag],
+    )
+
+    # add_second_tag is the last mutator that has modified a job
+    expected_location = Location.from_callable(add_second_tag.function)
+
+    assert not diagnostics.has_error()
+    assert new_resources._locations[("resources", "jobs", "job_0")] == expected_location
+    assert new_resources.jobs["job_0"].tags == {
+        "first": "tag",
+        "second": "tag",
+        "tag": "value",
+    }
+
+
+def test_mutators_unmodified():
+    bundle = Bundle(target="default")
+
+    resources = Resources()
+    resources.add_job("job_0", Job(description="description"))
+
+    @job_mutator
+    def mutator_1(job: Job) -> Job:
+        return replace(job, description="updated description")
+
+    @job_mutator
+    def mutator_2(job: Job) -> Job:
+        return job
+
+    new_resources, diagnostics = _apply_mutators(
+        bundle=bundle,
+        resources=resources,
+        mutator_functions=[mutator_1, mutator_2],
+    )
+
+    # despite mutator_2 being called last, it doesn't change the job, and we should use location of mutator_1
+    expected_location = Location.from_callable(mutator_1.function)
+
+    assert not diagnostics.has_error()
+    assert new_resources._locations[("resources", "jobs", "job_0")] == expected_location
+    assert new_resources.jobs["job_0"].description == "updated description"

From cea96083156a03b9fc2de60944040721fd4d48b2 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Tue, 18 Mar 2025 14:17:13 +0100
Subject: [PATCH 2/2] Fix fmt

---
 experimental/python/databricks_tests/test_build.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/experimental/python/databricks_tests/test_build.py b/experimental/python/databricks_tests/test_build.py
index 90b42e94aa..2d08e138a4 100644
--- a/experimental/python/databricks_tests/test_build.py
+++ b/experimental/python/databricks_tests/test_build.py
@@ -364,6 +364,7 @@ def mutator_2(job: Job) -> Job:
     assert new_resources._locations[("resources", "jobs", "job_0")] == expected_location
     assert new_resources.jobs["job_0"].description == "updated description"
 
+
 def test_load_resources():
     bundle = Bundle(target="default")