[Python] Add volumes support

kanterov · kanterov · commit f179705e903a · 2025-08-21T17:39:07.000+02:00
diff --git a/acceptance/bundle/python/volumes-support/databricks.yml b/acceptance/bundle/python/volumes-support/databricks.yml
@@ -0,0 +1,18 @@
+bundle:
+  name: my_project
+
+sync: {paths: []} # don't need to copy files
+
+experimental:
+  python:
+    resources:
+      - "resources:load_resources"
+    mutators:
+      - "mutators:update_volume"
+
+resources:
+  volumes:
+    my_volume_1:
+      name: "My Volume"
+      catalog_name: "my_catalog"
+      schema_name: "my_schema"
diff --git a/acceptance/bundle/python/volumes-support/mutators.py b/acceptance/bundle/python/volumes-support/mutators.py
@@ -0,0 +1,11 @@
+from dataclasses import replace
+
+from databricks.bundles.core import volume_mutator
+from databricks.bundles.catalog import Volume
+
+
+@volume_mutator
+def update_volume(volume: Volume) -> Volume:
+    assert isinstance(volume.name, str)
+
+    return replace(volume, name=f"{volume.name} (updated)")
diff --git a/acceptance/bundle/python/volumes-support/out.test.toml b/acceptance/bundle/python/volumes-support/out.test.toml
@@ -0,0 +1,6 @@
+Local = true
+Cloud = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]
+  UV_ARGS = ["--with-requirements requirements-latest.txt --no-cache"]
diff --git a/acceptance/bundle/python/volumes-support/output.txt b/acceptance/bundle/python/volumes-support/output.txt
@@ -0,0 +1,30 @@
+
+>>> uv run --with-requirements requirements-latest.txt --no-cache -q [CLI] bundle validate --output json
+{
+  "experimental": {
+    "python": {
+      "mutators": [
+        "mutators:update_volume"
+      ],
+      "resources": [
+        "resources:load_resources"
+      ]
+    }
+  },
+  "resources": {
+    "volumes": {
+      "my_volume_1": {
+        "catalog_name": "my_catalog",
+        "name": "My Volume (updated)",
+        "schema_name": "my_schema",
+        "volume_type": "MANAGED"
+      },
+      "my_volume_2": {
+        "catalog_name": "my_catalog_2",
+        "name": "My Volume (2) (updated)",
+        "schema_name": "my_schema_2",
+        "volume_type": "MANAGED"
+      }
+    }
+  }
+}
diff --git a/acceptance/bundle/python/volumes-support/resources.py b/acceptance/bundle/python/volumes-support/resources.py
@@ -0,0 +1,16 @@
+from databricks.bundles.core import Resources
+
+
+def load_resources() -> Resources:
+    resources = Resources()
+
+    resources.add_volume(
+        "my_volume_2",
+        {
+            "name": "My Volume (2)",
+            "catalog_name": "my_catalog_2",
+            "schema_name": "my_schema_2",
+        },
+    )
+
+    return resources
diff --git a/acceptance/bundle/python/volumes-support/script b/acceptance/bundle/python/volumes-support/script
@@ -0,0 +1,6 @@
+echo "$DATABRICKS_BUNDLES_WHEEL" > "requirements-latest.txt"
+
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq "pick(.experimental.python, .resources)"
+
+rm -fr .databricks __pycache__
diff --git a/acceptance/bundle/python/volumes-support/test.toml b/acceptance/bundle/python/volumes-support/test.toml
@@ -0,0 +1,8 @@
+Local = true
+Cloud = false # tests don't interact with APIs
+
+[EnvMatrix]
+UV_ARGS = [
+    # pipelines are only supported in the latest version of the wheel
+    "--with-requirements requirements-latest.txt --no-cache",
+]
diff --git a/experimental/python/codegen/codegen/packages.py b/experimental/python/codegen/codegen/packages.py
@@ -5,6 +5,7 @@
 RESOURCE_NAMESPACE = {
     "resources.Job": "jobs",
     "resources.Pipeline": "pipelines",
+    "resources.Volume": "volumes",
 }
 
 RESOURCE_TYPES = list(RESOURCE_NAMESPACE.keys())
@@ -17,6 +18,7 @@
     "jobs",
     "pipelines",
     "resources",
+    "catalog",
 ]
 
 RENAMES = {
diff --git a/experimental/python/databricks/bundles/core/__init__.py b/experimental/python/databricks/bundles/core/__init__.py
@@ -13,12 +13,13 @@
     "VariableOrList",
     "VariableOrOptional",
     "job_mutator",
-    "pipeline_mutator",
     "load_resources_from_current_package_module",
     "load_resources_from_module",
     "load_resources_from_modules",
     "load_resources_from_package_module",
+    "pipeline_mutator",
     "variables",
+    "volume_mutator",
 ]
 
 from databricks.bundles.core._bundle import Bundle
@@ -39,6 +40,7 @@
     ResourceMutator,
     job_mutator,
     pipeline_mutator,
+    volume_mutator,
 )
 from databricks.bundles.core._resources import Resources
 from databricks.bundles.core._variable import (
diff --git a/experimental/python/databricks/bundles/core/_resource_mutator.py b/experimental/python/databricks/bundles/core/_resource_mutator.py
@@ -6,6 +6,7 @@
 from databricks.bundles.core._resource import Resource
 
 if TYPE_CHECKING:
+    from databricks.bundles.catalog._models.volume import Volume
     from databricks.bundles.jobs._models.job import Job
     from databricks.bundles.pipelines._models.pipeline import Pipeline
 
@@ -127,3 +128,35 @@ def my_pipeline_mutator(bundle: Bundle, pipeline: Pipeline) -> Pipeline:
     from databricks.bundles.pipelines._models.pipeline import Pipeline
 
     return ResourceMutator(resource_type=Pipeline, function=function)
+
+
+@overload
+def volume_mutator(
+    function: Callable[[Bundle, "Volume"], "Volume"],
+) -> ResourceMutator["Volume"]: ...
+
+
+@overload
+def volume_mutator(
+    function: Callable[["Volume"], "Volume"],
+) -> ResourceMutator["Volume"]: ...
+
+
+def volume_mutator(function: Callable) -> ResourceMutator["Volume"]:
+    """
+    Decorator for defining a volume mutator. Function should return a new instance of the volume with the desired changes,
+    instead of mutating the input volume.
+
+    Example:
+
+    .. code-block:: python
+
+        @volume_mutator
+        def my_volume_mutator(bundle: Bundle, volume: Volume) -> Volume:
+            return replace(volume, name="my_volume")
+
+    :param function: Function that mutates a volume.
+    """
+    from databricks.bundles.catalog._models.volume import Volume
+
+    return ResourceMutator(resource_type=Volume, function=function)
diff --git a/experimental/python/databricks/bundles/core/_resource_type.py b/experimental/python/databricks/bundles/core/_resource_type.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Type
 
+from databricks.bundles.catalog._models.volume import Volume
 from databricks.bundles.core._resource import Resource
 
 
@@ -45,4 +46,9 @@ def all(cls) -> tuple["_ResourceType", ...]:
                 plural_name="pipelines",
                 singular_name="pipeline",
             ),
+            _ResourceType(
+                resource_type=Volume,
+                plural_name="volumes",
+                singular_name="volume",
+            ),
         )
diff --git a/experimental/python/databricks/bundles/core/_resources.py b/experimental/python/databricks/bundles/core/_resources.py
@@ -6,6 +6,7 @@
 from databricks.bundles.core._transform import _transform
 
 if TYPE_CHECKING:
+    from databricks.bundles.catalog._models.volume import Volume, VolumeParam
     from databricks.bundles.jobs._models.job import Job, JobParam
     from databricks.bundles.pipelines._models.pipeline import Pipeline, PipelineParam
 
@@ -57,6 +58,7 @@ def load_resources(bundle: Bundle) -> Resources:
     def __init__(self):
         self._jobs = dict[str, "Job"]()
         self._pipelines = dict[str, "Pipeline"]()
+        self._volumes = dict[str, "Volume"]()
         self._locations = dict[tuple[str, ...], Location]()
         self._diagnostics = Diagnostics()
 
@@ -68,6 +70,10 @@ def jobs(self) -> dict[str, "Job"]:
     def pipelines(self) -> dict[str, "Pipeline"]:
         return self._pipelines
 
+    @property
+    def volumes(self) -> dict[str, "Volume"]:
+        return self._volumes
+
     @property
     def diagnostics(self) -> Diagnostics:
         """
@@ -91,6 +97,7 @@ def add_resource(
         :param location: optional location of the resource in the source code
         """
 
+        from databricks.bundles.catalog import Volume
         from databricks.bundles.jobs import Job
         from databricks.bundles.pipelines import Pipeline
 
@@ -101,6 +108,8 @@ def add_resource(
                 self.add_job(resource_name, resource, location=location)
             case Pipeline():
                 self.add_pipeline(resource_name, resource, location=location)
+            case Volume():
+                self.add_volume(resource_name, resource, location=location)
             case _:
                 raise ValueError(f"Unsupported resource type: {type(resource)}")
 
@@ -168,6 +177,38 @@ def add_pipeline(
 
             self._pipelines[resource_name] = pipeline
 
+    def add_volume(
+        self,
+        resource_name: str,
+        volume: "VolumeParam",
+        *,
+        location: Optional[Location] = None,
+    ) -> None:
+        """
+        Adds a volume to the collection of resources. Resource name must be unique across all volumes.
+
+        :param resource_name: unique identifier for the volume
+        :param volume: the volume to add, can be Volume or dict
+        :param location: optional location of the volume in the source code
+        """
+        from databricks.bundles.catalog import Volume
+
+        volume = _transform(Volume, volume)
+        path = ("resources", "volumes", resource_name)
+        location = location or Location.from_stack_frame(depth=1)
+
+        if self._volumes.get(resource_name):
+            self.add_diagnostic_error(
+                msg=f"Duplicate resource name '{resource_name}' for a volume. Resource names must be unique.",
+                location=location,
+                path=path,
+            )
+        else:
+            if location:
+                self.add_location(path, location)
+
+            self._volumes[resource_name] = volume
+
     def add_location(self, path: tuple[str, ...], location: Location) -> None:
         """
         Associate source code location with a path in the bundle configuration.
@@ -244,6 +285,9 @@ def add_resources(self, other: "Resources") -> None:
         for name, pipeline in other.pipelines.items():
             self.add_pipeline(name, pipeline)
 
+        for name, volume in other.volumes.items():
+            self.add_volume(name, volume)
+
         for path, location in other._locations.items():
             self.add_location(path, location)
 
diff --git a/experimental/python/databricks/bundles/volumes/__init__.py b/experimental/python/databricks/bundles/volumes/__init__.py
@@ -0,0 +1,25 @@
+__all__ = [
+    "Volume",
+    "VolumeDict",
+    "VolumeGrant",
+    "VolumeGrantDict",
+    "VolumeGrantParam",
+    "VolumeGrantPrivilege",
+    "VolumeGrantPrivilegeParam",
+    "VolumeParam",
+    "VolumeType",
+    "VolumeTypeParam",
+]
+
+
+from databricks.bundles.volumes._models.volume import Volume, VolumeDict, VolumeParam
+from databricks.bundles.volumes._models.volume_grant import (
+    VolumeGrant,
+    VolumeGrantDict,
+    VolumeGrantParam,
+)
+from databricks.bundles.volumes._models.volume_grant_privilege import (
+    VolumeGrantPrivilege,
+    VolumeGrantPrivilegeParam,
+)
+from databricks.bundles.volumes._models.volume_type import VolumeType, VolumeTypeParam
diff --git a/experimental/python/databricks/bundles/volumes/_models/volume.py b/experimental/python/databricks/bundles/volumes/_models/volume.py
diff --git a/experimental/python/databricks/bundles/volumes/_models/volume_grant.py b/experimental/python/databricks/bundles/volumes/_models/volume_grant.py
diff --git a/experimental/python/databricks/bundles/volumes/_models/volume_grant_privilege.py b/experimental/python/databricks/bundles/volumes/_models/volume_grant_privilege.py
diff --git a/experimental/python/databricks/bundles/volumes/_models/volume_type.py b/experimental/python/databricks/bundles/volumes/_models/volume_type.py
diff --git a/experimental/python/databricks_tests/core/test_resources.py b/experimental/python/databricks_tests/core/test_resources.py

Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@`
`5`	`5`	`RESOURCE_NAMESPACE = {`
`6`	`6`	`"resources.Job": "jobs",`
`7`	`7`	`"resources.Pipeline": "pipelines",`
	`8`	`+ "resources.Volume": "volumes",`
`8`	`9`	`}`
`9`	`10`
`10`	`11`	`RESOURCE_TYPES = list(RESOURCE_NAMESPACE.keys())`
`@@ -17,6 +18,7 @@`
`17`	`18`	`"jobs",`
`18`	`19`	`"pipelines",`
`19`	`20`	`"resources",`
	`21`	`+ "catalog",`
`20`	`22`	`]`
`21`	`23`
`22`	`24`	`RENAMES = {`