[Python] Support schemas resource type

Author: kanterov

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

@@ -13,11 +13,12 @@
     "VariableOrList",
     "VariableOrOptional",
     "job_mutator",
+    "pipeline_mutator",
+    "schema_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",
 ]
@@ -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"],