deployCore: use Finalize return value instead of re-opening state

Finalize now returns (ExportedResourcesMap, error), capturing the merged
state before clearing. This lets deploy.go use LoadFromState directly
instead of closing and re-opening the state file from disk.

LoadFromState is a new statemgmt constructor that accepts pre-computed
state and skips the engine dispatch in Load.Apply.

Co-authored-by: Denis Bilenko <denis.bilenko@databricks.com>

Author: denik

bundle/direct/bind.go

@@ -64,7 +64,9 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac
 	var checkStateDB dstate.DeploymentState
 	if err := checkStateDB.Open(ctx, statePath, dstate.WithRecovery(true), dstate.WithWrite(false)); err == nil {
 		existingID := checkStateDB.GetResourceID(resourceKey)
-		_ = checkStateDB.Finalize(ctx)
+		if _, err := checkStateDB.Finalize(ctx); err != nil {
+			log.Warnf(ctx, "failed to finalize state: %v", err)
+		}
 		if existingID != "" {
 			return nil, ErrResourceAlreadyBound{
 				ResourceKey: resourceKey,
@@ -98,7 +100,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac
 	}
 
 	// Finalize to persist temp state to disk
-	err = b.StateDB.Finalize(ctx)
+	_, err = b.StateDB.Finalize(ctx)
 	if err != nil {
 		os.Remove(tmpStatePath)
 		return nil, err
@@ -117,7 +119,9 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac
 		os.Remove(tmpStatePath)
 		return nil, err
 	}
-	_ = b.StateDB.Finalize(ctx)
+	if _, err := b.StateDB.Finalize(ctx); err != nil {
+		log.Warnf(ctx, "failed to finalize state: %v", err)
+	}
 
 	// Populate the state with the resolved config
 	entry := plan.Plan[resourceKey]
@@ -152,7 +156,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac
 			return nil, err
 		}
 
-		err = b.StateDB.Finalize(ctx)
+		_, err = b.StateDB.Finalize(ctx)
 		if err != nil {
 			os.Remove(tmpStatePath)
 			return nil, err
@@ -166,7 +170,9 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac
 		return nil, err
 	}
 	plan, err = b.CalculatePlan(ctx, client, configRoot)
-	_ = b.StateDB.Finalize(ctx)
+	if _, ferr := b.StateDB.Finalize(ctx); ferr != nil {
+		log.Warnf(ctx, "failed to finalize state: %v", ferr)
+	}
 	if err != nil {
 		os.Remove(tmpStatePath)
 		return nil, err
@@ -236,5 +242,6 @@ func (b *DeploymentBundle) Unbind(ctx context.Context, statePath, resourceKey st
 		}
 	}
 
-	return b.StateDB.Finalize(ctx)
+	_, err = b.StateDB.Finalize(ctx)
+	return err
 }

bundle/direct/dstate/state.go

@@ -32,19 +32,25 @@ const (
 var errStaleWAL = errors.New("stale WAL")
 
 type DeploymentState struct {
-	Path     string
-	Data     Database
-	mu       sync.Mutex
-	walFile  *os.File
+	Path    string
+	Data    Database
+	mu      sync.Mutex
+	walFile *os.File
+
+	// Maps resource key to ID. Unlike Data.State, this is up to during writes (deploys).
 	stateIDs map[string]string
 }
 
 type Database struct {
-	StateVersion int                      `json:"state_version"`
-	CLIVersion   string                   `json:"cli_version"`
-	Lineage      string                   `json:"lineage"`
-	Serial       int                      `json:"serial"`
-	State        map[string]ResourceEntry `json:"state"`
+	StateVersion int    `json:"state_version"`
+	CLIVersion   string `json:"cli_version"`
+	Lineage      string `json:"lineage"`
+	Serial       int    `json:"serial"`
+
+	// Maps resource key to ResourceEntry which includes ID + full serialized state.
+	// This is not updated during write/deploy, those writes go to WAL instead.
+	// The State is then reconstructed from WAL.
+	State map[string]ResourceEntry `json:"state"`
 }
 
 type ResourceEntry struct {
@@ -346,14 +352,15 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error)
 	return lineNumber > 1, nil
 }
 
-// Finalize replays the WAL (if open for write) and resets the state.
+// Finalize replays the WAL (if open for write), captures the resulting state, and resets.
 // Safe to call multiple times or on an already-finalized state.
-func (db *DeploymentState) Finalize(ctx context.Context) error {
+// Returns the exported state as of the end of this operation.
+func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.ExportedResourcesMap, error) {
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
 	if db.Path == "" {
-		return nil
+		return nil, nil
 	}
 
 	var err error
@@ -364,11 +371,13 @@ func (db *DeploymentState) Finalize(ctx context.Context) error {
 		err = db.replayWAL(ctx)
 	}
 
+	state := ExportStateFromData(db.Data)
+
 	db.Path = ""
 	db.Data = Database{}
-	db.stateIDs = make(map[string]string)
+	db.stateIDs = nil
 
-	return err
+	return state, err
 }
 
 // UpgradeToWrite transitions from read mode to write mode without re-reading state.
@@ -427,9 +436,10 @@ func (db *DeploymentState) AssertOpenedForWrite() {
 	}
 }
 
-func (db *DeploymentState) ExportState(ctx context.Context) resourcestate.ExportedResourcesMap {
+// ExportStateFromData extracts resource IDs and ETags from a database snapshot.
+func ExportStateFromData(data Database) resourcestate.ExportedResourcesMap {
 	result := make(resourcestate.ExportedResourcesMap)
-	for key, entry := range db.Data.State {
+	for key, entry := range data.State {
 		var etag string
 		// Extract etag for dashboards.
 		// covered by test case: bundle/deploy/dashboard/detect-change
@@ -450,6 +460,10 @@ func (db *DeploymentState) ExportState(ctx context.Context) resourcestate.Export
 	return result
 }
 
+func (db *DeploymentState) ExportState(ctx context.Context) resourcestate.ExportedResourcesMap {
+	return ExportStateFromData(db.Data)
+}
+
 func (db *DeploymentState) unlockedSave() error {
 	data, err := json.MarshalIndent(db.Data, "", " ")
 	if err != nil {

bundle/direct/dstate/state_test.go

@@ -9,29 +9,35 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func mustFinalize(t *testing.T, db *DeploymentState) {
+	t.Helper()
+	_, err := db.Finalize(t.Context())
+	require.NoError(t, err)
+}
+
 func TestOpenSaveFinalizeRoundTrip(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)))
 
 	require.NoError(t, db.SaveState("jobs.my_job", "123", map[string]string{"key": "val"}, nil))
-	require.NoError(t, db.Finalize(t.Context()))
+	mustFinalize(t, &db)
 
 	// Re-open and verify persisted data.
 	var db2 DeploymentState
 	require.NoError(t, db2.Open(t.Context(), path, WithRecovery(false), WithWrite(false)))
 	assert.Equal(t, 1, db2.Data.Serial)
 	assert.Equal(t, "123", db2.GetResourceID("jobs.my_job"))
-	require.NoError(t, db2.Finalize(t.Context()))
+	mustFinalize(t, &db2)
 }
 
 func TestFinalizeWithNoEntriesDoesNotWriteStateFile(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)))
-	require.NoError(t, db.Finalize(t.Context()))
+	mustFinalize(t, &db)
 
 	_, err := os.Stat(path)
 	assert.ErrorIs(t, err, os.ErrNotExist)
@@ -46,7 +52,7 @@ func TestPanicOnDoubleOpen(t *testing.T) {
 	assert.Panics(t, func() {
 		_ = db.Open(t.Context(), path, WithRecovery(true), WithWrite(true))
 	})
-	require.NoError(t, db.Finalize(t.Context()))
+	mustFinalize(t, &db)
 }
 
 func TestDeleteState(t *testing.T) {
@@ -55,16 +61,16 @@ func TestDeleteState(t *testing.T) {
 	var db DeploymentState
 	require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true)))
 	require.NoError(t, db.SaveState("jobs.my_job", "123", map[string]string{}, nil))
-	require.NoError(t, db.Finalize(t.Context()))
+	mustFinalize(t, &db)
 
 	var db2 DeploymentState
 	require.NoError(t, db2.Open(t.Context(), path, WithRecovery(true), WithWrite(true)))
 	require.NoError(t, db2.DeleteState("jobs.my_job"))
-	require.NoError(t, db2.Finalize(t.Context()))
+	mustFinalize(t, &db2)
 
 	var db3 DeploymentState
 	require.NoError(t, db3.Open(t.Context(), path, WithRecovery(false), WithWrite(false)))
 	assert.Equal(t, 2, db3.Data.Serial)
 	assert.Equal(t, "", db3.GetResourceID("jobs.my_job"))
-	require.NoError(t, db3.Finalize(t.Context()))
+	mustFinalize(t, &db3)
 }

bundle/phases/deploy.go

@@ -15,7 +15,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/libraries"
 	"github.com/databricks/cli/bundle/metrics"
 	"github.com/databricks/cli/bundle/permissions"
@@ -75,14 +74,12 @@ func deployCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan, ta
 		bundle.ApplyContext(ctx, b, terraform.Apply())
 	}
 
-	// Close state to replay WAL into state file, then reopen for read.
-	// PushResourcesState needs the file on disk, Load needs the state in memory.
+	// Flush WAL to state file on disk; capture the resulting state for Load below.
+	var directState statemgmt.ExportedResourcesMap
 	if targetEngine.IsDirect() {
-		if err := b.DeploymentBundle.StateDB.Finalize(ctx); err != nil {
-			logdiag.LogError(ctx, err)
-		}
-		_, localPath := b.StateFilenameDirect(ctx)
-		if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false)); err != nil {
+		var err error
+		directState, err = b.DeploymentBundle.StateDB.Finalize(ctx)
+		if err != nil {
 			logdiag.LogError(ctx, err)
 		}
 	}
@@ -93,8 +90,15 @@ func deployCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan, ta
 		return
 	}
 
+	var loadMutator bundle.Mutator
+	if targetEngine.IsDirect() {
+		loadMutator = statemgmt.LoadFromState(directState)
+	} else {
+		loadMutator = statemgmt.Load(targetEngine)
+	}
+
 	bundle.ApplySeqContext(ctx, b,
-		statemgmt.Load(targetEngine),
+		loadMutator,
 		metadata.Compute(),
 		metadata.Upload(),
 		statemgmt.UploadStateForYamlSync(targetEngine),

bundle/phases/destroy.go

@@ -86,7 +86,7 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan, e
 	// Warn instead of hard-error: resources are already deleted, so proceed
 	// with file cleanup regardless of whether state flush succeeds.
 	if engine.IsDirect() {
-		if err := b.DeploymentBundle.StateDB.Finalize(ctx); err != nil {
+		if _, err := b.DeploymentBundle.StateDB.Finalize(ctx); err != nil {
 			diags := diag.WarningFromErr(err)
 			if len(diags) > 0 {
 				logdiag.LogDiag(ctx, diags[0])

bundle/statemgmt/state_load.go

@@ -35,7 +35,6 @@ func (l *load) Name() string {
 }
 
 func (l *load) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-	var err error
 	var state ExportedResourcesMap
 
 	if l.engine.IsDirect() {
@@ -48,14 +47,29 @@ func (l *load) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 		}
 	}
 
-	err = l.validateState(state)
-	if err != nil {
+	return applyState(ctx, b, state, l.modes)
+}
+
+type loadFromState struct {
+	state ExportedResourcesMap
+	modes []LoadMode
+}
+
+func (l *loadFromState) Name() string {
+	return "statemgmt.Load"
+}
+
+func (l *loadFromState) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	return applyState(ctx, b, l.state, l.modes)
+}
+
+// applyState merges the exported resource state into the bundle configuration.
+func applyState(ctx context.Context, b *bundle.Bundle, state ExportedResourcesMap, modes []LoadMode) diag.Diagnostics {
+	if err := validateLoadedState(state, modes); err != nil {
 		return diag.FromErr(err)
 	}
 
-	// Merge state into configuration.
-	err = StateToBundle(ctx, state, &b.Config)
-	if err != nil {
+	if err := StateToBundle(ctx, state, &b.Config); err != nil {
 		return diag.FromErr(err)
 	}
 
@@ -160,14 +174,19 @@ func StateToBundle(ctx context.Context, state ExportedResourcesMap, config *conf
 	})
 }
 
-func (l *load) validateState(state ExportedResourcesMap) error {
-	if len(state) == 0 && slices.Contains(l.modes, ErrorOnEmptyState) {
+func validateLoadedState(state ExportedResourcesMap, modes []LoadMode) error {
+	if len(state) == 0 && slices.Contains(modes, ErrorOnEmptyState) {
 		return errors.New("resource not found or not yet deployed. Did you forget to run 'databricks bundle deploy'?")
 	}
-
 	return nil
 }
 
 func Load(engine engine.EngineType, modes ...LoadMode) bundle.Mutator {
 	return &load{modes: modes, engine: engine}
 }
+
+// LoadFromState returns a mutator that loads the provided pre-computed state into the bundle,
+// skipping the engine-specific state retrieval step.
+func LoadFromState(state ExportedResourcesMap, modes ...LoadMode) bundle.Mutator {
+	return &loadFromState{state: state, modes: modes}
+}

bundle/statemgmt/upload_state_for_yaml_sync.go

@@ -198,7 +198,7 @@ func (m *uploadStateForYamlSync) convertState(ctx context.Context, b *bundle.Bun
 	}
 
 	deploymentBundle.Apply(ctx, b.WorkspaceClient(ctx), plan, direct.MigrateMode(true))
-	if err := deploymentBundle.StateDB.Finalize(ctx); err != nil {
+	if _, err := deploymentBundle.StateDB.Finalize(ctx); err != nil {
 		return false, err
 	}
 

cmd/bundle/deployment/migrate.go

@@ -282,7 +282,7 @@ To start using direct engine, set "engine: direct" under bundle in your databric
 		}
 
 		deploymentBundle.Apply(ctx, b.WorkspaceClient(ctx), plan, direct.MigrateMode(true))
-		if err := deploymentBundle.StateDB.Finalize(ctx); err != nil {
+		if _, err := deploymentBundle.StateDB.Finalize(ctx); err != nil {
 			logdiag.LogError(ctx, err)
 		}
 		if logdiag.HasError(ctx) {