From eadf08121af42227081d6fedbbc773b7f7c65828 Mon Sep 17 00:00:00 2001 From: Ankur Mittal Date: Fri, 22 May 2026 10:18:18 +0530 Subject: [PATCH] Stream database dump exports --- src/export/dump.test.ts | 60 ++++++++++++++-- src/export/dump.ts | 156 +++++++++++++++++++++++++++++++--------- 2 files changed, 175 insertions(+), 41 deletions(-) diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..10a113e 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -71,13 +71,13 @@ describe('Database Dump Module', () => { expect(dumpText).toContain( 'CREATE TABLE users (id INTEGER, name TEXT);' ) - expect(dumpText).toContain("INSERT INTO users VALUES (1, 'Alice');") - expect(dumpText).toContain("INSERT INTO users VALUES (2, 'Bob');") + expect(dumpText).toContain('INSERT INTO "users" VALUES (1, \'Alice\');') + expect(dumpText).toContain('INSERT INTO "users" VALUES (2, \'Bob\');') expect(dumpText).toContain( 'CREATE TABLE orders (id INTEGER, total REAL);' ) - expect(dumpText).toContain('INSERT INTO orders VALUES (1, 99.99);') - expect(dumpText).toContain('INSERT INTO orders VALUES (2, 49.5);') + expect(dumpText).toContain('INSERT INTO "orders" VALUES (1, 99.99);') + expect(dumpText).toContain('INSERT INTO "orders" VALUES (2, 49.5);') }) it('should handle empty databases (no tables)', async () => { @@ -108,7 +108,7 @@ describe('Database Dump Module', () => { expect(dumpText).toContain( 'CREATE TABLE users (id INTEGER, name TEXT);' ) - expect(dumpText).not.toContain('INSERT INTO users VALUES') + expect(dumpText).not.toContain('INSERT INTO "users" VALUES') }) it('should escape single quotes properly in string values', async () => { @@ -124,7 +124,55 @@ describe('Database Dump Module', () => { expect(response).toBeInstanceOf(Response) const dumpText = await response.text() expect(dumpText).toContain( - "INSERT INTO users VALUES (1, 'Alice''s adventure');" + "INSERT INTO \"users\" VALUES (1, 'Alice''s adventure');" + ) + }) + + it('should export table rows in batches without loading the whole table at once', async () => { + const firstBatch = Array.from({ length: 500 }, (_, index) => ({ + id: index + 1, + name: `User ${index + 1}`, + })) + const secondBatch = [{ id: 501, name: 'User 501' }] + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + ]) + .mockResolvedValueOnce(firstBatch) + .mockResolvedValueOnce(secondBatch) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await response.text() + + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (1, \'User 1\');' + ) + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (501, \'User 501\');' + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 3, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [500, 0], + }, + ], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 4, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [500, 500], + }, + ], + mockDataSource, + mockConfig ) }) diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..ecc3946 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -3,67 +3,153 @@ import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' -export async function dumpDatabaseRoute( +const DEFAULT_EXPORT_BATCH_SIZE = 500 + +function quoteIdentifier(identifier: string) { + return `"${identifier.replace(/"/g, '""')}"` +} + +function formatSqlValue(value: unknown) { + if (value === null || value === undefined) { + return 'NULL' + } + + if (typeof value === 'string') { + return `'${value.replace(/'/g, "''")}'` + } + + if (typeof value === 'boolean') { + return value ? '1' : '0' + } + + if (value instanceof ArrayBuffer) { + return `X'${arrayBufferToHex(value)}'` + } + + if (ArrayBuffer.isView(value)) { + const view = value as ArrayBufferView + return `X'${arrayBufferToHex( + view.buffer.slice( + view.byteOffset, + view.byteOffset + view.byteLength + ) + )}'` + } + + return String(value) +} + +function arrayBufferToHex(value: ArrayBuffer) { + return Array.from(new Uint8Array(value)) + .map((byte) => byte.toString(16).padStart(2, '0')) + .join('') +} + +async function pauseForQueuedWork() { + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +async function enqueueDatabaseDump( + controller: ReadableStreamDefaultController, + encoder: TextEncoder, + tables: string[], dataSource: DataSource, - config: StarbaseDBConfiguration -): Promise { - try { - // Get all table names - const tablesResult = await executeOperation( - [{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }], + config: StarbaseDBConfiguration, + batchSize = DEFAULT_EXPORT_BATCH_SIZE +) { + const enqueue = (chunk: string) => controller.enqueue(encoder.encode(chunk)) + + enqueue('SQLite format 3\0') + + for (const table of tables) { + const quotedTable = quoteIdentifier(table) + const schemaResult = await executeOperation( + [ + { + sql: "SELECT sql FROM sqlite_master WHERE type='table' AND name=?;", + params: [table], + }, + ], dataSource, config ) - const tables = tablesResult.map((row: any) => row.name) - let dumpContent = 'SQLite format 3\0' // SQLite file header + if (schemaResult.length) { + const schema = schemaResult[0].sql + enqueue(`\n-- Table: ${table}\n${schema};\n\n`) + } - // Iterate through all tables - for (const table of tables) { - // Get table schema - const schemaResult = await executeOperation( + for (let offset = 0; ; offset += batchSize) { + const rows = await executeOperation( [ { - sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name='${table}';`, + sql: `SELECT * FROM ${quotedTable} LIMIT ? OFFSET ?;`, + params: [batchSize, offset], }, ], dataSource, config ) - if (schemaResult.length) { - const schema = schemaResult[0].sql - dumpContent += `\n-- Table: ${table}\n${schema};\n\n` + for (const row of rows) { + const values = Object.values(row).map(formatSqlValue) + enqueue( + `INSERT INTO ${quotedTable} VALUES (${values.join(', ')});\n` + ) } - // Get table data - const dataResult = await executeOperation( - [{ sql: `SELECT * FROM ${table};` }], - dataSource, - config - ) - - for (const row of dataResult) { - const values = Object.values(row).map((value) => - typeof value === 'string' - ? `'${value.replace(/'/g, "''")}'` - : value - ) - dumpContent += `INSERT INTO ${table} VALUES (${values.join(', ')});\n` + if (rows.length < batchSize) { + break } - dumpContent += '\n' + await pauseForQueuedWork() } - // Create a Blob from the dump content - const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' }) + enqueue('\n') + } +} + +export async function dumpDatabaseRoute( + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + try { + const tablesResult = await executeOperation( + [ + { + sql: "SELECT name FROM sqlite_master WHERE type='table';", + }, + ], + dataSource, + config + ) + const tables = tablesResult.map((row: any) => String(row.name)) + + const encoder = new TextEncoder() + const stream = new ReadableStream({ + async start(controller) { + try { + await enqueueDatabaseDump( + controller, + encoder, + tables, + dataSource, + config + ) + controller.close() + } catch (error) { + console.error('Database Dump Stream Error:', error) + controller.error(error) + } + }, + }) const headers = new Headers({ 'Content-Type': 'application/x-sqlite3', 'Content-Disposition': 'attachment; filename="database_dump.sql"', }) - return new Response(blob, { headers }) + return new Response(stream, { headers }) } catch (error: any) { console.error('Database Dump Error:', error) return createResponse(undefined, 'Failed to create database dump', 500)