From b1fe040e43cd9705baf64e0da9d8cc5efe405b13 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Mon, 9 Jun 2025 21:32:03 -0700 Subject: [PATCH] Add tasks totalCount support Signed-off-by: Bryan Frimin --- pkg/coredata/task.go | 64 ++++++++++++ pkg/probo/task_service.go | 50 +++++++++ pkg/server/api/console/v1/schema.graphql | 6 +- pkg/server/api/console/v1/schema/schema.go | 114 +++++++++++++++++++-- pkg/server/api/console/v1/types/task.go | 21 +++- pkg/server/api/console/v1/types/types.go | 5 - pkg/server/api/console/v1/v1_resolver.go | 30 +++++- 7 files changed, 273 insertions(+), 17 deletions(-) diff --git a/pkg/coredata/task.go b/pkg/coredata/task.go index 0dc88a29a..08ccbce37 100644 --- a/pkg/coredata/task.go +++ b/pkg/coredata/task.go @@ -248,6 +248,38 @@ RETURNING return nil } +func (c *Tasks) CountByOrganizationID( + ctx context.Context, + conn pg.Conn, + scope Scoper, + organizationID gid.GID, +) (int, error) { + q := ` + SELECT + COUNT(id) + FROM + tasks + WHERE + %s + AND organization_id = @organization_id + ` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"organization_id": organizationID} + maps.Copy(args, scope.SQLArguments()) + + row := conn.QueryRow(ctx, q, args) + + var count int + err := row.Scan(&count) + if err != nil { + return 0, fmt.Errorf("cannot collect tasks: %w", err) + } + + return count, nil +} + func (c *Tasks) LoadByOrganizationID( ctx context.Context, conn pg.Conn, @@ -297,6 +329,38 @@ func (c *Tasks) LoadByOrganizationID( return nil } +func (c *Tasks) CountByMeasureID( + ctx context.Context, + conn pg.Conn, + scope Scoper, + measureID gid.GID, +) (int, error) { + q := ` +SELECT + COUNT(id) +FROM + tasks +WHERE + %s + AND measure_id = @measure_id +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"measure_id": measureID} + maps.Copy(args, scope.SQLArguments()) + + row := conn.QueryRow(ctx, q, args) + + var count int + err := row.Scan(&count) + if err != nil { + return 0, fmt.Errorf("cannot collect tasks: %w", err) + } + + return count, nil +} + func (c *Tasks) LoadByMeasureID( ctx context.Context, conn pg.Conn, diff --git a/pkg/probo/task_service.go b/pkg/probo/task_service.go index f8f9e77a4..719247934 100644 --- a/pkg/probo/task_service.go +++ b/pkg/probo/task_service.go @@ -245,6 +245,31 @@ func (s TaskService) Delete( return nil } +func (s TaskService) CountForOrganizationID( + ctx context.Context, + organizationID gid.GID, +) (int, error) { + var count int + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) (err error) { + tasks := coredata.Tasks{} + count, err = tasks.CountByOrganizationID(ctx, conn, s.svc.scope, organizationID) + if err != nil { + return fmt.Errorf("cannot count tasks: %w", err) + } + + return nil + }, + ) + if err != nil { + return 0, err + } + + return count, nil +} + func (s TaskService) ListForOrganizationID( ctx context.Context, organizationID gid.GID, @@ -265,6 +290,31 @@ func (s TaskService) ListForOrganizationID( return page.NewPage(tasks, cursor), nil } +func (s TaskService) CountForMeasureID( + ctx context.Context, + measureID gid.GID, +) (int, error) { + var count int + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) (err error) { + tasks := coredata.Tasks{} + count, err = tasks.CountByMeasureID(ctx, conn, s.svc.scope, measureID) + if err != nil { + return fmt.Errorf("cannot count tasks: %w", err) + } + + return nil + }, + ) + if err != nil { + return 0, err + } + + return count, nil +} + func (s TaskService) ListForMeasureID( ctx context.Context, measureID gid.GID, diff --git a/pkg/server/api/console/v1/schema.graphql b/pkg/server/api/console/v1/schema.graphql index 71b5f9320..eb9eaef63 100644 --- a/pkg/server/api/console/v1/schema.graphql +++ b/pkg/server/api/console/v1/schema.graphql @@ -1132,7 +1132,11 @@ type MeasureEdge { node: Measure! } -type TaskConnection { +type TaskConnection + @goModel( + model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.TaskConnection" + ) { + totalCount: Int! @goField(forceResolver: true) edges: [TaskEdge!]! pageInfo: PageInfo! } diff --git a/pkg/server/api/console/v1/schema/schema.go b/pkg/server/api/console/v1/schema/schema.go index fdb3a6ece..6d8c8f983 100644 --- a/pkg/server/api/console/v1/schema/schema.go +++ b/pkg/server/api/console/v1/schema/schema.go @@ -61,6 +61,7 @@ type ResolverRoot interface { Risk() RiskResolver RiskConnection() RiskConnectionResolver Task() TaskResolver + TaskConnection() TaskConnectionResolver User() UserResolver Vendor() VendorResolver VendorComplianceReport() VendorComplianceReportResolver @@ -690,8 +691,9 @@ type ComplexityRoot struct { } TaskConnection struct { - Edges func(childComplexity int) int - PageInfo func(childComplexity int) int + Edges func(childComplexity int) int + PageInfo func(childComplexity int) int + TotalCount func(childComplexity int) int } TaskEdge struct { @@ -1037,6 +1039,9 @@ type TaskResolver interface { Measure(ctx context.Context, obj *types.Task) (*types.Measure, error) Evidences(ctx context.Context, obj *types.Task, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.EvidenceOrderBy) (*types.EvidenceConnection, error) } +type TaskConnectionResolver interface { + TotalCount(ctx context.Context, obj *types.TaskConnection) (int, error) +} type UserResolver interface { People(ctx context.Context, obj *types.User, organizationID gid.GID) (*types.People, error) } @@ -3941,6 +3946,13 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin return e.complexity.TaskConnection.PageInfo(childComplexity), true + case "TaskConnection.totalCount": + if e.complexity.TaskConnection.TotalCount == nil { + break + } + + return e.complexity.TaskConnection.TotalCount(childComplexity), true + case "TaskEdge.cursor": if e.complexity.TaskEdge.Cursor == nil { break @@ -5889,7 +5901,11 @@ type MeasureEdge { node: Measure! } -type TaskConnection { +type TaskConnection + @goModel( + model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.TaskConnection" + ) { + totalCount: Int! @goField(forceResolver: true) edges: [TaskEdge!]! pageInfo: PageInfo! } @@ -21032,6 +21048,8 @@ func (ec *executionContext) fieldContext_Measure_tasks(ctx context.Context, fiel IsResolver: true, Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) { switch field.Name { + case "totalCount": + return ec.fieldContext_TaskConnection_totalCount(ctx, field) case "edges": return ec.fieldContext_TaskConnection_edges(ctx, field) case "pageInfo": @@ -25992,6 +26010,8 @@ func (ec *executionContext) fieldContext_Organization_tasks(ctx context.Context, IsResolver: true, Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) { switch field.Name { + case "totalCount": + return ec.fieldContext_TaskConnection_totalCount(ctx, field) case "edges": return ec.fieldContext_TaskConnection_edges(ctx, field) case "pageInfo": @@ -29760,6 +29780,50 @@ func (ec *executionContext) fieldContext_Task_updatedAt(_ context.Context, field return fc, nil } +func (ec *executionContext) _TaskConnection_totalCount(ctx context.Context, field graphql.CollectedField, obj *types.TaskConnection) (ret graphql.Marshaler) { + fc, err := ec.fieldContext_TaskConnection_totalCount(ctx, field) + if err != nil { + return graphql.Null + } + ctx = graphql.WithFieldContext(ctx, fc) + defer func() { + if r := recover(); r != nil { + ec.Error(ctx, ec.Recover(ctx, r)) + ret = graphql.Null + } + }() + resTmp, err := ec.ResolverMiddleware(ctx, func(rctx context.Context) (any, error) { + ctx = rctx // use context from middleware stack in children + return ec.resolvers.TaskConnection().TotalCount(rctx, obj) + }) + if err != nil { + ec.Error(ctx, err) + return graphql.Null + } + if resTmp == nil { + if !graphql.HasFieldError(ctx, fc) { + ec.Errorf(ctx, "must not be null") + } + return graphql.Null + } + res := resTmp.(int) + fc.Result = res + return ec.marshalNInt2int(ctx, field.Selections, res) +} + +func (ec *executionContext) fieldContext_TaskConnection_totalCount(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) { + fc = &graphql.FieldContext{ + Object: "TaskConnection", + Field: field, + IsMethod: true, + IsResolver: true, + Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) { + return nil, errors.New("field of type Int does not have child fields") + }, + } + return fc, nil +} + func (ec *executionContext) _TaskConnection_edges(ctx context.Context, field graphql.CollectedField, obj *types.TaskConnection) (ret graphql.Marshaler) { fc, err := ec.fieldContext_TaskConnection_edges(ctx, field) if err != nil { @@ -29836,9 +29900,9 @@ func (ec *executionContext) _TaskConnection_pageInfo(ctx context.Context, field } return graphql.Null } - res := resTmp.(*types.PageInfo) + res := resTmp.(types.PageInfo) fc.Result = res - return ec.marshalNPageInfo2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐPageInfo(ctx, field.Selections, res) + return ec.marshalNPageInfo2githubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐPageInfo(ctx, field.Selections, res) } func (ec *executionContext) fieldContext_TaskConnection_pageInfo(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) { @@ -46973,15 +47037,51 @@ func (ec *executionContext) _TaskConnection(ctx context.Context, sel ast.Selecti switch field.Name { case "__typename": out.Values[i] = graphql.MarshalString("TaskConnection") + case "totalCount": + field := field + + innerFunc := func(ctx context.Context, fs *graphql.FieldSet) (res graphql.Marshaler) { + defer func() { + if r := recover(); r != nil { + ec.Error(ctx, ec.Recover(ctx, r)) + } + }() + res = ec._TaskConnection_totalCount(ctx, field, obj) + if res == graphql.Null { + atomic.AddUint32(&fs.Invalids, 1) + } + return res + } + + if field.Deferrable != nil { + dfs, ok := deferred[field.Deferrable.Label] + di := 0 + if ok { + dfs.AddField(field) + di = len(dfs.Values) - 1 + } else { + dfs = graphql.NewFieldSet([]graphql.CollectedField{field}) + deferred[field.Deferrable.Label] = dfs + } + dfs.Concurrently(di, func(ctx context.Context) graphql.Marshaler { + return innerFunc(ctx, dfs) + }) + + // don't run the out.Concurrently() call below + out.Values[i] = graphql.Null + continue + } + + out.Concurrently(i, func(ctx context.Context) graphql.Marshaler { return innerFunc(ctx, out) }) case "edges": out.Values[i] = ec._TaskConnection_edges(ctx, field, obj) if out.Values[i] == graphql.Null { - out.Invalids++ + atomic.AddUint32(&out.Invalids, 1) } case "pageInfo": out.Values[i] = ec._TaskConnection_pageInfo(ctx, field, obj) if out.Values[i] == graphql.Null { - out.Invalids++ + atomic.AddUint32(&out.Invalids, 1) } default: panic("unknown field " + strconv.Quote(field.Name)) diff --git a/pkg/server/api/console/v1/types/task.go b/pkg/server/api/console/v1/types/task.go index 4757b1a71..2a1a334e3 100644 --- a/pkg/server/api/console/v1/types/task.go +++ b/pkg/server/api/console/v1/types/task.go @@ -16,14 +16,28 @@ package types import ( "github.com/getprobo/probo/pkg/coredata" + "github.com/getprobo/probo/pkg/gid" "github.com/getprobo/probo/pkg/page" ) type ( TaskOrderBy OrderBy[coredata.TaskOrderField] + + TaskConnection struct { + TotalCount int + Edges []*TaskEdge + PageInfo PageInfo + + Resolver any + ParentID gid.GID + } ) -func NewTaskConnection(p *page.Page[*coredata.Task, coredata.TaskOrderField]) *TaskConnection { +func NewTaskConnection( + p *page.Page[*coredata.Task, coredata.TaskOrderField], + parentType any, + parentID gid.GID, +) *TaskConnection { var edges = make([]*TaskEdge, len(p.Data)) for i := range edges { @@ -32,7 +46,10 @@ func NewTaskConnection(p *page.Page[*coredata.Task, coredata.TaskOrderField]) *T return &TaskConnection{ Edges: edges, - PageInfo: NewPageInfo(p), + PageInfo: *NewPageInfo(p), + + Resolver: parentType, + ParentID: parentID, } } diff --git a/pkg/server/api/console/v1/types/types.go b/pkg/server/api/console/v1/types/types.go index f47f576d5..205f8c3e1 100644 --- a/pkg/server/api/console/v1/types/types.go +++ b/pkg/server/api/console/v1/types/types.go @@ -914,11 +914,6 @@ type Task struct { func (Task) IsNode() {} func (this Task) GetID() gid.GID { return this.ID } -type TaskConnection struct { - Edges []*TaskEdge `json:"edges"` - PageInfo *PageInfo `json:"pageInfo"` -} - type TaskEdge struct { Cursor page.CursorKey `json:"cursor"` Node *Task `json:"node"` diff --git a/pkg/server/api/console/v1/v1_resolver.go b/pkg/server/api/console/v1/v1_resolver.go index 96c5c4127..3d2783996 100644 --- a/pkg/server/api/console/v1/v1_resolver.go +++ b/pkg/server/api/console/v1/v1_resolver.go @@ -675,7 +675,7 @@ func (r *measureResolver) Tasks(ctx context.Context, obj *types.Measure, first * return nil, fmt.Errorf("cannot list measure tasks: %w", err) } - return types.NewTaskConnection(page), nil + return types.NewTaskConnection(page, r, obj.ID), nil } // Risks is the resolver for the risks field. @@ -2293,7 +2293,7 @@ func (r *organizationResolver) Tasks(ctx context.Context, obj *types.Organizatio panic(fmt.Errorf("cannot list organization tasks: %w", err)) } - return types.NewTaskConnection(page), nil + return types.NewTaskConnection(page, r, obj.ID), nil } // Assets is the resolver for the assets field. @@ -2695,6 +2695,28 @@ func (r *taskResolver) Evidences(ctx context.Context, obj *types.Task, first *in return types.NewEvidenceConnection(page), nil } +// TotalCount is the resolver for the totalCount field. +func (r *taskConnectionResolver) TotalCount(ctx context.Context, obj *types.TaskConnection) (int, error) { + svc := GetTenantService(ctx, r.proboSvc, obj.ParentID.TenantID()) + + switch obj.Resolver.(type) { + case *measureResolver: + count, err := svc.Tasks.CountForMeasureID(ctx, obj.ParentID) + if err != nil { + return 0, fmt.Errorf("cannot count tasks: %w", err) + } + return count, nil + case *organizationResolver: + count, err := svc.Tasks.CountForOrganizationID(ctx, obj.ParentID) + if err != nil { + return 0, fmt.Errorf("cannot count tasks: %w", err) + } + return count, nil + } + + panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver)) +} + // People is the resolver for the people field. func (r *userResolver) People(ctx context.Context, obj *types.User, organizationID gid.GID) (*types.People, error) { svc := GetTenantService(ctx, r.proboSvc, organizationID.TenantID()) @@ -2962,6 +2984,9 @@ func (r *Resolver) RiskConnection() schema.RiskConnectionResolver { return &risk // Task returns schema.TaskResolver implementation. func (r *Resolver) Task() schema.TaskResolver { return &taskResolver{r} } +// TaskConnection returns schema.TaskConnectionResolver implementation. +func (r *Resolver) TaskConnection() schema.TaskConnectionResolver { return &taskConnectionResolver{r} } + // User returns schema.UserResolver implementation. func (r *Resolver) User() schema.UserResolver { return &userResolver{r} } @@ -3000,6 +3025,7 @@ type queryResolver struct{ *Resolver } type riskResolver struct{ *Resolver } type riskConnectionResolver struct{ *Resolver } type taskResolver struct{ *Resolver } +type taskConnectionResolver struct{ *Resolver } type userResolver struct{ *Resolver } type vendorResolver struct{ *Resolver } type vendorComplianceReportResolver struct{ *Resolver }