[Python] Add volumes support

kanterov · kanterov · commit 7eb092f5a6f0 · 2025-08-08T13:59:56.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,15 @@
+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
@@ -4,16 +4,19 @@
 RESOURCE_NAMESPACE_OVERRIDE = {
     "resources.Job": "jobs",
     "resources.Pipeline": "pipelines",
+    "resources.Volume": "catalog",
     "resources.JobPermission": "jobs",
     "resources.JobPermissionLevel": "jobs",
     "resources.PipelinePermission": "pipelines",
     "resources.PipelinePermissionLevel": "pipelines",
+    "resources.Grant": "catalog",
 }
 
 # All supported resource types
 RESOURCE_TYPES = [
     "resources.Job",
     "resources.Pipeline",
+    "resources.Volume",
 ]
 
 # Namespaces to load from OpenAPI spec.
@@ -24,6 +27,7 @@
     "jobs",
     "pipelines",
     "resources",
+    "catalog",
 ]
 
 RENAMES = {
diff --git a/experimental/python/databricks/bundles/catalog/__init__.py b/experimental/python/databricks/bundles/catalog/__init__.py
@@ -0,0 +1,15 @@
+__all__ = [
+    "Grant",
+    "GrantDict",
+    "GrantParam",
+    "Volume",
+    "VolumeDict",
+    "VolumeParam",
+    "VolumeType",
+    "VolumeTypeParam",
+]
+
+
+from databricks.bundles.catalog._models.grant import Grant, GrantDict, GrantParam
+from databricks.bundles.catalog._models.volume import Volume, VolumeDict, VolumeParam
+from databricks.bundles.catalog._models.volume_type import VolumeType, VolumeTypeParam
diff --git a/experimental/python/databricks/bundles/catalog/_models/grant.py b/experimental/python/databricks/bundles/catalog/_models/grant.py
@@ -0,0 +1,48 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, TypedDict
+
+from databricks.bundles.core._transform import _transform
+from databricks.bundles.core._transform_to_json import _transform_to_json_value
+from databricks.bundles.core._variable import VariableOr, VariableOrList
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+
+@dataclass(kw_only=True)
+class Grant:
+    """"""
+
+    principal: VariableOr[str]
+    """
+    The name of the principal that will be granted privileges
+    """
+
+    privileges: VariableOrList[str] = field(default_factory=list)
+    """
+    The privileges to grant to the specified entity
+    """
+
+    @classmethod
+    def from_dict(cls, value: "GrantDict") -> "Self":
+        return _transform(cls, value)
+
+    def as_dict(self) -> "GrantDict":
+        return _transform_to_json_value(self)  # type:ignore
+
+
+class GrantDict(TypedDict, total=False):
+    """"""
+
+    principal: VariableOr[str]
+    """
+    The name of the principal that will be granted privileges
+    """
+
+    privileges: VariableOrList[str]
+    """
+    The privileges to grant to the specified entity
+    """
+
+
+GrantParam = GrantDict | Grant
diff --git a/experimental/python/databricks/bundles/catalog/_models/volume.py b/experimental/python/databricks/bundles/catalog/_models/volume.py
@@ -0,0 +1,93 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, TypedDict
+
+from databricks.bundles.catalog._models.grant import Grant, GrantParam
+from databricks.bundles.catalog._models.volume_type import VolumeType, VolumeTypeParam
+from databricks.bundles.core._resource import Resource
+from databricks.bundles.core._transform import _transform
+from databricks.bundles.core._transform_to_json import _transform_to_json_value
+from databricks.bundles.core._variable import (
+    VariableOr,
+    VariableOrList,
+    VariableOrOptional,
+)
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+
+@dataclass(kw_only=True)
+class Volume(Resource):
+    """"""
+
+    catalog_name: VariableOr[str]
+    """
+    The name of the catalog where the schema and the volume are
+    """
+
+    name: VariableOr[str]
+    """
+    The name of the volume
+    """
+
+    schema_name: VariableOr[str]
+    """
+    The name of the schema where the volume is
+    """
+
+    comment: VariableOrOptional[str] = None
+    """
+    The comment attached to the volume
+    """
+
+    grants: VariableOrList[Grant] = field(default_factory=list)
+
+    storage_location: VariableOrOptional[str] = None
+    """
+    The storage location on the cloud
+    """
+
+    volume_type: VariableOrOptional[VolumeType] = None
+
+    @classmethod
+    def from_dict(cls, value: "VolumeDict") -> "Self":
+        return _transform(cls, value)
+
+    def as_dict(self) -> "VolumeDict":
+        return _transform_to_json_value(self)  # type:ignore
+
+
+class VolumeDict(TypedDict, total=False):
+    """"""
+
+    catalog_name: VariableOr[str]
+    """
+    The name of the catalog where the schema and the volume are
+    """
+
+    name: VariableOr[str]
+    """
+    The name of the volume
+    """
+
+    schema_name: VariableOr[str]
+    """
+    The name of the schema where the volume is
+    """
+
+    comment: VariableOrOptional[str]
+    """
+    The comment attached to the volume
+    """
+
+    grants: VariableOrList[GrantParam]
+
+    storage_location: VariableOrOptional[str]
+    """
+    The storage location on the cloud
+    """
+
+    volume_type: VariableOrOptional[VolumeTypeParam]
+
+
+VolumeParam = VolumeDict | Volume
diff --git a/experimental/python/databricks/bundles/catalog/_models/volume_type.py b/experimental/python/databricks/bundles/catalog/_models/volume_type.py
@@ -0,0 +1,14 @@
+from enum import Enum
+from typing import Literal
+
+
+class VolumeType(Enum):
+    """
+    The type of the volume. An external volume is located in the specified external location. A managed volume is located in the default location which is specified by the parent schema, or the parent catalog, or the Metastore. [Learn more](https://docs.databricks.com/aws/en/volumes/managed-vs-external)
+    """
+
+    EXTERNAL = "EXTERNAL"
+    MANAGED = "MANAGED"
+
+
+VolumeTypeParam = Literal["EXTERNAL", "MANAGED"] | VolumeType
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
@@ -8,6 +8,7 @@
 if TYPE_CHECKING:
     from databricks.bundles.jobs._models.job import Job
     from databricks.bundles.pipelines._models.pipeline import Pipeline
+    from databricks.bundles.catalog._models.volume import Volume
 
 _T = TypeVar("_T", bound=Resource)
 
@@ -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
diff --git a/experimental/python/databricks_tests/core/test_resources.py b/experimental/python/databricks_tests/core/test_resources.py
