diff --git a/pkg/coredata/asset.go b/pkg/coredata/asset.go index 19144b570..b7a80cd9a 100644 --- a/pkg/coredata/asset.go +++ b/pkg/coredata/asset.go @@ -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, diff --git a/pkg/probo/asset_service.go b/pkg/probo/asset_service.go index 316f20e42..75190158a 100644 --- a/pkg/probo/asset_service.go +++ b/pkg/probo/asset_service.go @@ -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, diff --git a/pkg/server/api/console/v1/schema.graphql b/pkg/server/api/console/v1/schema.graphql index 77220892e..69db514f1 100644 --- a/pkg/server/api/console/v1/schema.graphql +++ b/pkg/server/api/console/v1/schema.graphql @@ -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! } diff --git a/pkg/server/api/console/v1/schema/schema.go b/pkg/server/api/console/v1/schema/schema.go index 5d1e65d1c..cbec66d0e 100644 --- a/pkg/server/api/console/v1/schema/schema.go +++ b/pkg/server/api/console/v1/schema/schema.go @@ -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)) diff --git a/pkg/server/api/console/v1/types/asset.go b/pkg/server/api/console/v1/types/asset.go index 48106a36d..dba3295d1 100644 --- a/pkg/server/api/console/v1/types/asset.go +++ b/pkg/server/api/console/v1/types/asset.go @@ -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, } } diff --git a/pkg/server/api/console/v1/types/types.go b/pkg/server/api/console/v1/types/types.go index 58686e9b8..0ad071888 100644 --- a/pkg/server/api/console/v1/types/types.go +++ b/pkg/server/api/console/v1/types/types.go @@ -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"` diff --git a/pkg/server/api/console/v1/v1_resolver.go b/pkg/server/api/console/v1/v1_resolver.go index 3b3df1f53..5fd58ba30 100644 --- a/pkg/server/api/console/v1/v1_resolver.go +++ b/pkg/server/api/console/v1/v1_resolver.go @@ -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 }