[FS2524] Add Feature Spec as a Bundle Resource

.nextchanges/bundles/feature-specs.md

@@ -0,0 +1 @@
+* Add support for the `feature_specs` resource type, including UC grants on a feature spec. ([#7364](https://github.com/databricks/cli/pull/7364))

acceptance/bin/create_feature_spec_fixtures.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""
+Creates the feature tables and Python UDFs that the feature_specs acceptance tests look up, in an
+existing schema. Creating a feature spec resolves each of them in Unity Catalog.
+
+Usage: create_feature_spec_fixtures.py CATALOG.SCHEMA
+"""
+
+import json
+import os
+import subprocess
+import sys
+import time
+
+schema_full_name = sys.argv[1]
+cli = os.environ.get("CLI", "databricks")
+warehouse_id = os.environ["TEST_DEFAULT_WAREHOUSE_ID"]
+
+statements = [
+    f"CREATE TABLE {schema_full_name}.user_features "
+    "(user_id INT NOT NULL, age INT, country STRING, CONSTRAINT user_features_pk PRIMARY KEY (user_id))",
+    f"CREATE TABLE {schema_full_name}.tenant_user_features "
+    "(tenant_id INT NOT NULL, user_id INT NOT NULL, age INT, country STRING, "
+    "CONSTRAINT tenant_user_features_pk PRIMARY KEY (tenant_id, user_id))",
+    f"CREATE FUNCTION {schema_full_name}.score(x INT) RETURNS DOUBLE LANGUAGE PYTHON AS $$ return float(x) $$",
+    f"CREATE FUNCTION {schema_full_name}.country_score(x INT, y STRING) RETURNS DOUBLE LANGUAGE PYTHON "
+    "AS $$ return float(x) + len(y) $$",
+]
+
+
+def call_api(*args):
+    result = subprocess.run([cli, "api", *args], capture_output=True, text=True)
+    if result.returncode != 0:
+        print(f"databricks api {' '.join(args)} failed: {result.stderr}", file=sys.stderr)
+        sys.exit(1)
+    return json.loads(result.stdout)
+
+
+for statement in statements:
+    payload = json.dumps({"warehouse_id": warehouse_id, "statement": statement, "wait_timeout": "50s"})
+    response = call_api("post", "/api/2.0/sql/statements/", "--json", payload)
+    # A warehouse that is still starting outlasts wait_timeout.
+    deadline = time.monotonic() + 600
+    while response["status"]["state"] in ("PENDING", "RUNNING") and time.monotonic() < deadline:
+        time.sleep(5)
+        response = call_api("get", f"/api/2.0/sql/statements/{response['statement_id']}")
+    if response["status"]["state"] != "SUCCEEDED":
+        print(f"Statement {statement!r} ended in {response['status']}", file=sys.stderr)
+        sys.exit(1)

acceptance/bundle/deployment/bind/feature_spec/databricks.yml.tmpl

@@ -0,0 +1,12 @@
+bundle:
+  name: test-bind-feature-spec-$UNIQUE_NAME
+
+resources:
+  feature_specs:
+    spec1:
+      name: $SCHEMA_FULL_NAME.boundspec
+      features:
+        - feature_lookup:
+            table: $SCHEMA_FULL_NAME.user_features
+            lookup_keys: [user_id]
+            feature_names: [age]

acceptance/bundle/deployment/bind/feature_spec/out.test.toml

@@ -0,0 +1,2 @@
+Cloud = true
+EnvMatrix.DMS = ["", "true"]

acceptance/bundle/deployment/bind/feature_spec/output.txt

@@ -0,0 +1,53 @@
+
+>>> [CLI] schemas create fspec_[UNIQUE_NAME] main -o json
+{
+  "full_name": "main.fspec_[UNIQUE_NAME]"
+}
+
+>>> create_feature_spec_fixtures.py main.fspec_[UNIQUE_NAME]
+
+=== Create the feature spec outside the bundle
+>>> [CLI] api post /api/2.0/feature-engineering/feature-specs --json {
+  "name": "main.fspec_[UNIQUE_NAME].boundspec",
+  "features": [{"feature_lookup": {"table": "main.fspec_[UNIQUE_NAME].user_features", "lookup_keys": ["user_id"], "feature_names": ["age"]}}]
+}
+{
+  "name": "main.fspec_[UNIQUE_NAME].boundspec"
+}
+
+>>> [CLI] bundle deployment bind spec1 main.fspec_[UNIQUE_NAME].boundspec --auto-approve
+Successfully bound feature_spec with an id 'main.fspec_[UNIQUE_NAME].boundspec'
+Run 'bundle deploy' to deploy changes to your workspace
+
+>>> [CLI] bundle summary
+Name: test-bind-feature-spec-[UNIQUE_NAME]
+Target: default
+Workspace:
+  User: [USERNAME]
+  Path: /Workspace/Users/[USERNAME]/.bundle/test-bind-feature-spec-[UNIQUE_NAME]/default
+Resources:
+  Feature Specs:
+    spec1:
+      Name: main.fspec_[UNIQUE_NAME].boundspec
+      URL:  [DATABRICKS_URL]/explore/data/functions/main/fspec_[UNIQUE_NAME]/boundspec?w=[WORKSPACE_ID]
+
+=== The bound feature spec matches the config, so there is nothing to change
+>>> [CLI] bundle plan
+Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged
+
+>>> [CLI] bundle deployment unbind spec1
+
+>>> [CLI] bundle summary
+Name: test-bind-feature-spec-[UNIQUE_NAME]
+Target: default
+Workspace:
+  User: [USERNAME]
+  Path: /Workspace/Users/[USERNAME]/.bundle/test-bind-feature-spec-[UNIQUE_NAME]/default
+Resources:
+  Feature Specs:
+    spec1:
+      Name: main.fspec_[UNIQUE_NAME].boundspec
+      URL:  (not deployed)
+
+>>> [CLI] api delete /api/2.0/feature-engineering/feature-specs/main.fspec_[UNIQUE_NAME].boundspec
+{}

acceptance/bundle/deployment/bind/feature_spec/script

@@ -0,0 +1,29 @@
+# Creating a feature spec resolves its lookup tables in Unity Catalog, so the test creates them in
+# a schema of its own. Dropping the schema also drops the feature spec.
+export SCHEMA_FULL_NAME="main.fspec_$UNIQUE_NAME"
+trace $CLI schemas create "fspec_$UNIQUE_NAME" main -o json | jq '{full_name}'
+trace create_feature_spec_fixtures.py "$SCHEMA_FULL_NAME"
+
+cleanup() {
+    $CLI schemas delete "$SCHEMA_FULL_NAME" --force >> LOG.cleanup 2>&1 || true
+}
+trap cleanup EXIT
+
+envsubst < databricks.yml.tmpl > databricks.yml
+
+title "Create the feature spec outside the bundle"
+MSYS_NO_PATHCONV=1 trace $CLI api post /api/2.0/feature-engineering/feature-specs --json "{
+  \"name\": \"$SCHEMA_FULL_NAME.boundspec\",
+  \"features\": [{\"feature_lookup\": {\"table\": \"$SCHEMA_FULL_NAME.user_features\", \"lookup_keys\": [\"user_id\"], \"feature_names\": [\"age\"]}}]
+}" | jq '{name}'
+
+trace $CLI bundle deployment bind spec1 $SCHEMA_FULL_NAME.boundspec --auto-approve
+trace $CLI bundle summary
+
+title "The bound feature spec matches the config, so there is nothing to change"
+trace $CLI bundle plan
+
+trace $CLI bundle deployment unbind spec1
+trace $CLI bundle summary
+
+MSYS_NO_PATHCONV=1 trace $CLI api delete /api/2.0/feature-engineering/feature-specs/$SCHEMA_FULL_NAME.boundspec

acceptance/bundle/deployment/bind/feature_spec/test.toml

@@ -0,0 +1,8 @@
+Ignore = [
+    ".databricks",
+]
+
+# Fake SQL endpoint so create_feature_spec_fixtures.py works against the local test server too.
+[[Server]]
+Pattern = "POST /api/2.0/sql/statements/"
+Response.Body = '{"status": {"state": "SUCCEEDED"}, "manifest": {"schema": {"columns": []}}}'

acceptance/bundle/invariant/auto-migrate/test.toml

@@ -43,6 +43,8 @@ EnvMatrixExclude.no_model_service = ["INPUT_CONFIG=model_service.yml.tmpl"]
 EnvMatrixExclude.no_data_quality_monitor = ["INPUT_CONFIG=data_quality_monitor.yml.tmpl"]
 # Features are direct-only; the terraform deploy that seeds the migration fails for them.
 EnvMatrixExclude.no_feature = ["INPUT_CONFIG=feature.yml.tmpl"]
+# Feature specs are direct-only; the terraform deploy that seeds the migration fails for them.
+EnvMatrixExclude.no_feature_spec = ["INPUT_CONFIG=feature_spec.yml.tmpl"]
 # The Postgres Data API is direct-only; the terraform deploy that seeds the migration fails for it.
 EnvMatrixExclude.no_postgres_data_api = ["INPUT_CONFIG=postgres_data_api.yml.tmpl"]
 

acceptance/bundle/invariant/configs/feature_spec.yml.tmpl

@@ -0,0 +1,16 @@
+bundle:
+  name: test-bundle-$UNIQUE_NAME
+
+resources:
+  feature_specs:
+    foo:
+      name: main.fspec_inv_$UNIQUE_NAME.test_feature_spec
+      features:
+        - feature_lookup:
+            table: main.fspec_inv_$UNIQUE_NAME.user_features
+            lookup_keys: [user_id]
+            feature_names: [age]
+      grants:
+        - principal: account users
+          privileges:
+            - EXECUTE

acceptance/bundle/invariant/configs/feature_spec.yml.tmpl-cleanup.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+echo "Cleaning up temporary schema"
+$CLI schemas delete main.fspec_inv_$UNIQUE_NAME --force || true

acceptance/bundle/invariant/configs/feature_spec.yml.tmpl-init.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+# Creating a feature spec resolves its lookup table in Unity Catalog, and the invariant harness
+# provisions no tables of its own.
+echo "Creating temporary schema and feature table: main.fspec_inv_$UNIQUE_NAME"
+
+$CLI schemas create fspec_inv_$UNIQUE_NAME main -o json | jq '{full_name}'
+create_feature_spec_fixtures.py main.fspec_inv_$UNIQUE_NAME

acceptance/bundle/invariant/continue_293/out.test.toml

@@ -17,6 +17,7 @@ EnvMatrix.INPUT_CONFIG = [
   "experiment.yml.tmpl",
   "external_location.yml.tmpl",
   "feature.yml.tmpl",
+  "feature_spec.yml.tmpl",
   "job.yml.tmpl",
   "job_pydabs_10_tasks.yml.tmpl",
   "job_run_job_ref.yml.tmpl",

acceptance/bundle/invariant/delete_idempotent/out.test.toml

@@ -19,6 +19,7 @@ EnvMatrix.INPUT_CONFIG = [
   "experiment.yml.tmpl",
   "external_location.yml.tmpl",
   "feature.yml.tmpl",
+  "feature_spec.yml.tmpl",
   "genie_space.yml.tmpl",
   "instance_pool.yml.tmpl",
   "job.yml.tmpl",

acceptance/bundle/invariant/destroy_idempotent/out.test.toml

@@ -19,6 +19,7 @@ EnvMatrix.INPUT_CONFIG = [
   "experiment.yml.tmpl",
   "external_location.yml.tmpl",
   "feature.yml.tmpl",
+  "feature_spec.yml.tmpl",
   "genie_space.yml.tmpl",
   "instance_pool.yml.tmpl",
   "job.yml.tmpl",

acceptance/bundle/invariant/migrate/test.toml

@@ -43,6 +43,8 @@ EnvMatrixExclude.no_model_service = ["INPUT_CONFIG=model_service.yml.tmpl"]
 EnvMatrixExclude.no_data_quality_monitor = ["INPUT_CONFIG=data_quality_monitor.yml.tmpl"]
 # Features are direct-only; the terraform deploy that seeds the migration fails for them.
 EnvMatrixExclude.no_feature = ["INPUT_CONFIG=feature.yml.tmpl"]
+# Feature specs are direct-only; the terraform deploy that seeds the migration fails for them.
+EnvMatrixExclude.no_feature_spec = ["INPUT_CONFIG=feature_spec.yml.tmpl"]
 # The Postgres Data API is direct-only; the terraform deploy that seeds the migration fails for it.
 EnvMatrixExclude.no_postgres_data_api = ["INPUT_CONFIG=postgres_data_api.yml.tmpl"]
 

acceptance/bundle/invariant/no_drift/out.test.toml

@@ -19,6 +19,7 @@ EnvMatrix.INPUT_CONFIG = [
   "experiment.yml.tmpl",
   "external_location.yml.tmpl",
   "feature.yml.tmpl",
+  "feature_spec.yml.tmpl",
   "genie_space.yml.tmpl",
   "instance_pool.yml.tmpl",
   "job.yml.tmpl",

acceptance/bundle/invariant/test.toml

@@ -33,6 +33,7 @@ EnvMatrix.INPUT_CONFIG = [
     "experiment.yml.tmpl",
     "external_location.yml.tmpl",
     "feature.yml.tmpl",
+    "feature_spec.yml.tmpl",
     "genie_space.yml.tmpl",
     "instance_pool.yml.tmpl",
     "job.yml.tmpl",
@@ -125,6 +126,8 @@ EnvMatrixExclude.no_data_quality_monitor_on_cloud = ["CONFIG_Cloud=true", "INPUT
 # Deployment history has no "features." resource-key prefix, so the deploy fails before the feature
 # is created; see dmsResourceTypes in libs/testserver/bundledeployments.go.
 EnvMatrixExclude.no_feature_dms = ["DMS=true", "INPUT_CONFIG=feature.yml.tmpl"]
+# Same for "feature_specs.".
+EnvMatrixExclude.no_feature_spec_dms = ["DMS=true", "INPUT_CONFIG=feature_spec.yml.tmpl"]
 # Same for the "postgres_data_apis." prefix.
 EnvMatrixExclude.no_postgres_data_api_dms = ["DMS=true", "INPUT_CONFIG=postgres_data_api.yml.tmpl"]
 # Fake SQL endpoint for local tests

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

@@ -0,0 +1,23 @@
+bundle:
+  name: my_project
+
+sync: {paths: []} # don't need to copy files
+
+python:
+  resources:
+    - "resources:load_resources"
+  mutators:
+    - "mutators:update_feature_spec"
+
+resources:
+  feature_specs:
+    my_feature_spec_1:
+      name: main.default.my_feature_spec_1
+      features:
+        - feature_lookup:
+            table: main.default.user_features
+            lookup_keys: [user_id]
+            feature_names: [age]
+      grants:
+        - principal: account users
+          privileges: [EXECUTE]

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

@@ -0,0 +1,11 @@
+from dataclasses import replace
+
+from databricks.bundles.core import feature_spec_mutator
+from databricks.bundles.feature_specs import FeatureSpec
+
+
+@feature_spec_mutator
+def update_feature_spec(feature_spec: FeatureSpec) -> FeatureSpec:
+    assert isinstance(feature_spec.name, str)
+
+    return replace(feature_spec, name=f"{feature_spec.name}_updated")

acceptance/bundle/python/feature_specs-support/out.test.toml

@@ -0,0 +1,3 @@
+Cloud = true
+EnvMatrix.DMS = ["", "true"]
+EnvMatrix.PYDAB_VERSION = ["current"]

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

@@ -0,0 +1,57 @@
+
+>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json
+{
+  "experimental": {
+    "python": {
+      "mutators": [
+        "mutators:update_feature_spec"
+      ],
+      "resources": [
+        "resources:load_resources"
+      ]
+    }
+  },
+  "resources": {
+    "feature_specs": {
+      "my_feature_spec_1": {
+        "features": [
+          {
+            "feature_lookup": {
+              "feature_names": [
+                "age"
+              ],
+              "lookup_keys": [
+                "user_id"
+              ],
+              "table": "main.default.user_features"
+            }
+          }
+        ],
+        "grants": [
+          {
+            "principal": "account users",
+            "privileges": [
+              "EXECUTE"
+            ]
+          }
+        ],
+        "name": "main.default.my_feature_spec_1_updated"
+      },
+      "my_feature_spec_2": {
+        "extra_pip_requirements": {
+          "requirements": [
+            "numpy==2.1.0"
+          ]
+        },
+        "features": [
+          {
+            "declarative_feature": {
+              "feature": "main.default.my_feature"
+            }
+          }
+        ],
+        "name": "main.default.my_feature_spec_2_updated"
+      }
+    }
+  }
+}

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

@@ -0,0 +1,18 @@
+from databricks.bundles.core import Resources
+
+
+def load_resources() -> Resources:
+    resources = Resources()
+
+    resources.add_feature_spec(
+        "my_feature_spec_2",
+        {
+            "name": "main.default.my_feature_spec_2",
+            "features": [
+                {"declarative_feature": {"feature": "main.default.my_feature"}},
+            ],
+            "extra_pip_requirements": {"requirements": ["numpy==2.1.0"]},
+        },
+    )
+
+    return resources

acceptance/bundle/python/feature_specs-support/script

@@ -0,0 +1,4 @@
+trace uv run $UV_ARGS -q $CLI bundle validate --output json | \
+  jq "pick(.experimental.python, .resources)"
+
+rm -fr .databricks __pycache__

acceptance/bundle/python/feature_specs-support/test.toml

@@ -0,0 +1,4 @@
+Cloud = true
+
+# feature specs are only supported in the current version of the wheel
+EnvMatrix.PYDAB_VERSION = ["current"]