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,37 @@
+
+=== 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
+>>> [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

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

@@ -0,0 +1,19 @@
+# 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.
+
+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 $CLI_V1 bundle plan 2>&1 | contains.py "state version 3 is newer than supported version 2; upgrade the CLI"

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

@@ -10,16 +10,22 @@ 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.
 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,30 @@
+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")
+}

bundle/direct/dstate/migrate_test.go

@@ -21,10 +21,32 @@ import (
 )
 
 const (
+	// currentStateVersion is the schema version written for normal deployments
+	// and the target that older states are migrated up to on load.
+	//
+	// NOTE: the next bump to the baseline schema must go to 4, not 3 (which is
+	// reserved by dmsStateVersion below), and should delete dmsStateVersion,
+	// folding DMS state handling back into normal versioning.
 	currentStateVersion = 2
-	initialBufferSize   = 64 * 1024
-	maxWalEntrySize     = 10 * 1024 * 1024
-	walSuffix           = ".wal"
+
+	// dmsStateVersion is the schema version written when the bundle opts into the
+	// deployment metadata service via experimental.record_deployment_history. It
+	// is also the newest version this CLI understands; newer states are rejected.
+	//
+	// It is kept separate from currentStateVersion on purpose: previewing DMS
+	// must not force a state upgrade on everyone else. Non-DMS deploys stay at
+	// currentStateVersion while only DMS opt-in bumps the state to this version.
+	// Remove it once currentStateVersion is bumped (to 4) and the two reconcile.
+	dmsStateVersion = 3
+
+	// currentDmsVersion is stamped into DMS-upgraded state. Bump it when the DMS
+	// protocol changes in a way the CLI must record (e.g. a breaking change or a
+	// new minimum CLI version) so future deploys can gate behavior on it.
+	currentDmsVersion = 1
+
+	initialBufferSize = 64 * 1024
+	maxWalEntrySize   = 10 * 1024 * 1024
+	walSuffix         = ".wal"
 )
 
 // errStaleWAL is returned when the WAL serial is behind the expected serial.
@@ -42,10 +64,16 @@ type DeploymentState struct {
 }
 
 type Header struct {
-	StateVersion int    `json:"state_version"`
-	CLIVersion   string `json:"cli_version"`
-	Lineage      string `json:"lineage"`
-	Serial       int    `json:"serial"`
+	StateVersion int `json:"state_version"`
+
+	// CurrentDmsVersion records the deployment metadata service (DMS) protocol
+	// version this state was written with. Set only for states opted into DMS
+	// (see dmsStateVersion) and omitted otherwise.
+	CurrentDmsVersion int `json:"current_dms_version,omitempty"`
+
+	CLIVersion string `json:"cli_version"`
+	Lineage    string `json:"lineage"`
+	Serial     int    `json:"serial"`
 }
 
 type Database struct {
@@ -414,6 +442,17 @@ func (db *DeploymentState) UpgradeToWrite() error {
 	return appendJSONLine(db.walFile, walHead)
 }
 
+// UpgradeToDMS marks the state as opted into the deployment metadata service
+// (DMS): it bumps the schema to dmsStateVersion and stamps the current DMS
+// version. The change is persisted on the next save. State must be open for write.
+func (db *DeploymentState) UpgradeToDMS() {
+	db.AssertOpenedForWrite()
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	db.Data.StateVersion = dmsStateVersion
+	db.Data.CurrentDmsVersion = currentDmsVersion
+}
+
 func (db *DeploymentState) AssertOpenedForReadOrWrite() {
 	if db.Path == "" {
 		panic("internal error: DeploymentState must be opened first")