Cloud-SDK extraction — PR 6: add ctx context.Context to module.IaCStateStore

cmd/wfctl/infra_state_store.go

@@ -260,8 +260,8 @@ type spacesWfctlStateStore struct {
 	inner *module.SpacesIaCStateStore
 }
 
-func (s *spacesWfctlStateStore) ListResources(_ context.Context) ([]interfaces.ResourceState, error) {
-	records, err := s.inner.ListStates(nil)
+func (s *spacesWfctlStateStore) ListResources(ctx context.Context) ([]interfaces.ResourceState, error) {
+	records, err := s.inner.ListStates(ctx, nil)
 	if err != nil {
 		return nil, fmt.Errorf("list spaces state: %w", err)
 	}
@@ -272,12 +272,12 @@ func (s *spacesWfctlStateStore) ListResources(_ context.Context) ([]interfaces.R
 	return states, nil
 }
 
-func (s *spacesWfctlStateStore) SaveResource(_ context.Context, state interfaces.ResourceState) error {
-	return s.inner.SaveState(resourceStateToIaCState(state))
+func (s *spacesWfctlStateStore) SaveResource(ctx context.Context, state interfaces.ResourceState) error {
+	return s.inner.SaveState(ctx, resourceStateToIaCState(state))
 }
 
-func (s *spacesWfctlStateStore) DeleteResource(_ context.Context, name string) error {
-	return s.inner.DeleteState(name)
+func (s *spacesWfctlStateStore) DeleteResource(ctx context.Context, name string) error {
+	return s.inner.DeleteState(ctx, name)
 }
 
 // ── Postgres backend ───────────────────────────────────────────────────────────
@@ -306,8 +306,8 @@ type postgresWfctlStateStore struct {
 	inner *module.PostgresIaCStateStore
 }
 
-func (s *postgresWfctlStateStore) ListResources(_ context.Context) ([]interfaces.ResourceState, error) {
-	records, err := s.inner.ListStates(nil)
+func (s *postgresWfctlStateStore) ListResources(ctx context.Context) ([]interfaces.ResourceState, error) {
+	records, err := s.inner.ListStates(ctx, nil)
 	if err != nil {
 		return nil, fmt.Errorf("list postgres state: %w", err)
 	}
@@ -318,12 +318,12 @@ func (s *postgresWfctlStateStore) ListResources(_ context.Context) ([]interfaces
 	return states, nil
 }
 
-func (s *postgresWfctlStateStore) SaveResource(_ context.Context, state interfaces.ResourceState) error {
-	return s.inner.SaveState(resourceStateToIaCState(state))
+func (s *postgresWfctlStateStore) SaveResource(ctx context.Context, state interfaces.ResourceState) error {
+	return s.inner.SaveState(ctx, resourceStateToIaCState(state))
 }
 
-func (s *postgresWfctlStateStore) DeleteResource(_ context.Context, name string) error {
-	return s.inner.DeleteState(name)
+func (s *postgresWfctlStateStore) DeleteResource(ctx context.Context, name string) error {
+	return s.inner.DeleteState(ctx, name)
 }
 
 // ── Conversion helpers ─────────────────────────────────────────────────────────

module/benchmark_iac_state_backend_test.go

@@ -31,18 +31,19 @@ func oneMBState() *IaCState {
 func BenchmarkIaCStateBackend_InProcess(b *testing.B) {
 	store := NewMemoryIaCStateStore()
 	st := oneMBState()
+	ctx := context.Background()
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		if err := store.Lock(st.ResourceID); err != nil {
+		if err := store.Lock(ctx, st.ResourceID); err != nil {
 			b.Fatal(err)
 		}
-		if _, err := store.GetState(st.ResourceID); err != nil {
+		if _, err := store.GetState(ctx, st.ResourceID); err != nil {
 			b.Fatal(err)
 		}
-		if err := store.SaveState(st); err != nil {
+		if err := store.SaveState(ctx, st); err != nil {
 			b.Fatal(err)
 		}
-		if err := store.Unlock(st.ResourceID); err != nil {
+		if err := store.Unlock(ctx, st.ResourceID); err != nil {
 			b.Fatal(err)
 		}
 	}

module/iac_state.go

@@ -1,5 +1,7 @@
 package module
 
+import "context"
+
 // IaCState tracks the state of an infrastructure resource.
 type IaCState struct {
 	ResourceID   string         `json:"resource_id"`
@@ -20,22 +22,23 @@ type IaCState struct {
 // IaCStateStore is the interface for IaC state persistence backends.
 type IaCStateStore interface {
 	// GetState retrieves a state record by resource ID. Returns nil, nil when not found.
-	GetState(resourceID string) (*IaCState, error)
+	GetState(ctx context.Context, resourceID string) (*IaCState, error)
 
 	// SaveState inserts or replaces a state record.
-	SaveState(state *IaCState) error
+	SaveState(ctx context.Context, state *IaCState) error
 
 	// ListStates returns all state records matching the provided key=value filter.
-	// Pass an empty map to return all records.
-	ListStates(filter map[string]string) ([]*IaCState, error)
+	// Pass a nil or empty map to return all records — both are treated as "no
+	// filter" (ranging over a nil map is valid Go, and most call sites pass nil).
+	ListStates(ctx context.Context, filter map[string]string) ([]*IaCState, error)
 
 	// DeleteState removes a state record by resource ID.
-	DeleteState(resourceID string) error
+	DeleteState(ctx context.Context, resourceID string) error
 
 	// Lock acquires an exclusive lock for the given resource ID.
 	// Returns an error if the resource is already locked.
-	Lock(resourceID string) error
+	Lock(ctx context.Context, resourceID string) error
 
 	// Unlock releases the lock for the given resource ID.
-	Unlock(resourceID string) error
+	Unlock(ctx context.Context, resourceID string) error
 }

module/iac_state_azure.go

@@ -81,8 +81,8 @@ func (s *AzureBlobIaCStateStore) lockBlobName(resourceID string) string {
 }
 
 // GetState retrieves a state record by resource ID. Returns nil, nil when not found.
-func (s *AzureBlobIaCStateStore) GetState(resourceID string) (*IaCState, error) {
-	data, err := s.client.DownloadBlob(context.Background(), s.blobName(resourceID))
+func (s *AzureBlobIaCStateStore) GetState(ctx context.Context, resourceID string) (*IaCState, error) {
+	data, err := s.client.DownloadBlob(ctx, s.blobName(resourceID))
 	if err != nil {
 		if errors.Is(err, ErrAzureBlobNotFound) {
 			return nil, nil
@@ -97,7 +97,7 @@ func (s *AzureBlobIaCStateStore) GetState(resourceID string) (*IaCState, error)
 }
 
 // SaveState writes the state record as a JSON blob.
-func (s *AzureBlobIaCStateStore) SaveState(state *IaCState) error {
+func (s *AzureBlobIaCStateStore) SaveState(ctx context.Context, state *IaCState) error {
 	if state == nil {
 		return fmt.Errorf("iac azure state: SaveState: state must not be nil")
 	}
@@ -108,15 +108,15 @@ func (s *AzureBlobIaCStateStore) SaveState(state *IaCState) error {
 	if err != nil {
 		return fmt.Errorf("iac azure state: SaveState %q: marshal: %w", state.ResourceID, err)
 	}
-	if err := s.client.UploadBlob(context.Background(), s.blobName(state.ResourceID), data, "application/json"); err != nil {
+	if err := s.client.UploadBlob(ctx, s.blobName(state.ResourceID), data, "application/json"); err != nil {
 		return fmt.Errorf("iac azure state: SaveState %q: upload: %w", state.ResourceID, err)
 	}
 	return nil
 }
 
 // ListStates lists all state blobs and returns those matching the filter.
-func (s *AzureBlobIaCStateStore) ListStates(filter map[string]string) ([]*IaCState, error) {
-	names, err := s.client.ListBlobs(context.Background(), s.prefix)
+func (s *AzureBlobIaCStateStore) ListStates(ctx context.Context, filter map[string]string) ([]*IaCState, error) {
+	names, err := s.client.ListBlobs(ctx, s.prefix)
 	if err != nil {
 		return nil, fmt.Errorf("iac azure state: ListStates: %w", err)
 	}
@@ -125,8 +125,14 @@ func (s *AzureBlobIaCStateStore) ListStates(filter map[string]string) ([]*IaCSta
 		if !strings.HasSuffix(name, ".json") {
 			continue
 		}
-		data, err := s.client.DownloadBlob(context.Background(), name)
+		data, err := s.client.DownloadBlob(ctx, name)
 		if err != nil {
+			// A canceled / deadlined context must abort the listing rather
+			// than silently return partial results; only genuinely unreadable
+			// blobs are skipped.
+			if ctx.Err() != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return nil, fmt.Errorf("iac azure state: ListStates: %w", err)
+			}
 			continue
 		}
 		var st IaCState
@@ -141,8 +147,8 @@ func (s *AzureBlobIaCStateStore) ListStates(filter map[string]string) ([]*IaCSta
 }
 
 // DeleteState removes the state blob for resourceID.
-func (s *AzureBlobIaCStateStore) DeleteState(resourceID string) error {
-	if err := s.client.DeleteBlob(context.Background(), s.blobName(resourceID)); err != nil {
+func (s *AzureBlobIaCStateStore) DeleteState(ctx context.Context, resourceID string) error {
+	if err := s.client.DeleteBlob(ctx, s.blobName(resourceID)); err != nil {
 		if errors.Is(err, ErrAzureBlobNotFound) {
 			return fmt.Errorf("iac azure state: DeleteState %q: not found", resourceID)
 		}
@@ -152,12 +158,12 @@ func (s *AzureBlobIaCStateStore) DeleteState(resourceID string) error {
 }
 
 // Lock acquires a blob lease on the lock blob for resourceID (60-second duration).
-func (s *AzureBlobIaCStateStore) Lock(resourceID string) error {
+func (s *AzureBlobIaCStateStore) Lock(ctx context.Context, resourceID string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	lockBlob := s.lockBlobName(resourceID)
-	leaseID, err := s.client.AcquireLease(context.Background(), lockBlob, 60)
+	leaseID, err := s.client.AcquireLease(ctx, lockBlob, 60)
 	if err != nil {
 		if strings.Contains(err.Error(), "already leased") || strings.Contains(err.Error(), "leased") {
 			return fmt.Errorf("iac azure state: Lock %q: resource is already locked", resourceID)
@@ -169,7 +175,7 @@ func (s *AzureBlobIaCStateStore) Lock(resourceID string) error {
 }
 
 // Unlock releases the lease on the lock blob for resourceID.
-func (s *AzureBlobIaCStateStore) Unlock(resourceID string) error {
+func (s *AzureBlobIaCStateStore) Unlock(ctx context.Context, resourceID string) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -178,7 +184,7 @@ func (s *AzureBlobIaCStateStore) Unlock(resourceID string) error {
 		return fmt.Errorf("iac azure state: Unlock %q: not locked", resourceID)
 	}
 	lockBlob := s.lockBlobName(resourceID)
-	if err := s.client.ReleaseLease(context.Background(), lockBlob, leaseID); err != nil {
+	if err := s.client.ReleaseLease(ctx, lockBlob, leaseID); err != nil {
 		return fmt.Errorf("iac azure state: Unlock %q: %w", resourceID, err)
 	}
 	delete(s.leaseIDs, resourceID)

module/iac_state_azure_test.go

@@ -97,7 +97,7 @@ func newTestAzureStore(client module.AzureBlobClient) *module.AzureBlobIaCStateS
 
 func TestAzureBlobIaCStateStore_GetState_NotFound(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
-	st, err := store.GetState("nonexistent")
+	st, err := store.GetState(context.Background(), "nonexistent")
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -115,11 +115,11 @@ func TestAzureBlobIaCStateStore_SaveAndGetState(t *testing.T) {
 		Provider:     "azure",
 		Status:       "active",
 	}
-	if err := store.SaveState(state); err != nil {
+	if err := store.SaveState(context.Background(), state); err != nil {
 		t.Fatalf("SaveState: %v", err)
 	}
 
-	got, err := store.GetState("az-cluster")
+	got, err := store.GetState(context.Background(), "az-cluster")
 	if err != nil {
 		t.Fatalf("GetState: %v", err)
 	}
@@ -133,14 +133,14 @@ func TestAzureBlobIaCStateStore_SaveAndGetState(t *testing.T) {
 
 func TestAzureBlobIaCStateStore_SaveState_Nil(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
-	if err := store.SaveState(nil); err == nil {
+	if err := store.SaveState(context.Background(), nil); err == nil {
 		t.Fatal("expected error for nil state")
 	}
 }
 
 func TestAzureBlobIaCStateStore_SaveState_EmptyID(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
-	if err := store.SaveState(&module.IaCState{}); err == nil {
+	if err := store.SaveState(context.Background(), &module.IaCState{}); err == nil {
 		t.Fatal("expected error for empty resource_id")
 	}
 }
@@ -153,20 +153,20 @@ func TestAzureBlobIaCStateStore_ListStates(t *testing.T) {
 		{ResourceID: "r2", ResourceType: "db", Provider: "azure", Status: "active"},
 		{ResourceID: "r3", ResourceType: "k8s", Provider: "gcp", Status: "destroyed"},
 	} {
-		if err := store.SaveState(st); err != nil {
+		if err := store.SaveState(context.Background(), st); err != nil {
 			t.Fatalf("SaveState %q: %v", st.ResourceID, err)
 		}
 	}
 
-	all, err := store.ListStates(nil)
+	all, err := store.ListStates(context.Background(), nil)
 	if err != nil {
 		t.Fatalf("ListStates(nil): %v", err)
 	}
 	if len(all) != 3 {
 		t.Errorf("ListStates = %d, want 3", len(all))
 	}
 
-	filtered, err := store.ListStates(map[string]string{"provider": "azure"})
+	filtered, err := store.ListStates(context.Background(), map[string]string{"provider": "azure"})
 	if err != nil {
 		t.Fatalf("ListStates(provider=azure): %v", err)
 	}
@@ -178,14 +178,14 @@ func TestAzureBlobIaCStateStore_ListStates(t *testing.T) {
 func TestAzureBlobIaCStateStore_ListStates_SkipsLockBlobs(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
 
-	if err := store.SaveState(&module.IaCState{ResourceID: "r1", Status: "active"}); err != nil {
+	if err := store.SaveState(context.Background(), &module.IaCState{ResourceID: "r1", Status: "active"}); err != nil {
 		t.Fatalf("SaveState: %v", err)
 	}
-	if err := store.Lock("r1"); err != nil {
+	if err := store.Lock(context.Background(), "r1"); err != nil {
 		t.Fatalf("Lock: %v", err)
 	}
 
-	results, err := store.ListStates(nil)
+	results, err := store.ListStates(context.Background(), nil)
 	if err != nil {
 		t.Fatalf("ListStates: %v", err)
 	}
@@ -197,13 +197,13 @@ func TestAzureBlobIaCStateStore_ListStates_SkipsLockBlobs(t *testing.T) {
 func TestAzureBlobIaCStateStore_DeleteState(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
 
-	if err := store.SaveState(&module.IaCState{ResourceID: "del-me", Status: "active"}); err != nil {
+	if err := store.SaveState(context.Background(), &module.IaCState{ResourceID: "del-me", Status: "active"}); err != nil {
 		t.Fatalf("SaveState: %v", err)
 	}
-	if err := store.DeleteState("del-me"); err != nil {
+	if err := store.DeleteState(context.Background(), "del-me"); err != nil {
 		t.Fatalf("DeleteState: %v", err)
 	}
-	st, err := store.GetState("del-me")
+	st, err := store.GetState(context.Background(), "del-me")
 	if err != nil {
 		t.Fatalf("GetState after delete: %v", err)
 	}
@@ -214,31 +214,31 @@ func TestAzureBlobIaCStateStore_DeleteState(t *testing.T) {
 
 func TestAzureBlobIaCStateStore_DeleteState_NotFound(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
-	if err := store.DeleteState("nonexistent"); err == nil {
+	if err := store.DeleteState(context.Background(), "nonexistent"); err == nil {
 		t.Fatal("expected error deleting nonexistent state")
 	}
 }
 
 func TestAzureBlobIaCStateStore_LockUnlock(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
 
-	if err := store.Lock("res-1"); err != nil {
+	if err := store.Lock(context.Background(), "res-1"); err != nil {
 		t.Fatalf("Lock: %v", err)
 	}
-	if err := store.Lock("res-1"); err == nil {
+	if err := store.Lock(context.Background(), "res-1"); err == nil {
 		t.Fatal("expected error on double lock")
 	}
-	if err := store.Unlock("res-1"); err != nil {
+	if err := store.Unlock(context.Background(), "res-1"); err != nil {
 		t.Fatalf("Unlock: %v", err)
 	}
-	if err := store.Lock("res-1"); err != nil {
+	if err := store.Lock(context.Background(), "res-1"); err != nil {
 		t.Fatalf("Lock after unlock: %v", err)
 	}
 }
 
 func TestAzureBlobIaCStateStore_Unlock_NotLocked(t *testing.T) {
 	store := newTestAzureStore(newMockAzureClient())
-	if err := store.Unlock("not-locked"); err == nil {
+	if err := store.Unlock(context.Background(), "not-locked"); err == nil {
 		t.Fatal("expected error unlocking non-locked resource")
 	}
 }
@@ -250,15 +250,15 @@ func TestAzureBlobIaCStateStore_Unlock_PassesLeaseID(t *testing.T) {
 	client := newMockAzureClient()
 	store := newTestAzureStore(client)
 
-	if err := store.Lock("res-lease"); err != nil {
+	if err := store.Lock(context.Background(), "res-lease"); err != nil {
 		t.Fatalf("Lock: %v", err)
 	}
 	// Unlock must pass the correct leaseID — mock rejects wrong/empty leaseIDs.
-	if err := store.Unlock("res-lease"); err != nil {
+	if err := store.Unlock(context.Background(), "res-lease"); err != nil {
 		t.Fatalf("Unlock with leaseID: %v", err)
 	}
 	// After unlock, should be able to re-lock.
-	if err := store.Lock("res-lease"); err != nil {
+	if err := store.Lock(context.Background(), "res-lease"); err != nil {
 		t.Fatalf("Lock after Unlock: %v", err)
 	}
 }
@@ -272,10 +272,10 @@ func TestAzureBlobIaCStateStore_JSONRoundTrip(t *testing.T) {
 		Status:     "active",
 		Outputs:    map[string]any{"fqdn": "myapp.azurewebsites.net"},
 	}
-	if err := store.SaveState(state); err != nil {
+	if err := store.SaveState(context.Background(), state); err != nil {
 		t.Fatalf("SaveState: %v", err)
 	}
-	got, err := store.GetState("az-rt")
+	got, err := store.GetState(context.Background(), "az-rt")
 	if err != nil {
 		t.Fatalf("GetState: %v", err)
 	}