History

Batch serving-state copies in the postgres runner

Rhys Sullivan committed 4 months ago

Commit 77756bb1 parent b3660ee1 file · +97 −67Browse files

Showing the changes to scripts/migration/service-split-postgres-cli.ts (1 of 1 file). Show all

Unified · Split

scripts/migration/service-split-postgres-cli.ts+97−67
@@ -232,6 +232,8 @@ const stableId = (...parts: readonly string[]): string => { · 231 unmodified lines
232232 const operationKeyPrefix = (integration: string): string =>
233233 `${operationStorageKey(integration, "").split(".").slice(0, 2).join(".")}.%`;
234234
235+const rowKey = (...parts: readonly string[]): string => parts.join("\u0000");
236+
235237 export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => {
236238 if (org.completed) return;
237239 if (org.hardErrors.length > 0) {
@@ -323,74 +325,103 @@ export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => { · 85 unmodified lines
323325 );
324326 }
325327
328+ // Batched serving-state copy: the naive per-tool loop costs 2-3 WAN
329+ // round-trips per operation (~200k total across prod), which stalls the
330+ // per-org transaction for hours. Fetch the org's operation rows once,
331+ // group them in JS (storageDataRecord tolerates both jsonb-object and
332+ // json-string data shapes, matching the planner), bulk-insert the
333+ // re-keyed copies in chunks, and clone tool rows set-based.
334+ const operationRowsBySource = new Map<
335+ string,
336+ Map<string, { readonly data: unknown; readonly created_at: string }>
337+ >();
338+ {
339+ const orgOperationRows = await tx.unsafe<
340+ {
341+ readonly plugin_id: string;
342+ readonly data: unknown;
343+ readonly created_at: string;
344+ }[]
345+ >(
346+ `
347+ SELECT plugin_id, data, created_at::text
348+ FROM plugin_storage
349+ WHERE tenant = $1
350+ AND owner = 'org'
351+ AND subject = ''
352+ AND plugin_id IN ('google', 'microsoft')
353+ AND collection = 'operation'
354+ `,
355+ [org.tenant],
356+ );
357+ for (const row of orgOperationRows) {
358+ const record = storageDataRecord(row);
359+ const sourceIntegration = record.integration;
360+ const toolName = record.toolName;
361+ if (typeof sourceIntegration !== "string" || typeof toolName !== "string") continue;
362+ const key = rowKey(row.plugin_id, sourceIntegration);
363+ const byTool = operationRowsBySource.get(key) ?? new Map();
364+ byTool.set(toolName, row);
365+ operationRowsBySource.set(key, byTool);
366+ }
367+ }
368+
326369 for (const integration of org.integrations) {
327370 for (const contribution of integration.sourceContributions) {
328− for (const toolName of contribution.operationToolNames) {
329− const [operation] = await tx.unsafe<
330− {
331− readonly data: unknown;
332− readonly created_at: string;
333− readonly updated_at: string;
334− }[]
335− >(
336− `
337− SELECT data, created_at::text, updated_at::text
338− FROM plugin_storage
339− WHERE tenant = $1
340− AND owner = 'org'
341− AND subject = ''
342− AND plugin_id = $2
343− AND collection = 'operation'
344− AND (
345− key = $3
346− OR (
347− data::jsonb ->> 'integration' = $4
348− AND data::jsonb ->> 'toolName' = $5
349− )
350− )
351− LIMIT 1
352− `,
353− [
354− org.tenant,
355− contribution.source.plugin_id,
356− operationStorageKey(contribution.source.slug, toolName),
357− contribution.source.slug,
358− toolName,
359− ],
371+ if (contribution.operationToolNames.length === 0) continue;
372+ const byToolName =
373+ operationRowsBySource.get(
374+ rowKey(contribution.source.plugin_id, contribution.source.slug),
375+ ) ?? new Map<string, { readonly data: unknown; readonly created_at: string }>();
376+ const missing = contribution.operationToolNames.filter((name) => !byToolName.has(name));
377+ if (missing.length > 0) {
378+ throw new Error(
379+ `Missing operation row for ${tenantHash(org.tenant)}/${contribution.source.slug}/${missing[0]}`,
360380 );
361− if (!operation) {
362− throw new Error(
363− `Missing operation row for ${tenantHash(org.tenant)}/${contribution.source.slug}/${toolName}`,
364− );
365− }
381+ }
382+
383+ const CHUNK = 200;
384+ for (let start = 0; start < contribution.operationToolNames.length; start += CHUNK) {
385+ const chunk = contribution.operationToolNames.slice(start, start + CHUNK);
386+ const params: unknown[] = [];
387+ const rowsSql = chunk
388+ .map((toolName) => {
389+ const operation = byToolName.get(toolName)!;
390+ const base = params.length;
391+ params.push(
392+ integration.target.pluginId,
393+ operationStorageKey(integration.target.slug, toolName),
394+ JSON.stringify(
395+ scrubJson({
396+ ...storageDataRecord(operation),
397+ integration: integration.target.slug,
398+ toolName,
399+ }),
400+ ),
401+ operation.created_at,
402+ now,
403+ stableId("operation", org.tenant, integration.target.slug, toolName),
404+ org.tenant,
405+ );
406+ return `($${base + 1}, 'operation', $${base + 2}, $${base + 3}::json, $${base + 4}, $${base + 5}, $${base + 6}, $${base + 7}, 'org', '')`;
407+ })
408+ .join(", ");
366409 await tx.unsafe(
367410 `
368411 INSERT INTO plugin_storage (
369412 plugin_id, collection, key, data, created_at, updated_at, row_id,
370413 tenant, owner, subject
371414 )
372− VALUES ($1, 'operation', $2, $3::json, $4, $5, $6, $7, 'org', '')
415+ VALUES ${rowsSql}
373416 ON CONFLICT (tenant, owner, subject, plugin_id, collection, key)
374417 DO UPDATE SET data = EXCLUDED.data, updated_at = EXCLUDED.updated_at
375418 `,
376− [
377− integration.target.pluginId,
378− operationStorageKey(integration.target.slug, toolName),
379− JSON.stringify(
380− scrubJson({
381− ...storageDataRecord(operation),
382− integration: integration.target.slug,
383− toolName,
384− }),
385− ),
386− operation.created_at,
387− now,
388− stableId("operation", org.tenant, integration.target.slug, toolName),
389− org.tenant,
390− ],
419+ params,
391420 );
392− await tx.unsafe(
393− `
421+ }
422+
423+ await tx.unsafe(
424+ `
394425 INSERT INTO tool (
395426 integration, connection, plugin_id, name, description, input_schema, output_schema,
396427 annotations, created_at, updated_at, row_id, tenant, owner, subject
@@ -398,20 +429,19 @@ export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => { · 1 unmodified line
398429 SELECT $1, connection, $2, name, description, input_schema, output_schema,
399430 annotations, created_at, $3, concat($4::text, '_', row_id), tenant, owner, subject
400431 FROM tool
401− WHERE tenant = $5 AND integration = $6 AND name = $7
432+ WHERE tenant = $5 AND integration = $6 AND name = ANY($7::text[])
402433 ON CONFLICT (tenant, owner, subject, integration, connection, name) DO NOTHING
403434 `,
404− [
405− integration.target.slug,
406− integration.target.pluginId,
407− now,
408− stableId("tool", org.tenant, integration.target.slug),
409− org.tenant,
410− contribution.source.slug,
411− toolName,
412− ],
413− );
414− }
435+ [
436+ integration.target.slug,
437+ integration.target.pluginId,
438+ now,
439+ stableId("tool", org.tenant, integration.target.slug),
440+ org.tenant,
441+ contribution.source.slug,
442+ contribution.operationToolNames,
443+ ],
444+ );
415445 }
416446 }
417447