cmd/bundle: read state from DMS when record_deployment_history is set

Wire the StateReader into the direct-engine read path. NewStateReader selects
the DMS reader when experimental.record_deployment_history is enabled and a
prior deployment exists (lineage recorded in resources.json), otherwise the
local file reader. process.go replaces the direct StateDB.Open call with the
selector, so plan/deploy/summary read resource state from the deployment
metadata service under the flag.

The deployment ID is the state lineage (matching the lock package), read from
the local resources.json; with no lineage yet (first deploy) there is nothing
in DMS, so the local file reader is used.

Co-authored-by: Isaac

Author: shreyas-goenka

bundle/statemgmt/statereader.go

@@ -3,8 +3,13 @@ package statemgmt
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io/fs"
+	"os"
+	"path/filepath"
 
+	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/direct/dstate"
 	sdkbundle "github.com/databricks/databricks-sdk-go/service/bundle"
 )
@@ -97,3 +102,45 @@ func fetchDeploymentState(ctx context.Context, client sdkbundle.BundleInterface,
 	}
 	return data, nil
 }
+
+// NewStateReader selects the StateReader for the bundle: the DMS reader when
+// experimental.record_deployment_history is enabled and a prior deployment
+// exists, otherwise the local resources.json reader.
+//
+// The DMS deployment ID is the state lineage, which is recorded in the local
+// resources.json (see dstate.DeploymentState.GetOrInitLineage). When there is no
+// lineage yet — a first deploy that has not registered with DMS — there is
+// nothing to read from the service, so we read the (empty) local file instead.
+func NewStateReader(ctx context.Context, b *bundle.Bundle, path string) (StateReader, error) {
+	if b.Config.Experimental == nil || !b.Config.Experimental.RecordDeploymentHistory {
+		return NewFileStateReader(path), nil
+	}
+
+	lineage, err := readLineage(path)
+	if err != nil {
+		return nil, err
+	}
+	if lineage == "" {
+		return NewFileStateReader(path), nil
+	}
+
+	return NewDMSStateReader(b.WorkspaceClient(ctx).Bundle, lineage, path), nil
+}
+
+// readLineage reads the deployment lineage from the local resources.json state
+// file. It returns an empty string when the file does not exist yet.
+func readLineage(path string) (string, error) {
+	content, err := os.ReadFile(path)
+	if errors.Is(err, fs.ErrNotExist) {
+		return "", nil
+	}
+	if err != nil {
+		return "", err
+	}
+
+	var db dstate.Database
+	if err := json.Unmarshal(content, &db); err != nil {
+		return "", fmt.Errorf("parsing %s: %w", filepath.ToSlash(path), err)
+	}
+	return db.Lineage, nil
+}

bundle/statemgmt/statereader_test.go

@@ -8,6 +8,8 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/direct/dstate"
 	sdkbundle "github.com/databricks/databricks-sdk-go/service/bundle"
 	"github.com/stretchr/testify/assert"
@@ -130,3 +132,63 @@ func TestFileStateReaderMissingFileIsEmptyState(t *testing.T) {
 	_, ok := db.GetResourceEntry("resources.jobs.foo")
 	assert.False(t, ok)
 }
+
+func writeStateFile(t *testing.T, lineage string) string {
+	t.Helper()
+	path := filepath.Join(t.TempDir(), "resources.json")
+	content, err := json.Marshal(dstate.NewDatabase(lineage, 1))
+	require.NoError(t, err)
+	require.NoError(t, os.WriteFile(path, content, 0o600))
+	return path
+}
+
+func TestReadLineage(t *testing.T) {
+	t.Run("present", func(t *testing.T) {
+		lineage, err := readLineage(writeStateFile(t, "lineage-9"))
+		require.NoError(t, err)
+		assert.Equal(t, "lineage-9", lineage)
+	})
+
+	t.Run("missing file", func(t *testing.T) {
+		lineage, err := readLineage(filepath.Join(t.TempDir(), "absent.json"))
+		require.NoError(t, err)
+		assert.Empty(t, lineage)
+	})
+
+	t.Run("corrupt file", func(t *testing.T) {
+		path := filepath.Join(t.TempDir(), "resources.json")
+		require.NoError(t, os.WriteFile(path, []byte("not json"), 0o600))
+		_, err := readLineage(path)
+		assert.Error(t, err)
+	})
+}
+
+func TestNewStateReaderSelectsFileReader(t *testing.T) {
+	// Flag disabled: always the file reader, regardless of any recorded lineage.
+	b := &bundle.Bundle{}
+	reader, err := NewStateReader(t.Context(), b, writeStateFile(t, "lineage-1"))
+	require.NoError(t, err)
+	assert.IsType(t, &fileStateReader{}, reader)
+}
+
+func TestNewStateReaderFallsBackToFileWhenNoLineage(t *testing.T) {
+	// Flag enabled but no prior deployment (no lineage yet): nothing in DMS to
+	// read, so fall back to the local file reader.
+	b := &bundle.Bundle{}
+	b.Config.Experimental = &config.Experimental{RecordDeploymentHistory: true}
+
+	reader, err := NewStateReader(t.Context(), b, writeStateFile(t, ""))
+	require.NoError(t, err)
+	assert.IsType(t, &fileStateReader{}, reader)
+}
+
+func TestNewStateReaderPropagatesCorruptStateError(t *testing.T) {
+	b := &bundle.Bundle{}
+	b.Config.Experimental = &config.Experimental{RecordDeploymentHistory: true}
+
+	path := filepath.Join(t.TempDir(), "resources.json")
+	require.NoError(t, os.WriteFile(path, []byte("not json"), 0o600))
+
+	_, err := NewStateReader(t.Context(), b, path)
+	assert.Error(t, err)
+}

cmd/bundle/utils/process.go

@@ -14,7 +14,6 @@ import (
 	"github.com/databricks/cli/bundle/deploy/terraform"
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/direct"
-	"github.com/databricks/cli/bundle/direct/dstate"
 	"github.com/databricks/cli/bundle/phases"
 	"github.com/databricks/cli/bundle/statemgmt"
 	"github.com/databricks/cli/cmd/root"
@@ -189,7 +188,12 @@ func ProcessBundleRet(cmd *cobra.Command, opts ProcessOptions) (b *bundle.Bundle
 		needDirectState := stateDesc.Engine.IsDirect() && (opts.InitIDs || opts.ErrorOnEmptyState || opts.Deploy || opts.ReadPlanPath != "" || opts.PreDeployChecks || opts.PostStateFunc != nil)
 		if needDirectState {
 			_, localPath := b.StateFilenameDirect(ctx)
-			if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false)); err != nil {
+			reader, err := statemgmt.NewStateReader(ctx, b, localPath)
+			if err != nil {
+				logdiag.LogError(ctx, err)
+				return b, stateDesc, root.ErrAlreadyPrinted
+			}
+			if err := reader.Load(ctx, &b.DeploymentBundle.StateDB); err != nil {
 				logdiag.LogError(ctx, err)
 				return b, stateDesc, root.ErrAlreadyPrinted
 			}