Skip to content

Commit

Permalink
batch
Browse files Browse the repository at this point in the history
  • Loading branch information
n1ru4l committed Jan 24, 2025
1 parent c0b78c5 commit cb29617
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 45 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -327,20 +327,20 @@ export class OrganizationMembers {
async resolveGraphQLMemberResourceAssignment(
member: OrganizationMembership,
): Promise<GraphQLSchema.ResolversTypes['ResourceAssignment']> {
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),
);

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),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand Down
7 changes: 5 additions & 2 deletions packages/services/api/src/modules/shared/providers/storage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<
| {
Expand Down Expand Up @@ -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[]>;
Expand Down
97 changes: 68 additions & 29 deletions packages/services/storage/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1398,18 +1398,36 @@ 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 allProjectsLookupMap = new Map<string, Project>();

const map = new Map<string, Project>();
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<Slonik<projects>>(
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<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(
Expand Down Expand Up @@ -1701,29 +1719,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 */
Expand Down

0 comments on commit cb29617

Please sign in to comment.