rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add vendors totalCount support

Browse Source 
Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2025-06-09 22:01:25 -07:00
parent f7ee529b3b
commit d128e91c34
9 changed files with 407 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

110
pkg/coredata/vendor.go
View File

@@ -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,

21
pkg/probo/asset_service.go
View File

@@ -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
}

27
pkg/probo/datum_service.go
View File

@@ -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
}

126
pkg/probo/vendor_service.go
View File

@@ -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,

6
pkg/server/api/console/v1/schema.graphql
View File

@@ -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!
}

116
pkg/server/api/console/v1/schema/schema.go
View File

@@ -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))

5
pkg/server/api/console/v1/types/types.go
View File

@@ -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"`

21
pkg/server/api/console/v1/types/vendor.go
View File

@@ -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,
}
}

42
pkg/server/api/console/v1/v1_resolver.go
View File

@@ -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 }