bundle: record and read deployment state via DMS

acceptance/bin/dms_resources.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""
+Read resource ids and state from the deployment metadata service.
+
+While a bundle records deployment history the service owns the resource set, so the state file is
+not where ids and state come from. The service is asked instead, which takes two lookups: the CLI
+keeps no deployment id locally, and the id is the object id of the workspace node the service
+registers under <state path>/resources.deployment.json (see libs/dms/resolve.go).
+"""
+
+import functools
+import glob
+import json
+import os
+import posixpath
+import subprocess
+import sys
+
+sys.path.insert(0, os.path.dirname(__file__))
+from print_state import get_state_file
+
+CLI = os.environ.get("CLI", "databricks")
+
+# Must match dms.DeploymentNodeName.
+DEPLOYMENT_NODE_NAME = "resources.deployment.json"
+
+
+def run_json(cmd, allow_failure=False):
+    """Run cmd and parse its stdout, or return None if it fails and allow_failure is set. stderr is
+    captured rather than inherited: these lookups are plumbing, and a CLI warning like "no files to
+    sync" would otherwise land in the test output."""
+    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8")
+    if result.returncode != 0:
+        if allow_failure:
+            return None
+        raise SystemExit(f"{cmd} failed with code {result.returncode}\n{result.stdout}{result.stderr}".strip())
+    return json.loads(result.stdout)
+
+
+def records_deployment_history():
+    """Whether this run records deployment history, so the service is what to ask."""
+    return os.environ.get("DATABRICKS_BUNDLE_RECORD_DEPLOYMENT_HISTORY") == "true"
+
+
+def get_remote_state_path(target):
+    """The bundle's remote state directory.
+
+    Preferred source is the sync snapshot, because it needs no CLI call: re-running the config
+    load would need whatever --var and flags the test deployed with, which a helper cannot know.
+    A bundle with no files to sync writes no snapshot, so fall back to asking the CLI - those
+    bundles are the ones with nothing to parameterize."""
+    target_dir = os.path.dirname(get_state_file(target, False))
+    snapshots = glob.glob(f"{target_dir}/sync-snapshots/*.json")
+    if snapshots:
+        # One snapshot per remote path, so a test that moved its root leaves several: the newest
+        # is the one the last deploy used.
+        newest = max(snapshots, key=os.path.getmtime)
+        remote_path = json.loads(open(newest).read())["remote_path"]
+        # state and files are siblings under the bundle root.
+        return posixpath.join(posixpath.dirname(remote_path), "state")
+
+    args = [CLI, "bundle", "validate", "--output", "json"]
+    if target:
+        args += ["-t", target]
+    return run_json(args)["workspace"]["state_path"]
+
+
+@functools.cache
+def get_resources(target):
+    """Map every recorded resource key ("jobs.foo") to its {"id", "state"}.
+
+    Empty when the bundle has no deployment recorded yet. Cached because a lookup costs three
+    round trips and a script asks for one resource at a time.
+    """
+    state_path = get_remote_state_path(target)
+    if not state_path:
+        return {}
+
+    # No node means nothing has been recorded, the conclusion dms.resolveDeploymentID also draws
+    # from a 404 - the deployment is gone once the bundle is destroyed.
+    node = run_json([CLI, "workspace", "get-status", f"{state_path}/{DEPLOYMENT_NODE_NAME}"], allow_failure=True)
+    if not node or not node.get("object_id"):
+        return {}
+    deployment_id = node["object_id"]
+
+    result = {}
+    # The service pages at 50 resources; the local fake returns everything at once.
+    page_token = None
+    while True:
+        url = f"/api/2.0/bundle/deployments/{deployment_id}/resources"
+        if page_token:
+            url += f"?page_token={page_token}"
+        listed = run_json([CLI, "api", "get", url])
+        for resource in listed.get("resources") or []:
+            # The service stores state as the opaque envelope the CLI wrote (dstate.RecordedState),
+            # so unwrap it to the resource state itself.
+            envelope = json.loads(resource["state"]) if resource.get("state") else {}
+            result[resource["resource_key"]] = {
+                "id": resource.get("resource_id"),
+                "state": envelope.get("state") or {},
+                "depends_on": envelope.get("depends_on") or [],
+            }
+        page_token = listed.get("next_page_token")
+        if not page_token:
+            return result

acceptance/bin/print_state.py

@@ -8,9 +8,14 @@
 
 import argparse
 import glob
+import json
 import os
 
 
+def records_deployment_history():
+    return os.environ.get("DATABRICKS_BUNDLE_RECORD_DEPLOYMENT_HISTORY") == "true"
+
+
 def print_file(filename):
     data = open(filename).read()
     print(data, end="")
@@ -53,14 +58,38 @@ def get_state_file(target, backup):
     return filtered[0] if filtered else result[0]
 
 
+def print_recorded_state(filename, target):
+    """Print the state file with its resources filled in from the deployment metadata service.
+
+    While recording, the file itself carries only the header - the service holds the resources - so
+    printing it raw would show an empty state and differ from the same test's non-recording run.
+    """
+    # Imported here rather than at module level: dms_resources reads get_state_file from this module.
+    from dms_resources import get_resources
+
+    data = json.loads(open(filename).read())
+    state = {}
+    for key, value in sorted(get_resources(target).items()):
+        entry = {"__id__": value["id"], "state": value["state"]}
+        if value["depends_on"]:
+            entry["depends_on"] = value["depends_on"]
+        state[f"resources.{key}"] = entry
+    data["state"] = state
+    print(json.dumps(data, indent=1))
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("-t", "--target")
     parser.add_argument("--backup", action="store_true")
     args = parser.parse_args()
 
     for filename in get_state_files(args.target, args.backup):
-        if os.path.exists(filename):
+        if not os.path.exists(filename):
+            continue
+        if filename.endswith("resources.json") and records_deployment_history():
+            print_recorded_state(filename, args.target)
+        else:
             print_file(filename)
 
 

acceptance/bin/read_id.py

@@ -15,6 +15,7 @@
 
 sys.path.insert(0, str(Path(__file__).parent))
 from add_repl import add_repl
+from dms_resources import get_resources, records_deployment_history
 from print_state import get_state_file
 
 
@@ -34,6 +35,15 @@ def get_id_terraform(filename, name):
     print(f"Cannot find resource with {name=}. Available: {available}", file=sys.stderr)
 
 
+def get_id_recorded(target, name):
+    resources = get_resources(target)
+    for key, value in resources.items():
+        if key.split(".")[1] == name:
+            return value["id"]
+
+    print(f"Cannot find recorded resource with {name=}. Available: {list(resources)}", file=sys.stderr)
+
+
 def get_id_direct(filename, name):
     raw = open(filename).read()
     data = json.loads(raw)
@@ -53,11 +63,14 @@ def main():
     parser.add_argument("name")
     args = parser.parse_args()
 
-    filename = get_state_file(args.target, args.backup)
-    if filename.endswith(".tfstate"):
-        id = get_id_terraform(filename, args.name)
+    if records_deployment_history():
+        id = get_id_recorded(args.target, args.name)
     else:
-        id = get_id_direct(filename, args.name)
+        filename = get_state_file(args.target, args.backup)
+        if filename.endswith(".tfstate"):
+            id = get_id_terraform(filename, args.name)
+        else:
+            id = get_id_direct(filename, args.name)
 
     if id:
         print(id)

acceptance/bin/read_state.py

@@ -9,6 +9,9 @@
 import os
 import sys
 
+sys.path.insert(0, os.path.dirname(__file__))
+from dms_resources import get_resources, records_deployment_history
+
 
 def print_resource_terraform(group, name, *attrs):
     resource_type = "databricks_" + group[:-1]
@@ -50,7 +53,21 @@ def print_resource_direct(group, name, *attrs):
     print(group, name, " ".join(values))
 
 
-if os.environ.get("DATABRICKS_BUNDLE_ENGINE", "").startswith("direct"):
+def print_resource_recorded(group, name, *attrs):
+    result = get_resources(None).get(f"{group}.{name}")
+    if result is None:
+        print(f"State not found for {group}.{name}")
+        return
+
+    state = dict(result["state"])
+    state.setdefault("id", result["id"])
+    values = [f"{x}={state.get(x)!r}" for x in attrs]
+    print(group, name, " ".join(values))
+
+
+if records_deployment_history():
+    print_resource_recorded(*sys.argv[1:])
+elif os.environ.get("DATABRICKS_BUNDLE_ENGINE", "").startswith("direct"):
     print_resource_direct(*sys.argv[1:])
 else:
     print_resource_terraform(*sys.argv[1:])

acceptance/bin/replace_ids.py

@@ -10,6 +10,7 @@
 
 sys.path.insert(0, str(Path(__file__).parent))
 from add_repl import add_repl
+from dms_resources import get_resources, records_deployment_history
 from print_state import get_state_file
 
 
@@ -26,6 +27,12 @@ def iter_ids_terraform(filename):
             yield r_name, id
 
 
+def iter_ids_recorded(target):
+    for key, value in get_resources(target).items():
+        if value["id"]:
+            yield key.split(".")[1], value["id"]
+
+
 def iter_ids_direct(filename):
     raw = open(filename).read()
     data = json.loads(raw)
@@ -44,11 +51,14 @@ def main():
     parser.add_argument("--backup", action="store_true")
     args = parser.parse_args()
 
-    filename = get_state_file(args.target, args.backup)
-    if filename.endswith(".tfstate"):
-        it = iter_ids_terraform(filename)
+    if records_deployment_history():
+        it = iter_ids_recorded(args.target)
     else:
-        it = iter_ids_direct(filename)
+        filename = get_state_file(args.target, args.backup)
+        if filename.endswith(".tfstate"):
+            it = iter_ids_terraform(filename)
+        else:
+            it = iter_ids_direct(filename)
 
     for name, id in it:
         add_repl(id, name.upper() + "_ID")

acceptance/bundle/deploy/readplan/test.toml

@@ -1,3 +1,3 @@
-# Saved plans don't carry the deployment stamp. Applying one plans an update on the next run.
-# See dms_no_readplan in acceptance/bundle/test.toml.
+# Dumps post-deploy state/plan, which carry the DMS deployment stamp under recording (applied at
+# deploy, not saved in the plan). deploy --plan recording is covered by bundle/dms, so opt out here.
 EnvMatrix.DMS = [""]

acceptance/bundle/deploy/wal/header-only-wal/output.txt

@@ -28,7 +28,7 @@ Resources: 1 created, 0 changed, 0 deleted, 0 unchanged
 
 >>> errcode assert_not_exists.py .databricks/bundle/default/resources.json.wal
 
->>> errcode cat .databricks/bundle/default/resources.json
+>>> errcode print_state.py
 {
   "serial": 1,
   "state_keys": [

acceptance/bundle/deploy/wal/header-only-wal/script

@@ -19,4 +19,4 @@ title "Third deploy (must recover and succeed, not blocked by the leftover WAL)"
 trace errcode $CLI bundle deploy --force-lock
 
 trace errcode assert_not_exists.py .databricks/bundle/default/resources.json.wal
-trace errcode cat .databricks/bundle/default/resources.json | jq -S '{serial: .serial, state_keys: (.state | keys)}'
+trace errcode print_state.py | jq -S '{serial: .serial, state_keys: (.state | keys)}'

acceptance/bundle/dms/declined-deploy/databricks.yml.tmpl

@@ -0,0 +1,10 @@
+bundle:
+  name: dms-declined-deploy-$UNIQUE_NAME
+experimental:
+  record_deployment_history: true
+
+resources:
+  schemas:
+    foo:
+      name: dms_declined_deploy_schema_$UNIQUE_NAME
+      catalog_name: main

acceptance/bundle/dms/declined-deploy/out.requests.txt

@@ -0,0 +1,7 @@
+{
+  "method": "DELETE",
+  "path": "/api/2.1/unity-catalog/catalogs/dms_other_[UNIQUE_NAME]",
+  "q": {
+    "force": "true"
+  }
+}