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>
2026-03-25 15:04:51 +01:00
parent 2004bf6050
commit d24712344e
24 changed files with 1091 additions and 103 deletions
--- 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,
--- 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,
--- 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,
--- 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,
--- 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,
--- 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,
--- 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,
--- 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,