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/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,26 @@
+
+=== 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
+}

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

@@ -0,0 +1,10 @@
+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}'

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

@@ -13,13 +13,14 @@ import (
 // migrateState runs all necessary migrations on the database.
 // It is called after loading state from disk.
 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)
 	}
 
+	// Migrate legacy states forward to the baseline current version. A state at
+	// dmsStateVersion is already current and skipped here; opting into DMS is an
+	// explicit upgrade applied during deploy (see UpgradeToDMS), not an automatic
+	// migration.
 	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,23 @@ import (
 )
 
 const (
+	// currentStateVersion is the schema version written for normal deployments
+	// and the target that older states are migrated up to on load.
 	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.
+	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 +55,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 +433,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")

bundle/direct/dstate/state.go

@@ -32,6 +32,24 @@ func TestOpenSaveFinalizeRoundTrip(t *testing.T) {
 	mustFinalize(t, &db2)
 }
 
+func TestUpgradeToDMSPersistsVersions(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "state.json")
+
+	var db DeploymentState
+	require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true)))
+	db.UpgradeToDMS()
+	require.NoError(t, db.SaveState("jobs.my_job", "123", map[string]string{"key": "val"}, nil))
+	mustFinalize(t, &db)
+
+	// Re-open and verify the upgraded schema version and DMS version persisted,
+	// and that loading the upgraded state does not error or downgrade it.
+	var db2 DeploymentState
+	require.NoError(t, db2.Open(t.Context(), path, WithRecovery(false), WithWrite(false)))
+	assert.Equal(t, dmsStateVersion, db2.Data.StateVersion)
+	assert.Equal(t, currentDmsVersion, db2.Data.CurrentDmsVersion)
+	mustFinalize(t, &db2)
+}
+
 func TestFinalizeWithNoEntriesDoesNotWriteStateFile(t *testing.T) {
 	path := filepath.Join(t.TempDir(), "state.json")