From d2ec690585e713a967a4256b31e0cd48f12a01a3 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <kanterov@users.noreply.github.com>
Date: Tue, 18 Mar 2025 12:56:36 +0100
Subject: [PATCH] [Python] Fix issues with multiple resource loaders (#2511)

## Changes
Fix issues with multiple resource loaders.

Previously, we discarded outputs of all loaders except the last one.
That worked well if there is only a single resource loader.

## Tests
- Unit tests
- Acceptance tests in https://github.com/databricks/cli/pull/2493
---
 .../python/databricks/bundles/build.py        |  6 ++-
 .../python/databricks_tests/test_build.py     | 49 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/experimental/python/databricks/bundles/build.py b/experimental/python/databricks/bundles/build.py
index d9ddc626b..949ac39ae 100644
--- a/experimental/python/databricks/bundles/build.py
+++ b/experimental/python/databricks/bundles/build.py
@@ -241,9 +241,11 @@ def _load_resources(
 
     for function in functions:
         try:
-            resources, diagnostics = diagnostics.extend_tuple(
+            function_resources, diagnostics = diagnostics.extend_tuple(
                 _load_resources_from_function(bundle, function)
             )
+
+            resources.add_resources(function_resources)
         except Exception as exc:
             diagnostics = diagnostics.extend(
                 Diagnostics.from_exception(
@@ -342,7 +344,7 @@ def _explain_common_import_error(exc: Exception) -> Diagnostics:
     # a common case when default name of the module is not found
     # we can give a hint to the user how to fix it
     explanation = """Make sure to create a new Python file at resources/__init__.py with contents:
-    
+
 from databricks.bundles.core import load_resources_from_current_package_module
 
 
diff --git a/experimental/python/databricks_tests/test_build.py b/experimental/python/databricks_tests/test_build.py
index 980cbbe76..427d447b6 100644
--- a/experimental/python/databricks_tests/test_build.py
+++ b/experimental/python/databricks_tests/test_build.py
@@ -7,6 +7,7 @@ from databricks.bundles.build import (
     _Args,
     _Conf,
     _load_object,
+    _load_resources,
     _parse_args,
     _parse_bundle_info,
     _relativize_location,
@@ -296,3 +297,51 @@ def test_conf_from_dict():
         ],
         venv_path="venv",
     )
+
+
+def test_load_resources():
+    bundle = Bundle(target="default")
+
+    def load_resources_1() -> Resources:
+        resources = Resources()
+        resources.add_job(
+            resource_name="my_job_1",
+            job={"name": "Job 1"},
+            location=Location(file="my_job_1.py", line=42, column=1),
+        )
+
+        return resources
+
+    def load_resources_2() -> Resources:
+        resources = Resources()
+        resources.add_job(
+            resource_name="my_job_2",
+            job={"name": "Job 2"},
+            location=Location(file="my_job_2.py", line=42, column=1),
+        )
+
+        return resources
+
+    resources, diagnostics = _load_resources(
+        bundle=bundle,
+        functions=[
+            load_resources_1,
+            load_resources_2,
+        ],
+    )
+
+    assert not diagnostics.has_error()
+
+    assert resources.jobs == {
+        "my_job_1": Job(name="Job 1"),
+        "my_job_2": Job(name="Job 2"),
+    }
+
+    assert resources._locations == {
+        ("resources", "jobs", "my_job_1"): Location(
+            file="my_job_1.py", line=42, column=1
+        ),
+        ("resources", "jobs", "my_job_2"): Location(
+            file="my_job_2.py", line=42, column=1
+        ),
+    }