batch

Author: n1ru4l

packages/services/api/src/modules/organization/providers/organization-members.ts

@@ -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),
@@ -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) {

packages/services/api/src/modules/shared/providers/storage.ts

@@ -205,7 +205,7 @@ export interface Storage {
 
   getProjects(_: OrganizationSelector): Promise<Project[] | never>;
 
-  findProjectsByIds(projectIds: Array<string>): Promise<Map<string, Project>>;
+  findProjectsByIds(args: { projectIds: Array<string> }): Promise<Map<string, Project>>;
 
   createProject(_: Pick<Project, 'type'> & { slug: string } & OrganizationSelector): Promise<
     | {
@@ -304,7 +304,10 @@ export interface Storage {
 
   getTargets(_: ProjectSelector): Promise<readonly Target[]>;
 
-  findTargetsByIds(organizationId: string, targetIds: Array<string>): Promise<Map<string, Target>>;
+  findTargetsByIds(args: {
+    organizationId: string;
+    targetIds: Array<string>;
+  }): Promise<Map<string, Target>>;
 
   getTargetIdsOfOrganization(_: OrganizationSelector): Promise<readonly string[]>;
   getTargetIdsOfProject(_: ProjectSelector): Promise<readonly string[]>;

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,30 @@ export async function createStorage(
 
       return result.rows.map(transformProject);
     },
-    async findProjectsByIds(ids) {
-      const result = await pool.query<Slonik<projects>>(
-        sql`/* findProjectsByIds */ SELECT * FROM projects WHERE id = ANY(${sql.array(ids, 'uuid')}) AND type != 'CUSTOM'`,
-      );
+    findProjectsByIds: batch<{ projectIds: Array<string> }, Map<string, Project>>(
+      async function FindProjectByIdsBatchHandler(args) {
+        const allProjectIds = args.flatMap(args => args.projectIds);
+        const result = await pool.query<Slonik<projects>>(
+          sql`/* findProjectsByIds */ SELECT * FROM projects WHERE id = ANY(${sql.array(allProjectIds, 'uuid')}) AND type != 'CUSTOM'`,
+        );
 
-      const map = new Map<string, Project>();
-      result.rows.forEach(row => {
-        const project = transformProject(row);
-        map.set(project.id, project);
-      });
-      return map;
-    },
+        const allProjectsLookupMap = new Map<string, Project>();
+        result.rows.forEach(row => {
+          const project = transformProject(row);
+          allProjectsLookupMap.set(project.id, project);
+        });
+        return args.map(async arg => {
+          const map = new Map<string, Project>();
+          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 +1713,50 @@ export async function createStorage(
         orgId: organization,
       }));
     },
-    async findTargetsByIds(organizationId, targetIds) {
-      const map = new Map<string, Target>();
+    findTargetsByIds: batchBy<
+      {
+        organizationId: string;
+        targetIds: Array<string>;
+      },
+      Map<string, Target>
+    >(
+      org => org.organizationId,
+      async function FindTargetsByIdsBatchHandler(args) {
+        const resultLookupMap = new Map<string, Target>();
 
-      if (targetIds.length === 0) {
-        return map;
-      }
+        const allTargetIds = args.flatMap(arg => arg.targetIds);
 
-      const results = await pool.query<unknown>(sql`/* getTargets */
-        SELECT
+        if (allTargetIds.length === 0) {
+          return args.map(async () => resultLookupMap);
+        }
+
+        const orgId = args[0].organizationId;
+
+        const results = await pool.query<unknown>(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<string, Target>();
+          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<Slonik<Pick<targets, 'id'>>>(
         sql`/* getTargetIdsOfOrganization */