Add GraphQL dataloaders for batched record lookups

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>

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

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=

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,

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,

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,

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,

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,

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,

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,

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,

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

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,

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,

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,

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,

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,

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,

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,

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,

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,

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,

pkg/server/api/console/v1/dataloader/dataloader.go

@@ -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
+}

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)
 

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