Showing the changes to scripts/migration/service-split-postgres-cli.ts (1 of 1 file). Show all
Unified · Split
| @@ -232,6 +232,8 @@ const stableId = (...parts: readonly string[]): string => { · 231 unmodified lines | |||
| 232 | const operationKeyPrefix = (integration: string): string => | 232 | const operationKeyPrefix = (integration: string): string => |
| 233 | `${operationStorageKey(integration, "").split(".").slice(0, 2).join(".")}.%`; | 233 | `${operationStorageKey(integration, "").split(".").slice(0, 2).join(".")}.%`; |
| 234 | 234 | ||
| 235 | const rowKey = (...parts: readonly string[]): string => parts.join("\u0000"); | ||
| 236 | |||
| 235 | export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => { | 237 | export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => { |
| 236 | if (org.completed) return; | 238 | if (org.completed) return; |
| 237 | if (org.hardErrors.length > 0) { | 239 | if (org.hardErrors.length > 0) { |
| @@ -323,74 +325,103 @@ export const applyOrg = async (sql: SqlClient, org: OrgPlan): Promise<void> => { · 85 unmodified lines | |||
| 323 | ); | 325 | ); |
| 324 | } | 326 | } |
| 325 | 327 | ||
| 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 | |||
| 326 | for (const integration of org.integrations) { | 369 | for (const integration of org.integrations) { |
| 327 | for (const contribution of integration.sourceContributions) { | 370 | for (const contribution of integration.sourceContributions) { |
| 328 | for (const toolName of contribution.operationToolNames) { | 371 | if (contribution.operationToolNames.length === 0) continue; |
| 329 | const [operation] = await tx.unsafe< | 372 | const byToolName = |
| 330 | { | 373 | operationRowsBySource.get( |
| 331 | readonly data: unknown; | 374 | rowKey(contribution.source.plugin_id, contribution.source.slug), |
| 332 | readonly created_at: string; | 375 | ) ?? new Map<string, { readonly data: unknown; readonly created_at: string }>(); |
| 333 | readonly updated_at: string; | 376 | const missing = contribution.operationToolNames.filter((name) => !byToolName.has(name)); |
| 334 | }[] | 377 | if (missing.length > 0) { |
| 335 | >( | 378 | throw new Error( |
| 336 | ` | 379 | `Missing operation row for ${tenantHash(org.tenant)}/${contribution.source.slug}/${missing[0]}`, |
| 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 | ], | ||
| 360 | ); | 380 | ); |
| 361 | if (!operation) { | 381 | } |
| 362 | throw new Error( | 382 | |
| 363 | `Missing operation row for ${tenantHash(org.tenant)}/${contribution.source.slug}/${toolName}`, | 383 | const CHUNK = 200; |
| 364 | ); | 384 | for (let start = 0; start < contribution.operationToolNames.length; start += CHUNK) { |
| 365 | } | 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(", "); | ||
| 366 | await tx.unsafe( | 409 | await tx.unsafe( |
| 367 | ` | 410 | ` |
| 368 | INSERT INTO plugin_storage ( | 411 | INSERT INTO plugin_storage ( |
| 369 | plugin_id, collection, key, data, created_at, updated_at, row_id, | 412 | plugin_id, collection, key, data, created_at, updated_at, row_id, |
| 370 | tenant, owner, subject | 413 | tenant, owner, subject |
| 371 | ) | 414 | ) |
| 372 | VALUES ($1, 'operation', $2, $3::json, $4, $5, $6, $7, 'org', '') | 415 | VALUES ${rowsSql} |
| 373 | ON CONFLICT (tenant, owner, subject, plugin_id, collection, key) | 416 | ON CONFLICT (tenant, owner, subject, plugin_id, collection, key) |
| 374 | DO UPDATE SET data = EXCLUDED.data, updated_at = EXCLUDED.updated_at | 417 | DO UPDATE SET data = EXCLUDED.data, updated_at = EXCLUDED.updated_at |
| 375 | `, | 418 | `, |
| 376 | [ | 419 | params, |
| 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 | ], | ||
| 391 | ); | 420 | ); |
| 392 | await tx.unsafe( | 421 | } |
| 393 | ` | 422 | |
| 423 | await tx.unsafe( | ||
| 424 | ` | ||
| 394 | INSERT INTO tool ( | 425 | INSERT INTO tool ( |
| 395 | integration, connection, plugin_id, name, description, input_schema, output_schema, | 426 | integration, connection, plugin_id, name, description, input_schema, output_schema, |
| 396 | annotations, created_at, updated_at, row_id, tenant, owner, subject | 427 | 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 | |||
| 398 | SELECT $1, connection, $2, name, description, input_schema, output_schema, | 429 | SELECT $1, connection, $2, name, description, input_schema, output_schema, |
| 399 | annotations, created_at, $3, concat($4::text, '_', row_id), tenant, owner, subject | 430 | annotations, created_at, $3, concat($4::text, '_', row_id), tenant, owner, subject |
| 400 | FROM tool | 431 | FROM tool |
| 401 | WHERE tenant = $5 AND integration = $6 AND name = $7 | 432 | WHERE tenant = $5 AND integration = $6 AND name = ANY($7::text[]) |
| 402 | ON CONFLICT (tenant, owner, subject, integration, connection, name) DO NOTHING | 433 | ON CONFLICT (tenant, owner, subject, integration, connection, name) DO NOTHING |
| 403 | `, | 434 | `, |
| 404 | [ | 435 | [ |
| 405 | integration.target.slug, | 436 | integration.target.slug, |
| 406 | integration.target.pluginId, | 437 | integration.target.pluginId, |
| 407 | now, | 438 | now, |
| 408 | stableId("tool", org.tenant, integration.target.slug), | 439 | stableId("tool", org.tenant, integration.target.slug), |
| 409 | org.tenant, | 440 | org.tenant, |
| 410 | contribution.source.slug, | 441 | contribution.source.slug, |
| 411 | toolName, | 442 | contribution.operationToolNames, |
| 412 | ], | 443 | ], |
| 413 | ); | 444 | ); |
| 414 | } | ||
| 415 | } | 445 | } |
| 416 | } | 446 | } |
| 417 | 447 | ||