Add datum totalCount support

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2025-06-09 22:22:12 -07:00
parent 9932212c39
commit 771ae6dbc3
7 changed files with 237 additions and 80 deletions

View File

@@ -26,15 +26,19 @@ import (
"go.gearno.de/kit/pg"
)
type Data struct {
ID gid.GID `db:"id"`
Name string `db:"name"`
OrganizationID gid.GID `db:"organization_id"`
OwnerID gid.GID `db:"owner_id"`
DataClassification DataClassification `db:"data_classification"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
type (
Data struct {
ID gid.GID `db:"id"`
Name string `db:"name"`
OrganizationID gid.GID `db:"organization_id"`
OwnerID gid.GID `db:"owner_id"`
DataClassification DataClassification `db:"data_classification"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
DataList []*Data
)
func (d *Data) CursorKey(field DatumOrderField) page.CursorKey {
switch field {
@@ -49,8 +53,6 @@ func (d *Data) CursorKey(field DatumOrderField) page.CursorKey {
panic(fmt.Sprintf("unsupported order by: %s", field))
}
type DataList []*Data
func (d *Data) LoadByID(
ctx context.Context,
conn pg.Conn,
@@ -136,48 +138,36 @@ LIMIT 1;
return nil
}
func (dl *DataList) LoadByOwnerID(
func (dl *DataList) CountByOrganizationID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
ownerID gid.GID,
cursor *page.Cursor[DatumOrderField],
) error {
organizationID gid.GID,
) (int, error) {
q := `
SELECT
id,
name,
owner_id,
data_classification,
created_at,
updated_at
COUNT(id)
FROM
data
WHERE
%s
AND owner_id = @owner_id
AND %s
AND organization_id = @organization_id
`
q = fmt.Sprintf(q, scope.SQLFragment(), cursor.SQLFragment())
q = fmt.Sprintf(q, scope.SQLFragment())
args := pgx.StrictNamedArgs{"owner_id": ownerID}
args := pgx.StrictNamedArgs{"organization_id": organizationID}
maps.Copy(args, scope.SQLArguments())
maps.Copy(args, cursor.SQLArguments())
rows, err := conn.Query(ctx, q, args)
row := conn.QueryRow(ctx, q, args)
var count int
err := row.Scan(&count)
if err != nil {
return fmt.Errorf("cannot query data: %w", err)
return 0, fmt.Errorf("cannot count data: %w", err)
}
data, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Data])
if err != nil {
return fmt.Errorf("cannot collect data: %w", err)
}
*dl = data
return nil
return count, nil
}
func (dl *DataList) LoadByOrganizationID(

View File

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

View File

@@ -1236,7 +1236,11 @@ type DocumentVersionEdge {
node: DocumentVersion!
}
type DatumConnection {
type DatumConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DatumConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [DatumEdge!]!
pageInfo: PageInfo!
}
@@ -2219,7 +2223,10 @@ type Datum implements Node {
updatedAt: Datetime!
}
input DatumOrder {
input DatumOrder
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DatumOrderBy"
) {
direction: OrderDirection!
field: DatumOrderField!
}

View File

@@ -46,6 +46,7 @@ type ResolverRoot interface {
Control() ControlResolver
ControlConnection() ControlConnectionResolver
Datum() DatumResolver
DatumConnection() DatumConnectionResolver
Document() DocumentResolver
DocumentConnection() DocumentConnectionResolver
DocumentVersion() DocumentVersionResolver
@@ -242,8 +243,9 @@ type ComplexityRoot struct {
}
DatumConnection 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
}
DatumEdge struct {
@@ -560,7 +562,7 @@ type ComplexityRoot struct {
Connectors func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.ConnectorOrder) int
Controls func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.ControlOrderBy, filter *types.ControlFilter) int
CreatedAt func(childComplexity int) int
Data func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrder) int
Data func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrderBy) int
Documents func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DocumentOrderBy, filter *types.DocumentFilter) int
Frameworks func(childComplexity int, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.FrameworkOrderBy) int
ID func(childComplexity int) int
@@ -900,6 +902,9 @@ type DatumResolver interface {
Vendors(ctx context.Context, obj *types.Datum, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.VendorOrderBy) (*types.VendorConnection, error)
Organization(ctx context.Context, obj *types.Datum) (*types.Organization, error)
}
type DatumConnectionResolver interface {
TotalCount(ctx context.Context, obj *types.DatumConnection) (int, error)
}
type DocumentResolver interface {
Owner(ctx context.Context, obj *types.Document) (*types.People, error)
Organization(ctx context.Context, obj *types.Document) (*types.Organization, error)
@@ -1026,7 +1031,7 @@ type OrganizationResolver interface {
Risks(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.RiskOrderBy, filter *types.RiskFilter) (*types.RiskConnection, error)
Tasks(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.TaskOrderBy) (*types.TaskConnection, error)
Assets(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.AssetOrder) (*types.AssetConnection, error)
Data(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrder) (*types.DatumConnection, error)
Data(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrderBy) (*types.DatumConnection, error)
}
type PeopleConnectionResolver interface {
TotalCount(ctx context.Context, obj *types.PeopleConnection) (int, error)
@@ -1638,6 +1643,13 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
return e.complexity.DatumConnection.PageInfo(childComplexity), true
case "DatumConnection.totalCount":
if e.complexity.DatumConnection.TotalCount == nil {
break
}
return e.complexity.DatumConnection.TotalCount(childComplexity), true
case "DatumEdge.cursor":
if e.complexity.DatumEdge.Cursor == nil {
break
@@ -3327,7 +3339,7 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
return 0, false
}
return e.complexity.Organization.Data(childComplexity, args["first"].(*int), args["after"].(*page.CursorKey), args["last"].(*int), args["before"].(*page.CursorKey), args["orderBy"].(*types.DatumOrder)), true
return e.complexity.Organization.Data(childComplexity, args["first"].(*int), args["after"].(*page.CursorKey), args["last"].(*int), args["before"].(*page.CursorKey), args["orderBy"].(*types.DatumOrderBy)), true
case "Organization.documents":
if e.complexity.Organization.Documents == nil {
@@ -6041,7 +6053,11 @@ type DocumentVersionEdge {
node: DocumentVersion!
}
type DatumConnection {
type DatumConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DatumConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [DatumEdge!]!
pageInfo: PageInfo!
}
@@ -7024,7 +7040,10 @@ type Datum implements Node {
updatedAt: Datetime!
}
input DatumOrder {
input DatumOrder
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DatumOrderBy"
) {
direction: OrderDirection!
field: DatumOrderField!
}
@@ -10171,13 +10190,13 @@ func (ec *executionContext) field_Organization_data_argsBefore(
func (ec *executionContext) field_Organization_data_argsOrderBy(
ctx context.Context,
rawArgs map[string]any,
) (*types.DatumOrder, error) {
) (*types.DatumOrderBy, error) {
ctx = graphql.WithPathContext(ctx, graphql.NewPathWithField("orderBy"))
if tmp, ok := rawArgs["orderBy"]; ok {
return ec.unmarshalODatumOrder2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐDatumOrder(ctx, tmp)
return ec.unmarshalODatumOrder2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐDatumOrderBy(ctx, tmp)
}
var zeroVal *types.DatumOrder
var zeroVal *types.DatumOrderBy
return zeroVal, nil
}
@@ -15577,6 +15596,50 @@ func (ec *executionContext) fieldContext_Datum_updatedAt(_ context.Context, fiel
return fc, nil
}
func (ec *executionContext) _DatumConnection_totalCount(ctx context.Context, field graphql.CollectedField, obj *types.DatumConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_DatumConnection_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.DatumConnection().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_DatumConnection_totalCount(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
fc = &graphql.FieldContext{
Object: "DatumConnection",
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) _DatumConnection_edges(ctx context.Context, field graphql.CollectedField, obj *types.DatumConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_DatumConnection_edges(ctx, field)
if err != nil {
@@ -15653,9 +15716,9 @@ func (ec *executionContext) _DatumConnection_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_DatumConnection_pageInfo(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
@@ -26211,7 +26274,7 @@ func (ec *executionContext) _Organization_data(ctx context.Context, field graphq
}()
resTmp, err := ec.ResolverMiddleware(ctx, func(rctx context.Context) (any, error) {
ctx = rctx // use context from middleware stack in children
return ec.resolvers.Organization().Data(rctx, obj, fc.Args["first"].(*int), fc.Args["after"].(*page.CursorKey), fc.Args["last"].(*int), fc.Args["before"].(*page.CursorKey), fc.Args["orderBy"].(*types.DatumOrder))
return ec.resolvers.Organization().Data(rctx, obj, fc.Args["first"].(*int), fc.Args["after"].(*page.CursorKey), fc.Args["last"].(*int), fc.Args["before"].(*page.CursorKey), fc.Args["orderBy"].(*types.DatumOrderBy))
})
if err != nil {
ec.Error(ctx, err)
@@ -26236,6 +26299,8 @@ func (ec *executionContext) fieldContext_Organization_data(ctx context.Context,
IsResolver: true,
Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
switch field.Name {
case "totalCount":
return ec.fieldContext_DatumConnection_totalCount(ctx, field)
case "edges":
return ec.fieldContext_DatumConnection_edges(ctx, field)
case "pageInfo":
@@ -37929,8 +37994,8 @@ func (ec *executionContext) unmarshalInputCreateVendorRiskAssessmentInput(ctx co
return it, nil
}
func (ec *executionContext) unmarshalInputDatumOrder(ctx context.Context, obj any) (types.DatumOrder, error) {
var it types.DatumOrder
func (ec *executionContext) unmarshalInputDatumOrder(ctx context.Context, obj any) (types.DatumOrderBy, error) {
var it types.DatumOrderBy
asMap := map[string]any{}
for k, v := range obj.(map[string]any) {
asMap[k] = v
@@ -42254,15 +42319,51 @@ func (ec *executionContext) _DatumConnection(ctx context.Context, sel ast.Select
switch field.Name {
case "__typename":
out.Values[i] = graphql.MarshalString("DatumConnection")
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._DatumConnection_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._DatumConnection_edges(ctx, field, obj)
if out.Values[i] == graphql.Null {
out.Invalids++
atomic.AddUint32(&out.Invalids, 1)
}
case "pageInfo":
out.Values[i] = ec._DatumConnection_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))
@@ -53733,7 +53834,7 @@ func (ec *executionContext) marshalODatetime2ᚖtimeᚐTime(ctx context.Context,
return res
}
func (ec *executionContext) unmarshalODatumOrder2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐDatumOrder(ctx context.Context, v any) (*types.DatumOrder, error) {
func (ec *executionContext) unmarshalODatumOrder2ᚖgithubᚗcomᚋgetproboᚋproboᚋpkgᚋserverᚋapiᚋconsoleᚋv1ᚋtypesᚐDatumOrderBy(ctx context.Context, v any) (*types.DatumOrderBy, error) {
if v == nil {
return nil, nil
}

View File

@@ -16,9 +16,42 @@ package types
import (
"github.com/getprobo/probo/pkg/coredata"
"github.com/getprobo/probo/pkg/gid"
"github.com/getprobo/probo/pkg/page"
)
type (
DatumOrderBy OrderBy[coredata.DatumOrderField]
DatumConnection struct {
TotalCount int
Edges []*DatumEdge
PageInfo PageInfo
Resolver any
ParentID gid.GID
}
)
func NewDataConnection(
p *page.Page[*coredata.Data, coredata.DatumOrderField],
parentType any,
parentID gid.GID,
) *DatumConnection {
edges := make([]*DatumEdge, len(p.Data))
for i, data := range p.Data {
edges[i] = NewDatumEdge(data, p.Cursor.OrderBy.Field)
}
return &DatumConnection{
Edges: edges,
PageInfo: *NewPageInfo(p),
Resolver: parentType,
ParentID: parentID,
}
}
func NewDatum(d *coredata.Data) *Datum {
return &Datum{
ID: d.ID,
@@ -36,15 +69,3 @@ func NewDatumEdge(d *coredata.Data, orderField coredata.DatumOrderField) *DatumE
Cursor: d.CursorKey(orderField),
}
}
func NewDataConnection(page *page.Page[*coredata.Data, coredata.DatumOrderField]) *DatumConnection {
edges := make([]*DatumEdge, len(page.Data))
for i, data := range page.Data {
edges[i] = NewDatumEdge(data, page.Cursor.OrderBy.Field)
}
return &DatumConnection{
Edges: edges,
PageInfo: NewPageInfo(page),
}
}

View File

@@ -365,21 +365,11 @@ type Datum struct {
func (Datum) IsNode() {}
func (this Datum) GetID() gid.GID { return this.ID }
type DatumConnection struct {
Edges []*DatumEdge `json:"edges"`
PageInfo *PageInfo `json:"pageInfo"`
}
type DatumEdge struct {
Cursor page.CursorKey `json:"cursor"`
Node *Datum `json:"node"`
}
type DatumOrder struct {
Direction page.OrderDirection `json:"direction"`
Field coredata.DatumOrderField `json:"field"`
}
type DeleteAssetInput struct {
AssetID gid.GID `json:"assetId"`
}

View File

@@ -261,6 +261,22 @@ func (r *datumResolver) Organization(ctx context.Context, obj *types.Datum) (*ty
return types.NewOrganization(org), nil
}
// TotalCount is the resolver for the totalCount field.
func (r *datumConnectionResolver) TotalCount(ctx context.Context, obj *types.DatumConnection) (int, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ParentID.TenantID())
switch obj.Resolver.(type) {
case *organizationResolver:
count, err := svc.Data.CountForOrganizationID(ctx, obj.ParentID)
if err != nil {
return 0, fmt.Errorf("cannot count data: %w", err)
}
return count, nil
}
panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver))
}
// Owner is the resolver for the owner field.
func (r *documentResolver) Owner(ctx context.Context, obj *types.Document) (*types.People, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ID.TenantID())
@@ -2344,7 +2360,7 @@ func (r *organizationResolver) Assets(ctx context.Context, obj *types.Organizati
}
// Assets is the resolver for the assets field.
func (r *organizationResolver) Data(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrder) (*types.DatumConnection, error) {
func (r *organizationResolver) Data(ctx context.Context, obj *types.Organization, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DatumOrderBy) (*types.DatumConnection, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ID.TenantID())
pageOrderBy := page.OrderBy[coredata.DatumOrderField]{
@@ -2365,7 +2381,7 @@ func (r *organizationResolver) Data(ctx context.Context, obj *types.Organization
panic(fmt.Errorf("cannot list organization data: %w", err))
}
return types.NewDataConnection(page), nil
return types.NewDataConnection(page, r, obj.ID), nil
}
// TotalCount is the resolver for the totalCount field.
@@ -2995,6 +3011,11 @@ func (r *Resolver) ControlConnection() schema.ControlConnectionResolver {
// Datum returns schema.DatumResolver implementation.
func (r *Resolver) Datum() schema.DatumResolver { return &datumResolver{r} }
// DatumConnection returns schema.DatumConnectionResolver implementation.
func (r *Resolver) DatumConnection() schema.DatumConnectionResolver {
return &datumConnectionResolver{r}
}
// Document returns schema.DocumentResolver implementation.
func (r *Resolver) Document() schema.DocumentResolver { return &documentResolver{r} }
@@ -3091,6 +3112,7 @@ type assetResolver struct{ *Resolver }
type controlResolver struct{ *Resolver }
type controlConnectionResolver struct{ *Resolver }
type datumResolver struct{ *Resolver }
type datumConnectionResolver struct{ *Resolver }
type documentResolver struct{ *Resolver }
type documentConnectionResolver struct{ *Resolver }
type documentVersionResolver struct{ *Resolver }