From c8381ac5074fa196962770aa6af6820a94214f52 Mon Sep 17 00:00:00 2001 From: Marco Dinis Date: Fri, 31 Jan 2025 12:11:10 +0000 Subject: [PATCH] Integration Status endpoint: include pending UserTasks This PR changes the Integration Status endpoint to include the pending UserTasks for the integration. This will be used by the frontend to show the number of pending tasks. --- lib/web/integrations.go | 26 +++++++++++ lib/web/integrations_test.go | 85 ++++++++++++++++++++++++++++++++++++ lib/web/ui/integration.go | 2 + 3 files changed, 113 insertions(+) diff --git a/lib/web/integrations.go b/lib/web/integrations.go index 988543bf9d7b7..b0d3a64fd2d41 100644 --- a/lib/web/integrations.go +++ b/lib/web/integrations.go @@ -32,8 +32,10 @@ import ( discoveryconfigv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/discoveryconfig/v1" integrationv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/integration/v1" pluginspb "github.com/gravitational/teleport/api/gen/proto/go/teleport/plugins/v1" + usertasksv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/usertasks/v1" "github.com/gravitational/teleport/api/types" "github.com/gravitational/teleport/api/types/discoveryconfig" + "github.com/gravitational/teleport/api/types/usertasks" "github.com/gravitational/teleport/integrations/access/msteams" "github.com/gravitational/teleport/lib/defaults" "github.com/gravitational/teleport/lib/httplib" @@ -265,6 +267,7 @@ func (h *Handler) integrationStats(w http.ResponseWriter, r *http.Request, p htt discoveryConfigLister: clt.DiscoveryConfigClient(), databaseGetter: clt, awsOIDCClient: clt.IntegrationAWSOIDCClient(), + userTasksClient: clt.UserTasksServiceClient(), } summary, err := collectIntegrationStats(r.Context(), req) if err != nil { @@ -274,12 +277,17 @@ func (h *Handler) integrationStats(w http.ResponseWriter, r *http.Request, p htt return summary, nil } +type userTasksByIntegrationLister interface { + ListUserTasksByIntegration(ctx context.Context, pageSize int64, nextToken string, integration string) ([]*usertasksv1.UserTask, string, error) +} + type collectIntegrationStatsRequest struct { logger *slog.Logger integration types.Integration discoveryConfigLister discoveryConfigLister databaseGetter databaseGetter awsOIDCClient deployedDatabaseServiceLister + userTasksClient userTasksByIntegrationLister } func collectIntegrationStats(ctx context.Context, req collectIntegrationStatsRequest) (*ui.IntegrationWithSummary, error) { @@ -292,6 +300,24 @@ func collectIntegrationStats(ctx context.Context, req collectIntegrationStatsReq ret.Integration = uiIg var nextPage string + for { + userTasks, nextToken, err := req.userTasksClient.ListUserTasksByIntegration(ctx, 0, nextPage, req.integration.GetName()) + if err != nil { + return nil, err + } + for _, userTask := range userTasks { + if userTask.GetSpec().GetState() == usertasks.TaskStateOpen { + ret.UnresolvedUserTasks++ + } + } + + if nextToken == "" { + break + } + nextPage = nextToken + } + + nextPage = "" for { discoveryConfigs, nextToken, err := req.discoveryConfigLister.ListDiscoveryConfigs(ctx, 0, nextPage) if err != nil { diff --git a/lib/web/integrations_test.go b/lib/web/integrations_test.go index dba752211e6c9..fd59c099b2f84 100644 --- a/lib/web/integrations_test.go +++ b/lib/web/integrations_test.go @@ -21,6 +21,7 @@ package web import ( "context" "encoding/json" + "strconv" "testing" "time" @@ -32,9 +33,11 @@ import ( "github.com/gravitational/teleport/api/client/proto" discoveryconfigv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/discoveryconfig/v1" integrationv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/integration/v1" + usertasksv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/usertasks/v1" "github.com/gravitational/teleport/api/types" "github.com/gravitational/teleport/api/types/discoveryconfig" "github.com/gravitational/teleport/api/types/header" + "github.com/gravitational/teleport/api/types/usertasks" "github.com/gravitational/teleport/lib/auth/integration/credentials" "github.com/gravitational/teleport/lib/modules" "github.com/gravitational/teleport/lib/services" @@ -102,6 +105,44 @@ func TestIntegrationsCreateWithAudience(t *testing.T) { } } +type mockUserTasksLister struct { + defaultPageSize int64 + userTasks []*usertasksv1.UserTask +} + +func (m *mockUserTasksLister) ListUserTasksByIntegration(ctx context.Context, pageSize int64, nextToken string, integration string) ([]*usertasksv1.UserTask, string, error) { + var ret []*usertasksv1.UserTask + if pageSize == 0 { + pageSize = m.defaultPageSize + } + + if len(m.userTasks) == 0 { + return ret, "", nil + } + + var sliceStart int + if nextToken != "" { + nextTokenInt, err := strconv.Atoi(nextToken) + if err != nil { + return nil, "", trace.Wrap(err) + } + sliceStart = nextTokenInt + } + userTasksSlice := m.userTasks[sliceStart:] + + for i, userTask := range userTasksSlice { + if userTask.GetSpec().GetState() == "OPEN" { + ret = append(ret, userTask) + if len(ret) == int(pageSize) { + nextTokenInt := sliceStart + i + 1 + return ret, strconv.Itoa(nextTokenInt), nil + } + } + } + + return ret, "", nil +} + func TestCollectAWSOIDCAutoDiscoverStats(t *testing.T) { ctx := context.Background() logger := utils.NewSlogLoggerForTests() @@ -138,6 +179,47 @@ func TestCollectAWSOIDCAutoDiscoverStats(t *testing.T) { discoveryConfigLister: clt, databaseGetter: clt, awsOIDCClient: deployedDatabaseServicesClient, + userTasksClient: &mockUserTasksLister{}, + } + gotSummary, err := collectIntegrationStats(ctx, req) + require.NoError(t, err) + expectedSummary := &ui.IntegrationWithSummary{ + Integration: &ui.Integration{ + Name: integrationName, + SubKind: "aws-oidc", + AWSOIDC: &ui.IntegrationAWSOIDCSpec{RoleARN: "arn:role"}, + }, + } + require.Equal(t, expectedSummary, gotSummary) + }) + + t.Run("returns the number of unresolved user tasks", func(t *testing.T) { + clt := &mockRelevantAWSRegionsClient{ + databaseServices: &proto.ListResourcesResponse{ + Resources: []*proto.PaginatedResource{}, + }, + databases: make([]types.Database, 0), + discoveryConfigs: make([]*discoveryconfig.DiscoveryConfig, 0), + } + + var userTasksList []*usertasksv1.UserTask + for range 10 { + userTasksList = append(userTasksList, &usertasksv1.UserTask{Spec: &usertasksv1.UserTaskSpec{State: usertasks.TaskStateOpen}}) + userTasksList = append(userTasksList, &usertasksv1.UserTask{Spec: &usertasksv1.UserTaskSpec{State: usertasks.TaskStateResolved}}) + } + + userTasksClient := &mockUserTasksLister{ + defaultPageSize: 3, + userTasks: userTasksList, + } + + req := collectIntegrationStatsRequest{ + logger: logger, + integration: integration, + discoveryConfigLister: clt, + databaseGetter: clt, + awsOIDCClient: deployedDatabaseServicesClient, + userTasksClient: userTasksClient, } gotSummary, err := collectIntegrationStats(ctx, req) require.NoError(t, err) @@ -147,6 +229,7 @@ func TestCollectAWSOIDCAutoDiscoverStats(t *testing.T) { SubKind: "aws-oidc", AWSOIDC: &ui.IntegrationAWSOIDCSpec{RoleARN: "arn:role"}, }, + UnresolvedUserTasks: 10, } require.Equal(t, expectedSummary, gotSummary) }) @@ -217,6 +300,7 @@ func TestCollectAWSOIDCAutoDiscoverStats(t *testing.T) { discoveryConfigLister: clt, databaseGetter: clt, awsOIDCClient: deployedDatabaseServicesClient, + userTasksClient: &mockUserTasksLister{}, } gotSummary, err := collectIntegrationStats(ctx, req) require.NoError(t, err) @@ -286,6 +370,7 @@ func TestCollectAWSOIDCAutoDiscoverStats(t *testing.T) { discoveryConfigLister: clt, databaseGetter: clt, awsOIDCClient: deployedDatabaseServicesClient, + userTasksClient: &mockUserTasksLister{}, } gotSummary, err := collectIntegrationStats(ctx, req) require.NoError(t, err) diff --git a/lib/web/ui/integration.go b/lib/web/ui/integration.go index 1d823d979cb88..142105ae323f5 100644 --- a/lib/web/ui/integration.go +++ b/lib/web/ui/integration.go @@ -71,6 +71,8 @@ func (r *IntegrationAWSOIDCSpec) CheckAndSetDefaults() error { // IntegrationWithSummary describes Integration fields and the fields required to return the summary. type IntegrationWithSummary struct { *Integration + // UnresolvedUserTasks contains the count of unresolved user tasks related to this integration. + UnresolvedUserTasks int `json:"unresolvedUserTasks,omitempty"` // AWSEC2 contains the summary for the AWS EC2 resources for this integration. AWSEC2 ResourceTypeSummary `json:"awsec2,omitempty"` // AWSRDS contains the summary for the AWS RDS resources and agents for this integration.