bundle/generate: read recorded state from DMS for in-place generate

acceptance/bundle/generate/dashboard-inplace/test.toml

@@ -1,7 +1,3 @@
-# bundle generate reading recorded state from DMS is a fast followup; until then this
-# deploy-then-generate flow can't run under recording (local state is a tombstone).
-EnvMatrix.DMS = [""]
-
 [[Repls]]
 Old = "[0-9a-f]{32}"
 New = "[DASHBOARD_ID]"

acceptance/bundle/generate/genie_space_inplace/test.toml

@@ -1,8 +1,5 @@
 # Genie spaces are only deployed via the direct deployment engine.
 EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
-# bundle generate reading recorded state from DMS is a fast followup; until then this
-# deploy-then-generate flow can't run under recording (local state is a tombstone).
-EnvMatrix.DMS = [""]
 
 [[Repls]]
 Old = "[0-9a-f]{32}"

acceptance/bundle/resources/dashboards/generate_inplace/test.toml

@@ -3,6 +3,3 @@ RecordRequests = false
 RunsOnDbr = true
 
 EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"]
-
-# bundle generate reading recorded state from DMS is a fast followup; opt out of recording.
-EnvMatrix.DMS = [""]

cmd/bundle/generate/dashboard.go

@@ -17,7 +17,6 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/terraform"
-	"github.com/databricks/cli/bundle/direct/dstate"
 	"github.com/databricks/cli/bundle/generate"
 	"github.com/databricks/cli/bundle/phases"
 	"github.com/databricks/cli/bundle/resources"
@@ -403,8 +402,7 @@ func (d *dashboard) runForResource(ctx context.Context, b *bundle.Bundle) {
 
 	var state statemgmt.ExportedResourcesMap
 	if stateDesc.Engine.IsDirect() {
-		_, localPath := b.StateFilenameDirect(ctx)
-		if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil {
+		if err := utils.OpenDirectStateForRead(ctx, b); err != nil {
 			logdiag.LogError(ctx, err)
 			return
 		}

cmd/bundle/generate/genie_space.go

@@ -14,7 +14,6 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/terraform"
-	"github.com/databricks/cli/bundle/direct/dstate"
 	"github.com/databricks/cli/bundle/generate"
 	"github.com/databricks/cli/bundle/phases"
 	"github.com/databricks/cli/bundle/resources"
@@ -321,8 +320,7 @@ func (g *genieSpace) runForResource(ctx context.Context, b *bundle.Bundle) {
 
 	var state statemgmt.ExportedResourcesMap
 	if stateDesc.Engine.IsDirect() {
-		_, localPath := b.StateFilenameDirect(ctx)
-		if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil {
+		if err := utils.OpenDirectStateForRead(ctx, b); err != nil {
 			logdiag.LogError(ctx, err)
 			return
 		}

cmd/bundle/utils/process.go

@@ -258,15 +258,10 @@ func ProcessBundleRet(cmd *cobra.Command, opts ProcessOptions) (b *bundle.Bundle
 				// the plan carries them. version_id is always known (last recorded + 1); deployment_id
 				// does not exist until a first deploy creates it, so it is left off here and the deploy
 				// phase stamps the created id.
-				// The service reports the version as a string; parse it here so everything below
-				// carries a number.
-				lastVersionID := 0
-				if dmsDeployment != nil && dmsDeployment.LastVersionId != "" {
-					lastVersionID, err = strconv.Atoi(dmsDeployment.LastVersionId)
-					if err != nil {
-						logdiag.LogError(ctx, fmt.Errorf("failed to parse last_version_id %q: %w", dmsDeployment.LastVersionId, err))
-						return b, stateDesc, root.ErrAlreadyPrinted
-					}
+				lastVersionID, err := parseLastVersionID(dmsDeployment)
+				if err != nil {
+					logdiag.LogError(ctx, err)
+					return b, stateDesc, root.ErrAlreadyPrinted
 				}
 				nextVersion := lastVersionID + 1
 				muts := []bundle.Mutator{metadata.AnnotateDeploymentVersion(nextVersion)}
@@ -514,6 +509,43 @@ func fetchDeploymentFromStatePath(ctx context.Context, w *databricks.WorkspaceCl
 	return deploymentID, deployment, nil
 }
 
+// parseLastVersionID parses the deployment's last recorded version, which the service reports
+// as a string. It returns 0 when the deployment does not exist yet or has no recorded version.
+func parseLastVersionID(dmsDeployment *bundledeployments.Deployment) (int, error) {
+	if dmsDeployment == nil || dmsDeployment.LastVersionId == "" {
+		return 0, nil
+	}
+	v, err := strconv.Atoi(dmsDeployment.LastVersionId)
+	if err != nil {
+		return 0, fmt.Errorf("failed to parse last_version_id %q: %w", dmsDeployment.LastVersionId, err)
+	}
+	return v, nil
+}
+
+// OpenDirectStateForRead opens the direct-engine state database read-only. When the bundle
+// records deployment history the local state file is only a tombstone, so the resources are
+// read from the deployment metadata service instead.
+func OpenDirectStateForRead(ctx context.Context, b *bundle.Bundle) error {
+	_, localPath := b.StateFilenameDirect(ctx)
+	if !b.ConfiguresDeploymentHistory(ctx) {
+		return b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{})
+	}
+
+	dmsDeploymentID, dmsDeployment, err := fetchDeploymentFromStatePath(ctx, b.WorkspaceClient(ctx), b.Config.Workspace.StatePath)
+	if err != nil {
+		return err
+	}
+	lastVersionID, err := parseLastVersionID(dmsDeployment)
+	if err != nil {
+		return err
+	}
+	// StateDB.Open builds the DMS client from the workspace client on the context, so ensure one is set.
+	if !cmdctx.HasWorkspaceClient(ctx) {
+		ctx = cmdctx.SetWorkspaceClient(ctx, b.WorkspaceClient(ctx))
+	}
+	return b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(false), dstate.WithWrite(false), dstate.WithDeploymentHistory(true), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID})
+}
+
 // isNewerVersion reports whether the state's recorded CLI version is strictly
 // newer than the running build. Both are bare versions without a leading "v".
 // An empty stateVersion (state not written by any CLI yet) or an unparseable

cmd/bundle/utils/process_test.go

@@ -3,7 +3,9 @@ package utils
 import (
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestIsNewerVersion(t *testing.T) {
@@ -42,3 +44,30 @@ func TestIsNewerVersion(t *testing.T) {
 		})
 	}
 }
+
+func TestParseLastVersionID(t *testing.T) {
+	tests := []struct {
+		name       string
+		deployment *bundledeployments.Deployment
+		want       int
+		wantErr    bool
+	}{
+		{"nil deployment", nil, 0, false},
+		{"empty version", &bundledeployments.Deployment{}, 0, false},
+		{"version zero", &bundledeployments.Deployment{LastVersionId: "0"}, 0, false},
+		{"valid version", &bundledeployments.Deployment{LastVersionId: "7"}, 7, false},
+		{"invalid version", &bundledeployments.Deployment{LastVersionId: "abc"}, 0, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseLastVersionID(tt.deployment)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}