acc: read recorded ids and state from the deployment metadata service

acceptance/bin/dms_resources.py

@@ -0,0 +1,64 @@
+#!/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 json
+import os
+import subprocess
+
+CLI = os.environ.get("CLI", "databricks")
+
+# Must match dms.DeploymentNodeName.
+DEPLOYMENT_NODE_NAME = "resources.deployment.json"
+
+
+def run_json(cmd):
+    """Run cmd and parse its stdout. 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:
+        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"
+
+
+@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.
+    """
+    args = [CLI, "bundle", "validate", "--output", "json"]
+    if target:
+        args += ["-t", target]
+    state_path = run_json(args)["workspace"]["state_path"]
+
+    node = run_json([CLI, "workspace", "get-status", f"{state_path}/{DEPLOYMENT_NODE_NAME}"])
+    deployment_id = node.get("object_id")
+    if not deployment_id:
+        return {}
+
+    listed = run_json([CLI, "api", "get", f"/api/2.0/bundle/deployments/{deployment_id}/resources"])
+
+    result = {}
+    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 {},
+        }
+    return result

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")