perf(sdk): upsert plugin storage bulk writes

Author: aryasaatvik

.changeset/plugin-storage-bulk-upserts.md

@@ -0,0 +1,7 @@
+---
+"@executor-js/fumadb": patch
+"@executor-js/sdk": patch
+---
+
+Add a FumaDB bulk upsert query path and route plugin-storage bulk writes through
+it so existing rows are updated without delete/reinsert churn.

packages/core/sdk/src/executor.ts

@@ -143,7 +143,6 @@ import { connectionIdentifier } from "./connection-name-identifier";
 import { annotateToolResultOutcome } from "./tool-result";
 
 const PLUGIN_STORAGE_DELETE_KEY_BATCH_SIZE = 90;
-const PLUGIN_STORAGE_CREATE_ROW_BATCH_SIZE = 90;
 const MAX_APPROVAL_ARGUMENT_PREVIEW_CHARS = 4_000;
 
 // ---------------------------------------------------------------------------
@@ -596,6 +595,14 @@ type LooseStorageDb = {
     tableName: string,
     rows: readonly Record<string, unknown>[],
   ) => Promise<readonly unknown[]>;
+  readonly upsertMany: (
+    tableName: string,
+    options: {
+      readonly target: readonly string[];
+      readonly update: readonly string[];
+      readonly values: readonly Record<string, unknown>[];
+    },
+  ) => Promise<void>;
   readonly deleteMany: (tableName: string, options?: unknown) => Promise<void>;
   readonly findFirst: (
     tableName: string,
@@ -632,6 +639,19 @@ const makeCoreDb = (fuma: ReturnType<typeof makeFumaClient>) => ({
       : fuma
           .use(`${tableName}.createMany`, (db) => asLooseStorageDb(db).createMany(tableName, rows))
           .pipe(Effect.asVoid),
+  upsertMany: <TName extends CoreTableName>(
+    tableName: TName,
+    options: {
+      readonly target: readonly string[];
+      readonly update: readonly string[];
+      readonly values: readonly Record<string, unknown>[];
+    },
+  ): Effect.Effect<void, StorageFailure> =>
+    options.values.length === 0
+      ? Effect.void
+      : fuma.use(`${tableName}.upsertMany`, (db) =>
+          asLooseStorageDb(db).upsertMany(tableName, options),
+        ),
   deleteMany: <TName extends CoreTableName>(
     tableName: TName,
     options: { readonly where?: CoreWhere } = {},
@@ -981,33 +1001,22 @@ const makePluginStorageFacade = (input: {
       const uniqueEntries = [...entriesById.values()];
       if (uniqueEntries.length === 0) return;
 
-      yield* deleteManyImpl(owner, os.subject, uniqueEntries);
-
       const now = new Date();
-      for (
-        let offset = 0;
-        offset < uniqueEntries.length;
-        offset += PLUGIN_STORAGE_CREATE_ROW_BATCH_SIZE
-      ) {
-        const batchEntries = uniqueEntries.slice(
-          offset,
-          offset + PLUGIN_STORAGE_CREATE_ROW_BATCH_SIZE,
-        );
-        yield* input.core.createMany(
-          "plugin_storage",
-          batchEntries.map((entry) => ({
-            tenant,
-            owner: os.owner,
-            subject: os.subject,
-            plugin_id: input.pluginId,
-            collection: entry.collection,
-            key: entry.key,
-            data: entry.data,
-            created_at: now,
-            updated_at: now,
-          })),
-        );
-      }
+      yield* input.core.upsertMany("plugin_storage", {
+        target: ["tenant", "owner", "subject", "plugin_id", "collection", "key"],
+        update: ["data", "updated_at"],
+        values: uniqueEntries.map((entry) => ({
+          tenant,
+          owner: os.owner,
+          subject: os.subject,
+          plugin_id: input.pluginId,
+          collection: entry.collection,
+          key: entry.key,
+          data: entry.data,
+          created_at: now,
+          updated_at: now,
+        })),
+      });
     });
 
   const removeManyImpl = (
@@ -1104,10 +1113,24 @@ const makePluginStorageFacade = (input: {
           PluginStorageEntry<PluginStorageCollectionData<typeof definition>>,
           StorageFailure
         >,
+      putMany: (storageInput) =>
+        putManyImpl(
+          storageInput.owner,
+          storageInput.entries.map((entry) => ({
+            collection: definition.name,
+            key: entry.key,
+            data: entry.data,
+          })),
+        ),
       query: (storageInput) => queryCollection(definition, storageInput),
       count: (storageInput) =>
         queryCollection(definition, storageInput).pipe(Effect.map((rows) => rows.length)),
       remove: (storageInput) => removeImpl(storageInput.owner, definition.name, storageInput.key),
+      removeMany: (storageInput) =>
+        removeManyImpl(
+          storageInput.owner,
+          storageInput.keys.map((key) => ({ collection: definition.name, key })),
+        ),
     }),
     get: (storageInput) => getVisible(storageInput.collection, storageInput.key),
     getForOwner: (storageInput) =>

packages/core/sdk/src/plugin-storage.test.ts

@@ -65,18 +65,14 @@ const executionHistoryPlugin = definePlugin(() => ({
       owner: Owner,
       rows: readonly { readonly key: string; readonly data: ToolCall }[],
     ) =>
-      ctx.pluginStorage.putMany({
+      ctx.storage.toolCalls.putMany({
         owner,
-        entries: rows.map((row) => ({
-          collection: toolCalls.name,
-          key: row.key,
-          data: row.data,
-        })),
+        entries: rows,
       }),
     removeMany: (owner: Owner, keys: readonly string[]) =>
-      ctx.pluginStorage.removeMany({
+      ctx.storage.toolCalls.removeMany({
         owner,
-        entries: keys.map((key) => ({ collection: toolCalls.name, key })),
+        keys,
       }),
     get: (key: string) => ctx.storage.toolCalls.get({ key }),
     getForOwner: (owner: Owner, key: string) => ctx.storage.toolCalls.getForOwner({ owner, key }),

packages/core/sdk/src/plugin-storage.ts

@@ -135,6 +135,21 @@ export interface PluginStorageCollectionScopedKeyInput extends PluginStorageColl
   readonly owner: Owner;
 }
 
+export interface PluginStorageCollectionPutManyEntry<TData extends object> {
+  readonly key: string;
+  readonly data: TData;
+}
+
+export interface PluginStorageCollectionPutManyInput<TData extends object> {
+  readonly owner: Owner;
+  readonly entries: readonly PluginStorageCollectionPutManyEntry<TData>[];
+}
+
+export interface PluginStorageCollectionRemoveManyInput {
+  readonly owner: Owner;
+  readonly keys: readonly string[];
+}
+
 export interface PluginStorageCollectionListInput {
   readonly keyPrefix?: string;
 }
@@ -188,6 +203,9 @@ export interface PluginStorageCollectionFacade<
   readonly put: (
     input: PluginStorageCollectionPutInput<PluginStorageCollectionData<TDefinition>>,
   ) => Effect.Effect<PluginStorageEntry<PluginStorageCollectionData<TDefinition>>, StorageFailure>;
+  readonly putMany: (
+    input: PluginStorageCollectionPutManyInput<PluginStorageCollectionData<TDefinition>>,
+  ) => Effect.Effect<void, StorageFailure>;
   readonly query: (
     input?: PluginStorageCollectionQueryInput<TDefinition>,
   ) => Effect.Effect<
@@ -200,6 +218,9 @@ export interface PluginStorageCollectionFacade<
   readonly remove: (
     input: PluginStorageCollectionScopedKeyInput,
   ) => Effect.Effect<void, StorageFailure>;
+  readonly removeMany: (
+    input: PluginStorageCollectionRemoveManyInput,
+  ) => Effect.Effect<void, StorageFailure>;
 }
 
 export interface PluginStorageFacade {