bundle: record and read deployment state via DMS

acceptance/bin/dms_resources.py

@@ -0,0 +1,117 @@
+#!/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
+
+
+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

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_RECORD_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 fields deployment_id, next_version_id, last_version_id.
+"""
+
+import argparse
+import json
+import sys
+
+_STAMP_KEYS = ("deployment_id", "version_id")
+_CHANGE_KEYS = ("deployment.deployment_id", "deployment.version_id")
+_PLAN_HEADER_KEYS = ("deployment_id", "next_version_id", "last_version_id")
+
+
+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

@@ -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,32 @@ 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_recorded_state
+
+    data = json.loads(open(filename).read())
+    data["state"] = get_recorded_state(target)
+    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:])