Handle SpacetimeDB migration imports with chunked uploads
Some checks failed
CI / verify (push) Has been cancelled
Some checks failed
CI / verify (push) Has been cancelled
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
#!/usr/bin/env node
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { readFile } from 'node:fs/promises';
|
||||
import path from 'node:path';
|
||||
import {
|
||||
@@ -11,6 +12,8 @@ import {
|
||||
parseArgs,
|
||||
} from './spacetime-migration-common.mjs';
|
||||
|
||||
const DEFAULT_MIGRATION_IMPORT_CHUNK_SIZE = 512 * 1024;
|
||||
|
||||
try {
|
||||
const options = parseArgs(process.argv.slice(2));
|
||||
if (!options.in) {
|
||||
@@ -30,7 +33,7 @@ try {
|
||||
const webOptions = await prepareWebImportOptions(options);
|
||||
let result;
|
||||
try {
|
||||
result = await importMigrationJsonDirect(webOptions, migrationJson);
|
||||
result = await importMigrationJsonWithFallback(webOptions, migrationJson);
|
||||
} finally {
|
||||
await revokeTemporaryWebIdentity(webOptions);
|
||||
}
|
||||
@@ -86,6 +89,25 @@ async function prepareWebImportOptions(options) {
|
||||
};
|
||||
}
|
||||
|
||||
async function importMigrationJsonWithFallback(options, migrationJson) {
|
||||
const chunkSize = resolveChunkSize(options);
|
||||
if (Buffer.byteLength(migrationJson, 'utf8') > chunkSize) {
|
||||
return importMigrationJsonChunked(options, migrationJson, chunkSize);
|
||||
}
|
||||
|
||||
try {
|
||||
return await importMigrationJsonDirect(options, migrationJson);
|
||||
} catch (error) {
|
||||
if (!isRequestBodyTooLargeError(error)) {
|
||||
throw error;
|
||||
}
|
||||
console.warn(
|
||||
`[spacetime:migration:import] 直接导入触发 HTTP 413,改用 ${chunkSize} bytes 分片上传。`,
|
||||
);
|
||||
return importMigrationJsonChunked(options, migrationJson, chunkSize);
|
||||
}
|
||||
}
|
||||
|
||||
async function importMigrationJsonDirect(options, migrationJson) {
|
||||
const includeTables = resolveImportIncludeTables(options, migrationJson);
|
||||
const procedureName =
|
||||
@@ -108,6 +130,60 @@ async function importMigrationJsonDirect(options, migrationJson) {
|
||||
return callSpacetimeProcedure(options, procedureName, input);
|
||||
}
|
||||
|
||||
async function importMigrationJsonChunked(options, migrationJson, chunkSize) {
|
||||
const includeTables = resolveImportIncludeTables(options, migrationJson);
|
||||
const procedureName =
|
||||
options.incremental === true
|
||||
? 'import_database_migration_incremental_from_chunks'
|
||||
: 'import_database_migration_from_chunks';
|
||||
const uploadId = `migration-${Date.now()}-${randomUUID()}`;
|
||||
const chunks = splitStringByUtf8Bytes(migrationJson, chunkSize);
|
||||
console.log(
|
||||
`[spacetime:migration:import] 使用分片导入: upload_id=${uploadId}, chunks=${chunks.length}, chunk_size=${chunkSize}`,
|
||||
);
|
||||
if (options.replaceExisting === true) {
|
||||
console.log(
|
||||
`[spacetime:migration:import] replace-existing 仅覆盖本次文件内的表: ${includeTables.join(', ') || '无'}`,
|
||||
);
|
||||
} else if (options.incremental === true) {
|
||||
console.log(`[spacetime:migration:import] 使用增量模式,已存在或冲突的行会跳过`);
|
||||
}
|
||||
|
||||
let committed = false;
|
||||
try {
|
||||
for (let index = 0; index < chunks.length; index += 1) {
|
||||
const chunkResult = await callSpacetimeProcedure(
|
||||
options,
|
||||
'put_database_migration_import_chunk',
|
||||
{
|
||||
upload_id: uploadId,
|
||||
chunk_index: index,
|
||||
chunk_count: chunks.length,
|
||||
chunk: chunks[index],
|
||||
},
|
||||
);
|
||||
ensureProcedureOk(chunkResult);
|
||||
console.log(
|
||||
`[spacetime:migration:import] 已上传迁移分片 ${index + 1}/${chunks.length}`,
|
||||
);
|
||||
}
|
||||
|
||||
const result = await callSpacetimeProcedure(options, procedureName, {
|
||||
upload_id: uploadId,
|
||||
include_tables: includeTables,
|
||||
replace_existing: options.replaceExisting === true,
|
||||
dry_run: options.dryRun === true,
|
||||
});
|
||||
ensureProcedureOk(result);
|
||||
committed = true;
|
||||
return result;
|
||||
} finally {
|
||||
if (!committed) {
|
||||
await clearMigrationChunksBestEffort(options, uploadId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function resolveImportIncludeTables(options, migrationJson) {
|
||||
if (options.replaceExisting !== true) {
|
||||
return options.includeTables;
|
||||
@@ -152,6 +228,63 @@ function readMigrationTableNames(migrationJson) {
|
||||
return tableNames;
|
||||
}
|
||||
|
||||
function resolveChunkSize(options) {
|
||||
const chunkSize = options.chunkSize || DEFAULT_MIGRATION_IMPORT_CHUNK_SIZE;
|
||||
if (chunkSize > 1024 * 1024) {
|
||||
throw new Error('--chunk-size 不能超过 1048576,避免触发迁移分片 procedure 单片限制。');
|
||||
}
|
||||
return chunkSize;
|
||||
}
|
||||
|
||||
function splitStringByUtf8Bytes(value, maxBytes) {
|
||||
const chunks = [];
|
||||
let current = '';
|
||||
let currentBytes = 0;
|
||||
for (const character of value) {
|
||||
const characterBytes = Buffer.byteLength(character, 'utf8');
|
||||
if (characterBytes > maxBytes) {
|
||||
throw new Error(`单个字符超过 chunk-size,当前 chunk-size: ${maxBytes}`);
|
||||
}
|
||||
if (currentBytes + characterBytes > maxBytes && current) {
|
||||
chunks.push(current);
|
||||
current = '';
|
||||
currentBytes = 0;
|
||||
}
|
||||
current += character;
|
||||
currentBytes += characterBytes;
|
||||
}
|
||||
if (current) {
|
||||
chunks.push(current);
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
function isRequestBodyTooLargeError(error) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
return (
|
||||
message.includes('HTTP 413') ||
|
||||
message.toLowerCase().includes('length limit exceeded')
|
||||
);
|
||||
}
|
||||
|
||||
async function clearMigrationChunksBestEffort(options, uploadId) {
|
||||
try {
|
||||
const result = await callSpacetimeProcedure(
|
||||
options,
|
||||
'clear_database_migration_import_chunks',
|
||||
{ upload_id: uploadId },
|
||||
);
|
||||
ensureProcedureOk(result);
|
||||
console.warn(`[spacetime:migration:import] 已清理失败导入的临时分片: ${uploadId}`);
|
||||
} catch (error) {
|
||||
console.warn(
|
||||
`[spacetime:migration:import] 清理临时迁移分片失败: ${
|
||||
error instanceof Error ? error.message : String(error)
|
||||
}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function revokeTemporaryWebIdentity(options) {
|
||||
if (!options.temporaryWebIdentity) {
|
||||
return;
|
||||
|
||||
Reference in New Issue
Block a user