[Python] Fix issues with multiple mutators (#2509)

## Changes
Fix issues with bundles using multiple mutators. 

Previously, we didn't correctly chain output between them.

## Tests
- Unit tests
- Acceptance tests in https://github.com/databricks/cli/pull/2493
This commit is contained in:
Gleb Kanterov 2025-03-18 19:41:14 +01:00 committed by GitHub
parent ace37a8822
commit fc9544e6d9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 67 additions and 0 deletions

View File

@ -128,6 +128,7 @@ def _apply_mutators(
("resources", "jobs", resource_name), location ("resources", "jobs", resource_name), location
) )
resources.jobs[resource_name] = new_job resources.jobs[resource_name] = new_job
job = new_job
except Exception as exc: except Exception as exc:
mutator_name = mutator.function.__name__ mutator_name = mutator.function.__name__

View File

@ -1,9 +1,11 @@
import sys import sys
from dataclasses import replace
from io import StringIO from io import StringIO
from pathlib import Path from pathlib import Path
from databricks.bundles.build import ( from databricks.bundles.build import (
_append_resources, _append_resources,
_apply_mutators,
_Args, _Args,
_Conf, _Conf,
_load_object, _load_object,
@ -21,6 +23,7 @@ from databricks.bundles.core import (
Location, Location,
Resources, Resources,
Severity, Severity,
job_mutator,
) )
from databricks.bundles.jobs import Job from databricks.bundles.jobs import Job
@ -299,6 +302,69 @@ def test_conf_from_dict():
) )
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"
def test_load_resources(): def test_load_resources():
bundle = Bundle(target="default") bundle = Bundle(target="default")