Persist state to .wal immediately after resource update

NEXT_CHANGELOG.md

@@ -13,5 +13,6 @@
 * Make sure warnings asking for approval are understood by agents ([#5239](https://github.com/databricks/cli/pull/5239))
 * Support `replace_existing: true` on `postgres_branches` and `postgres_endpoints` so bundles can manage the implicitly-created production branch and primary read-write endpoint of a Lakebase project.
 * Add `postgres_catalogs` resource to bind a Unity Catalog catalog to a Postgres database on a Lakebase Autoscaling branch ([#5265](https://github.com/databricks/cli/pull/5265)).
+* engine/direct: Changes to state file now persisted to .wal file right away instead of being saved in the end ([#5149](https://github.com/databricks/cli/pull/5149))
 
 ### Dependency updates

acceptance/bin/assert_exists.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+import os, sys
+
+errors = 0
+
+for filename in sys.argv[1:]:
+    if not os.path.exists(filename):
+        sys.stderr.write(f"Unexpected: {filename} does not exist.\n")
+        errors += 1
+
+if errors:
+    sys.exit(1)

acceptance/bin/assert_not_exists.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+import os, sys
+
+errors = 0
+
+for filename in sys.argv[1:]:
+    if os.path.exists(filename):
+        sys.stderr.write(f"Unexpected: {filename} exists.\n")
+        errors += 1
+
+if errors:
+    sys.exit(1)

acceptance/bin/kill_after.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+"""Set up a kill rule on the testserver for the current test token.
+
+Usage: kill_after.py PATTERN OFFSET TIMES
+
+  PATTERN   HTTP method and path, e.g. "POST /api/2.2/jobs/create"
+  OFFSET    number of requests to let through before killing starts
+  TIMES     number of times to kill the caller
+
+The rule is scoped to the current DATABRICKS_TOKEN so it only affects
+the test that registers it, even when tests share a server.
+"""
+
+import json
+import os
+import sys
+import urllib.request
+
+host = os.environ.get("DATABRICKS_HOST", "")
+token = os.environ.get("DATABRICKS_TOKEN", "")
+
+if not host:
+    print("DATABRICKS_HOST not set", file=sys.stderr)
+    sys.exit(1)
+
+if len(sys.argv) != 4:
+    print(f"usage: {sys.argv[0]} PATTERN OFFSET TIMES", file=sys.stderr)
+    sys.exit(1)
+
+pattern, offset, times = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
+
+data = json.dumps({"pattern": pattern, "offset": offset, "times": times}).encode()
+req = urllib.request.Request(
+    f"{host}/__testserver/kill",
+    data=data,
+    headers={"Content-Type": "application/json", "Authorization": f"Bearer {token}"},
+    method="POST",
+)
+urllib.request.urlopen(req)

acceptance/bundle/deploy/wal/chain-3-jobs/databricks.yml

@@ -0,0 +1,37 @@
+bundle:
+  name: wal-chain-test
+
+resources:
+  jobs:
+    # Linear chain: job_01 -> job_02 -> job_03
+    # Execution order: job_01 first, job_03 last
+    job_01:
+      name: "job-01"
+      description: "first in chain"
+      tasks:
+        - task_key: "task"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge
+    job_02:
+      name: "job-02"
+      description: "depends on ${resources.jobs.job_01.id}"
+      tasks:
+        - task_key: "task"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge
+    job_03:
+      name: "job-03"
+      description: "depends on ${resources.jobs.job_02.id}"
+      tasks:
+        - task_key: "task"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge

acceptance/bundle/deploy/wal/chain-3-jobs/output.txt

@@ -0,0 +1,110 @@
+=== First deploy (crashes on job_03) ===
+
+>>> errcode [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default/files...
+Deploying resources...
+[PROCESS_KILLED]
+
+Exit code: [KILLED]
+
+=== WAL content after crash ===
+{
+  "cli_version": "[DEV_VERSION]",
+  "lineage": "[UUID]",
+  "serial": 1,
+  "state_version": 2
+}
+{
+  "k": "resources.jobs.job_01",
+  "v": {
+    "__id__": "[JOB_01_ID]",
+    "state": {
+      "deployment": {
+        "kind": "BUNDLE",
+        "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default/state/metadata.json"
+      },
+      "description": "first in chain",
+      "edit_mode": "UI_LOCKED",
+      "format": "MULTI_TASK",
+      "max_concurrent_runs": 1,
+      "name": "job-01",
+      "queue": {
+        "enabled": true
+      },
+      "tasks": [
+        {
+          "new_cluster": {
+            "node_type_id": "[NODE_TYPE_ID]",
+            "spark_version": "15.4.x-scala2.12"
+          },
+          "spark_python_task": {
+            "python_file": "/Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default/files/test.py"
+          },
+          "task_key": "task"
+        }
+      ]
+    }
+  }
+}
+{
+  "k": "resources.jobs.job_02",
+  "v": {
+    "__id__": "[JOB_02_ID]",
+    "depends_on": [
+      {
+        "label": "${resources.jobs.job_01.id}",
+        "node": "resources.jobs.job_01"
+      }
+    ],
+    "state": {
+      "deployment": {
+        "kind": "BUNDLE",
+        "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default/state/metadata.json"
+      },
+      "description": "depends on [JOB_01_ID]",
+      "edit_mode": "UI_LOCKED",
+      "format": "MULTI_TASK",
+      "max_concurrent_runs": 1,
+      "name": "job-02",
+      "queue": {
+        "enabled": true
+      },
+      "tasks": [
+        {
+          "new_cluster": {
+            "node_type_id": "[NODE_TYPE_ID]",
+            "spark_version": "15.4.x-scala2.12"
+          },
+          "spark_python_task": {
+            "python_file": "/Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default/files/test.py"
+          },
+          "task_key": "task"
+        }
+      ]
+    }
+  }
+}
+
+=== Number of jobs saved in WAL ===
+2
+
+=== Bundle summary (reads from WAL) ===
+Name: wal-chain-test
+Target: default
+Workspace:
+  User: [USERNAME]
+  Path: /Workspace/Users/[USERNAME]/.bundle/wal-chain-test/default
+Resources:
+  Jobs:
+    job_01:
+      Name: job-01
+      URL:  [DATABRICKS_URL]/jobs/[JOB_01_ID]?o=[NUMID]
+    job_02:
+      Name: job-02
+      URL:  [DATABRICKS_URL]/jobs/[JOB_02_ID]?o=[NUMID]
+    job_03:
+      Name: job-03
+      URL:  (not deployed)
+
+=== WAL after successful deploy ===
+WAL deleted (expected)

acceptance/bundle/deploy/wal/chain-3-jobs/script

@@ -0,0 +1,24 @@
+# Linear chain: job_01 -> job_02 -> job_03
+# Let first 2 jobs/create succeed, then kill on the 3rd
+kill_after.py "POST /api/2.2/jobs/create" 2 1
+
+echo "=== First deploy (crashes on job_03) ==="
+trace errcode $CLI bundle deploy
+
+echo ""
+echo "=== WAL content after crash ==="
+jq -S . .databricks/bundle/default/resources.json.wal 2>/dev/null || echo "No WAL file"
+
+echo ""
+echo "=== Number of jobs saved in WAL ==="
+grep -c '"k":"resources.jobs' .databricks/bundle/default/resources.json.wal 2>/dev/null || echo "0"
+
+echo ""
+echo "=== Bundle summary (reads from WAL) ==="
+$CLI bundle summary
+
+echo ""
+echo "=== WAL after successful deploy ==="
+cat .databricks/bundle/default/resources.json.wal 2>/dev/null || echo "WAL deleted (expected)"
+
+replace_ids.py

acceptance/bundle/deploy/wal/chain-3-jobs/test.py

@@ -0,0 +1 @@
+print("test")

acceptance/bundle/deploy/wal/corrupted-wal-entry/databricks.yml

@@ -0,0 +1,23 @@
+bundle:
+  name: wal-corrupted-test
+
+resources:
+  jobs:
+    valid_job:
+      name: "valid-job"
+      tasks:
+        - task_key: "task-a"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge
+    another_valid:
+      name: "another-valid"
+      tasks:
+        - task_key: "task-b"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge

acceptance/bundle/deploy/wal/corrupted-wal-entry/output.txt

@@ -0,0 +1,34 @@
+
+>>> cat .databricks/bundle/default/resources.json.wal
+{"lineage":"test-lineage-123","serial":6}
+{"k":"resources.jobs.valid_job","v":{"__id__":"","state":{"name":"valid-job"}}}
+{"k":"resources.jobs.another_valid","v":{"__id__":"","state":{"name":"another-valid"}}}
+{"k":"resources.jobs.partial_write","v":{"__id__":"33","state":{"name":"partial-
+
+>>> [CLI] bundle deploy
+Warn: Skipping corrupted WAL entry at [TEST_TMP_DIR]/.databricks/bundle/default/resources.json.wal:4: unexpected end of JSON input
+Warn: Saved 1 corrupted WAL entries to [TEST_TMP_DIR]/.databricks/bundle/default/resources.json.wal.corrupted
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/wal-corrupted-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] bundle summary
+Name: wal-corrupted-test
+Target: default
+Workspace:
+  User: [USERNAME]
+  Path: /Workspace/Users/[USERNAME]/.bundle/wal-corrupted-test/default
+Resources:
+  Jobs:
+    another_valid:
+      Name: another-valid
+      URL:  [DATABRICKS_URL]/jobs/[NUMID]?o=[NUMID]
+    valid_job:
+      Name: valid-job
+      URL:  [DATABRICKS_URL]/jobs/[NUMID]?o=[NUMID]
+
+>>> cat .databricks/bundle/default/resources.json.wal.corrupted
+{"k":"resources.jobs.partial_write","v":{"__id__":"33","state":{"name":"partial-
+=== WAL after successful deploy ===
+WAL deleted (expected)

acceptance/bundle/deploy/wal/corrupted-wal-entry/resources.json

@@ -0,0 +1,7 @@
+{
+  "state_version": 1,
+  "cli_version": "0.0.0",
+  "lineage": "test-lineage-123",
+  "serial": 5,
+  "state": {}
+}

acceptance/bundle/deploy/wal/corrupted-wal-entry/resources.json.wal.tmpl

@@ -0,0 +1,4 @@
+{"lineage":"test-lineage-123","serial":6}
+{"k":"resources.jobs.valid_job","v":{"__id__":"$JOB1","state":{"name":"valid-job"}}}
+{"k":"resources.jobs.another_valid","v":{"__id__":"$JOB2","state":{"name":"another-valid"}}}
+{"k":"resources.jobs.partial_write","v":{"__id__":"33","state":{"name":"partial-

acceptance/bundle/deploy/wal/corrupted-wal-entry/script

@@ -0,0 +1,22 @@
+# Create pre-existing jobs in the testserver so WAL recovery triggers DoUpdate (reset) instead of DoCreate
+JOB1=$($CLI jobs create --json '{"name":"valid-job"}' | jq -r '.job_id')
+JOB2=$($CLI jobs create --json '{"name":"another-valid"}' | jq -r '.job_id')
+echo "$JOB1:JOB1_ID" >> ACC_REPLS
+echo "$JOB2:JOB2_ID" >> ACC_REPLS
+
+mkdir -p .databricks/bundle/default
+cp resources.json .databricks/bundle/default/
+
+envsubst < resources.json.wal.tmpl > .databricks/bundle/default/resources.json.wal
+
+trace cat .databricks/bundle/default/resources.json.wal
+trace $CLI bundle deploy
+trace $CLI bundle summary
+trace cat .databricks/bundle/default/resources.json.wal.corrupted
+
+printf "\n=== WAL after successful deploy ===\n"
+if [ -f ".databricks/bundle/default/resources.json.wal" ]; then
+    echo "WAL exists (unexpected)"
+else
+    echo "WAL deleted (expected)"
+fi

acceptance/bundle/deploy/wal/corrupted-wal-entry/test.py

@@ -0,0 +1 @@
+print("test")

acceptance/bundle/deploy/wal/crash-after-create/databricks.yml

@@ -0,0 +1,25 @@
+bundle:
+  name: wal-crash-test
+
+resources:
+  jobs:
+    job_a:
+      name: "test-job-a"
+      description: "first job"
+      tasks:
+        - task_key: "task-a"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge
+    job_b:
+      name: "test-job-b"
+      description: "depends on ${resources.jobs.job_a.id}"
+      tasks:
+        - task_key: "task-b"
+          spark_python_task:
+            python_file: ./test.py
+          new_cluster:
+            spark_version: 15.4.x-scala2.12
+            node_type_id: i3.xlarge