diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..c89c3e4 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -38,8 +38,12 @@ beforeEach(() => { } }) +async function collectStream(response: Response): Promise { + return await response.text() +} + describe('Database Dump Module', () => { - it('should return a database dump when tables exist', async () => { + it('should return a streaming database dump when tables exist', async () => { vi.mocked(executeOperation) .mockResolvedValueOnce([{ name: 'users' }, { name: 'orders' }]) .mockResolvedValueOnce([ @@ -60,24 +64,24 @@ describe('Database Dump Module', () => { const response = await dumpDatabaseRoute(mockDataSource, mockConfig) expect(response).toBeInstanceOf(Response) - expect(response.headers.get('Content-Type')).toBe( - 'application/x-sqlite3' - ) + expect(response.headers.get('Content-Type')).toBe('text/sql') expect(response.headers.get('Content-Disposition')).toBe( 'attachment; filename="database_dump.sql"' ) - const dumpText = await response.text() + const dumpText = await collectStream(response) + expect(dumpText).toContain('BEGIN TRANSACTION;') + expect(dumpText).toContain('COMMIT;') 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 () => { @@ -86,11 +90,11 @@ describe('Database Dump Module', () => { const response = await dumpDatabaseRoute(mockDataSource, mockConfig) expect(response).toBeInstanceOf(Response) - expect(response.headers.get('Content-Type')).toBe( - 'application/x-sqlite3' - ) - const dumpText = await response.text() - expect(dumpText).toBe('SQLite format 3\0') + expect(response.headers.get('Content-Type')).toBe('text/sql') + const dumpText = await collectStream(response) + expect(dumpText).toContain('BEGIN TRANSACTION;') + expect(dumpText).toContain('COMMIT;') + expect(dumpText).not.toContain('INSERT INTO') }) it('should handle databases with tables but no data', async () => { @@ -104,11 +108,11 @@ describe('Database Dump Module', () => { const response = await dumpDatabaseRoute(mockDataSource, mockConfig) expect(response).toBeInstanceOf(Response) - const dumpText = await response.text() + const dumpText = await collectStream(response) 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 () => { @@ -122,12 +126,55 @@ describe('Database Dump Module', () => { const response = await dumpDatabaseRoute(mockDataSource, mockConfig) expect(response).toBeInstanceOf(Response) - const dumpText = await response.text() + const dumpText = await collectStream(response) expect(dumpText).toContain( - "INSERT INTO users VALUES (1, 'Alice''s adventure');" + "INSERT INTO \"users\" VALUES (1, 'Alice''s adventure');" ) }) + it('should handle NULL values', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + ]) + .mockResolvedValueOnce([{ id: 1, name: null }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await collectStream(response) + expect(dumpText).toContain('INSERT INTO "users" VALUES (1, NULL);') + }) + + it('should handle boolean values', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'flags' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE flags (id INTEGER, active INTEGER);' }, + ]) + .mockResolvedValueOnce([ + { id: 1, active: true }, + { id: 2, active: false }, + ]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await collectStream(response) + expect(dumpText).toContain('INSERT INTO "flags" VALUES (1, 1);') + expect(dumpText).toContain('INSERT INTO "flags" VALUES (2, 0);') + }) + + it('should escape table names with special characters', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'my-table' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE "my-table" (id INTEGER);' }, + ]) + .mockResolvedValueOnce([{ id: 1 }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await collectStream(response) + expect(dumpText).toContain('INSERT INTO "my-table" VALUES (1);') + }) + it('should return a 500 response when an error occurs', async () => { const consoleErrorMock = vi .spyOn(console, 'error') @@ -141,5 +188,39 @@ describe('Database Dump Module', () => { expect(response.status).toBe(500) const jsonResponse: { error: string } = await response.json() expect(jsonResponse.error).toBe('Failed to create database dump') + consoleErrorMock.mockRestore() + }) + + it('should fetch data in batches for large tables', async () => { + const firstBatch = Array.from({ length: 500 }, (_, i) => ({ + id: i + 1, + name: `user${i + 1}`, + })) + const secondBatch = Array.from({ length: 100 }, (_, i) => ({ + id: i + 501, + name: `user${i + 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 collectStream(response) + + expect(dumpText).toContain('INSERT INTO "users" VALUES (1, \'user1\');') + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (500, \'user500\');' + ) + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (501, \'user501\');' + ) + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (600, \'user600\');' + ) }) }) diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..b1e1862 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -3,67 +3,170 @@ import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +const BATCH_SIZE = 500 + +function escapeSqlIdentifier(name: string): string { + return `"${name.replace(/"/g, '""')}"` +} + +function escapeSqlValue(value: unknown): string { + if (value === null || value === undefined) { + return 'NULL' + } + if (typeof value === 'number') { + return String(value) + } + if (typeof value === 'boolean') { + return value ? '1' : '0' + } + if (value instanceof ArrayBuffer || ArrayBuffer.isView(value)) { + const bytes = new Uint8Array( + value instanceof ArrayBuffer ? value : value.buffer + ) + const hex = Array.from(bytes) + .map((b) => b.toString(16).padStart(2, '0')) + .join('') + return `X'${hex}'` + } + return `'${String(value).replace(/'/g, "''")}'` +} + +async function* streamTableData( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): AsyncGenerator { + const escapedName = escapeSqlIdentifier(tableName) + + const schemaResult = await executeOperation( + [ + { + sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name=?;`, + params: [tableName], + }, + ], + dataSource, + config + ) + + if (schemaResult.length) { + const schema = schemaResult[0].sql + yield `\n-- Table: ${tableName}\n${schema};\n\n` + } + + let offset = 0 + let hasMore = true + + while (hasMore) { + const dataResult = await executeOperation( + [ + { + sql: `SELECT * FROM ${escapedName} LIMIT ? OFFSET ?;`, + params: [BATCH_SIZE, offset], + }, + ], + dataSource, + config + ) + + if (!dataResult || dataResult.length === 0) { + hasMore = false + continue + } + + for (const row of dataResult) { + const values = Object.values(row).map((v) => escapeSqlValue(v)) + yield `INSERT INTO ${escapedName} VALUES (${values.join(', ')});\n` + } + + if (dataResult.length < BATCH_SIZE) { + hasMore = false + } else { + offset += BATCH_SIZE + } + } + + yield '\n' +} + +function createSqlDumpStream( + tableNames: string[], + dataSource: DataSource, + config: StarbaseDBConfiguration +): ReadableStream { + const encoder = new TextEncoder() + let tableIterator: AsyncGenerator | null = null + let tableIndex = 0 + let chunkBuffer = 'BEGIN TRANSACTION;\n' + + return new ReadableStream({ + async pull(controller) { + try { + while (tableIndex < tableNames.length) { + if (!tableIterator) { + tableIterator = streamTableData( + tableNames[tableIndex], + dataSource, + config + ) + } + + const result = await tableIterator.next() + if (result.done) { + tableIterator = null + tableIndex++ + continue + } + + chunkBuffer += result.value + + if (chunkBuffer.length >= 8192) { + controller.enqueue(encoder.encode(chunkBuffer)) + chunkBuffer = '' + return + } + } + + chunkBuffer += 'COMMIT;\n' + if (chunkBuffer.length > 0) { + controller.enqueue(encoder.encode(chunkBuffer)) + chunkBuffer = '' + } + controller.close() + } catch (error: any) { + console.error('Database Dump Stream Error:', error) + controller.error(error) + } + }, + }) +} + export async function dumpDatabaseRoute( dataSource: DataSource, config: StarbaseDBConfiguration ): Promise { try { - // Get all table names const tablesResult = await executeOperation( - [{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }], + [ + { + sql: "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'tmp_%' ORDER BY name;", + }, + ], dataSource, config ) - const tables = tablesResult.map((row: any) => row.name) - let dumpContent = 'SQLite format 3\0' // SQLite file header - - // Iterate through all tables - for (const table of tables) { - // Get table schema - const schemaResult = await executeOperation( - [ - { - sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name='${table}';`, - }, - ], - dataSource, - config - ) - - if (schemaResult.length) { - const schema = schemaResult[0].sql - dumpContent += `\n-- Table: ${table}\n${schema};\n\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` - } - - dumpContent += '\n' - } + const tableNames = tablesResult.map((row: any) => row.name) - // Create a Blob from the dump content - const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' }) + const stream = createSqlDumpStream(tableNames, dataSource, config) const headers = new Headers({ - 'Content-Type': 'application/x-sqlite3', - 'Content-Disposition': 'attachment; filename="database_dump.sql"', + 'Content-Type': 'text/sql', + 'Content-Disposition': `attachment; filename="database_dump.sql"`, + 'Transfer-Encoding': 'chunked', }) - 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)