diff --git a/pkg/coredata/vendor.go b/pkg/coredata/vendor.go index 7d9196e50..3952fa40c 100644 --- a/pkg/coredata/vendor.go +++ b/pkg/coredata/vendor.go @@ -235,6 +235,38 @@ DELETE FROM vendors WHERE %s AND id = @vendor_id return err } +func (v *Vendors) CountByOrganizationID( + ctx context.Context, + conn pg.Conn, + scope Scoper, + organizationID gid.GID, +) (int, error) { + q := ` +SELECT + COUNT(id) +FROM + vendors +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 count vendors: %w", err) + } + + return count, nil +} + func (v *Vendors) LoadByOrganizationID( ctx context.Context, conn pg.Conn, @@ -390,6 +422,45 @@ func (v Vendor) ExpireNonExpiredRiskAssessments( return nil } +func (v *Vendors) CountByAssetID( + ctx context.Context, + conn pg.Conn, + scope Scoper, + assetID gid.GID, +) (int, error) { + q := ` +WITH vend AS ( + SELECT + v.id + FROM + vendors v + INNER JOIN + asset_vendors av ON v.id = av.vendor_id + WHERE + av.asset_id = @asset_id +) +SELECT + COUNT(id) +FROM + vend +WHERE %s +` + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"asset_id": assetID} + 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 count vendors: %w", err) + } + + return count, nil +} + func (v *Vendors) LoadByAssetID( ctx context.Context, conn pg.Conn, @@ -480,6 +551,45 @@ WHERE %s return nil } +func (v *Vendors) CountByDatumID( + ctx context.Context, + conn pg.Conn, + scope Scoper, + datumID gid.GID, +) (int, error) { + q := ` +WITH vend AS ( + SELECT + v.id + FROM + vendors v + INNER JOIN + data_vendors dv ON v.id = dv.vendor_id + WHERE + dv.datum_id = @datum_id +) +SELECT + COUNT(id) +FROM + vend +WHERE %s +` + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"datum_id": datumID} + 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 count vendors: %w", err) + } + + return count, nil +} + func (vs *Vendors) LoadByDatumID( ctx context.Context, conn pg.Conn, diff --git a/pkg/probo/asset_service.go b/pkg/probo/asset_service.go index b9b697ebc..316f20e42 100644 --- a/pkg/probo/asset_service.go +++ b/pkg/probo/asset_service.go @@ -205,24 +205,3 @@ func (s AssetService) Delete( }, ) } - -func (s AssetService) ListVendors( - ctx context.Context, - assetID gid.GID, - cursor *page.Cursor[coredata.VendorOrderField], -) (*page.Page[*coredata.Vendor, coredata.VendorOrderField], error) { - var vendors coredata.Vendors - - err := s.svc.pg.WithConn( - ctx, - func(conn pg.Conn) error { - return vendors.LoadByAssetID(ctx, conn, s.svc.scope, assetID, cursor) - }, - ) - - if err != nil { - return nil, err - } - - return page.NewPage(vendors, cursor), nil -} diff --git a/pkg/probo/datum_service.go b/pkg/probo/datum_service.go index cddebc639..1845f4be7 100644 --- a/pkg/probo/datum_service.go +++ b/pkg/probo/datum_service.go @@ -218,30 +218,3 @@ func (s DatumService) ListVendors( return page.NewPage(vendors, cursor), nil } - -func (s VendorService) ListForDatumID( - ctx context.Context, - datumID gid.GID, - cursor *page.Cursor[coredata.VendorOrderField], -) (*page.Page[*coredata.Vendor, coredata.VendorOrderField], error) { - var vendors coredata.Vendors - - err := s.svc.pg.WithConn( - ctx, - func(conn pg.Conn) error { - return vendors.LoadByDatumID( - ctx, - conn, - s.svc.scope, - datumID, - cursor, - ) - }, - ) - - if err != nil { - return nil, err - } - - return page.NewPage(vendors, cursor), nil -} diff --git a/pkg/probo/vendor_service.go b/pkg/probo/vendor_service.go index cb83a24c1..526335dee 100644 --- a/pkg/probo/vendor_service.go +++ b/pkg/probo/vendor_service.go @@ -89,6 +89,32 @@ type ( } ) +func (s VendorService) CountForOrganizationID( + ctx context.Context, + organizationID gid.GID, +) (int, error) { + var count int + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) (err error) { + vendors := coredata.Vendors{} + count, err = vendors.CountByOrganizationID(ctx, conn, s.svc.scope, organizationID) + if err != nil { + return fmt.Errorf("cannot count vendors: %w", err) + } + + return nil + }, + ) + + if err != nil { + return 0, err + } + + return count, nil +} + func (s VendorService) ListForOrganizationID( ctx context.Context, organizationID gid.GID, @@ -121,6 +147,59 @@ func (s VendorService) ListForOrganizationID( return page.NewPage(vendors, cursor), nil } +func (s VendorService) CountForDatumID( + ctx context.Context, + datumID gid.GID, +) (int, error) { + var count int + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) (err error) { + vendors := coredata.Vendors{} + count, err = vendors.CountByDatumID(ctx, conn, s.svc.scope, datumID) + if err != nil { + return fmt.Errorf("cannot count vendors: %w", err) + } + + return nil + }, + ) + + if err != nil { + return 0, err + } + + return count, nil +} + +func (s VendorService) ListForDatumID( + ctx context.Context, + datumID gid.GID, + cursor *page.Cursor[coredata.VendorOrderField], +) (*page.Page[*coredata.Vendor, coredata.VendorOrderField], error) { + var vendors coredata.Vendors + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + return vendors.LoadByDatumID( + ctx, + conn, + s.svc.scope, + datumID, + cursor, + ) + }, + ) + + if err != nil { + return nil, err + } + + return page.NewPage(vendors, cursor), nil +} + func (s VendorService) Update( ctx context.Context, req UpdateVendorRequest, @@ -355,6 +434,53 @@ func (s VendorService) Create( return vendor, nil } +func (s VendorService) CountForAssetID( + ctx context.Context, + assetID gid.GID, +) (int, error) { + var count int + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) (err error) { + vendors := coredata.Vendors{} + count, err = vendors.CountByAssetID(ctx, conn, s.svc.scope, assetID) + if err != nil { + return fmt.Errorf("cannot count vendors: %w", err) + } + + return nil + }, + ) + + if err != nil { + return 0, err + } + + return count, nil +} + +func (s VendorService) ListForAssetID( + ctx context.Context, + assetID gid.GID, + cursor *page.Cursor[coredata.VendorOrderField], +) (*page.Page[*coredata.Vendor, coredata.VendorOrderField], error) { + var vendors coredata.Vendors + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + return vendors.LoadByAssetID(ctx, conn, s.svc.scope, assetID, cursor) + }, + ) + + if err != nil { + return nil, err + } + + return page.NewPage(vendors, cursor), nil +} + func (s VendorService) ListRiskAssessments( ctx context.Context, vendorID gid.GID, diff --git a/pkg/server/api/console/v1/schema.graphql b/pkg/server/api/console/v1/schema.graphql index a83981b98..7d67a59c8 100644 --- a/pkg/server/api/console/v1/schema.graphql +++ b/pkg/server/api/console/v1/schema.graphql @@ -1080,7 +1080,11 @@ type PeopleEdge { node: People! } -type VendorConnection { +type VendorConnection + @goModel( + model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.VendorConnection" + ) { + totalCount: Int! @goField(forceResolver: true) edges: [VendorEdge!]! pageInfo: PageInfo! } diff --git a/pkg/server/api/console/v1/schema/schema.go b/pkg/server/api/console/v1/schema/schema.go index 84319d707..2a3c3f4e0 100644 --- a/pkg/server/api/console/v1/schema/schema.go +++ b/pkg/server/api/console/v1/schema/schema.go @@ -66,6 +66,7 @@ type ResolverRoot interface { User() UserResolver Vendor() VendorResolver VendorComplianceReport() VendorComplianceReportResolver + VendorConnection() VendorConnectionResolver VendorRiskAssessment() VendorRiskAssessmentResolver Viewer() ViewerResolver } @@ -836,8 +837,9 @@ type ComplexityRoot struct { } VendorConnection 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 } VendorEdge struct { @@ -1062,6 +1064,9 @@ type VendorComplianceReportResolver interface { FileURL(ctx context.Context, obj *types.VendorComplianceReport) (string, error) } +type VendorConnectionResolver interface { + TotalCount(ctx context.Context, obj *types.VendorConnection) (int, error) +} type VendorRiskAssessmentResolver interface { Vendor(ctx context.Context, obj *types.VendorRiskAssessment) (*types.Vendor, error) @@ -4449,6 +4454,13 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin return e.complexity.VendorConnection.PageInfo(childComplexity), true + case "VendorConnection.totalCount": + if e.complexity.VendorConnection.TotalCount == nil { + break + } + + return e.complexity.VendorConnection.TotalCount(childComplexity), true + case "VendorEdge.cursor": if e.complexity.VendorEdge.Cursor == nil { break @@ -5861,7 +5873,11 @@ type PeopleEdge { node: People! } -type VendorConnection { +type VendorConnection + @goModel( + model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.VendorConnection" + ) { + totalCount: Int! @goField(forceResolver: true) edges: [VendorEdge!]! pageInfo: PageInfo! } @@ -12182,6 +12198,8 @@ func (ec *executionContext) fieldContext_Asset_vendors(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_VendorConnection_totalCount(ctx, field) case "edges": return ec.fieldContext_VendorConnection_edges(ctx, field) case "pageInfo": @@ -15351,6 +15369,8 @@ func (ec *executionContext) fieldContext_Datum_vendors(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_VendorConnection_totalCount(ctx, field) case "edges": return ec.fieldContext_VendorConnection_edges(ctx, field) case "pageInfo": @@ -25761,6 +25781,8 @@ func (ec *executionContext) fieldContext_Organization_vendors(ctx context.Contex IsResolver: true, Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) { switch field.Name { + case "totalCount": + return ec.fieldContext_VendorConnection_totalCount(ctx, field) case "edges": return ec.fieldContext_VendorConnection_edges(ctx, field) case "pageInfo": @@ -33462,6 +33484,50 @@ func (ec *executionContext) fieldContext_VendorComplianceReportEdge_node(_ conte return fc, nil } +func (ec *executionContext) _VendorConnection_totalCount(ctx context.Context, field graphql.CollectedField, obj *types.VendorConnection) (ret graphql.Marshaler) { + fc, err := ec.fieldContext_VendorConnection_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.VendorConnection().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_VendorConnection_totalCount(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) { + fc = &graphql.FieldContext{ + Object: "VendorConnection", + 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) _VendorConnection_edges(ctx context.Context, field graphql.CollectedField, obj *types.VendorConnection) (ret graphql.Marshaler) { fc, err := ec.fieldContext_VendorConnection_edges(ctx, field) if err != nil { @@ -33538,9 +33604,9 @@ func (ec *executionContext) _VendorConnection_pageInfo(ctx context.Context, fiel } 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_VendorConnection_pageInfo(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) { @@ -48555,15 +48621,51 @@ func (ec *executionContext) _VendorConnection(ctx context.Context, sel ast.Selec switch field.Name { case "__typename": out.Values[i] = graphql.MarshalString("VendorConnection") + 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._VendorConnection_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._VendorConnection_edges(ctx, field, obj) if out.Values[i] == graphql.Null { - out.Invalids++ + atomic.AddUint32(&out.Invalids, 1) } case "pageInfo": out.Values[i] = ec._VendorConnection_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/types.go b/pkg/server/api/console/v1/types/types.go index 9bf04fe31..96a8d982a 100644 --- a/pkg/server/api/console/v1/types/types.go +++ b/pkg/server/api/console/v1/types/types.go @@ -1193,11 +1193,6 @@ type VendorComplianceReportEdge struct { Node *VendorComplianceReport `json:"node"` } -type VendorConnection struct { - Edges []*VendorEdge `json:"edges"` - PageInfo *PageInfo `json:"pageInfo"` -} - type VendorEdge struct { Cursor page.CursorKey `json:"cursor"` Node *Vendor `json:"node"` diff --git a/pkg/server/api/console/v1/types/vendor.go b/pkg/server/api/console/v1/types/vendor.go index 17b615ea9..70ebbe47e 100644 --- a/pkg/server/api/console/v1/types/vendor.go +++ b/pkg/server/api/console/v1/types/vendor.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 ( VendorOrderBy OrderBy[coredata.VendorOrderField] + + VendorConnection struct { + TotalCount int + Edges []*VendorEdge + PageInfo PageInfo + + Resolver any + ParentID gid.GID + } ) -func NewVendorConnection(p *page.Page[*coredata.Vendor, coredata.VendorOrderField]) *VendorConnection { +func NewVendorConnection( + p *page.Page[*coredata.Vendor, coredata.VendorOrderField], + parentType any, + parentID gid.GID, +) *VendorConnection { var edges = make([]*VendorEdge, len(p.Data)) for i := range edges { @@ -32,7 +46,10 @@ func NewVendorConnection(p *page.Page[*coredata.Vendor, coredata.VendorOrderFiel return &VendorConnection{ Edges: edges, - PageInfo: NewPageInfo(p), + PageInfo: *NewPageInfo(p), + + Resolver: parentType, + ParentID: parentID, } } diff --git a/pkg/server/api/console/v1/v1_resolver.go b/pkg/server/api/console/v1/v1_resolver.go index 735e4f935..44b8a5031 100644 --- a/pkg/server/api/console/v1/v1_resolver.go +++ b/pkg/server/api/console/v1/v1_resolver.go @@ -54,12 +54,12 @@ func (r *assetResolver) Vendors(ctx context.Context, obj *types.Asset, first *in cursor := types.NewCursor(first, after, last, before, pageOrderBy) - page, err := svc.Assets.ListVendors(ctx, obj.ID, cursor) + page, err := svc.Vendors.ListForAssetID(ctx, obj.ID, cursor) if err != nil { panic(fmt.Errorf("cannot list asset vendors: %w", err)) } - return types.NewVendorConnection(page), nil + return types.NewVendorConnection(page, r, obj.ID), nil } // AssetType is the resolver for the assetType field. @@ -246,7 +246,7 @@ func (r *datumResolver) Vendors(ctx context.Context, obj *types.Datum, first *in panic(fmt.Errorf("cannot list data vendors: %w", err)) } - return types.NewVendorConnection(page), nil + return types.NewVendorConnection(page, r, obj.ID), nil } // Organization is the resolver for the organization field. @@ -2175,7 +2175,7 @@ func (r *organizationResolver) Vendors(ctx context.Context, obj *types.Organizat panic(fmt.Errorf("cannot list organization vendors: %w", err)) } - return types.NewVendorConnection(page), nil + return types.NewVendorConnection(page, r, obj.ID), nil } // Peoples is the resolver for the peoples field. @@ -2883,6 +2883,34 @@ func (r *vendorComplianceReportResolver) FileURL(ctx context.Context, obj *types return fileURL, nil } +// TotalCount is the resolver for the totalCount field. +func (r *vendorConnectionResolver) TotalCount(ctx context.Context, obj *types.VendorConnection) (int, error) { + svc := GetTenantService(ctx, r.proboSvc, obj.ParentID.TenantID()) + + switch obj.Resolver.(type) { + case *organizationResolver: + count, err := svc.Vendors.CountForOrganizationID(ctx, obj.ParentID) + if err != nil { + return 0, fmt.Errorf("cannot count vendors: %w", err) + } + return count, nil + case *assetResolver: + count, err := svc.Vendors.CountForAssetID(ctx, obj.ParentID) + if err != nil { + return 0, fmt.Errorf("cannot count vendors: %w", err) + } + return count, nil + case *datumResolver: + count, err := svc.Vendors.CountForDatumID(ctx, obj.ParentID) + if err != nil { + return 0, fmt.Errorf("cannot count vendors: %w", err) + } + return count, nil + } + + panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver)) +} + // Vendor is the resolver for the vendor field. func (r *vendorRiskAssessmentResolver) Vendor(ctx context.Context, obj *types.VendorRiskAssessment) (*types.Vendor, error) { svc := GetTenantService(ctx, r.proboSvc, obj.ID.TenantID()) @@ -3025,6 +3053,11 @@ func (r *Resolver) VendorComplianceReport() schema.VendorComplianceReportResolve return &vendorComplianceReportResolver{r} } +// VendorConnection returns schema.VendorConnectionResolver implementation. +func (r *Resolver) VendorConnection() schema.VendorConnectionResolver { + return &vendorConnectionResolver{r} +} + // VendorRiskAssessment returns schema.VendorRiskAssessmentResolver implementation. func (r *Resolver) VendorRiskAssessment() schema.VendorRiskAssessmentResolver { return &vendorRiskAssessmentResolver{r} @@ -3057,5 +3090,6 @@ type taskConnectionResolver struct{ *Resolver } type userResolver struct{ *Resolver } type vendorResolver struct{ *Resolver } type vendorComplianceReportResolver struct{ *Resolver } +type vendorConnectionResolver struct{ *Resolver } type vendorRiskAssessmentResolver struct{ *Resolver } type viewerResolver struct{ *Resolver }