Add assets totalCount support

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

pkg/coredata/asset.go

@@ -26,7 +26,8 @@ import (
 	"go.gearno.de/kit/pg"
 )
 
-type Asset struct {
+type (
+	Asset struct {
 		ID              gid.GID        `db:"id"`
 		Name            string         `db:"name"`
 		Amount          int            `db:"amount"`
@@ -39,6 +40,9 @@ type Asset struct {
 		UpdatedAt       time.Time      `db:"updated_at"`
 	}
 
+	Assets []*Asset
+)
+
 func (a *Asset) CursorKey(field AssetOrderField) page.CursorKey {
 	switch field {
 	case AssetOrderFieldCreatedAt:
@@ -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,

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,

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

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
@@ -98,6 +99,7 @@ type ComplexityRoot struct {
 	AssetConnection struct {
 		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))

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 {
+func NewAssetConnection(
-	edges := make([]*AssetEdge, len(page.Data))
+	p *page.Page[*coredata.Asset, coredata.AssetOrderField],
-	for i, asset := range page.Data {
+	resolver any,
-		edges[i] = NewAssetEdge(asset, page.Cursor.OrderBy.Field)
+	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,
 	}
 }
 

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

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 }