bundle: record and read deployment state via DMS

acceptance/acceptance_test.go

@@ -443,6 +443,10 @@ func testAccept(t *testing.T, inprocessMode bool, selectedTests []string, skipTo
 	if base, _, found := strings.Cut(cliVersion, "+"); found {
 		repls.Set(base, "[CLI_VERSION]")
 	}
+	// A dev build may embed a +<git-sha> that the base-version replacement above leaves
+	// behind (e.g. "[CLI_VERSION]+abc123def456"), which would otherwise bake into a
+	// regenerated golden. Strip any such trailing suffix so goldens stay sha-independent.
+	repls.Repls = append(repls.Repls, testdiff.Replacement{Old: regexp.MustCompile(`\[CLI_VERSION\]\+[0-9a-f]{7,40}`), New: "[CLI_VERSION]"})
 	testdiff.PrepareReplacementSdkVersion(t, &repls)
 	testdiff.PrepareReplacementTfProviderVersion(t, &repls)
 	testdiff.PrepareReplacementsGoVersion(t, &repls)
@@ -1078,7 +1082,7 @@ func checkEnvFilters(t *testing.T, testEnv, envFilters []string) {
 // matrix key ends up in the variant's test name, so a long one makes every name that carries
 // it hard to read. Tests may still name the variable itself; the alias is only shorter.
 var envAliases = map[string]string{
-	"DMS": "DATABRICKS_BUNDLE_RECORD_DEPLOYMENT_HISTORY",
+	"DMS": "DATABRICKS_BUNDLE_DEPLOYMENT_HISTORY",
 }
 
 // buildTestEnv builds the test environment from config.Env and customEnv.

acceptance/bin/nostamp.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Read JSON on stdin, write it back with the DMS deployment stamp removed.
+
+Deployment history recording (DATABRICKS_BUNDLE_DEPLOYMENT_HISTORY=true; see
+bundle/test.toml) adds a deployment stamp to plans, states, and resource payloads. Pipe any of
+those through this so an acceptance golden compares equal whether or not recording is on. Output
+is a 2-space indent, keys in input order, <>& left unescaped, and integers at full precision.
+Tests under bundle/dms assert the stamp itself and must not use it.
+
+Removed in three shapes plus the plan header:
+  1. deployment_id/version_id nested in a deployment block, recognized by the
+     neighbouring "kind" and "metadata_file_path", which are kept; version_id
+     == "" is kept (a Terraform state dump carries that for an unstamped job).
+  2. changes entries keyed "deployment.deployment_id" / "deployment.version_id".
+  3. a changes object emptied by (2), or already empty.
+  4. the recorded plan header field features.
+"""
+
+import argparse
+import json
+import sys
+
+_STAMP_KEYS = ("deployment_id", "version_id")
+_CHANGE_KEYS = ("deployment.deployment_id", "deployment.version_id")
+_PLAN_HEADER_KEYS = ("features",)
+
+
+def scrub(node):
+    if isinstance(node, dict):
+        # 1. A deployment block is the only object carrying both of these.
+        if "kind" in node and "metadata_file_path" in node:
+            for k in _STAMP_KEYS:
+                if node.get(k, "") != "":
+                    node.pop(k, None)
+        out = {}
+        for k, v in node.items():
+            if k == "changes" and isinstance(v, dict):
+                v = {ck: scrub(cv) for ck, cv in v.items() if ck not in _CHANGE_KEYS}
+                if not v:
+                    continue  # 3. drop a changes object left (or already) empty
+                out[k] = v
+            else:
+                out[k] = scrub(v)
+        return out
+    if isinstance(node, list):
+        return [scrub(x) for x in node]
+    return node
+
+
+def render(data, indent):
+    data = scrub(data)
+    if isinstance(data, dict):
+        for k in _PLAN_HEADER_KEYS:
+            data.pop(k, None)  # 4. plan header, present only at the root
+    return json.dumps(data, indent=indent, ensure_ascii=False, separators=(",", ": "))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    # A state dump is printed with a single-space indent; plans use two.
+    parser.add_argument("--indent", type=int, default=2)
+    args = parser.parse_args()
+
+    # Input is one JSON value (a plan or state dump) or a whitespace-separated stream of them (each
+    # request emitted by a print_requests filter) - jq accepted both here, so this must too.
+    text = sys.stdin.read()
+    decoder = json.JSONDecoder()
+    idx, n = 0, len(text)
+    while idx < n:
+        while idx < n and text[idx].isspace():
+            idx += 1
+        if idx >= n:
+            break
+        data, idx = decoder.raw_decode(text, idx)
+        sys.stdout.write(render(data, args.indent) + "\n")
+
+
+if __name__ == "__main__":
+    main()

acceptance/bin/print_state.py

@@ -1,14 +1,20 @@
 #!/usr/bin/env python3
 """
-Print resources state from default target.
+Print resources state from default target, and handle DMS (deployment metadata service) interactions.
 
 Note, this intentionally has no logic on guessing what is the right state file (e.g. via DATABRICKS_BUNDLE_ENGINE),
 the goal is to record all states that are available.
 """
 
 import argparse
+import functools
 import glob
+import json
 import os
+import posixpath
+import subprocess
+
+from nostamp import scrub
 
 
 def print_file(filename):
@@ -53,14 +59,173 @@ def get_state_file(target, backup):
     return filtered[0] if filtered else result[0]
 
 
+# DMS (deployment metadata service) interaction helpers
+# https://go.databricks.com/dms-service
+
+CLI = os.environ.get("CLI", "databricks")
+DEPLOYMENT_NODE_NAME = "resources.deployment.json"  # Must match dms.DeploymentNodeName
+
+
+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 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_deployment_id(target):
+    """The recorded deployment's id, or None when nothing is recorded.
+
+    Cached and shared: both the resource listing and the version lookup need it, and resolving it
+    costs a workspace round trip. 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.
+    """
+    state_path = get_remote_state_path(target)
+    if not state_path:
+        return None
+    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 None
+    return node["object_id"]
+
+
+@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.
+    """
+    deployment_id = get_deployment_id(target)
+    if not deployment_id:
+        return {}
+
+    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
+
+
+def get_recorded_state(target):
+    """The recorded resources in the on-disk state file's `state` shape (resources.<key> ->
+    {__id__, state, depends_on}), so a recording run prints the same shape as a non-recording one."""
+    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
+    return state
+
+
+def get_last_version_id(target):
+    """The version the service has recorded for this deployment, or None when nothing is recorded."""
+    deployment_id = get_deployment_id(target)
+    if not deployment_id:
+        return None
+    deployment = run_json([CLI, "api", "get", f"/api/2.0/bundle/deployments/{deployment_id}"])
+    return (deployment or {}).get("last_version_id")
+
+
+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.
+    """
+    data = json.loads(open(filename).read())
+    # Recording stamps each resource payload with the deployment and version; drop it here so the
+    # printed state matches a non-recording run without every caller piping through nostamp.
+    data["state"] = scrub(get_recorded_state(target))
+
+    # The service owns the version and the file persists no serial, so take it from the deployment.
+    # Rebuilt in header order, since the file has no serial key to overwrite in place.
+    last_version_id = get_last_version_id(target)
+    serial = int(last_version_id) if last_version_id else 0
+
+    # Recording itself is not what these tests assert, so drop the feature that marks it.
+    features = {k: v for k, v in (data.get("features") or {}).items() if k != "deployment_history"}
+
+    rebuilt = {}
+    for key in ("state_version", "cli_version", "lineage"):
+        if key in data:
+            rebuilt[key] = data[key]
+    rebuilt["serial"] = serial
+    if features:
+        rebuilt["features"] = features
+    for key, value in data.items():
+        if key not in rebuilt and key != "features":
+            rebuilt[key] = value
+    print(json.dumps(rebuilt, indent=1))
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("-t", "--target")
     parser.add_argument("--backup", action="store_true")
+    parser.add_argument(
+        "--no-dms",
+        action="store_true",
+        help="Print resources.json as-is without fetching state from DMS. "
+        "Use in bundle/dms tests to assert the tombstone content of the file.",
+    )
     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
+        # Recording only applies to the direct engine, so a terraform run prints the file as-is.
+        recording = os.environ.get("DATABRICKS_BUNDLE_DEPLOYMENT_HISTORY") == "true"
+        terraform = os.environ.get("DATABRICKS_BUNDLE_ENGINE") == "terraform"
+        if recording and not terraform and not args.no_dms:
+            print_recorded_state(filename, args.target)
+        else:
             print_file(filename)
 
 

acceptance/bin/read_id.py

@@ -10,12 +10,13 @@
 
 import argparse
 import json
+import os
 import sys
 from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent))
 from add_repl import add_repl
-from print_state import get_state_file
+from print_state import get_resources, get_state_file
 
 
 def get_id_terraform(filename, name):
@@ -34,6 +35,23 @@ def get_id_terraform(filename, name):
     print(f"Cannot find resource with {name=}. Available: {available}", file=sys.stderr)
 
 
+def get_id_recorded(target, name):
+    """Find a recorded resource's id by its bare name ("foo"), scanning the resource list.
+
+    The GET resource API would be the direct lookup, but it needs the full resource key
+    ("jobs.foo"), and callers pass only the leaf name - matching how user-visible commands
+    take a bare name ("bundle run my_job"), which DABs allows because it enforces name
+    uniqueness across types. So list and match on the leaf instead. Switching to GET is a
+    reasonable follow-up; it likely needs the service to key resources by name as well.
+    """
+    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 +71,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 os.environ.get("DATABRICKS_BUNDLE_DEPLOYMENT_HISTORY") == "true":
+        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)