Move LoadStateFromDMS from lock package to statemgmt package

LoadStateFromDMS is a state-loading function, not a lock function.
Moving it to statemgmt where it belongs alongside other state
management code.

Co-authored-by: Isaac

Author: shreyas-goenka

bundle/deploy/lock/deployment_metadata_service.go

@@ -17,7 +17,6 @@ import (
 	"github.com/databricks/cli/bundle/deploy"
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/direct"
-	"github.com/databricks/cli/bundle/direct/dstate"
 	"github.com/databricks/cli/bundle/statemgmt"
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/filer"
@@ -257,58 +256,6 @@ func makeOperationReporter(svc *tmpdms.DeploymentMetadataAPI, deploymentID, vers
 	}
 }
 
-// LoadStateFromDMS loads resource state from the deployment metadata service
-// into the state DB. It first opens the local state file (which contains the
-// deployment ID pointer), then populates the resource state from the server.
-func LoadStateFromDMS(ctx context.Context, b *bundle.Bundle) error {
-	if b.DeploymentID == "" {
-		return nil
-	}
-
-	// Open the local state file first so the state DB path is set.
-	// The local file contains {"deployment_id":"..."} with no resource state.
-	db := &b.DeploymentBundle.StateDB
-	_, localPath := b.StateFilenameDirect(ctx)
-	if err := db.Open(localPath); err != nil {
-		return fmt.Errorf("opening local state: %w", err)
-	}
-
-	svc, err := tmpdms.NewDeploymentMetadataAPI(b.WorkspaceClient())
-	if err != nil {
-		return fmt.Errorf("failed to create metadata service client: %w", err)
-	}
-
-	resources, err := svc.ListResources(ctx, tmpdms.ListResourcesRequest{
-		DeploymentID: b.DeploymentID,
-	})
-	if err != nil {
-		return fmt.Errorf("failed to list resources from deployment metadata service: %w", err)
-	}
-
-	// Populate resource state from the server.
-	db.Data.State = make(map[string]dstate.ResourceEntry)
-
-	for _, r := range resources {
-		// The DMS stores keys without the "resources." prefix (e.g., "jobs.foo").
-		// The state DB expects the full key (e.g., "resources.jobs.foo").
-		resourceKey := "resources." + r.ResourceKey
-
-		var stateBytes json.RawMessage
-		if r.State != nil {
-			stateBytes, err = json.Marshal(r.State)
-			if err != nil {
-				return fmt.Errorf("marshaling state for %s: %w", resourceKey, err)
-			}
-		}
-
-		db.Data.State[resourceKey] = dstate.ResourceEntry{
-			ID:    r.ResourceID,
-			State: stateBytes,
-		}
-	}
-
-	return nil
-}
 
 // planActionToOperationAction maps a deploy plan action to a metadata service
 // operation action type. No-op actions like Skip return ("", nil) and should

bundle/statemgmt/state_dms.go

@@ -0,0 +1,64 @@
+package statemgmt
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/direct/dstate"
+	"github.com/databricks/cli/libs/tmpdms"
+)
+
+// LoadStateFromDMS loads resource state from the deployment metadata service
+// into the state DB. It first opens the local state file (which contains the
+// deployment ID pointer), then populates the resource state from the server.
+func LoadStateFromDMS(ctx context.Context, b *bundle.Bundle) error {
+	if b.DeploymentID == "" {
+		return nil
+	}
+
+	// Open the local state file first so the state DB path is set.
+	// The local file contains {"deployment_id":"..."} with no resource state.
+	db := &b.DeploymentBundle.StateDB
+	_, localPath := b.StateFilenameDirect(ctx)
+	if err := db.Open(localPath); err != nil {
+		return fmt.Errorf("opening local state: %w", err)
+	}
+
+	svc, err := tmpdms.NewDeploymentMetadataAPI(b.WorkspaceClient())
+	if err != nil {
+		return fmt.Errorf("failed to create metadata service client: %w", err)
+	}
+
+	resources, err := svc.ListResources(ctx, tmpdms.ListResourcesRequest{
+		DeploymentID: b.DeploymentID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to list resources from deployment metadata service: %w", err)
+	}
+
+	// Populate resource state from the server.
+	db.Data.State = make(map[string]dstate.ResourceEntry)
+
+	for _, r := range resources {
+		// The DMS stores keys without the "resources." prefix (e.g., "jobs.foo").
+		// The state DB expects the full key (e.g., "resources.jobs.foo").
+		resourceKey := "resources." + r.ResourceKey
+
+		var stateBytes json.RawMessage
+		if r.State != nil {
+			stateBytes, err = json.Marshal(r.State)
+			if err != nil {
+				return fmt.Errorf("marshaling state for %s: %w", resourceKey, err)
+			}
+		}
+
+		db.Data.State[resourceKey] = dstate.ResourceEntry{
+			ID:    r.ResourceID,
+			State: stateBytes,
+		}
+	}
+
+	return nil
+}

cmd/bundle/utils/process.go

@@ -11,7 +11,6 @@ import (
 	"github.com/databricks/cli/bundle/config/engine"
 	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/bundle/config/validate"
-	"github.com/databricks/cli/bundle/deploy/lock"
 	"github.com/databricks/cli/bundle/deployplan"
 	"github.com/databricks/cli/bundle/direct"
 	"github.com/databricks/cli/bundle/phases"
@@ -189,7 +188,7 @@ 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 {
 			if b.DeploymentID != "" {
-				if err := lock.LoadStateFromDMS(ctx, b); err != nil {
+				if err := statemgmt.LoadStateFromDMS(ctx, b); err != nil {
 					logdiag.LogError(ctx, err)
 					return b, stateDesc, root.ErrAlreadyPrinted
 				}