From cb296177a402afc9ef6bae263f1f393c97012a77 Mon Sep 17 00:00:00 2001 From: Laurin Quast Date: Fri, 24 Jan 2025 16:12:34 +0100 Subject: [PATCH] batch --- .../providers/organization-members.ts | 28 +++--- .../src/modules/shared/providers/storage.ts | 7 +- packages/services/storage/src/index.ts | 97 +++++++++++++------ 3 files changed, 87 insertions(+), 45 deletions(-) diff --git a/packages/services/api/src/modules/organization/providers/organization-members.ts b/packages/services/api/src/modules/organization/providers/organization-members.ts index 76a2a51d59..bdb11bc1b1 100644 --- a/packages/services/api/src/modules/organization/providers/organization-members.ts +++ b/packages/services/api/src/modules/organization/providers/organization-members.ts @@ -327,9 +327,9 @@ export class OrganizationMembers { async resolveGraphQLMemberResourceAssignment( member: OrganizationMembership, ): Promise { - const projects = await this.storage.findProjectsByIds( - member.assignedRole.resources.projects.map(project => project.id), - ); + const projects = await this.storage.findProjectsByIds({ + projectIds: member.assignedRole.resources.projects.map(project => project.id), + }); const filteredProjects = member.assignedRole.resources.projects.filter(row => projects.get(row.id), @@ -337,10 +337,10 @@ export class OrganizationMembers { const targetAssignments = filteredProjects.flatMap(project => project.targets.targets); - const targets = await this.storage.findTargetsByIds( - member.organizationId, - targetAssignments.map(target => target.id), - ); + const targets = await this.storage.findTargetsByIds({ + organizationId: member.organizationId, + targetIds: targetAssignments.map(target => target.id), + }); return { mode: member.assignedRole.resources.mode === '*' ? ('all' as const) : ('granular' as const), @@ -418,9 +418,9 @@ export class OrganizationMembers { const sanitizedProjects = input.projects.filter(project => isUUID(project.projectId)); - const projects = await this.storage.findProjectsByIds( - sanitizedProjects.map(record => record.projectId), - ); + const projects = await this.storage.findProjectsByIds({ + projectIds: sanitizedProjects.map(record => record.projectId), + }); // In case we are not assigning all targets to the project, // we need to load all the targets/projects that would be assigned @@ -468,10 +468,10 @@ export class OrganizationMembers { } } - const targets = await this.storage.findTargetsByIds( - organization.id, - Array.from(targetLookupIds), - ); + const targets = await this.storage.findTargetsByIds({ + organizationId: organization.id, + targetIds: Array.from(targetLookupIds), + }); for (const record of projectTargetAssignments) { for (const targetRecord of record.targets) { diff --git a/packages/services/api/src/modules/shared/providers/storage.ts b/packages/services/api/src/modules/shared/providers/storage.ts index f2e8f73fa0..3701c38227 100644 --- a/packages/services/api/src/modules/shared/providers/storage.ts +++ b/packages/services/api/src/modules/shared/providers/storage.ts @@ -205,7 +205,7 @@ export interface Storage { getProjects(_: OrganizationSelector): Promise; - findProjectsByIds(projectIds: Array): Promise>; + findProjectsByIds(args: { projectIds: Array }): Promise>; createProject(_: Pick & { slug: string } & OrganizationSelector): Promise< | { @@ -304,7 +304,10 @@ export interface Storage { getTargets(_: ProjectSelector): Promise; - findTargetsByIds(organizationId: string, targetIds: Array): Promise>; + findTargetsByIds(args: { + organizationId: string; + targetIds: Array; + }): Promise>; getTargetIdsOfOrganization(_: OrganizationSelector): Promise; getTargetIdsOfProject(_: ProjectSelector): Promise; diff --git a/packages/services/storage/src/index.ts b/packages/services/storage/src/index.ts index 0fe9417ac0..0fa4d00c07 100644 --- a/packages/services/storage/src/index.ts +++ b/packages/services/storage/src/index.ts @@ -32,7 +32,7 @@ import { type SchemaLog, type SchemaPolicy, } from '../../api/src/shared/entities'; -import { batch } from '../../api/src/shared/helpers'; +import { batch, batchBy } from '../../api/src/shared/helpers'; import { alert_channels, alerts, @@ -1398,18 +1398,36 @@ export async function createStorage( return result.rows.map(transformProject); }, - async findProjectsByIds(ids) { - const result = await pool.query>( - sql`/* findProjectsByIds */ SELECT * FROM projects WHERE id = ANY(${sql.array(ids, 'uuid')}) AND type != 'CUSTOM'`, - ); + findProjectsByIds: batch<{ projectIds: Array }, Map>( + async function FindProjectByIdsBatchHandler(args) { + const allProjectIds = args.flatMap(args => args.projectIds); + const allProjectsLookupMap = new Map(); - const map = new Map(); - result.rows.forEach(row => { - const project = transformProject(row); - map.set(project.id, project); - }); - return map; - }, + if (allProjectIds.length === 0) { + return args.map(async () => allProjectsLookupMap); + } + + const result = await pool.query>( + sql`/* findProjectsByIds */ SELECT * FROM projects WHERE id = ANY(${sql.array(allProjectIds, 'uuid')}) AND type != 'CUSTOM'`, + ); + + result.rows.forEach(row => { + const project = transformProject(row); + allProjectsLookupMap.set(project.id, project); + }); + + return args.map(async arg => { + const map = new Map(); + for (const projectId of arg.projectIds) { + const project = allProjectsLookupMap.get(projectId); + if (!project) continue; + map.set(projectId, project); + } + + return map; + }); + }, + ), async updateProjectSlug({ slug, organizationId: organization, projectId: project }) { return pool.transaction(async t => { const projectSlugExists = await t.exists( @@ -1701,29 +1719,50 @@ export async function createStorage( orgId: organization, })); }, - async findTargetsByIds(organizationId, targetIds) { - const map = new Map(); + findTargetsByIds: batchBy< + { + organizationId: string; + targetIds: Array; + }, + Map + >( + org => org.organizationId, + async function FindTargetsByIdsBatchHandler(args) { + const resultLookupMap = new Map(); - if (targetIds.length === 0) { - return map; - } + const allTargetIds = args.flatMap(arg => arg.targetIds); - const results = await pool.query(sql`/* getTargets */ - SELECT + if (allTargetIds.length === 0) { + return args.map(async () => resultLookupMap); + } + + const orgId = args[0].organizationId; + + const results = await pool.query(sql`/* getTargets */ + SELECT ${targetSQLFields} - FROM + FROM "targets" - WHERE - "id" = ANY(${sql.array(targetIds, 'uuid')}) - `); + WHERE + "id" = ANY(${sql.array(allTargetIds, 'uuid')}) + `); - for (const row of results.rows) { - const target: Target = { ...TargetModel.parse(row), orgId: organizationId }; - map.set(target.id, target); - } + for (const row of results.rows) { + const target: Target = { ...TargetModel.parse(row), orgId }; + resultLookupMap.set(target.id, target); + } - return map; - }, + return args.map(async arg => { + const map = new Map(); + for (const targetId of arg.targetIds) { + const target = resultLookupMap.get(targetId); + if (!target) continue; + map.set(targetId, target); + } + return map; + }); + }, + ), async getTargetIdsOfOrganization({ organizationId: organization }) { const results = await pool.query>>( sql`/* getTargetIdsOfOrganization */