History

tools: collapse annotation N+1 + push sourceId filter to DB

openapi/graphql resolveAnnotations ran one listOperationsBySource per distinct scope in the toolRows. Store the owning scopeId on each StoredOperation instead and issue a single query, partitioning by op.scopeId so shadowed tools still resolve against the right binding. Also push executor.tools.list's sourceId filter down to storage so `/sources/:id/tools` stops pulling every row in the scope stack.

Rhys Sullivan committed 6 months ago

Commit 400a8061 parent aa220d75 files · +44 −40Browse files

Unified · Split

5 files
  1. packages/core/sdk/src/executor.ts +10 −1
  2. packages/plugins/graphql/src/sdk/plugin.ts +13 −15
  3. packages/plugins/graphql/src/sdk/store.ts +4 −6
  4. packages/plugins/openapi/src/sdk/plugin.ts +13 −12
  5. packages/plugins/openapi/src/sdk/store.ts +4 −6
packages/core/sdk/src/executor.ts+10−1
@@ -1762,7 +1762,16 @@ export const createExecutor = < · 1761 unmodified lines
1762 1762
1763 const listTools = (filter?: ToolListFilter) =>1763 const listTools = (filter?: ToolListFilter) =>
1764 Effect.gen(function* () {1764 Effect.gen(function* () {
1765 const dynamic = yield* core.findMany({ model: "tool" });1765 // Push sourceId down to storage so we don't pull every tool row
1766 // in the stack when the caller only wants one source (e.g. the
1767 // `/sources/:id/tools` endpoint). `query` is a free-text match
1768 // that still needs the row contents, so we apply it in memory.
1769 const dynamic = yield* core.findMany({
1770 model: "tool",
1771 where: filter?.sourceId
1772 ? [{ field: "source_id", value: filter.sourceId }]
1773 : undefined,
1774 });
1766 // Dedup by tool id, innermost scope winning — same reason as1775 // Dedup by tool id, innermost scope winning — same reason as
1767 // `listSources` above: a shadowed id must surface as one entry1776 // `listSources` above: a shadowed id must surface as one entry
1768 // (the inner one), not two.1777 // (the inner one), not two.
packages/plugins/graphql/src/sdk/plugin.ts+13−15
@@ -366,6 +366,7 @@ export const graphqlPlugin = definePlugin( · 365 unmodified lines
366 const storedOps: StoredOperation[] = prepared.map((p) => ({366 const storedOps: StoredOperation[] = prepared.map((p) => ({
367 toolId: `${namespace}.${p.toolPath}`,367 toolId: `${namespace}.${p.toolPath}`,
368 sourceId: namespace,368 sourceId: namespace,
369 scopeId: config.scope,
369 binding: p.binding,370 binding: p.binding,
370 }));371 }));
371 372
@@ -527,22 +528,19 @@ export const graphqlPlugin = definePlugin( · 155 unmodified lines
527 // toolRows for a single (plugin_id, source_id) group can still528 // toolRows for a single (plugin_id, source_id) group can still
528 // straddle multiple scopes when the source is shadowed (e.g. an529 // straddle multiple scopes when the source is shadowed (e.g. an
529 // org-level GraphQL source plus a per-user override that530 // org-level GraphQL source plus a per-user override that
530 // re-registers the same tool ids). Run one listOperationsBySource531 // re-registers the same tool ids). One query returns every op
531 // per distinct scope so each lookup pins {source_id, scope_id}532 // row in the caller's stack (the scoped adapter pins
532 // and we don't fall through to the wrong scope's bindings.533 // `scope_id IN (stack)`); partition by `op.scopeId` so each
533 const scopes = new Set<string>();534 // tool row resolves against its own scope's binding.
534 for (const row of toolRows as readonly ToolRow[]) {535 const ops = yield* ctx.storage.listOperationsBySource(sourceId);
535 scopes.add(row.scope_id as string);
536 }
537 const byScope = new Map<string, Map<string, OperationBinding>>();536 const byScope = new Map<string, Map<string, OperationBinding>>();
538 for (const scope of scopes) {537 for (const op of ops) {
539 const ops = yield* ctx.storage.listOperationsBySource(538 let byId = byScope.get(op.scopeId);
540 sourceId,539 if (!byId) {
541 scope,540 byId = new Map<string, OperationBinding>();
542 );541 byScope.set(op.scopeId, byId);
543 const byId = new Map<string, OperationBinding>();542 }
544 for (const op of ops) byId.set(op.toolId, op.binding);543 byId.set(op.toolId, op.binding);
545 byScope.set(scope, byId);
546 }544 }
547 545
548 const out: Record<string, ToolAnnotations> = {};546 const out: Record<string, ToolAnnotations> = {};
packages/plugins/graphql/src/sdk/store.ts+4−6
@@ -50,6 +50,7 @@ export interface StoredGraphqlSource { · 49 unmodified lines
50export interface StoredOperation {50export interface StoredOperation {
51 readonly toolId: string;51 readonly toolId: string;
52 readonly sourceId: string;52 readonly sourceId: string;
53 readonly scopeId: string;
53 readonly binding: OperationBinding;54 readonly binding: OperationBinding;
54}55}
55 56
@@ -126,7 +127,6 @@ export interface GraphqlStore { · 70 unmodified lines
126 127
127 readonly listOperationsBySource: (128 readonly listOperationsBySource: (
128 sourceId: string,129 sourceId: string,
129 scope: string,
130 ) => Effect.Effect<readonly StoredOperation[], StorageFailure>;130 ) => Effect.Effect<readonly StoredOperation[], StorageFailure>;
131 131
132 readonly removeSource: (132 readonly removeSource: (
@@ -153,6 +153,7 @@ export const makeDefaultGraphqlStore = ({ · 20 unmodified lines
153 const rowToOperation = (row: Record<string, unknown>): StoredOperation => ({153 const rowToOperation = (row: Record<string, unknown>): StoredOperation => ({
154 toolId: row.id as string,154 toolId: row.id as string,
155 sourceId: row.source_id as string,155 sourceId: row.source_id as string,
156 scopeId: row.scope_id as string,
156 binding: decodeBinding(row.binding),157 binding: decodeBinding(row.binding),
157 });158 });
158 159
@@ -257,14 +258,11 @@ export const makeDefaultGraphqlStore = ({ · 98 unmodified lines
257 })258 })
258 .pipe(Effect.map((row) => (row ? rowToOperation(row) : null))),259 .pipe(Effect.map((row) => (row ? rowToOperation(row) : null))),
259 260
260 listOperationsBySource: (sourceId, scope) =>261 listOperationsBySource: (sourceId) =>
261 db262 db
262 .findMany({263 .findMany({
263 model: "graphql_operation",264 model: "graphql_operation",
264 where: [265 where: [{ field: "source_id", value: sourceId }],
265 { field: "source_id", value: sourceId },
266 { field: "scope_id", value: scope },
267 ],
268 })266 })
269 .pipe(Effect.map((rows) => rows.map(rowToOperation))),267 .pipe(Effect.map((rows) => rows.map(rowToOperation))),
270 268
packages/plugins/openapi/src/sdk/plugin.ts+13−12
@@ -408,6 +408,7 @@ export const openApiPlugin = definePlugin( · 407 unmodified lines
408 const storedOps: StoredOperation[] = definitions.map((def) => ({408 const storedOps: StoredOperation[] = definitions.map((def) => ({
409 toolId: `${namespace}.${def.toolPath}`,409 toolId: `${namespace}.${def.toolPath}`,
410 sourceId: namespace,410 sourceId: namespace,
411 scopeId: input.scope,
411 binding: toBinding(def),412 binding: toBinding(def),
412 }));413 }));
413 414
@@ -949,19 +950,19 @@ export const openApiPlugin = definePlugin( · 535 unmodified lines
949 // toolRows for a single (plugin_id, source_id) group can still950 // toolRows for a single (plugin_id, source_id) group can still
950 // straddle multiple scopes when the source is shadowed (e.g. an951 // straddle multiple scopes when the source is shadowed (e.g. an
951 // org-level openapi source plus a per-user override that952 // org-level openapi source plus a per-user override that
952 // re-registers the same tool ids). Run one listOperationsBySource953 // re-registers the same tool ids). One query returns every op
953 // per distinct scope so each lookup pins {source_id, scope_id}954 // row in the caller's stack (the scoped adapter pins
954 // and we don't fall through to the wrong scope's bindings.955 // `scope_id IN (stack)`); partition by `op.scopeId` so each
955 const scopes = new Set<string>();956 // tool row resolves against its own scope's binding.
956 for (const row of toolRows as readonly ToolRow[]) {957 const ops = yield* ctx.storage.listOperationsBySource(sourceId);
957 scopes.add(row.scope_id as string);
958 }
959 const byScope = new Map<string, Map<string, OperationBinding>>();958 const byScope = new Map<string, Map<string, OperationBinding>>();
960 for (const scope of scopes) {959 for (const op of ops) {
961 const ops = yield* ctx.storage.listOperationsBySource(sourceId, scope);960 let byId = byScope.get(op.scopeId);
962 const byId = new Map<string, OperationBinding>();961 if (!byId) {
963 for (const op of ops) byId.set(op.toolId, op.binding);962 byId = new Map<string, OperationBinding>();
964 byScope.set(scope, byId);963 byScope.set(op.scopeId, byId);
964 }
965 byId.set(op.toolId, op.binding);
965 }966 }
966 967
967 const out: Record<string, ToolAnnotations> = {};968 const out: Record<string, ToolAnnotations> = {};
packages/plugins/openapi/src/sdk/store.ts+4−6
@@ -119,6 +119,7 @@ export type StoredSourceSchemaType = typeof StoredSourceSchema.Type; · 118 unmodified lines
119export interface StoredOperation {119export interface StoredOperation {
120 readonly toolId: string;120 readonly toolId: string;
121 readonly sourceId: string;121 readonly sourceId: string;
122 readonly scopeId: string;
122 readonly binding: OperationBinding;123 readonly binding: OperationBinding;
123}124}
124 125
@@ -202,7 +203,6 @@ export interface OpenapiStore { · 77 unmodified lines
202 203
203 readonly listOperationsBySource: (204 readonly listOperationsBySource: (
204 sourceId: string,205 sourceId: string,
205 scope: string,
206 ) => Effect.Effect<readonly StoredOperation[], StorageFailure>;206 ) => Effect.Effect<readonly StoredOperation[], StorageFailure>;
207 207
208 readonly removeSource: (208 readonly removeSource: (
@@ -255,6 +255,7 @@ export const makeDefaultOpenapiStore = ({ · 46 unmodified lines
255 const rowToOperation = (row: Record<string, unknown>): StoredOperation => ({255 const rowToOperation = (row: Record<string, unknown>): StoredOperation => ({
256 toolId: row.id as string,256 toolId: row.id as string,
257 sourceId: row.source_id as string,257 sourceId: row.source_id as string,
258 scopeId: row.scope_id as string,
258 binding: decodeBinding(259 binding: decodeBinding(
259 typeof row.binding === "string" ? JSON.parse(row.binding) : row.binding,260 typeof row.binding === "string" ? JSON.parse(row.binding) : row.binding,
260 ),261 ),
@@ -388,14 +389,11 @@ export const makeDefaultOpenapiStore = ({ · 127 unmodified lines
388 })389 })
389 .pipe(Effect.map((row) => (row ? rowToOperation(row) : null))),390 .pipe(Effect.map((row) => (row ? rowToOperation(row) : null))),
390 391
391 listOperationsBySource: (sourceId, scope) =>392 listOperationsBySource: (sourceId) =>
392 adapter393 adapter
393 .findMany({394 .findMany({
394 model: "openapi_operation",395 model: "openapi_operation",
395 where: [396 where: [{ field: "source_id", value: sourceId }],
396 { field: "source_id", value: sourceId },
397 { field: "scope_id", value: scope },
398 ],
399 })397 })
400 .pipe(Effect.map((rows) => rows.map(rowToOperation))),398 .pipe(Effect.map((rows) => rows.map(rowToOperation))),
401 399