diff --git a/acceptance/bin/dms_resources.py b/acceptance/bin/dms_resources.py new file mode 100644 index 00000000000..019ee57709c --- /dev/null +++ b/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 /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 diff --git a/acceptance/bin/read_id.py b/acceptance/bin/read_id.py index 87cd2954bdf..06e7b69281e 100755 --- a/acceptance/bin/read_id.py +++ b/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) diff --git a/acceptance/bin/read_state.py b/acceptance/bin/read_state.py index 0166bf9abbb..4d9bf84186e 100755 --- a/acceptance/bin/read_state.py +++ b/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:]) diff --git a/acceptance/bin/replace_ids.py b/acceptance/bin/replace_ids.py index e0ebf1cd0e6..2f5d8165032 100755 --- a/acceptance/bin/replace_ids.py +++ b/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")