diff --git a/go.mod b/go.mod index d37829459..b5288661a 100644 --- a/go.mod +++ b/go.mod @@ -31,6 +31,7 @@ require ( github.com/scim2/filter-parser/v2 v2.2.0 github.com/stretchr/testify v1.11.1 github.com/vektah/gqlparser/v2 v2.5.32 + github.com/vikstrous/dataloadgen v0.0.10 github.com/yuin/goldmark v1.7.16 go.abhg.dev/goldmark/mermaid v0.6.0 go.gearno.de/crypto/uuid v0.1.1-0.20251208105319-3f587312a712 diff --git a/go.sum b/go.sum index 28fefd5c5..1e24ef6ae 100644 --- a/go.sum +++ b/go.sum @@ -316,6 +316,8 @@ github.com/urfave/cli/v3 v3.6.2 h1:lQuqiPrZ1cIz8hz+HcrG0TNZFxU70dPZ3Yl+pSrH9A8= github.com/urfave/cli/v3 v3.6.2/go.mod h1:ysVLtOEmg2tOy6PknnYVhDoouyC/6N42TMeoMzskhso= github.com/vektah/gqlparser/v2 v2.5.32 h1:k9QPJd4sEDTL+qB4ncPLflqTJ3MmjB9SrVzJrawpFSc= github.com/vektah/gqlparser/v2 v2.5.32/go.mod h1:c1I28gSOVNzlfc4WuDlqU7voQnsqI6OG2amkBAFmgts= +github.com/vikstrous/dataloadgen v0.0.10 h1:x07XAeEjIWXohvcjRvE72KY8pV5A3sTbKEFmxcj9RNM= +github.com/vikstrous/dataloadgen v0.0.10/go.mod h1:8vuQVpBH0ODbMKAPUdCAPcOGezoTIhgAjgex51t4vbg= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= diff --git a/pkg/coredata/control.go b/pkg/coredata/control.go index 82d305cb2..f0536e830 100644 --- a/pkg/coredata/control.go +++ b/pkg/coredata/control.go @@ -720,6 +720,51 @@ LIMIT 1; return nil } +func (c *Controls) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + controlIDs []gid.GID, +) error { + q := ` +SELECT + id, + section_title, + framework_id, + organization_id, + name, + description, + best_practice, + implemented, + not_implemented_justification, + created_at, + updated_at +FROM + controls +WHERE + %s + AND id = ANY(@control_ids) +` + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"control_ids": controlIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query controls: %w", err) + } + + controls, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Control]) + if err != nil { + return fmt.Errorf("cannot collect controls: %w", err) + } + + *c = controls + + return nil +} + func (c Control) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/file.go b/pkg/coredata/file.go index ad2c181d5..3f90d59f2 100644 --- a/pkg/coredata/file.go +++ b/pkg/coredata/file.go @@ -131,6 +131,52 @@ LIMIT 1; return nil } +func (f *Files) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + fileIDs []gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + bucket_name, + mime_type, + file_name, + file_key, + file_size, + visibility, + created_at, + updated_at, + deleted_at +FROM + files +WHERE + %s + AND id = ANY(@file_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"file_ids": fileIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query files: %w", err) + } + + files, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[File]) + if err != nil { + return fmt.Errorf("cannot collect files: %w", err) + } + + *f = files + + return nil +} + func (f File) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/framework.go b/pkg/coredata/framework.go index 558b72483..7d7dd4d5f 100644 --- a/pkg/coredata/framework.go +++ b/pkg/coredata/framework.go @@ -241,6 +241,50 @@ LIMIT 1; return nil } +func (f *Frameworks) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + frameworkIDs []gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + reference_id, + name, + description, + light_logo_file_id, + dark_logo_file_id, + created_at, + updated_at +FROM + frameworks +WHERE + %s + AND id = ANY(@framework_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"framework_ids": frameworkIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query frameworks: %w", err) + } + + frameworks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Framework]) + if err != nil { + return fmt.Errorf("cannot collect frameworks: %w", err) + } + + *f = frameworks + + return nil +} + func (f Framework) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/measure.go b/pkg/coredata/measure.go index c641d3f1c..8560bc9d0 100644 --- a/pkg/coredata/measure.go +++ b/pkg/coredata/measure.go @@ -461,6 +461,50 @@ LIMIT 1; return nil } +func (m *Measures) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + measureIDs []gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + category, + name, + description, + state, + reference_id, + created_at, + updated_at +FROM + measures +WHERE + %s + AND id = ANY(@measure_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"measure_ids": measureIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query measures: %w", err) + } + + measures, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Measure]) + if err != nil { + return fmt.Errorf("cannot collect measures: %w", err) + } + + *m = measures + + return nil +} + func (m *Measure) Upsert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/organization.go b/pkg/coredata/organization.go index 3a7e59996..0b053daf7 100644 --- a/pkg/coredata/organization.go +++ b/pkg/coredata/organization.go @@ -125,6 +125,53 @@ LIMIT 1; return nil } +func (o *Organizations) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + organizationIDs []gid.GID, +) error { + q := ` +SELECT + tenant_id, + id, + name, + logo_file_id, + horizontal_logo_file_id, + description, + website_url, + email, + headquarter_address, + custom_domain_id, + created_at, + updated_at +FROM + organizations +WHERE + %s + AND id = ANY(@organization_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"organization_ids": organizationIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query organizations: %w", err) + } + + organizations, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Organization]) + if err != nil { + return fmt.Errorf("cannot collect organizations: %w", err) + } + + *o = organizations + + return nil +} + func (o *Organizations) LoadByIdentityID( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/risk.go b/pkg/coredata/risk.go index 09846612d..494c76e50 100644 --- a/pkg/coredata/risk.go +++ b/pkg/coredata/risk.go @@ -401,6 +401,57 @@ LIMIT 1; return nil } +func (r *Risks) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + riskIDs []gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + name, + description, + category, + owner_profile_id, + NULL as owner_full_name, + treatment, + note, + inherent_likelihood, + inherent_impact, + inherent_risk_score, + residual_likelihood, + residual_impact, + residual_risk_score, + snapshot_id, + source_id, + created_at, + updated_at +FROM risks +WHERE %s + AND id = ANY(@risk_ids) +` + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"risk_ids": riskIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query risks: %w", err) + } + + risks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Risk]) + if err != nil { + return fmt.Errorf("cannot collect risks: %w", err) + } + + *r = risks + + return nil +} + func (r *Risk) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/task.go b/pkg/coredata/task.go index 04b893859..8024fed5a 100644 --- a/pkg/coredata/task.go +++ b/pkg/coredata/task.go @@ -123,6 +123,53 @@ LIMIT 1; return nil } +func (t *Tasks) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + taskIDs []gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + measure_id, + name, + description, + state, + reference_id, + time_estimate, + assigned_to_profile_id, + deadline, + created_at, + updated_at +FROM + tasks +WHERE + %s + AND id = ANY(@task_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"task_ids": taskIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query tasks: %w", err) + } + + tasks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Task]) + if err != nil { + return fmt.Errorf("cannot collect tasks: %w", err) + } + + *t = tasks + + return nil +} + func (c Task) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/coredata/vendor.go b/pkg/coredata/vendor.go index 436e514bc..1918bb8b8 100644 --- a/pkg/coredata/vendor.go +++ b/pkg/coredata/vendor.go @@ -160,6 +160,68 @@ LIMIT 1; return nil } +func (v *Vendors) LoadByIDs( + ctx context.Context, + conn pg.Conn, + scope Scoper, + vendorIDs []gid.GID, +) error { + q := ` +SELECT + id, + tenant_id, + organization_id, + name, + description, + category, + headquarter_address, + legal_name, + website_url, + privacy_policy_url, + service_level_agreement_url, + data_processing_agreement_url, + business_associate_agreement_url, + subprocessors_list_url, + certifications, + countries, + business_owner_profile_id, + security_owner_profile_id, + status_page_url, + terms_of_service_url, + security_page_url, + trust_page_url, + show_on_trust_center, + snapshot_id, + source_id, + created_at, + updated_at +FROM + vendors +WHERE + %s + AND id = ANY(@vendor_ids) +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"vendor_ids": vendorIDs} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query vendors: %w", err) + } + + vendors, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Vendor]) + if err != nil { + return fmt.Errorf("cannot collect vendors: %w", err) + } + + *v = vendors + + return nil +} + func (v Vendor) Insert( ctx context.Context, conn pg.Conn, diff --git a/pkg/iam/organization_service.go b/pkg/iam/organization_service.go index 1c9fad9d6..5abe903f7 100644 --- a/pkg/iam/organization_service.go +++ b/pkg/iam/organization_service.go @@ -1073,6 +1073,35 @@ func (s *OrganizationService) GetProfile(ctx context.Context, profileID gid.GID) return profile, nil } +func (s *OrganizationService) GetProfilesByIDs( + ctx context.Context, + scope coredata.Scoper, + profileIDs ...gid.GID, +) (coredata.MembershipProfiles, error) { + var profiles coredata.MembershipProfiles + + err := s.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := profiles.LoadByIDs( + ctx, + conn, + scope, + profileIDs, + ); err != nil { + return fmt.Errorf("cannot load profiles by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return profiles, nil +} + func (s *OrganizationService) GetProfileForIdentityAndOrganization(ctx context.Context, identityID gid.GID, organizationID gid.GID) (*coredata.MembershipProfile, error) { profile := &coredata.MembershipProfile{} diff --git a/pkg/probo/control_service.go b/pkg/probo/control_service.go index 3085913ae..e6e7739a7 100644 --- a/pkg/probo/control_service.go +++ b/pkg/probo/control_service.go @@ -913,6 +913,34 @@ func (s ControlService) Get( return control, nil } +func (s ControlService) GetByIDs( + ctx context.Context, + controlIDs ...gid.GID, +) (coredata.Controls, error) { + var controls coredata.Controls + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := controls.LoadByIDs( + ctx, + conn, + s.svc.scope, + controlIDs, + ); err != nil { + return fmt.Errorf("cannot load controls by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return controls, nil +} + func (s ControlService) Update( ctx context.Context, req UpdateControlRequest, diff --git a/pkg/probo/document_service.go b/pkg/probo/document_service.go index 10df4658a..3338873cc 100644 --- a/pkg/probo/document_service.go +++ b/pkg/probo/document_service.go @@ -202,6 +202,34 @@ func (s *DocumentService) Get( return document, nil } +func (s *DocumentService) GetByIDs( + ctx context.Context, + documentIDs ...gid.GID, +) (coredata.Documents, error) { + var documents coredata.Documents + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := documents.LoadByIDs( + ctx, + conn, + s.svc.scope, + documentIDs, + ); err != nil { + return fmt.Errorf("cannot load documents by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return documents, nil +} + func (s *DocumentService) ListApprovers( ctx context.Context, documentID gid.GID, diff --git a/pkg/probo/file_service.go b/pkg/probo/file_service.go index bbf74fbc4..fbc48c48b 100644 --- a/pkg/probo/file_service.go +++ b/pkg/probo/file_service.go @@ -73,6 +73,34 @@ func (s FileService) Get( return file, nil } +func (s FileService) GetByIDs( + ctx context.Context, + fileIDs ...gid.GID, +) (coredata.Files, error) { + var files coredata.Files + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := files.LoadByIDs( + ctx, + conn, + s.svc.scope, + fileIDs, + ); err != nil { + return fmt.Errorf("cannot load files by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return files, nil +} + func (s FileService) UploadAndSaveFile( ctx context.Context, fileValidator *filevalidation.FileValidator, diff --git a/pkg/probo/framework_service.go b/pkg/probo/framework_service.go index 599c03124..51d81c4b2 100644 --- a/pkg/probo/framework_service.go +++ b/pkg/probo/framework_service.go @@ -447,6 +447,34 @@ func (s FrameworkService) Get( return framework, nil } +func (s FrameworkService) GetByIDs( + ctx context.Context, + frameworkIDs ...gid.GID, +) (coredata.Frameworks, error) { + var frameworks coredata.Frameworks + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := frameworks.LoadByIDs( + ctx, + conn, + s.svc.scope, + frameworkIDs, + ); err != nil { + return fmt.Errorf("cannot load frameworks by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return frameworks, nil +} + func (s FrameworkService) Update( ctx context.Context, req UpdateFrameworkRequest, diff --git a/pkg/probo/measure_service.go b/pkg/probo/measure_service.go index 18678526a..86fd47f7f 100644 --- a/pkg/probo/measure_service.go +++ b/pkg/probo/measure_service.go @@ -333,6 +333,34 @@ func (s MeasureService) Get( return measure, nil } +func (s MeasureService) GetByIDs( + ctx context.Context, + measureIDs ...gid.GID, +) (coredata.Measures, error) { + var measures coredata.Measures + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := measures.LoadByIDs( + ctx, + conn, + s.svc.scope, + measureIDs, + ); err != nil { + return fmt.Errorf("cannot load measures by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return measures, nil +} + func (s MeasureService) Import( ctx context.Context, organizationID gid.GID, diff --git a/pkg/probo/organization_service.go b/pkg/probo/organization_service.go index 87b1bcb02..213b13a3d 100644 --- a/pkg/probo/organization_service.go +++ b/pkg/probo/organization_service.go @@ -110,6 +110,34 @@ func (s OrganizationService) Get( return organization, nil } +func (s OrganizationService) GetByIDs( + ctx context.Context, + organizationIDs ...gid.GID, +) (coredata.Organizations, error) { + var organizations coredata.Organizations + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := organizations.LoadByIDs( + ctx, + conn, + s.svc.scope, + organizationIDs, + ); err != nil { + return fmt.Errorf("cannot load organizations by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return organizations, nil +} + func (s OrganizationService) GetContext( ctx context.Context, organizationID gid.GID, diff --git a/pkg/probo/report_service.go b/pkg/probo/report_service.go index da7807187..09c87761c 100644 --- a/pkg/probo/report_service.go +++ b/pkg/probo/report_service.go @@ -54,6 +54,34 @@ func (s ReportService) Get( return report, nil } +func (s ReportService) GetByIDs( + ctx context.Context, + reportIDs ...gid.GID, +) (coredata.Reports, error) { + var reports coredata.Reports + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := reports.LoadByIDs( + ctx, + conn, + s.svc.scope, + reportIDs, + ); err != nil { + return fmt.Errorf("cannot load reports by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return reports, nil +} + func (s ReportService) Delete( ctx context.Context, reportID gid.GID, diff --git a/pkg/probo/risk_service.go b/pkg/probo/risk_service.go index 88670a7bc..7f8f1cdad 100644 --- a/pkg/probo/risk_service.go +++ b/pkg/probo/risk_service.go @@ -498,6 +498,34 @@ func (s RiskService) Get( return risk, nil } +func (s RiskService) GetByIDs( + ctx context.Context, + riskIDs ...gid.GID, +) (coredata.Risks, error) { + var risks coredata.Risks + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := risks.LoadByIDs( + ctx, + conn, + s.svc.scope, + riskIDs, + ); err != nil { + return fmt.Errorf("cannot load risks by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return risks, nil +} + func (s RiskService) Update( ctx context.Context, req UpdateRiskRequest, diff --git a/pkg/probo/task_service.go b/pkg/probo/task_service.go index b2709f363..7ad450526 100644 --- a/pkg/probo/task_service.go +++ b/pkg/probo/task_service.go @@ -162,6 +162,34 @@ func (s TaskService) Get( return task, nil } +func (s TaskService) GetByIDs( + ctx context.Context, + taskIDs ...gid.GID, +) (coredata.Tasks, error) { + var tasks coredata.Tasks + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := tasks.LoadByIDs( + ctx, + conn, + s.svc.scope, + taskIDs, + ); err != nil { + return fmt.Errorf("cannot load tasks by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return tasks, nil +} + func (s TaskService) Assign( ctx context.Context, taskID gid.GID, diff --git a/pkg/probo/vendor_service.go b/pkg/probo/vendor_service.go index b71808c2e..5c42b95f9 100644 --- a/pkg/probo/vendor_service.go +++ b/pkg/probo/vendor_service.go @@ -429,6 +429,34 @@ func (s VendorService) Get( return vendor, nil } +func (s VendorService) GetByIDs( + ctx context.Context, + vendorIDs ...gid.GID, +) (coredata.Vendors, error) { + var vendors coredata.Vendors + + err := s.svc.pg.WithConn( + ctx, + func(conn pg.Conn) error { + if err := vendors.LoadByIDs( + ctx, + conn, + s.svc.scope, + vendorIDs, + ); err != nil { + return fmt.Errorf("cannot load vendors by ids: %w", err) + } + + return nil + }, + ) + if err != nil { + return nil, err + } + + return vendors, nil +} + func (s VendorService) Delete( ctx context.Context, vendorID gid.GID, diff --git a/pkg/server/api/console/v1/dataloader/dataloader.go b/pkg/server/api/console/v1/dataloader/dataloader.go new file mode 100644 index 000000000..0ea0dca48 --- /dev/null +++ b/pkg/server/api/console/v1/dataloader/dataloader.go @@ -0,0 +1,248 @@ +// Copyright (c) 2025 Probo Inc . +// +// 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 +} diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go index bf27f783c..40dc93250 100644 --- a/pkg/server/api/console/v1/resolver.go +++ b/pkg/server/api/console/v1/resolver.go @@ -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) diff --git a/pkg/server/api/console/v1/v1_resolver.go b/pkg/server/api/console/v1/v1_resolver.go index ae9c47735..cd0fa0b6e 100644 --- a/pkg/server/api/console/v1/v1_resolver.go +++ b/pkg/server/api/console/v1/v1_resolver.go @@ -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) }