bundle/direct: conditionally upgrade state version when opted into DMS

Normal deploys keep writing state_version 2. When the bundle opts into
experimental.record_deployment_history, the deploy upgrades the local
state to dmsStateVersion (3) and stamps a new CurrentDmsVersion header
field, so future DMS breaking changes or minimum-CLI requirements can be
gated on what the state was written with.

migrateState now reads up to dmsStateVersion but still auto-migrates only
up to the baseline currentStateVersion, so a v3 state is neither
auto-applied nor downgraded. The upgrade is applied explicitly in
deployCore, never as an automatic migration.

Co-authored-by: Isaac

Author: shreyas-goenka

acceptance/acceptance_test.go

@@ -277,6 +277,13 @@ func testAccept(t *testing.T, inprocessMode bool, singleTest string) int {
 		cli293Path := DownloadCLI(t, buildDir, "0.293.0")
 		t.Setenv("CLI_293", cli293Path)
 		repls.SetPath(cli293Path, "[CLI_293]")
+
+		// v1.0.0 understands state schema versions only up to 2. Used by tests
+		// asserting that an older CLI rejects newer state (e.g. the v3 state
+		// written when a bundle opts into the deployment metadata service).
+		cliV1Path := DownloadCLI(t, buildDir, "1.0.0")
+		t.Setenv("CLI_V1", cliV1Path)
+		repls.SetPath(cliV1Path, "[CLI_V1]")
 	}
 
 	paths := []string{

acceptance/bundle/deploy/record-deployment-history/state-upgrade/databricks.dms.yml

@@ -0,0 +1,10 @@
+bundle:
+  name: test-rdh-state-upgrade
+
+experimental:
+  record_deployment_history: true
+
+resources:
+  jobs:
+    foo:
+      name: foo-job-renamed

acceptance/bundle/deploy/record-deployment-history/state-upgrade/databricks.yml

@@ -0,0 +1,7 @@
+bundle:
+  name: test-rdh-state-upgrade
+
+resources:
+  jobs:
+    foo:
+      name: foo-job

acceptance/bundle/deploy/record-deployment-history/state-upgrade/out.test.toml

@@ -0,0 +1,52 @@
+
+=== without the flag: state stays at the baseline schema version
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-rdh-state-upgrade/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_state.py
+{
+  "state_version": 2,
+  "current_dms_version": null
+}
+
+=== with experimental.record_deployment_history: state upgrades and records the DMS version
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-rdh-state-upgrade/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_state.py
+{
+  "state_version": 3,
+  "current_dms_version": 1
+}
+
+=== an older CLI (v1.0.0, max state version 2) rejects operations on the upgraded state
+>>> errcode [CLI_V1] bundle plan
+Warning: unknown field: record_deployment_history
+  at experimental
+  in databricks.yml:5:3
+
+Error: migrating state [TEST_TMP_DIR]/.databricks/bundle/default/resources.json: state version 3 is newer than supported version 2; upgrade the CLI
+
+
+Exit code: 1
+
+=== with the flag, a state written by a newer record_deployment_history version is rejected
+>>> update_file.py .databricks/bundle/default/resources.json "current_dms_version": 1 "current_dms_version": 2
+
+>>> errcode [CLI] bundle plan
+Error: record_deployment_history state version 2 is newer than supported version 1; upgrade the CLI
+
+
+Exit code: 1
+
+=== without the flag, the same state is accepted (the check is gated on opt-in)
+>>> update_file.py databricks.yml record_deployment_history: true record_deployment_history: false
+
+>>> errcode [CLI] bundle plan
+Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged

acceptance/bundle/deploy/record-deployment-history/state-upgrade/output.txt

@@ -0,0 +1,35 @@
+# State-version lifecycle for the DMS preview (experimental.record_deployment_history):
+#   1. Without the flag, deploys keep writing the baseline state version.
+#   2. Opting in upgrades the state version and records the DMS version.
+#   3. An older CLI that predates the DMS state version refuses to operate on the
+#      upgraded state and tells the user to upgrade, so it cannot mishandle it.
+#   4. With the flag set, a state stamped with a newer record_deployment_history
+#      version than this CLI knows is also refused (the state schema version is
+#      one we support). Without the flag the check is skipped.
+
+title "without the flag: state stays at the baseline schema version"
+trace $CLI bundle deploy
+trace print_state.py | jq '{state_version, current_dms_version}'
+
+title "with experimental.record_deployment_history: state upgrades and records the DMS version"
+# Also renames the job so the deploy writes state (a no-op deploy would not
+# rewrite the state file and the upgrade would not be observable yet).
+cp databricks.dms.yml databricks.yml
+trace $CLI bundle deploy
+trace print_state.py | jq '{state_version, current_dms_version}'
+
+title "an older CLI (v1.0.0, max state version 2) rejects operations on the upgraded state"
+trace errcode $CLI_V1 bundle plan 2>&1 | contains.py "state version 3 is newer than supported version 2; upgrade the CLI"
+
+title "with the flag, a state written by a newer record_deployment_history version is rejected"
+# The state must hold a record_deployment_history version greater than this CLI's
+# supported version (currentDmsVersion) to be rejected, so we set supported+1.
+# TODO: when currentDmsVersion is bumped, change the value written below to the new
+# supported+1. The error's version numbers are not asserted here, so the assertion
+# needs no change.
+trace update_file.py .databricks/bundle/default/resources.json '"current_dms_version": 1' '"current_dms_version": 2'
+trace errcode $CLI bundle plan 2>&1 | contains.py "record_deployment_history state version" "is newer than supported version" "upgrade the CLI"
+
+title "without the flag, the same state is accepted (the check is gated on opt-in)"
+trace update_file.py databricks.yml "record_deployment_history: true" "record_deployment_history: false"
+trace errcode $CLI bundle plan 2>&1 | contains.py "!is newer than supported version"

acceptance/bundle/deploy/record-deployment-history/state-upgrade/script

@@ -0,0 +1,10 @@
+Local = true
+Cloud = false
+
+# databricks.yml is rewritten in-place by the script (flag added in the second step).
+Ignore = [".databricks", "databricks.yml"]
+
+# The state-version upgrade only applies to the direct engine; terraform stores
+# state differently and would diverge.
+[EnvMatrix]
+DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/bundle/deploy/record-deployment-history/state-upgrade/test.toml

@@ -1,3 +1,3 @@
-state version 999 is newer than supported version 2; upgrade the CLI
+state version 999 is newer than supported version 3; upgrade the CLI
 
 Exit code: 1

acceptance/bundle/state/future_version/output.txt

@@ -10,16 +10,25 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/iam"
 )
 
-// migrateState runs all necessary migrations on the database.
+// migrateState brings a freshly-loaded state up to a version this CLI can use.
 // It is called after loading state from disk.
+//
+// Two versions are "current" and left untouched: the baseline currentStateVersion
+// and the opt-in dmsStateVersion (written when a bundle previews DMS). Legacy
+// states below the baseline are migrated forward via the migrations map. A state
+// newer than dmsStateVersion was written by a newer CLI, so we refuse it rather
+// than risk mishandling a format we don't understand.
+//
+// The DMS protocol version (current_dms_version) is enforced separately and only
+// when the bundle has opted into DMS; see DeploymentState.EnsureSupportedDmsVersion.
 func migrateState(db *Database) error {
-	if db.StateVersion == currentStateVersion {
-		return nil
-	}
-	if db.StateVersion > currentStateVersion {
-		return fmt.Errorf("state version %d is newer than supported version %d; upgrade the CLI", db.StateVersion, currentStateVersion)
+	if db.StateVersion > dmsStateVersion {
+		return fmt.Errorf("state version %d is newer than supported version %d; upgrade the CLI", db.StateVersion, dmsStateVersion)
 	}
 
+	// Only legacy states (below the baseline) migrate here. The DMS upgrade is an
+	// explicit deploy-time step (see UpgradeToDMS), never an automatic migration,
+	// so a dmsStateVersion state falls through this loop unchanged.
 	for version := db.StateVersion; version < currentStateVersion; version++ {
 		fn, ok := migrations[version]
 		if !ok {

bundle/direct/dstate/migrate.go

@@ -0,0 +1,95 @@
+package dstate
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMigrateStateLeavesCurrentUntouched(t *testing.T) {
+	db := &Database{Header: Header{StateVersion: currentStateVersion}}
+	require.NoError(t, migrateState(db))
+	assert.Equal(t, currentStateVersion, db.StateVersion)
+}
+
+func TestMigrateStateLeavesDMSStateUntouched(t *testing.T) {
+	// A DMS-upgraded state is already current; it must not be downgraded to the
+	// baseline version, and the recorded DMS version must be preserved.
+	db := &Database{Header: Header{StateVersion: dmsStateVersion, CurrentDmsVersion: currentDmsVersion}}
+	require.NoError(t, migrateState(db))
+	assert.Equal(t, dmsStateVersion, db.StateVersion)
+	assert.Equal(t, currentDmsVersion, db.CurrentDmsVersion)
+}
+
+func TestMigrateStateRejectsNewerThanSupported(t *testing.T) {
+	db := &Database{Header: Header{StateVersion: dmsStateVersion + 1}}
+	err := migrateState(db)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "upgrade the CLI")
+}
+
+// TestStateSchemaVersions pins the state schema version constants. They are part
+// of the on-disk format and a contract with older CLIs, so changing them must be
+// deliberate. If you are intentionally bumping the baseline schema, this test is
+// your checklist: follow the steps below, then update the assertions last.
+//
+// HOW TO BUMP THE BASELINE STATE VERSION (2 -> 4) AND RETIRE THE DMS SPECIAL-CASING
+//
+// Version 3 only exists because previewing DMS bumped the schema out-of-band via
+// UpgradeToDMS, so non-preview bundles weren't forced to upgrade. The next baseline
+// bump is when you delete all of that custom code and make 3 an ordinary version
+// in the linear migration chain. Go to 4, not 3 (3 is already in the wild):
+//
+//  1. state.go: set currentStateVersion = 4. Delete the dmsStateVersion constant
+//     and the UpgradeToDMS method. Version 3 is now reached only by migration,
+//     like any other version. (Leave currentDmsVersion and Header.CurrentDmsVersion
+//     alone; they are a separate concern, see step 4.)
+//
+//  2. migrate.go: change the upper-bound guard from "> dmsStateVersion" back to
+//     "> currentStateVersion" (there is again only one current version). Add
+//     stepwise migrations so every old state climbs to 4:
+//     migrations[2] = v2 (non-DMS baseline)   -> v3
+//     migrations[3] = v3 (former DMS preview) -> v4
+//     A v2 state climbs 2->3->4 and a v3 state climbs 3->4, exactly like any
+//     other version. Write a real transform for whatever the v4 change is.
+//
+//  3. deploy.go: delete the `if RecordDeploymentHistory { UpgradeToDMS() }` block.
+//     The version is no longer bumped conditionally; every deploy writes the
+//     baseline through the normal path.
+//
+//  4. current_dms_version is NOT part of this bump. It tracks the DMS *protocol*
+//     version (currentDmsVersion), independent of the state schema version, and is
+//     enforced by EnsureSupportedDmsVersion (called from cmd/bundle/utils/process.go
+//     only when the bundle has opted into DMS). Leave the field, the constant, and
+//     that check in place; if UpgradeToDMS was the only place stamping it, move
+//     that stamping into the normal write path.
+//
+//  5. Tests: update the assertions below (the dmsStateVersion assertion and the
+//     DMS-schema cases go away); add 2->3->4 and 3->4 migration tests.
+//     TestMigrationsCoverBaseline fails until migrations[2] and migrations[3] exist.
+//
+// RELATED COVERAGE
+//   - acceptance/bundle/state/permission_level_migration is a golden migration
+//     fixture: it commits a real v1 state and asserts the migrated v2 output,
+//     catching migration-correctness bugs (not just "did you mean to change this").
+//   - acceptance/bundle/deploy/record-deployment-history/state-upgrade drives the
+//     full lifecycle, including both rejections (older CLI vs newer state, and a
+//     newer DMS version vs this CLI).
+//   - bundle/invariant/continue_293 asserts the current CLI reads state written by
+//     an older released CLI, so we never break reading older state.
+func TestStateSchemaVersions(t *testing.T) {
+	assert.Equal(t, 2, currentStateVersion)
+	assert.Equal(t, 3, dmsStateVersion)
+	assert.Equal(t, 1, currentDmsVersion)
+}
+
+// TestMigrationsCoverBaseline guards a baseline bump: every state version below
+// currentStateVersion must have a migration to the next version, so migrateState
+// can always climb a legacy state up to the baseline. A bump that forgets a
+// migration fails here instead of at a user's deploy.
+func TestMigrationsCoverBaseline(t *testing.T) {
+	for v := range currentStateVersion {
+		assert.Containsf(t, migrations, v, "missing migration for state version %d", v)
+	}
+}