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

Add assets totalCount support

Browse Source 
Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2025-06-09 22:36:13 -07:00
parent 82effa84df
commit 60d41c0d1c
7 changed files with 228 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
pkg/coredata/asset.go
View File

@@ -26,18 +26,22 @@ import (
"go.gearno.de/kit/pg"
)
type Asset struct {
ID gid.GID `db:"id"`
Name string `db:"name"`
Amount int `db:"amount"`
OwnerID gid.GID `db:"owner_id"`
OrganizationID gid.GID `db:"organization_id"`
Criticity CriticityLevel `db:"criticity"`
AssetType AssetType `db:"asset_type"`
DataTypesStored string `db:"data_types_stored"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
type (
Asset struct {
ID gid.GID `db:"id"`
Name string `db:"name"`
Amount int `db:"amount"`
OwnerID gid.GID `db:"owner_id"`
OrganizationID gid.GID `db:"organization_id"`
Criticity CriticityLevel `db:"criticity"`
AssetType AssetType `db:"asset_type"`
DataTypesStored string `db:"data_types_stored"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
Assets []*Asset
)
func (a *Asset) CursorKey(field AssetOrderField) page.CursorKey {
switch field {
@@ -52,8 +56,6 @@ func (a *Asset) CursorKey(field AssetOrderField) page.CursorKey {
panic(fmt.Sprintf("unsupported order by: %s", field))
}
type Assets []*Asset
func (a *Asset) LoadByID(
ctx context.Context,
conn pg.Conn,
@@ -145,6 +147,37 @@ LIMIT 1;
return nil
}
func (a *Assets) CountByOrganizationID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
organizationID gid.GID,
) (int, error) {
q := `
SELECT
COUNT(id)
FROM
assets
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
if err := row.Scan(&count); err != nil {
return 0, fmt.Errorf("cannot scan count: %w", err)
}
return count, nil
}
func (a *Assets) LoadByOrganizationID(
ctx context.Context,
conn pg.Conn,

26
pkg/probo/asset_service.go
View File

@@ -77,6 +77,32 @@ func (s AssetService) GetByOwnerID(
return asset, nil
}
func (s AssetService) CountForOrganizationID(
ctx context.Context,
organizationID gid.GID,
) (int, error) {
var count int
err := s.svc.pg.WithConn(
ctx,
func(conn pg.Conn) (err error) {
assets := coredata.Assets{}
count, err = assets.CountByOrganizationID(ctx, conn, s.svc.scope, organizationID)
if err != nil {
return fmt.Errorf("cannot count assets: %w", err)
}
return nil
},
)
if err != nil {
return 0, err
}
return count, nil
}
func (s AssetService) ListForOrganizationID(
ctx context.Context,
organizationID gid.GID,

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

@@ -2153,7 +2153,11 @@ type Asset implements Node {
updatedAt: Datetime!
}
type AssetConnection {
type AssetConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.AssetConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [AssetEdge!]!
pageInfo: PageInfo!
}

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

@@ -43,6 +43,7 @@ type Config struct {
type ResolverRoot interface {
Asset() AssetResolver
AssetConnection() AssetConnectionResolver
Control() ControlResolver
ControlConnection() ControlConnectionResolver
Datum() DatumResolver
@@ -96,8 +97,9 @@ type ComplexityRoot struct {
}
AssetConnection 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
}
AssetEdge struct {
@@ -889,6 +891,9 @@ type AssetResolver interface {
Organization(ctx context.Context, obj *types.Asset) (*types.Organization, error)
}
type AssetConnectionResolver interface {
TotalCount(ctx context.Context, obj *types.AssetConnection) (int, error)
}
type ControlResolver interface {
Framework(ctx context.Context, obj *types.Control) (*types.Framework, error)
Measures(ctx context.Context, obj *types.Control, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.MeasureOrderBy, filter *types.MeasureFilter) (*types.MeasureConnection, error)
@@ -1208,6 +1213,13 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
return e.complexity.AssetConnection.PageInfo(childComplexity), true
case "AssetConnection.totalCount":
if e.complexity.AssetConnection.TotalCount == nil {
break
}
return e.complexity.AssetConnection.TotalCount(childComplexity), true
case "AssetEdge.cursor":
if e.complexity.AssetEdge.Cursor == nil {
break
@@ -6970,7 +6982,11 @@ type Asset implements Node {
updatedAt: Datetime!
}
type AssetConnection {
type AssetConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.AssetConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [AssetEdge!]!
pageInfo: PageInfo!
}
@@ -12560,6 +12576,50 @@ func (ec *executionContext) fieldContext_Asset_updatedAt(_ context.Context, fiel
return fc, nil
}
func (ec *executionContext) _AssetConnection_totalCount(ctx context.Context, field graphql.CollectedField, obj *types.AssetConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_AssetConnection_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.AssetConnection().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_AssetConnection_totalCount(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
fc = &graphql.FieldContext{
Object: "AssetConnection",
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) _AssetConnection_edges(ctx context.Context, field graphql.CollectedField, obj *types.AssetConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_AssetConnection_edges(ctx, field)
if err != nil {
@@ -12636,9 +12696,9 @@ func (ec *executionContext) _AssetConnection_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_AssetConnection_pageInfo(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
@@ -26241,6 +26301,8 @@ func (ec *executionContext) fieldContext_Organization_assets(ctx context.Context
IsResolver: true,
Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
switch field.Name {
case "totalCount":
return ec.fieldContext_AssetConnection_totalCount(ctx, field)
case "edges":
return ec.fieldContext_AssetConnection_edges(ctx, field)
case "pageInfo":
@@ -40819,15 +40881,51 @@ func (ec *executionContext) _AssetConnection(ctx context.Context, sel ast.Select
switch field.Name {
case "__typename":
out.Values[i] = graphql.MarshalString("AssetConnection")
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._AssetConnection_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._AssetConnection_edges(ctx, field, obj)
if out.Values[i] == graphql.Null {
out.Invalids++
atomic.AddUint32(&out.Invalids, 1)
}
case "pageInfo":
out.Values[i] = ec._AssetConnection_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))

27
pkg/server/api/console/v1/types/asset.go
View File

@@ -2,22 +2,39 @@ package types
import (
"github.com/getprobo/probo/pkg/coredata"
"github.com/getprobo/probo/pkg/gid"
"github.com/getprobo/probo/pkg/page"
)
type (
AssetOrderBy OrderBy[coredata.AssetOrderField]
AssetConnection struct {
TotalCount int
Edges []*AssetEdge
PageInfo PageInfo
Resolver any
ParentID gid.GID
}
)
func NewAssetConnection(page *page.Page[*coredata.Asset, coredata.AssetOrderField]) *AssetConnection {
edges := make([]*AssetEdge, len(page.Data))
for i, asset := range page.Data {
edges[i] = NewAssetEdge(asset, page.Cursor.OrderBy.Field)
func NewAssetConnection(
p *page.Page[*coredata.Asset, coredata.AssetOrderField],
resolver any,
parentID gid.GID,
) *AssetConnection {
edges := make([]*AssetEdge, len(p.Data))
for i, asset := range p.Data {
edges[i] = NewAssetEdge(asset, p.Cursor.OrderBy.Field)
}
return &AssetConnection{
Edges: edges,
PageInfo: NewPageInfo(page),
PageInfo: *NewPageInfo(p),
Resolver: resolver,
ParentID: parentID,
}
}

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

@@ -42,11 +42,6 @@ type Asset struct {
func (Asset) IsNode() {}
func (this Asset) GetID() gid.GID { return this.ID }
type AssetConnection struct {
Edges []*AssetEdge `json:"edges"`
PageInfo *PageInfo `json:"pageInfo"`
}
type AssetEdge struct {
Cursor page.CursorKey `json:"cursor"`
Node *Asset `json:"node"`

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

@@ -90,6 +90,22 @@ func (r *assetResolver) Organization(ctx context.Context, obj *types.Asset) (*ty
return types.NewOrganization(org), nil
}
// TotalCount is the resolver for the totalCount field.
func (r *assetConnectionResolver) TotalCount(ctx context.Context, obj *types.AssetConnection) (int, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ParentID.TenantID())
switch obj.Resolver.(type) {
case *organizationResolver:
count, err := svc.Assets.CountForOrganizationID(ctx, obj.ParentID)
if err != nil {
return 0, fmt.Errorf("cannot count assets: %w", err)
}
return count, nil
}
panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver))
}
// Framework is the resolver for the framework field.
func (r *controlResolver) Framework(ctx context.Context, obj *types.Control) (*types.Framework, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ID.TenantID())
@@ -2356,7 +2372,7 @@ func (r *organizationResolver) Assets(ctx context.Context, obj *types.Organizati
panic(fmt.Errorf("cannot list organization assets: %w", err))
}
return types.NewAssetConnection(page), nil
return types.NewAssetConnection(page, r, obj.ID), nil
}
// Assets is the resolver for the assets field.
@@ -3000,6 +3016,11 @@ func (r *viewerResolver) Organizations(ctx context.Context, obj *types.Viewer, f
// Asset returns schema.AssetResolver implementation.
func (r *Resolver) Asset() schema.AssetResolver { return &assetResolver{r} }
// AssetConnection returns schema.AssetConnectionResolver implementation.
func (r *Resolver) AssetConnection() schema.AssetConnectionResolver {
return &assetConnectionResolver{r}
}
// Control returns schema.ControlResolver implementation.
func (r *Resolver) Control() schema.ControlResolver { return &controlResolver{r} }
@@ -3109,6 +3130,7 @@ func (r *Resolver) VendorRiskAssessment() schema.VendorRiskAssessmentResolver {
func (r *Resolver) Viewer() schema.ViewerResolver { return &viewerResolver{r} }
type assetResolver struct{ *Resolver }
type assetConnectionResolver struct{ *Resolver }
type controlResolver struct{ *Resolver }
type controlConnectionResolver struct{ *Resolver }
type datumResolver struct{ *Resolver }