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

Add GraphQL dataloaders for batched record lookups

Browse Source 
Introduce dataloadgen-based dataloaders to batch individual
record-by-ID fetches in GraphQL resolvers into single SQL
queries. Each entity type (organization, framework, control,
vendor, document, risk, measure, task, file, report, profile)
gets a LoadByIDs method in coredata and a GetByIDs service
method with variadic arguments and dedicated collection return
types. Resolvers now use dataloader.FromContext instead of
direct service calls for single-record lookups.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-03-25 15:04:51 +01:00
parent 2004bf6050
commit d24712344e
24 changed files with 1091 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

248
pkg/server/api/console/v1/dataloader/dataloader.go Normal file
View File

@@ -0,0 +1,248 @@
// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package dataloader
import (
"context"
"fmt"
"net/http"
"github.com/vikstrous/dataloadgen"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
"go.probo.inc/probo/pkg/iam"
"go.probo.inc/probo/pkg/probo"
)
type (
ctxKey struct{ name string }
Loaders struct {
Organization *dataloadgen.Loader[gid.GID, *coredata.Organization]
Framework *dataloadgen.Loader[gid.GID, *coredata.Framework]
Control *dataloadgen.Loader[gid.GID, *coredata.Control]
Vendor *dataloadgen.Loader[gid.GID, *coredata.Vendor]
Document *dataloadgen.Loader[gid.GID, *coredata.Document]
Profile *dataloadgen.Loader[gid.GID, *coredata.MembershipProfile]
Risk *dataloadgen.Loader[gid.GID, *coredata.Risk]
Measure *dataloadgen.Loader[gid.GID, *coredata.Measure]
Task *dataloadgen.Loader[gid.GID, *coredata.Task]
File *dataloadgen.Loader[gid.GID, *coredata.File]
Report *dataloadgen.Loader[gid.GID, *coredata.Report]
}
batchFetcher struct {
probo *probo.Service
iam *iam.Service
}
)
var loadersKey = &ctxKey{name: "dataloaders"}
func FromContext(ctx context.Context) *Loaders {
return ctx.Value(loadersKey).(*Loaders)
}
func NewMiddleware(proboSvc *probo.Service, iamSvc *iam.Service) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
f := &batchFetcher{probo: proboSvc, iam: iamSvc}
loaders := f.newLoaders()
ctx := context.WithValue(r.Context(), loadersKey, loaders)
next.ServeHTTP(w, r.WithContext(ctx))
})
}
}
func (f *batchFetcher) newLoaders() *Loaders {
return &Loaders{
Organization: dataloadgen.NewMappedLoader(f.fetchOrganizations),
Framework: dataloadgen.NewMappedLoader(f.fetchFrameworks),
Control: dataloadgen.NewMappedLoader(f.fetchControls),
Vendor: dataloadgen.NewMappedLoader(f.fetchVendors),
Document: dataloadgen.NewMappedLoader(f.fetchDocuments),
Profile: dataloadgen.NewMappedLoader(f.fetchProfiles),
Risk: dataloadgen.NewMappedLoader(f.fetchRisks),
Measure: dataloadgen.NewMappedLoader(f.fetchMeasures),
Task: dataloadgen.NewMappedLoader(f.fetchTasks),
File: dataloadgen.NewMappedLoader(f.fetchFiles),
Report: dataloadgen.NewMappedLoader(f.fetchReports),
}
}
func (f *batchFetcher) fetchOrganizations(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Organization, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
orgs, err := tenantSvc.Organizations.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load organizations: %w", err)
}
result := make(map[gid.GID]*coredata.Organization, len(orgs))
for _, org := range orgs {
result[org.ID] = org
}
return result, nil
}
func (f *batchFetcher) fetchFrameworks(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Framework, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
frameworks, err := tenantSvc.Frameworks.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load frameworks: %w", err)
}
result := make(map[gid.GID]*coredata.Framework, len(frameworks))
for _, v := range frameworks {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchControls(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Control, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
controls, err := tenantSvc.Controls.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load controls: %w", err)
}
result := make(map[gid.GID]*coredata.Control, len(controls))
for _, v := range controls {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchVendors(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Vendor, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
vendors, err := tenantSvc.Vendors.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load vendors: %w", err)
}
result := make(map[gid.GID]*coredata.Vendor, len(vendors))
for _, v := range vendors {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchDocuments(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Document, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
documents, err := tenantSvc.Documents.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load documents: %w", err)
}
result := make(map[gid.GID]*coredata.Document, len(documents))
for _, v := range documents {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchProfiles(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.MembershipProfile, error) {
scope := coredata.NewScopeFromObjectID(keys[0])
profiles, err := f.iam.OrganizationService.GetProfilesByIDs(ctx, scope, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load profiles: %w", err)
}
result := make(map[gid.GID]*coredata.MembershipProfile, len(profiles))
for _, v := range profiles {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchRisks(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Risk, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
risks, err := tenantSvc.Risks.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load risks: %w", err)
}
result := make(map[gid.GID]*coredata.Risk, len(risks))
for _, v := range risks {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchMeasures(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Measure, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
measures, err := tenantSvc.Measures.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load measures: %w", err)
}
result := make(map[gid.GID]*coredata.Measure, len(measures))
for _, v := range measures {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchTasks(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Task, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
tasks, err := tenantSvc.Tasks.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load tasks: %w", err)
}
result := make(map[gid.GID]*coredata.Task, len(tasks))
for _, v := range tasks {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchFiles(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.File, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
files, err := tenantSvc.Files.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load files: %w", err)
}
result := make(map[gid.GID]*coredata.File, len(files))
for _, v := range files {
result[v.ID] = v
}
return result, nil
}
func (f *batchFetcher) fetchReports(ctx context.Context, keys []gid.GID) (map[gid.GID]*coredata.Report, error) {
tenantSvc := f.probo.WithTenant(keys[0].TenantID())
reports, err := tenantSvc.Reports.GetByIDs(ctx, keys...)
if err != nil {
return nil, fmt.Errorf("cannot batch load reports: %w", err)
}
result := make(map[gid.GID]*coredata.Report, len(reports))
for _, v := range reports {
result[v.ID] = v
}
return result, nil
}

2
pkg/server/api/console/v1/resolver.go
View File

@@ -37,6 +37,7 @@ import (
"go.probo.inc/probo/pkg/securecookie"
"go.probo.inc/probo/pkg/server/api/authn"
"go.probo.inc/probo/pkg/server/api/authz"
"go.probo.inc/probo/pkg/server/api/console/v1/dataloader"
"go.probo.inc/probo/pkg/server/api/console/v1/types"
)
@@ -74,6 +75,7 @@ func NewMux(
r.Use(authn.NewSessionMiddleware(iamSvc, cookieConfig))
r.Use(authn.NewAPIKeyMiddleware(iamSvc, tokenSecret))
r.Use(authn.NewIdentityPresenceMiddleware())
r.Use(dataloader.NewMiddleware(proboSvc, iamSvc))
r.Handle("/graphql", graphqlHandler)

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

@@ -14,6 +14,7 @@ import (
"time"
pgx "github.com/jackc/pgx/v5"
"github.com/vikstrous/dataloadgen"
"go.gearno.de/kit/log"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
@@ -22,6 +23,7 @@ import (
"go.probo.inc/probo/pkg/page"
"go.probo.inc/probo/pkg/probo"
"go.probo.inc/probo/pkg/server/api/authn"
"go.probo.inc/probo/pkg/server/api/console/v1/dataloader"
"go.probo.inc/probo/pkg/server/api/console/v1/schema"
"go.probo.inc/probo/pkg/server/api/console/v1/types"
"go.probo.inc/probo/pkg/server/gqlutils"
@@ -52,10 +54,14 @@ func (r *applicabilityStatementResolver) Control(ctx context.Context, obj *types
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
control, err := prb.Controls.Get(ctx, obj.Control.ID)
control, err := loaders.Control.Load(ctx, obj.Control.ID)
if err != nil {
if errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot get control", log.Error(err))
return nil, gqlutils.Internal(ctx)
}
@@ -96,9 +102,11 @@ func (r *assetResolver) Owner(ctx context.Context, obj *types.Asset) (*types.Pro
return nil, err
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -205,11 +213,11 @@ func (r *auditResolver) Organization(ctx context.Context, obj *types.Audit) (*ty
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -226,11 +234,11 @@ func (r *auditResolver) Framework(ctx context.Context, obj *types.Audit) (*types
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
framework, err := prb.Frameworks.Get(ctx, obj.Framework.ID)
framework, err := loaders.Framework.Load(ctx, obj.Framework.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -247,14 +255,18 @@ func (r *auditResolver) Report(ctx context.Context, obj *types.Audit) (*types.Re
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
if obj.Report == nil {
return nil, nil
}
report, err := prb.Reports.Get(ctx, obj.Report.ID)
loaders := dataloader.FromContext(ctx)
report, err := loaders.Report.Load(ctx, obj.Report.ID)
if err != nil {
if errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot load report", log.Error(err))
return nil, gqlutils.Internal(ctx)
}
@@ -444,10 +456,14 @@ func (r *complianceFrameworkResolver) Framework(ctx context.Context, obj *types.
return nil, err
}
prb := r.ProboService(ctx, obj.FrameworkID.TenantID())
loaders := dataloader.FromContext(ctx)
framework, err := prb.Frameworks.Get(ctx, obj.FrameworkID)
framework, err := loaders.Framework.Load(ctx, obj.FrameworkID)
if err != nil {
if errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot load framework", log.Error(err))
return nil, gqlutils.Internal(ctx)
}
@@ -461,11 +477,11 @@ func (r *controlResolver) Organization(ctx context.Context, obj *types.Control)
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -520,11 +536,11 @@ func (r *controlResolver) Framework(ctx context.Context, obj *types.Control) (*t
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
framework, err := prb.Frameworks.Get(ctx, obj.Framework.ID)
framework, err := loaders.Framework.Load(ctx, obj.Framework.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -851,9 +867,11 @@ func (r *datumResolver) Owner(ctx context.Context, obj *types.Datum) (*types.Pro
return nil, err
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -899,11 +917,11 @@ func (r *datumResolver) Organization(ctx context.Context, obj *types.Datum) (*ty
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
org, err := prb.Organizations.Get(ctx, obj.OrganizationID)
org, err := loaders.Organization.Load(ctx, obj.OrganizationID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -987,11 +1005,11 @@ func (r *documentResolver) Organization(ctx context.Context, obj *types.Document
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1116,11 +1134,11 @@ func (r *documentVersionResolver) Document(ctx context.Context, obj *types.Docum
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
document, err := prb.Documents.Get(ctx, obj.Document.ID)
document, err := loaders.Document.Load(ctx, obj.Document.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1281,9 +1299,11 @@ func (r *documentVersionSignatureResolver) SignedBy(ctx context.Context, obj *ty
return nil, err
}
signatory, err := r.iam.OrganizationService.GetProfile(ctx, obj.SignedBy.ID)
loaders := dataloader.FromContext(ctx)
signatory, err := loaders.Profile.Load(ctx, obj.SignedBy.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1365,15 +1385,15 @@ func (r *evidenceResolver) File(ctx context.Context, obj *types.Evidence) (*type
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
if obj.File == nil {
return nil, nil
}
file, err := prb.Files.Get(ctx, obj.File.ID)
loaders := dataloader.FromContext(ctx)
file, err := loaders.File.Load(ctx, obj.File.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1390,16 +1410,16 @@ func (r *evidenceResolver) Task(ctx context.Context, obj *types.Evidence) (*type
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
if obj.Task == nil {
r.logger.ErrorCtx(ctx, "evidence is not associated with a task")
return nil, gqlutils.Internal(ctx)
}
task, err := prb.Tasks.Get(ctx, obj.Task.ID)
loaders := dataloader.FromContext(ctx)
task, err := loaders.Task.Load(ctx, obj.Task.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1416,11 +1436,11 @@ func (r *evidenceResolver) Measure(ctx context.Context, obj *types.Evidence) (*t
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
measure, err := prb.Measures.Get(ctx, obj.Measure.ID)
measure, err := loaders.Measure.Load(ctx, obj.Measure.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1488,11 +1508,11 @@ func (r *findingResolver) Organization(ctx context.Context, obj *types.Finding)
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1543,9 +1563,11 @@ func (r *findingResolver) Owner(ctx context.Context, obj *types.Finding) (*types
return nil, err
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1566,11 +1588,11 @@ func (r *findingResolver) Risk(ctx context.Context, obj *types.Finding) (*types.
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
risk, err := prb.Risks.Get(ctx, obj.Risk.ID)
risk, err := loaders.Risk.Load(ctx, obj.Risk.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -1639,11 +1661,11 @@ func (r *frameworkResolver) Organization(ctx context.Context, obj *types.Framewo
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -2022,11 +2044,11 @@ func (r *meetingResolver) Organization(ctx context.Context, obj *types.Meeting)
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -6299,11 +6321,11 @@ func (r *obligationResolver) Organization(ctx context.Context, obj *types.Obliga
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -6320,9 +6342,11 @@ func (r *obligationResolver) Owner(ctx context.Context, obj *types.Obligation) (
return nil, err
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -7328,11 +7352,11 @@ func (r *processingActivityResolver) Organization(ctx context.Context, obj *type
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -7353,8 +7377,14 @@ func (r *processingActivityResolver) DataProtectionOfficer(ctx context.Context,
return nil, nil
}
dpo, err := r.iam.OrganizationService.GetProfile(ctx, obj.DataProtectionOfficer.ID)
loaders := dataloader.FromContext(ctx)
dpo, err := loaders.Profile.Load(ctx, obj.DataProtectionOfficer.ID)
if err != nil {
if errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot get data protection officer", log.Error(err))
return nil, gqlutils.Internal(ctx)
}
@@ -7937,9 +7967,11 @@ func (r *riskResolver) Owner(ctx context.Context, obj *types.Risk) (*types.Profi
return nil, nil
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -7956,11 +7988,11 @@ func (r *riskResolver) Organization(ctx context.Context, obj *types.Risk) (*type
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -8293,11 +8325,11 @@ func (r *stateOfApplicabilityResolver) Organization(ctx context.Context, obj *ty
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot load organization", log.Error(err))
@@ -8313,9 +8345,11 @@ func (r *stateOfApplicabilityResolver) Owner(ctx context.Context, obj *types.Sta
return nil, err
}
owner, err := r.iam.OrganizationService.GetProfile(ctx, obj.Owner.ID)
loaders := dataloader.FromContext(ctx)
owner, err := loaders.Profile.Load(ctx, obj.Owner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
r.logger.ErrorCtx(ctx, "cannot load owner", log.Error(err))
@@ -8388,9 +8422,11 @@ func (r *taskResolver) AssignedTo(ctx context.Context, obj *types.Task) (*types.
return nil, nil
}
assignee, err := r.iam.OrganizationService.GetProfile(ctx, obj.AssignedTo.ID)
loaders := dataloader.FromContext(ctx)
assignee, err := loaders.Profile.Load(ctx, obj.AssignedTo.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -8407,11 +8443,11 @@ func (r *taskResolver) Organization(ctx context.Context, obj *types.Task) (*type
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -8428,15 +8464,15 @@ func (r *taskResolver) Measure(ctx context.Context, obj *types.Task) (*types.Mea
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
if obj.Measure == nil {
return nil, nil
}
measure, err := prb.Measures.Get(ctx, obj.Measure.ID)
loaders := dataloader.FromContext(ctx)
measure, err := loaders.Measure.Load(ctx, obj.Measure.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -8533,11 +8569,11 @@ func (r *transferImpactAssessmentResolver) Organization(ctx context.Context, obj
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -9115,11 +9151,11 @@ func (r *vendorResolver) Organization(ctx context.Context, obj *types.Vendor) (*
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -9302,9 +9338,11 @@ func (r *vendorResolver) BusinessOwner(ctx context.Context, obj *types.Vendor) (
return nil, nil
}
businessOwner, err := r.iam.OrganizationService.GetProfile(ctx, obj.BusinessOwner.ID)
loaders := dataloader.FromContext(ctx)
businessOwner, err := loaders.Profile.Load(ctx, obj.BusinessOwner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -9325,9 +9363,11 @@ func (r *vendorResolver) SecurityOwner(ctx context.Context, obj *types.Vendor) (
return nil, nil
}
securityOwner, err := r.iam.OrganizationService.GetProfile(ctx, obj.SecurityOwner.ID)
loaders := dataloader.FromContext(ctx)
securityOwner, err := loaders.Profile.Load(ctx, obj.SecurityOwner.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -9586,11 +9626,11 @@ func (r *vendorServiceResolver) Vendor(ctx context.Context, obj *types.VendorSer
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
vendor, err := prb.Vendors.Get(ctx, obj.Vendor.ID)
vendor, err := loaders.Vendor.Load(ctx, obj.Vendor.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}
@@ -9708,11 +9748,11 @@ func (r *webhookSubscriptionResolver) Organization(ctx context.Context, obj *typ
return nil, err
}
prb := r.ProboService(ctx, obj.ID.TenantID())
loaders := dataloader.FromContext(ctx)
organization, err := prb.Organizations.Get(ctx, obj.Organization.ID)
organization, err := loaders.Organization.Load(ctx, obj.Organization.ID)
if err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
if errors.Is(err, coredata.ErrResourceNotFound) || errors.Is(err, dataloadgen.ErrNotFound) {
return nil, gqlutils.NotFound(ctx, err)
}