From 01c381ffdef1cab969aa207e9c4057f97facaff7 Mon Sep 17 00:00:00 2001
From: Lennart Kats <lennart.kats@databricks.com>
Date: Mon, 9 Sep 2024 09:51:33 +0200
Subject: [PATCH] Hoist state loading below prompt

---
 cmd/bundle/run.go | 33 ++++++++-------------------------
 1 file changed, 8 insertions(+), 25 deletions(-)

diff --git a/cmd/bundle/run.go b/cmd/bundle/run.go
index 07fa1b746..19d2f287d 100644
--- a/cmd/bundle/run.go
+++ b/cmd/bundle/run.go
@@ -2,10 +2,7 @@ package bundle
 
 import (
 	"encoding/json"
-	"errors"
 	"fmt"
-	"os"
-	"path/filepath"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/terraform"
@@ -43,9 +40,6 @@ task or a Python wheel task, the second example applies.
 `,
 	}
 
-	var forcePull bool
-	cmd.Flags().BoolVar(&forcePull, "force-pull", false, "Skip local cache and load the state from the remote workspace")
-
 	var runOptions run.Options
 	runOptions.Define(cmd)
 
@@ -66,25 +60,6 @@ task or a Python wheel task, the second example applies.
 			return err
 		}
 
-		cacheDir, err := terraform.Dir(ctx, b)
-		if err != nil {
-			return err
-		}
-		_, stateFileErr := os.Stat(filepath.Join(cacheDir, terraform.TerraformStateFileName))
-		_, configFileErr := os.Stat(filepath.Join(cacheDir, terraform.TerraformConfigFileName))
-		noCache := errors.Is(stateFileErr, os.ErrNotExist) || errors.Is(configFileErr, os.ErrNotExist)
-
-		if forcePull || noCache {
-			diags = bundle.Apply(ctx, b, bundle.Seq(
-				terraform.StatePull(),
-				terraform.Interpolate(),
-				terraform.Write(),
-			))
-			if err := diags.Error(); err != nil {
-				return err
-			}
-		}
-
 		// If no arguments are specified, prompt the user to select something to run.
 		if len(args) == 0 && cmdio.IsPromptSupported(ctx) {
 			// Invert completions from KEY -> NAME, to NAME -> KEY.
@@ -103,6 +78,14 @@ task or a Python wheel task, the second example applies.
 			return fmt.Errorf("expected a KEY of the resource to run")
 		}
 
+		diags = bundle.Apply(ctx, b, bundle.Seq(
+			terraform.Interpolate(),
+			terraform.Write(),
+			terraform.StatePull(),
+			terraform.Load(terraform.ErrorOnEmptyState),
+		))
+		return diags.Error()
+
 		runner, err := run.Find(b, args[0])
 		if err != nil {
 			return err