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

Add documents totalCount support

Browse Source 
Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2025-06-09 21:20:09 -07:00
parent c4e6de922c
commit 5bc65d3834
7 changed files with 373 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
pkg/coredata/document.go
View File

@@ -96,6 +96,39 @@ LIMIT 1;
return nil
}
func (p *Documents) CountByOrganizationID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
organizationID gid.GID,
filter *DocumentFilter,
) (int, error) {
q := `
SELECT
COUNT(id)
FROM
documents
WHERE
%s
AND organization_id = @organization_id
AND %s
`
q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
args := pgx.NamedArgs{"organization_id": organizationID}
maps.Copy(args, scope.SQLArguments())
maps.Copy(args, filter.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 (p *Documents) LoadByOrganizationID(
ctx context.Context,
conn pg.Conn,
@@ -246,6 +279,47 @@ WHERE %s
return nil
}
func (p *Documents) CountByControlID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
controlID gid.GID,
filter *DocumentFilter,
) (int, error) {
q := `
WITH plcs AS (
SELECT
p.id
FROM
documents p
INNER JOIN
controls_documents cp ON p.id = cp.document_id
WHERE
cp.control_id = @control_id
)
SELECT
COUNT(id)
FROM
plcs
WHERE %s
AND %s
`
q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
args := pgx.NamedArgs{"control_id": controlID}
maps.Copy(args, scope.SQLArguments())
maps.Copy(args, filter.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 (p *Documents) LoadByControlID(
ctx context.Context,
conn pg.Conn,
@@ -310,6 +384,47 @@ WHERE %s
return nil
}
func (p *Documents) CountByRiskID(
ctx context.Context,
conn pg.Conn,
scope Scoper,
riskID gid.GID,
filter *DocumentFilter,
) (int, error) {
q := `
WITH plcs AS (
SELECT
p.id
FROM
documents p
INNER JOIN
risks_documents rp ON p.id = rp.document_id
WHERE
rp.risk_id = @risk_id
)
SELECT
COUNT(id)
FROM
plcs
WHERE %s
AND %s
`
q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment())
args := pgx.NamedArgs{"risk_id": riskID}
maps.Copy(args, scope.SQLArguments())
maps.Copy(args, filter.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 (p *Documents) LoadByRiskID(
ctx context.Context,
conn pg.Conn,

81
pkg/probo/document_service.go
View File

@@ -656,6 +656,33 @@ func (s *DocumentService) GetVersion(
return documentVersion, nil
}
func (s *DocumentService) CountForOrganizationID(
ctx context.Context,
organizationID gid.GID,
filter *coredata.DocumentFilter,
) (int, error) {
var count int
err := s.svc.pg.WithConn(
ctx,
func(conn pg.Conn) (err error) {
documents := &coredata.Documents{}
count, err = documents.CountByOrganizationID(ctx, conn, s.svc.scope, organizationID, filter)
if err != nil {
return fmt.Errorf("cannot count documents: %w", err)
}
return nil
},
)
if err != nil {
return 0, fmt.Errorf("cannot count documents: %w", err)
}
return count, nil
}
func (s *DocumentService) ListByOrganizationID(
ctx context.Context,
organizationID gid.GID,
@@ -685,6 +712,33 @@ func (s *DocumentService) ListByOrganizationID(
return page.NewPage(documents, cursor), nil
}
func (s *DocumentService) CountForControlID(
ctx context.Context,
controlID gid.GID,
filter *coredata.DocumentFilter,
) (int, error) {
var count int
err := s.svc.pg.WithConn(
ctx,
func(conn pg.Conn) (err error) {
documents := &coredata.Documents{}
count, err = documents.CountByControlID(ctx, conn, s.svc.scope, controlID, filter)
if err != nil {
return fmt.Errorf("cannot count documents: %w", err)
}
return nil
},
)
if err != nil {
return 0, fmt.Errorf("cannot count documents: %w", err)
}
return count, nil
}
func (s *DocumentService) ListForControlID(
ctx context.Context,
controlID gid.GID,
@@ -707,6 +761,33 @@ func (s *DocumentService) ListForControlID(
return page.NewPage(documents, cursor), nil
}
func (s *DocumentService) CountForRiskID(
ctx context.Context,
riskID gid.GID,
filter *coredata.DocumentFilter,
) (int, error) {
var count int
err := s.svc.pg.WithConn(
ctx,
func(conn pg.Conn) (err error) {
documents := &coredata.Documents{}
count, err = documents.CountByRiskID(ctx, conn, s.svc.scope, riskID, filter)
if err != nil {
return fmt.Errorf("cannot count documents: %w", err)
}
return nil
},
)
if err != nil {
return 0, fmt.Errorf("cannot count documents: %w", err)
}
return count, nil
}
func (s *DocumentService) ListForRiskID(
ctx context.Context,
riskID gid.GID,

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

@@ -1152,7 +1152,11 @@ type EvidenceEdge {
node: Evidence!
}
type DocumentConnection {
type DocumentConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DocumentConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [DocumentEdge!]!
pageInfo: PageInfo!
}

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

@@ -47,6 +47,7 @@ type ResolverRoot interface {
ControlConnection() ControlConnectionResolver
Datum() DatumResolver
Document() DocumentResolver
DocumentConnection() DocumentConnectionResolver
DocumentVersion() DocumentVersionResolver
DocumentVersionSignature() DocumentVersionSignatureResolver
Evidence() EvidenceResolver
@@ -329,8 +330,9 @@ type ComplexityRoot struct {
}
DocumentConnection 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
}
DocumentEdge struct {
@@ -896,6 +898,9 @@ type DocumentResolver interface {
Versions(ctx context.Context, obj *types.Document, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.DocumentVersionOrderBy, filter *types.DocumentVersionFilter) (*types.DocumentVersionConnection, error)
Controls(ctx context.Context, obj *types.Document, first *int, after *page.CursorKey, last *int, before *page.CursorKey, orderBy *types.ControlOrderBy, filter *types.ControlFilter) (*types.ControlConnection, error)
}
type DocumentConnectionResolver interface {
TotalCount(ctx context.Context, obj *types.DocumentConnection) (int, error)
}
type DocumentVersionResolver interface {
Document(ctx context.Context, obj *types.DocumentVersion) (*types.Document, error)
@@ -1868,6 +1873,13 @@ func (e *executableSchema) Complexity(ctx context.Context, typeName, field strin
return e.complexity.DocumentConnection.PageInfo(childComplexity), true
case "DocumentConnection.totalCount":
if e.complexity.DocumentConnection.TotalCount == nil {
break
}
return e.complexity.DocumentConnection.TotalCount(childComplexity), true
case "DocumentEdge.cursor":
if e.complexity.DocumentEdge.Cursor == nil {
break
@@ -5897,7 +5909,11 @@ type EvidenceEdge {
node: Evidence!
}
type DocumentConnection {
type DocumentConnection
@goModel(
model: "github.com/getprobo/probo/pkg/server/api/console/v1/types.DocumentConnection"
) {
totalCount: Int! @goField(forceResolver: true)
edges: [DocumentEdge!]!
pageInfo: PageInfo!
}
@@ -13548,6 +13564,8 @@ func (ec *executionContext) fieldContext_Control_documents(ctx context.Context,
IsResolver: true,
Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
switch field.Name {
case "totalCount":
return ec.fieldContext_DocumentConnection_totalCount(ctx, field)
case "edges":
return ec.fieldContext_DocumentConnection_edges(ctx, field)
case "pageInfo":
@@ -17156,6 +17174,50 @@ func (ec *executionContext) fieldContext_Document_updatedAt(_ context.Context, f
return fc, nil
}
func (ec *executionContext) _DocumentConnection_totalCount(ctx context.Context, field graphql.CollectedField, obj *types.DocumentConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_DocumentConnection_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.DocumentConnection().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_DocumentConnection_totalCount(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
fc = &graphql.FieldContext{
Object: "DocumentConnection",
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) _DocumentConnection_edges(ctx context.Context, field graphql.CollectedField, obj *types.DocumentConnection) (ret graphql.Marshaler) {
fc, err := ec.fieldContext_DocumentConnection_edges(ctx, field)
if err != nil {
@@ -17232,9 +17294,9 @@ func (ec *executionContext) _DocumentConnection_pageInfo(ctx context.Context, fi
}
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_DocumentConnection_pageInfo(_ context.Context, field graphql.CollectedField) (fc *graphql.FieldContext, err error) {
@@ -25741,6 +25803,8 @@ func (ec *executionContext) fieldContext_Organization_documents(ctx context.Cont
IsResolver: true,
Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
switch field.Name {
case "totalCount":
return ec.fieldContext_DocumentConnection_totalCount(ctx, field)
case "edges":
return ec.fieldContext_DocumentConnection_edges(ctx, field)
case "pageInfo":
@@ -28498,6 +28562,8 @@ func (ec *executionContext) fieldContext_Risk_documents(ctx context.Context, fie
IsResolver: true,
Child: func(ctx context.Context, field graphql.CollectedField) (*graphql.FieldContext, error) {
switch field.Name {
case "totalCount":
return ec.fieldContext_DocumentConnection_totalCount(ctx, field)
case "edges":
return ec.fieldContext_DocumentConnection_edges(ctx, field)
case "pageInfo":
@@ -42874,15 +42940,51 @@ func (ec *executionContext) _DocumentConnection(ctx context.Context, sel ast.Sel
switch field.Name {
case "__typename":
out.Values[i] = graphql.MarshalString("DocumentConnection")
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._DocumentConnection_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._DocumentConnection_edges(ctx, field, obj)
if out.Values[i] == graphql.Null {
out.Invalids++
atomic.AddUint32(&out.Invalids, 1)
}
case "pageInfo":
out.Values[i] = ec._DocumentConnection_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))

31
pkg/server/api/console/v1/types/document.go
View File

@@ -16,22 +16,43 @@ package types
import (
"github.com/getprobo/probo/pkg/coredata"
"github.com/getprobo/probo/pkg/gid"
"github.com/getprobo/probo/pkg/page"
)
type (
DocumentOrderBy OrderBy[coredata.DocumentOrderField]
DocumentConnection struct {
TotalCount int
Edges []*DocumentEdge
PageInfo PageInfo
Resolver any
ParentID gid.GID
Filters *coredata.DocumentFilter
}
)
func NewDocumentConnection(page *page.Page[*coredata.Document, coredata.DocumentOrderField]) *DocumentConnection {
edges := make([]*DocumentEdge, len(page.Data))
for i, document := range page.Data {
edges[i] = NewDocumentEdge(document, page.Cursor.OrderBy.Field)
func NewDocumentConnection(
p *page.Page[*coredata.Document, coredata.DocumentOrderField],
parentType any,
parentID gid.GID,
filters *coredata.DocumentFilter,
) *DocumentConnection {
var edges = make([]*DocumentEdge, len(p.Data))
for i := range edges {
edges[i] = NewDocumentEdge(p.Data[i], p.Cursor.OrderBy.Field)
}
return &DocumentConnection{
Edges: edges,
PageInfo: NewPageInfo(page),
PageInfo: *NewPageInfo(p),
Resolver: parentType,
ParentID: parentID,
Filters: filters,
}
}

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

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

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

@@ -164,7 +164,7 @@ func (r *controlResolver) Documents(ctx context.Context, obj *types.Control, fir
return nil, fmt.Errorf("cannot list documents: %w", err)
}
return types.NewDocumentConnection(page), nil
return types.NewDocumentConnection(page, r, obj.ID, documentFilter), nil
}
// TotalCount is the resolver for the totalCount field.
@@ -351,6 +351,34 @@ func (r *documentResolver) Controls(ctx context.Context, obj *types.Document, fi
return types.NewControlConnection(page, r, obj.ID, controlFilter), nil
}
// TotalCount is the resolver for the totalCount field.
func (r *documentConnectionResolver) TotalCount(ctx context.Context, obj *types.DocumentConnection) (int, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ParentID.TenantID())
switch obj.Resolver.(type) {
case *controlResolver:
count, err := svc.Documents.CountForControlID(ctx, obj.ParentID, obj.Filters)
if err != nil {
return 0, fmt.Errorf("cannot count controls: %w", err)
}
return count, nil
case *organizationResolver:
count, err := svc.Documents.CountForOrganizationID(ctx, obj.ParentID, obj.Filters)
if err != nil {
return 0, fmt.Errorf("cannot count documents: %w", err)
}
return count, nil
case *riskResolver:
count, err := svc.Documents.CountForRiskID(ctx, obj.ParentID, obj.Filters)
if err != nil {
return 0, fmt.Errorf("cannot count risks: %w", err)
}
return count, nil
}
panic(fmt.Errorf("unsupported resolver: %T", obj.Resolver))
}
// Document is the resolver for the document field.
func (r *documentVersionResolver) Document(ctx context.Context, obj *types.DocumentVersion) (*types.Document, error) {
svc := GetTenantService(ctx, r.proboSvc, obj.ID.TenantID())
@@ -2180,7 +2208,7 @@ func (r *organizationResolver) Documents(ctx context.Context, obj *types.Organiz
panic(fmt.Errorf("cannot list organization documents: %w", err))
}
return types.NewDocumentConnection(page), nil
return types.NewDocumentConnection(page, r, obj.ID, documentFilter), nil
}
// Measures is the resolver for the measures field.
@@ -2533,7 +2561,7 @@ func (r *riskResolver) Documents(ctx context.Context, obj *types.Risk, first *in
panic(fmt.Errorf("cannot list risk documents: %w", err))
}
return types.NewDocumentConnection(page), nil
return types.NewDocumentConnection(page, r, obj.ID, documentFilter), nil
}
// Controls is the resolver for the controls field.
@@ -2882,6 +2910,11 @@ func (r *Resolver) Datum() schema.DatumResolver { return &datumResolver{r} }
// Document returns schema.DocumentResolver implementation.
func (r *Resolver) Document() schema.DocumentResolver { return &documentResolver{r} }
// DocumentConnection returns schema.DocumentConnectionResolver implementation.
func (r *Resolver) DocumentConnection() schema.DocumentConnectionResolver {
return &documentConnectionResolver{r}
}
// DocumentVersion returns schema.DocumentVersionResolver implementation.
func (r *Resolver) DocumentVersion() schema.DocumentVersionResolver {
return &documentVersionResolver{r}
@@ -2953,6 +2986,7 @@ type controlResolver struct{ *Resolver }
type controlConnectionResolver struct{ *Resolver }
type datumResolver struct{ *Resolver }
type documentResolver struct{ *Resolver }
type documentConnectionResolver struct{ *Resolver }
type documentVersionResolver struct{ *Resolver }
type documentVersionSignatureResolver struct{ *Resolver }
type evidenceResolver struct{ *Resolver }