diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a3bfb3..f0278b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ All notable changes to this project will be documented in this file. ### Changed - `grant status` now fetches sessions and eligibility data concurrently, reducing wall-clock time by ~2s +- `grant revoke` interactive mode now fetches workspace names concurrently across CSPs ### Fixed diff --git a/cmd/helpers.go b/cmd/helpers.go index 04b5712..6b25928 100644 --- a/cmd/helpers.go +++ b/cmd/helpers.go @@ -96,8 +96,9 @@ func fetchStatusData( } // buildWorkspaceNameMap fetches eligibility for each unique CSP in sessions -// and builds a workspaceID -> workspaceName map. Errors are silently ignored -// (graceful degradation — the raw workspace ID is shown as fallback). +// concurrently and builds a workspaceID -> workspaceName map. Errors are +// silently ignored (graceful degradation — the raw workspace ID is shown +// as fallback). func buildWorkspaceNameMap(ctx context.Context, eligLister eligibilityLister, sessions []scamodels.SessionInfo, errWriter io.Writer) map[string]string { nameMap := make(map[string]string) @@ -106,22 +107,46 @@ func buildWorkspaceNameMap(ctx context.Context, eligLister eligibilityLister, se for _, s := range sessions { csps[s.CSP] = true } + if len(csps) == 0 { + return nameMap + } + + type eligResult struct { + csp scamodels.CSP + targets []scamodels.EligibleTarget + err error + } - // Fetch eligibility for each CSP + // Fetch eligibility for each CSP concurrently + results := make(chan eligResult, len(csps)) + var wg sync.WaitGroup for csp := range csps { - if ctx.Err() != nil { - break - } - resp, err := eligLister.ListEligibility(ctx, csp) - if err != nil || resp == nil { - if verbose && err != nil { - fmt.Fprintf(errWriter, "Warning: failed to fetch names for %s: %v\n", csp, err) + wg.Add(1) + go func(csp scamodels.CSP) { + defer wg.Done() + resp, err := eligLister.ListEligibility(ctx, csp) + if err != nil || resp == nil { + results <- eligResult{csp: csp, err: err} + return + } + results <- eligResult{csp: csp, targets: resp.Response} + }(csp) + } + go func() { + wg.Wait() + close(results) + }() + + for r := range results { + if r.err != nil { + if verbose { + fmt.Fprintf(errWriter, "Warning: failed to fetch names for %s: %v\n", r.csp, r.err) } continue } - for _, target := range resp.Response { - if target.WorkspaceName != "" { - nameMap[target.WorkspaceID] = target.WorkspaceName + for _, t := range r.targets { + if t.WorkspaceName != "" { + nameMap[t.WorkspaceID] = t.WorkspaceName } } } diff --git a/cmd/helpers_test.go b/cmd/helpers_test.go index 11675c0..728aea8 100644 --- a/cmd/helpers_test.go +++ b/cmd/helpers_test.go @@ -271,6 +271,167 @@ func TestFetchStatusData_VerboseWarning(t *testing.T) { } } +func TestBuildWorkspaceNameMap(t *testing.T) { + tests := []struct { + name string + sessions []scamodels.SessionInfo + setupEligibility func() *mockEligibilityLister + wantNameMapKeys []string + wantNameMapVals map[string]string + }{ + { + name: "multiple CSPs fetched concurrently", + sessions: []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + {CSP: scamodels.CSPAWS, WorkspaceID: "arn:aws:iam::123:role/Admin"}, + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + switch csp { + case scamodels.CSPAzure: + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, nil + case scamodels.CSPAWS: + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "arn:aws:iam::123:role/Admin", WorkspaceName: "Admin Account"}, + }, + }, nil + } + return &scamodels.EligibilityResponse{}, nil + }, + } + }, + wantNameMapKeys: []string{"/subscriptions/sub-1", "arn:aws:iam::123:role/Admin"}, + wantNameMapVals: map[string]string{ + "/subscriptions/sub-1": "Dev Sub", + "arn:aws:iam::123:role/Admin": "Admin Account", + }, + }, + { + name: "no sessions returns empty map", + sessions: []scamodels.SessionInfo{}, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + t.Error("ListEligibility should not be called with no sessions") + return nil, nil + }, + } + }, + }, + { + name: "duplicate CSPs deduplicated", + sessions: []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-2"}, + }, + setupEligibility: func() *mockEligibilityLister { + var callCount int + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + callCount++ + if callCount > 1 { + t.Error("ListEligibility called more than once for same CSP") + } + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Sub 1"}, + {WorkspaceID: "/subscriptions/sub-2", WorkspaceName: "Sub 2"}, + }, + }, nil + }, + } + }, + wantNameMapKeys: []string{"/subscriptions/sub-1", "/subscriptions/sub-2"}, + }, + { + name: "partial failure - graceful degradation", + sessions: []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + {CSP: scamodels.CSPAWS, WorkspaceID: "arn:aws:iam::123:role/Admin"}, + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + if csp == scamodels.CSPAWS { + return nil, errors.New("AWS unavailable") + } + return &scamodels.EligibilityResponse{ + Response: []scamodels.EligibleTarget{ + {WorkspaceID: "/subscriptions/sub-1", WorkspaceName: "Dev Sub"}, + }, + }, nil + }, + } + }, + wantNameMapKeys: []string{"/subscriptions/sub-1"}, + }, + { + name: "all failures - empty map", + sessions: []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listErr: errors.New("eligibility unavailable"), + } + }, + wantNameMapKeys: nil, + }, + { + name: "context cancellation - no hang", + sessions: []scamodels.SessionInfo{ + {CSP: scamodels.CSPAzure, WorkspaceID: "/subscriptions/sub-1"}, + {CSP: scamodels.CSPAWS, WorkspaceID: "arn:aws:iam::123:role/Admin"}, + }, + setupEligibility: func() *mockEligibilityLister { + return &mockEligibilityLister{ + listFunc: func(ctx context.Context, csp scamodels.CSP) (*scamodels.EligibilityResponse, error) { + return nil, ctx.Err() + }, + } + }, + wantNameMapKeys: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + if tt.name == "context cancellation - no hang" { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(ctx) + cancel() + } + + eligLister := tt.setupEligibility() + var errBuf bytes.Buffer + nameMap := buildWorkspaceNameMap(ctx, eligLister, tt.sessions, &errBuf) + + for _, key := range tt.wantNameMapKeys { + if _, ok := nameMap[key]; !ok { + t.Errorf("nameMap missing key %q, got: %v", key, nameMap) + } + } + + for k, wantV := range tt.wantNameMapVals { + if gotV := nameMap[k]; gotV != wantV { + t.Errorf("nameMap[%q] = %q, want %q", k, gotV, wantV) + } + } + + if tt.wantNameMapKeys == nil && len(nameMap) > 0 { + t.Errorf("expected empty nameMap, got: %v", nameMap) + } + }) + } +} + func TestBuildWorkspaceNameMap_VerboseWarning(t *testing.T) { oldVerbose := verbose verbose = true