Move Resource.tools to loader, so multiple resources can load their tools in one query

Author: CarsonF

src/components/tools/tool-usage/resource-tools.resolver.ts

@@ -1,16 +1,20 @@
 import { Parent, ResolveField, Resolver } from '@nestjs/graphql';
+import { Loader, type LoaderOf } from '@seedcompany/data-loader';
 import { Resource } from '~/common';
 import { ToolUsage } from './dto';
-import { ToolUsageService } from './tool-usage.service';
+import { ToolUsageByContainerLoader } from './tool-usage-by-container.loader';
 
 @Resolver(() => Resource)
 export class ResourceToolsResolver {
-  constructor(private readonly toolUsageService: ToolUsageService) {}
-
   @ResolveField(() => [ToolUsage], {
     description: 'Tools used in this resource',
   })
-  async tools(@Parent() resource: Resource): Promise<readonly ToolUsage[]> {
-    return await this.toolUsageService.readByContainer(resource);
+  async tools(
+    @Parent() resource: Resource,
+    @Loader(() => ToolUsageByContainerLoader)
+    loader: LoaderOf<ToolUsageByContainerLoader>,
+  ): Promise<readonly ToolUsage[]> {
+    const { usages } = await loader.load(resource);
+    return usages;
   }
 }

src/components/tools/tool-usage/tool-usage-by-container.loader.ts

@@ -0,0 +1,33 @@
+import { type ID, type Resource } from '~/common';
+import {
+  type DataLoaderStrategy,
+  LoaderFactory,
+  type LoaderOptionsOf,
+} from '~/core/data-loader';
+import { type ToolUsage } from './dto';
+import { ToolUsageService } from './tool-usage.service';
+
+export interface UsagesByContainer {
+  container: Resource;
+  usages: readonly ToolUsage[];
+}
+
+@LoaderFactory()
+export class ToolUsageByContainerLoader
+  implements DataLoaderStrategy<UsagesByContainer, Resource, ID>
+{
+  constructor(private readonly usages: ToolUsageService) {}
+
+  getOptions() {
+    return {
+      propertyKey: ({ container }) => container,
+      cacheKeyFn: (container) => container.id,
+    } satisfies LoaderOptionsOf<ToolUsageByContainerLoader>;
+  }
+
+  async loadMany(
+    ids: readonly Resource[],
+  ): Promise<readonly UsagesByContainer[]> {
+    return await this.usages.readManyForContainers(ids);
+  }
+}

src/components/tools/tool-usage/tool-usage.gel.repository.ts

@@ -25,14 +25,17 @@ export class ToolUsageRepository
   })
   implements PublicOf<Neo4jRepository>
 {
-  async listForContainer(container: ID) {
-    return await this.db.run(this.listForContainerQuery, { container });
+  async listForContainers(containers: readonly ID[]) {
+    return await this.db.run(this.listForContainersQuery, { containers });
   }
-  private readonly listForContainerQuery = e.params(
-    { container: e.uuid },
+  private readonly listForContainersQuery = e.params(
+    { containers: e.array(e.uuid) },
     ($) => {
-      const container = e.cast(e.Resource, $.container);
-      return e.select(container.tools, this.hydrate);
+      const containers = e.cast(e.Resource, e.array_unpack($.containers));
+      return e.select(containers, (container) => ({
+        container: e.select(container, (c) => ({ id: c.id })),
+        usages: e.select(container.tools, this.hydrate),
+      }));
     },
   );
 }

src/components/tools/tool-usage/tool-usage.module.ts

@@ -2,6 +2,7 @@ import { Module } from '@nestjs/common';
 import { splitDb } from '~/core/database';
 import { ToolCoreModule } from '../tool/tool.module';
 import { ResourceToolsResolver } from './resource-tools.resolver';
+import { ToolUsageByContainerLoader } from './tool-usage-by-container.loader';
 import { ToolUsageRepository as GelRepository } from './tool-usage.gel.repository';
 import { ToolUsageLoader } from './tool-usage.loader';
 import { ToolUsageRepository as Neo4jRepository } from './tool-usage.neo4j.repository';
@@ -14,6 +15,7 @@ import { ToolUsageService } from './tool-usage.service';
     ToolUsageResolver,
     ResourceToolsResolver,
     ToolUsageLoader,
+    ToolUsageByContainerLoader,
     ToolUsageService,
     splitDb(Neo4jRepository, GelRepository),
   ],

src/components/tools/tool-usage/tool-usage.neo4j.repository.ts

@@ -4,12 +4,14 @@ import { CreationFailed, type ID, type UnsecuredDto } from '~/common';
 import { DtoRepository } from '~/core/database';
 import {
   ACTIVE,
+  collect,
   createNode,
   createRelationships,
   currentUser,
   filter,
   matchProps,
   merge,
+  variable,
 } from '~/core/database/query';
 import { toolFilters } from '../tool/tool.neo4j.repository';
 import {
@@ -44,16 +46,25 @@ export class ToolUsageRepository extends DtoRepository(ToolUsage) {
         );
   }
 
-  async listForContainer(containerId: ID) {
+  async listForContainers(containers: readonly ID[]) {
     const result = await this.db
       .query()
-      .match([
-        node('', 'BaseNode', { id: containerId }),
-        relation('out', '', 'uses', ACTIVE),
-        node('node', 'ToolUsage'),
-      ])
-      .apply(this.hydrate())
-      .map('dto')
+      .unwind([...containers], 'containerId')
+      .match(node('container', 'BaseNode', { id: variable('containerId') }))
+      .subQuery('container', (sub) =>
+        sub
+          .match([
+            node('container'),
+            relation('out', '', 'uses', ACTIVE),
+            node('node', 'ToolUsage'),
+          ])
+          .subQuery('node', this.hydrate())
+          .return(collect('dto').as('usages')),
+      )
+      .return<{
+        container: { id: ID };
+        usages: ReadonlyArray<UnsecuredDto<ToolUsage>>;
+      }>(['container { .id }', 'usages'])
       .run();
     return result;
   }

src/components/tools/tool-usage/tool-usage.service.ts

@@ -1,4 +1,5 @@
 import { Injectable } from '@nestjs/common';
+import { mapKeys } from '@seedcompany/common';
 import {
   type ID,
   InputException,
@@ -13,6 +14,7 @@ import { type BaseNode, isBaseNode } from '~/core/database/results';
 import { Privileges } from '../../authorization';
 import { Tool } from '../tool/dto';
 import { type CreateToolUsage, ToolUsage, type UpdateToolUsage } from './dto';
+import { type UsagesByContainer } from './tool-usage-by-container.loader';
 import { ToolUsageRepository } from './tool-usage.neo4j.repository';
 
 type TypedResource = Resource & { __typename: string };
@@ -48,9 +50,16 @@ export class ToolUsageService {
     return secured.flat();
   }
 
-  async readByContainer(container: Resource): Promise<readonly ToolUsage[]> {
-    const dtos = await this.repo.listForContainer(container.id);
-    return dtos.flatMap((dto) => this.secure(dto, container) ?? []);
+  async readManyForContainers(containers: readonly Resource[]) {
+    const containersById = mapKeys.fromList(containers, (r) => r.id).asMap;
+    const rows = await this.repo.listForContainers(containers.map((r) => r.id));
+    return rows.map((row): UsagesByContainer => {
+      const container = containersById.get(row.container.id)!;
+      return {
+        container,
+        usages: row.usages.flatMap((dto) => this.secure(dto, container) ?? []),
+      };
+    });
   }
 
   private secure(