filestate: Plumb context through constructor

Instead of using context.TODO everywhere,
plumb context.Context into filestate.New and filestate.Login,
and through that into other operations.

Extracted from #12134

Co-authored-by: Abhinav Gupta <abhinav@pulumi.com>

pkg/backend/filestate/backend.go

@@ -118,7 +118,7 @@ func IsFileStateBackendURL(urlstr string) bool {
 
 const FilePathPrefix = "file://"
 
-func New(d diag.Sink, originalURL string, project *workspace.Project) (Backend, error) {
+func New(ctx context.Context, d diag.Sink, originalURL string, project *workspace.Project) (Backend, error) {
 	if !IsFileStateBackendURL(originalURL) {
 		return nil, fmt.Errorf("local URL %s has an illegal prefix; expected one of: %s",
 			originalURL, strings.Join(blob.DefaultURLMux().BucketSchemes(), ", "))
@@ -139,13 +139,13 @@ func New(d diag.Sink, originalURL string, project *workspace.Project) (Backend,
 	// for gcp we want to support additional credentials
 	// schemes on top of go-cloud's default credentials mux.
 	if p.Scheme == gcsblob.Scheme {
-		blobmux, err = authhelpers.GoogleCredentialsMux(context.TODO())
+		blobmux, err = authhelpers.GoogleCredentialsMux(ctx)
 		if err != nil {
 			return nil, err
 		}
 	}
 
-	bucket, err := blobmux.OpenBucket(context.TODO(), u)
+	bucket, err := blobmux.OpenBucket(ctx, u)
 	if err != nil {
 		return nil, fmt.Errorf("unable to open bucket %s: %w", u, err)
 	}
@@ -161,7 +161,7 @@ func New(d diag.Sink, originalURL string, project *workspace.Project) (Backend,
 		}
 	}
 
-	isAcc, err := bucket.IsAccessible(context.TODO())
+	isAcc, err := bucket.IsAccessible(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("unable to check if bucket %s is accessible: %w", u, err)
 	}
@@ -233,8 +233,8 @@ func massageBlobPath(path string) (string, error) {
 	return FilePathPrefix + path, nil
 }
 
-func Login(d diag.Sink, url string, project *workspace.Project) (Backend, error) {
-	be, err := New(d, url, project)
+func Login(ctx context.Context, d diag.Sink, url string, project *workspace.Project) (Backend, error) {
+	be, err := New(ctx, d, url, project)
 	if err != nil {
 		return nil, err
 	}
@@ -704,7 +704,7 @@ func (b *localBackend) apply(
 			u.Path = filepath.ToSlash(path.Join(u.Path, b.stackPath(stackName)))
 			link = u.String()
 		} else {
-			link, err = b.bucket.SignedURL(context.TODO(), b.stackPath(stackName), nil)
+			link, err = b.bucket.SignedURL(ctx, b.stackPath(stackName), nil)
 			if err != nil {
 				// set link to be empty to when there is an error to hide use of Permalinks
 				link = ""

pkg/backend/filestate/backend_test.go

@@ -144,9 +144,9 @@ func makeUntypedDeployment(name tokens.QName, phrase, state string) (*apitype.Un
 func TestListStacksWithMultiplePassphrases(t *testing.T) {
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Create stack "a" and import a checkpoint with a secret
 	aStackRef, err := b.ParseStackReference("a")
@@ -204,9 +204,9 @@ func TestDrillError(t *testing.T) {
 
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Get a non-existent stack and expect a nil error because it won't be found.
 	stackRef, err := b.ParseStackReference("dev")
@@ -222,9 +222,9 @@ func TestCancel(t *testing.T) {
 
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Check that trying to cancel a stack that isn't created yet doesn't error
 	aStackRef, err := b.ParseStackReference("a")
@@ -260,7 +260,7 @@ func TestCancel(t *testing.T) {
 	assert.False(t, lockExists)
 
 	// Make another filestate backend which will have a different lockId
-	ob, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	ob, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
 	assert.NoError(t, err)
 	otherBackend, ok := ob.(*localBackend)
 	assert.True(t, ok)
@@ -283,9 +283,9 @@ func TestRemoveMakesBackups(t *testing.T) {
 
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Grab the bucket interface to test with
 	lb, ok := b.(*localBackend)
@@ -326,9 +326,9 @@ func TestRenameWorks(t *testing.T) {
 
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Grab the bucket interface to test with
 	lb, ok := b.(*localBackend)
@@ -393,7 +393,8 @@ func TestLoginToNonExistingFolderFails(t *testing.T) {
 	t.Parallel()
 
 	fakeDir := "file://" + filepath.ToSlash(os.TempDir()) + "/non-existing"
-	b, err := New(diagtest.LogSink(t), fakeDir, nil)
+	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), fakeDir, nil)
 	assert.Error(t, err)
 	assert.Nil(t, b)
 }
@@ -444,9 +445,9 @@ func TestHtmlEscaping(t *testing.T) {
 
 	// Login to a temp dir filestate backend
 	tmpDir := t.TempDir()
-	b, err := New(diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
-	assert.NoError(t, err)
 	ctx := context.Background()
+	b, err := New(ctx, diagtest.LogSink(t), "file://"+filepath.ToSlash(tmpDir), nil)
+	assert.NoError(t, err)
 
 	// Create stack "a" and import a checkpoint with a secret
 	aStackRef, err := b.ParseStackReference("a")

pkg/backend/filestate/bucket_test.go

@@ -32,7 +32,8 @@ func TestWrappedBucket(t *testing.T) {
 
 	// Initialize a filestate backend, using the default Pulumi directory.
 	cloudURL := FilePathPrefix + "~"
-	b, err := New(nil, cloudURL, nil)
+	ctx := context.Background()
+	b, err := New(ctx, nil, cloudURL, nil)
 	if err != nil {
 		t.Fatalf("Initializing new filestate backend: %v", err)
 	}
@@ -46,7 +47,6 @@ func TestWrappedBucket(t *testing.T) {
 		t.Fatalf("localBackend.bucket wasn't of type wrappedBucket?")
 	}
 
-	ctx := context.Background()
 	// Perform basic file operations using wrappedBucket and verify that it will
 	// successfully handle both "/" and "\" as file separators. (And probably fail in
 	// exciting ways if you try to give it a file on a system that supports "\" or "/" as

pkg/cmd/pulumi/login.go

@@ -142,7 +142,7 @@ func newLoginCmd() *cobra.Command {
 
 			var be backend.Backend
 			if filestate.IsFileStateBackendURL(cloudURL) {
-				be, err = filestate.Login(cmdutil.Diag(), cloudURL, project)
+				be, err = filestate.Login(ctx, cmdutil.Diag(), cloudURL, project)
 				if defaultOrg != "" {
 					return fmt.Errorf("unable to set default org for this type of backend")
 				}

pkg/cmd/pulumi/util.go

@@ -129,7 +129,7 @@ func nonInteractiveCurrentBackend(ctx context.Context, project *workspace.Projec
 	}
 
 	if filestate.IsFileStateBackendURL(url) {
-		return filestate.New(cmdutil.Diag(), url, project)
+		return filestate.New(ctx, cmdutil.Diag(), url, project)
 	}
 	return httpstate.NewLoginManager().Current(ctx, cmdutil.Diag(), url, project, workspace.GetCloudInsecure(url))
 }
@@ -145,7 +145,7 @@ func currentBackend(ctx context.Context, project *workspace.Project, opts displa
 	}
 
 	if filestate.IsFileStateBackendURL(url) {
-		return filestate.New(cmdutil.Diag(), url, project)
+		return filestate.New(ctx, cmdutil.Diag(), url, project)
 	}
 	return httpstate.NewLoginManager().Login(ctx, cmdutil.Diag(), url, project, workspace.GetCloudInsecure(url), opts)
 }