[Python] Add schemas support

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.267.0
 
 ### Notable Changes
+* Add support for schemas in Python support ([#3389])(https://github.com/databricks/cli/pull/3389))
 
 ### CLI
 * Introduce retries to `databricks psql` command ([#3492](https://github.com/databricks/cli/pull/3492))

acceptance/bundle/python/schemas-support/databricks.yml

@@ -0,0 +1,17 @@
+bundle:
+  name: my_project
+
+sync: {paths: []} # don't need to copy files
+
+experimental:
+  python:
+    resources:
+      - "resources:load_resources"
+    mutators:
+      - "mutators:update_schema"
+
+resources:
+  schemas:
+    my_schema_1:
+      name: "My Schema"
+      catalog_name: "my_catalog"

acceptance/bundle/python/schemas-support/mutators.py

@@ -0,0 +1,11 @@
+from dataclasses import replace
+
+from databricks.bundles.core import schema_mutator
+from databricks.bundles.schemas import Schema
+
+
+@schema_mutator
+def update_schema(schema: Schema) -> Schema:
+    assert isinstance(schema.name, str)
+
+    return replace(schema, name=f"{schema.name} (updated)")

acceptance/bundle/python/schemas-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"]

acceptance/bundle/python/schemas-support/output.txt

@@ -0,0 +1,26 @@
+
+>>> uv run --with-requirements requirements-latest.txt --no-cache -q [CLI] bundle validate --output json
+{
+  "experimental": {
+    "python": {
+      "mutators": [
+        "mutators:update_schema"
+      ],
+      "resources": [
+        "resources:load_resources"
+      ]
+    }
+  },
+  "resources": {
+    "schemas": {
+      "my_schema_1": {
+        "catalog_name": "my_catalog",
+        "name": "My Schema (updated)"
+      },
+      "my_schema_2": {
+        "catalog_name": "my_catalog_2",
+        "name": "My Schema (2) (updated)"
+      }
+    }
+  }
+}

acceptance/bundle/python/schemas-support/resources.py

@@ -0,0 +1,15 @@
+from databricks.bundles.core import Resources
+
+
+def load_resources() -> Resources:
+    resources = Resources()
+
+    resources.add_schema(
+        "my_schema_2",
+        {
+            "name": "My Schema (2)",
+            "catalog_name": "my_catalog_2",
+        },
+    )
+
+    return resources

acceptance/bundle/python/schemas-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__

acceptance/bundle/python/schemas-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",
+]

experimental/python/codegen/codegen/packages.py

@@ -5,6 +5,7 @@
 RESOURCE_NAMESPACE = {
     "resources.Job": "jobs",
     "resources.Pipeline": "pipelines",
+    "resources.Schema": "schemas",
     "resources.Volume": "volumes",
 }
 

experimental/python/databricks/bundles/core/__init__.py

@@ -18,6 +18,7 @@
     "load_resources_from_modules",
     "load_resources_from_package_module",
     "pipeline_mutator",
+    "schema_mutator",
     "variables",
     "volume_mutator",
 ]
@@ -40,6 +41,7 @@
     ResourceMutator,
     job_mutator,
     pipeline_mutator,
+    schema_mutator,
     volume_mutator,
 )
 from databricks.bundles.core._resources import Resources

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.schemas._models.schema import Schema
     from databricks.bundles.volumes._models.volume import Volume
 
 _T = TypeVar("_T", bound=Resource)
@@ -130,6 +131,38 @@ def my_pipeline_mutator(bundle: Bundle, pipeline: Pipeline) -> Pipeline:
     return ResourceMutator(resource_type=Pipeline, function=function)
 
 
+@overload
+def schema_mutator(
+    function: Callable[[Bundle, "Schema"], "Schema"],
+) -> ResourceMutator["Schema"]: ...
+
+
+@overload
+def schema_mutator(
+    function: Callable[["Schema"], "Schema"],
+) -> ResourceMutator["Schema"]: ...
+
+
+def schema_mutator(function: Callable) -> ResourceMutator["Schema"]:
+    """
+    Decorator for defining a schema mutator. Function should return a new instance of the schema with the desired changes,
+    instead of mutating the input schema.
+
+    Example:
+
+    .. code-block:: python
+
+        @schema_mutator
+        def my_schema_mutator(bundle: Bundle, schema: Schema) -> Schema:
+            return replace(schema, name="my_schema")
+
+    :param function: Function that mutates a schema.
+    """
+    from databricks.bundles.schemas._models.schema import Schema
+
+    return ResourceMutator(resource_type=Schema, function=function)
+
+
 @overload
 def volume_mutator(
     function: Callable[[Bundle, "Volume"], "Volume"],

experimental/python/databricks/bundles/core/_resource_type.py

@@ -2,7 +2,6 @@
 from typing import Type
 
 from databricks.bundles.core._resource import Resource
-from databricks.bundles.volumes._models.volume import Volume
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -34,6 +33,8 @@ def all(cls) -> tuple["_ResourceType", ...]:
 
         from databricks.bundles.jobs._models.job import Job
         from databricks.bundles.pipelines._models.pipeline import Pipeline
+        from databricks.bundles.schemas._models.schema import Schema
+        from databricks.bundles.volumes._models.volume import Volume
 
         return (
             _ResourceType(
@@ -51,4 +52,9 @@ def all(cls) -> tuple["_ResourceType", ...]:
                 plural_name="volumes",
                 singular_name="volume",
             ),
+            _ResourceType(
+                resource_type=Schema,
+                plural_name="schemas",
+                singular_name="schema",
+            ),
         )

experimental/python/databricks/bundles/core/_resources.py

@@ -8,6 +8,7 @@
 if TYPE_CHECKING:
     from databricks.bundles.jobs._models.job import Job, JobParam
     from databricks.bundles.pipelines._models.pipeline import Pipeline, PipelineParam
+    from databricks.bundles.schemas._models.schema import Schema, SchemaParam
     from databricks.bundles.volumes._models.volume import Volume, VolumeParam
 
 __all__ = ["Resources"]
@@ -58,6 +59,7 @@ def load_resources(bundle: Bundle) -> Resources:
     def __init__(self):
         self._jobs = dict[str, "Job"]()
         self._pipelines = dict[str, "Pipeline"]()
+        self._schemas = dict[str, "Schema"]()
         self._volumes = dict[str, "Volume"]()
         self._locations = dict[tuple[str, ...], Location]()
         self._diagnostics = Diagnostics()
@@ -70,6 +72,10 @@ def jobs(self) -> dict[str, "Job"]:
     def pipelines(self) -> dict[str, "Pipeline"]:
         return self._pipelines
 
+    @property
+    def schemas(self) -> dict[str, "Schema"]:
+        return self._schemas
+
     @property
     def volumes(self) -> dict[str, "Volume"]:
         return self._volumes
@@ -99,6 +105,7 @@ def add_resource(
 
         from databricks.bundles.jobs import Job
         from databricks.bundles.pipelines import Pipeline
+        from databricks.bundles.schemas import Schema
         from databricks.bundles.volumes import Volume
 
         location = location or Location.from_stack_frame(depth=1)
@@ -108,6 +115,8 @@ def add_resource(
                 self.add_job(resource_name, resource, location=location)
             case Pipeline():
                 self.add_pipeline(resource_name, resource, location=location)
+            case Schema():
+                self.add_schema(resource_name, resource, location=location)
             case Volume():
                 self.add_volume(resource_name, resource, location=location)
             case _:
@@ -177,6 +186,38 @@ def add_pipeline(
 
             self._pipelines[resource_name] = pipeline
 
+    def add_schema(
+        self,
+        resource_name: str,
+        schema: "SchemaParam",
+        *,
+        location: Optional[Location] = None,
+    ) -> None:
+        """
+        Adds a schema to the collection of resources. Resource name must be unique across all schemas.
+
+        :param resource_name: unique identifier for the schema
+        :param schema: the schema to add, can be Schema or dict
+        :param location: optional location of the schema in the source code
+        """
+        from databricks.bundles.schemas import Schema
+
+        schema = _transform(Schema, schema)
+        path = ("resources", "schemas", resource_name)
+        location = location or Location.from_stack_frame(depth=1)
+
+        if self._schemas.get(resource_name):
+            self.add_diagnostic_error(
+                msg=f"Duplicate resource name '{resource_name}' for a schema. Resource names must be unique.",
+                location=location,
+                path=path,
+            )
+        else:
+            if location:
+                self.add_location(path, location)
+
+            self._schemas[resource_name] = schema
+
     def add_volume(
         self,
         resource_name: str,
@@ -285,6 +326,9 @@ def add_resources(self, other: "Resources") -> None:
         for name, pipeline in other.pipelines.items():
             self.add_pipeline(name, pipeline)
 
+        for name, schema in other.schemas.items():
+            self.add_schema(name, schema)
+
         for name, volume in other.volumes.items():
             self.add_volume(name, volume)
 

experimental/python/databricks/bundles/schemas/__init__.py

@@ -0,0 +1,22 @@
+__all__ = [
+    "Schema",
+    "SchemaDict",
+    "SchemaGrant",
+    "SchemaGrantDict",
+    "SchemaGrantParam",
+    "SchemaGrantPrivilege",
+    "SchemaGrantPrivilegeParam",
+    "SchemaParam",
+]
+
+
+from databricks.bundles.schemas._models.schema import Schema, SchemaDict, SchemaParam
+from databricks.bundles.schemas._models.schema_grant import (
+    SchemaGrant,
+    SchemaGrantDict,
+    SchemaGrantParam,
+)
+from databricks.bundles.schemas._models.schema_grant_privilege import (
+    SchemaGrantPrivilege,
+    SchemaGrantPrivilegeParam,
+)