Add vendors totalCount support

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

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,

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

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

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,

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

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
+		Edges      func(childComplexity int) int
-		PageInfo 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))

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

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

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 }