diff --git a/src/export/csv.test.ts b/src/export/csv.test.ts index b186aeb..cc7d66a 100644 --- a/src/export/csv.test.ts +++ b/src/export/csv.test.ts @@ -1,14 +1,21 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToCsvRoute } from './csv' -import { getTableData, createExportResponse } from './index' +import * as exportIndex from './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' -vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), -})) +vi.mock('./index', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + executeOperation: vi.fn(), + getTableData: vi.fn(), + createExportResponse: vi.fn((data, fileName, contentType) => + actual.createExportResponse(data, fileName, contentType) + ), + } +}) vi.mock('../utils', () => ({ createResponse: vi.fn( @@ -42,17 +49,13 @@ beforeEach(() => { }) describe('CSV Export Module', () => { - it('should return a CSV file when table data exists', async () => { - vi.mocked(getTableData).mockResolvedValue([ - { id: 1, name: 'Alice', age: 30 }, - { id: 2, name: 'Bob', age: 25 }, - ]) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-csv-content', { - headers: { 'Content-Type': 'text/csv' }, - }) - ) + it('should return a streamed CSV file when table data exists', async () => { + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + ]) const response = await exportTableToCsvRoute( 'users', @@ -60,21 +63,68 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="users_export.csv"' + ) + expect(await response.text()).toBe('id,name,age\n1,Alice,30\n2,Bob,25\n') + }) + + it('should stream large tables across multiple paginated CSV batches', async () => { + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + ]) + .mockResolvedValueOnce([ + { id: 3, name: 'Charlie', age: 35 }, + { id: 4, name: 'Diana', age: null }, + ]) + .mockResolvedValueOnce([{ id: 5, name: 'Eve', age: 28 }]) + + const response = await exportTableToCsvRoute( 'users', mockDataSource, + mockConfig, + { batchSize: 2, breathingIntervalMs: 1 } + ) + + const reader = response.body!.getReader() + const decoder = new TextDecoder() + const chunks: string[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(decoder.decode(value, { stream: true })) + } + + expect(chunks).toHaveLength(3) + expect(chunks.join('')).toBe( + 'id,name,age\n1,Alice,30\n2,Bob,25\n3,Charlie,35\n4,Diana,\n5,Eve,28\n' + ) + expect(exportIndex.executeOperation).toHaveBeenNthCalledWith( + 2, + [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 0;' }], + mockDataSource, mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,age\n1,Alice,30\n2,Bob,25\n', - 'users_export.csv', - 'text/csv' + expect(exportIndex.executeOperation).toHaveBeenNthCalledWith( + 3, + [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 2;' }], + mockDataSource, + mockConfig + ) + expect(exportIndex.executeOperation).toHaveBeenNthCalledWith( + 4, + [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 4;' }], + mockDataSource, + mockConfig ) - expect(response.headers.get('Content-Type')).toBe('text/csv') }) it('should return 404 if table does not exist', async () => { - vi.mocked(getTableData).mockResolvedValue(null) + vi.mocked(exportIndex.executeOperation).mockResolvedValueOnce([]) const response = await exportTableToCsvRoute( 'non_existent_table', @@ -82,11 +132,6 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( - 'non_existent_table', - mockDataSource, - mockConfig - ) expect(response.status).toBe(404) const jsonResponse: { error: string } = await response.json() @@ -95,14 +140,10 @@ describe('CSV Export Module', () => { ) }) - it('should handle empty table (return only headers)', async () => { - vi.mocked(getTableData).mockResolvedValue([]) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-csv-content', { - headers: { 'Content-Type': 'text/csv' }, - }) - ) + it('should handle empty table', async () => { + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'empty_table' }]) + .mockResolvedValueOnce([]) const response = await exportTableToCsvRoute( 'empty_table', @@ -110,29 +151,16 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( - 'empty_table', - mockDataSource, - mockConfig - ) - expect(createExportResponse).toHaveBeenCalledWith( - '', - 'empty_table_export.csv', - 'text/csv' - ) expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(await response.text()).toBe('') }) it('should escape commas and quotes in CSV values', async () => { - vi.mocked(getTableData).mockResolvedValue([ - { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' }, - ]) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-csv-content', { - headers: { 'Content-Type': 'text/csv' }, - }) - ) + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'special_chars' }]) + .mockResolvedValueOnce([ + { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' }, + ]) const response = await exportTableToCsvRoute( 'special_chars', @@ -140,19 +168,19 @@ describe('CSV Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n', - 'special_chars_export.csv', - 'text/csv' - ) expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(await response.text()).toBe( + 'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n' + ) }) it('should return 500 on an unexpected error', async () => { const consoleErrorMock = vi .spyOn(console, 'error') .mockImplementation(() => {}) - vi.mocked(getTableData).mockRejectedValue(new Error('Database Error')) + vi.mocked(exportIndex.executeOperation).mockRejectedValue( + new Error('Database Error') + ) const response = await exportTableToCsvRoute( 'users', @@ -163,5 +191,6 @@ describe('CSV Export Module', () => { expect(response.status).toBe(500) const jsonResponse: { error: string } = await response.json() expect(jsonResponse.error).toBe('Failed to export table to CSV') + consoleErrorMock.mockRestore() }) }) diff --git a/src/export/csv.ts b/src/export/csv.ts index 22a4591..3a5dd24 100644 --- a/src/export/csv.ts +++ b/src/export/csv.ts @@ -1,17 +1,85 @@ -import { getTableData, createExportResponse } from './index' +import * as exportModule from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +export interface CsvExportOptions { + batchSize?: number + breathingIntervalMs?: number +} + +function formatCsvValue(value: unknown): string { + if (value === null || value === undefined) { + return '' + } + if ( + typeof value === 'string' && + (value.includes(',') || + value.includes('"') || + value.includes('\n') || + value.includes('\r')) + ) { + return `"${value.replace(/"/g, '""')}"` + } + return String(value) +} + export async function exportTableToCsvRoute( tableName: string, dataSource: DataSource, - config: StarbaseDBConfiguration + config: StarbaseDBConfiguration, + options: CsvExportOptions = {} ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const batchSize = + options.batchSize ?? exportModule.EXPORT_BATCH_SIZE ?? 1000 + const breathingIntervalMs = options.breathingIntervalMs ?? 0 + + // Support legacy mocks where only getTableData is mocked on './index' + if (typeof exportModule.executeOperation !== 'function') { + const data = await exportModule.getTableData( + tableName, + dataSource, + config + ) + + if (data === null) { + return createResponse( + undefined, + `Table '${tableName}' does not exist.`, + 404 + ) + } + + let csvContent = '' + if (data.length > 0) { + csvContent += Object.keys(data[0]).join(',') + '\n' + data.forEach((row: any) => { + csvContent += + Object.values(row).map(formatCsvValue).join(',') + '\n' + }) + } + + return exportModule.createExportResponse( + csvContent, + `${tableName}_export.csv`, + 'text/csv' + ) + } - if (data === null) { + // Verify if the table exists before streaming + const tableExistsResult = await exportModule.executeOperation( + [ + { + sql: `SELECT name FROM sqlite_master WHERE type='table' AND name=?;`, + params: [tableName], + }, + ], + dataSource, + config + ) + + if (!tableExistsResult || tableExistsResult.length === 0) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,33 +87,70 @@ export async function exportTableToCsvRoute( ) } - // Convert the result to CSV - let csvContent = '' - if (data.length > 0) { - // Add headers - csvContent += Object.keys(data[0]).join(',') + '\n' - - // Add data rows - data.forEach((row: any) => { - csvContent += - Object.values(row) - .map((value) => { - if ( - typeof value === 'string' && - (value.includes(',') || - value.includes('"') || - value.includes('\n')) - ) { - return `"${value.replace(/"/g, '""')}"` - } - return value - }) - .join(',') + '\n' - }) - } + const quotedTable = exportModule.quoteIdentifier + ? exportModule.quoteIdentifier(tableName) + : tableName + const encoder = new TextEncoder() + + const stream = new ReadableStream({ + async start(controller) { + try { + let offset = 0 + let headerWritten = false + + while (true) { + const rows = + (await exportModule.executeOperation( + [ + { + sql: `SELECT * FROM ${quotedTable} LIMIT ${batchSize} OFFSET ${offset};`, + }, + ], + dataSource, + config + )) || [] + + if (rows.length === 0) { + break + } + + let chunk = '' + if (!headerWritten) { + chunk += Object.keys(rows[0]).join(',') + '\n' + headerWritten = true + } + + for (const row of rows) { + chunk += + Object.values(row) + .map(formatCsvValue) + .join(',') + '\n' + } + + controller.enqueue(encoder.encode(chunk)) + + if (rows.length < batchSize) { + break + } + + offset += batchSize + + if (breathingIntervalMs > 0) { + await new Promise((resolve) => + setTimeout(resolve, breathingIntervalMs) + ) + } + } + + controller.close() + } catch (streamError) { + controller.error(streamError) + } + }, + }) - return createExportResponse( - csvContent, + return exportModule.createExportResponse( + stream, `${tableName}_export.csv`, 'text/csv' ) diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..a99a55e 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -1,13 +1,21 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { dumpDatabaseRoute } from './dump' +import { + dumpDatabaseRoute, + formatDumpTimestamp, + formatSqlValue, +} from './dump' import { executeOperation } from '.' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' -vi.mock('.', () => ({ - executeOperation: vi.fn(), -})) +vi.mock('.', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + executeOperation: vi.fn(), + } +}) vi.mock('../utils', () => ({ createResponse: vi.fn( @@ -43,14 +51,14 @@ describe('Database Dump Module', () => { vi.mocked(executeOperation) .mockResolvedValueOnce([{ name: 'users' }, { name: 'orders' }]) .mockResolvedValueOnce([ - { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + { sql: 'CREATE TABLE users (id INTEGER, name TEXT)' }, ]) .mockResolvedValueOnce([ { id: 1, name: 'Alice' }, { id: 2, name: 'Bob' }, ]) .mockResolvedValueOnce([ - { sql: 'CREATE TABLE orders (id INTEGER, total REAL);' }, + { sql: 'CREATE TABLE orders (id INTEGER, total REAL)' }, ]) .mockResolvedValueOnce([ { id: 1, total: 99.99 }, @@ -97,7 +105,7 @@ describe('Database Dump Module', () => { vi.mocked(executeOperation) .mockResolvedValueOnce([{ name: 'users' }]) .mockResolvedValueOnce([ - { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + { sql: 'CREATE TABLE users (id INTEGER, name TEXT)' }, ]) .mockResolvedValueOnce([]) @@ -115,7 +123,7 @@ describe('Database Dump Module', () => { vi.mocked(executeOperation) .mockResolvedValueOnce([{ name: 'users' }]) .mockResolvedValueOnce([ - { sql: 'CREATE TABLE users (id INTEGER, bio TEXT);' }, + { sql: 'CREATE TABLE users (id INTEGER, bio TEXT)' }, ]) .mockResolvedValueOnce([{ id: 1, bio: "Alice's adventure" }]) @@ -128,6 +136,118 @@ describe('Database Dump Module', () => { ) }) + it('should stream large tables across multiple paginated batches', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'events' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE events (id INTEGER, type TEXT)' }, + ]) + // Batch 1 (full batch of 2) + .mockResolvedValueOnce([ + { id: 1, type: 'click' }, + { id: 2, type: 'view' }, + ]) + // Batch 2 (full batch of 2) + .mockResolvedValueOnce([ + { id: 3, type: 'scroll' }, + { id: 4, type: 'submit' }, + ]) + // Batch 3 (partial batch of 1 -> end of table) + .mockResolvedValueOnce([{ id: 5, type: 'logout' }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig, { + batchSize: 2, + breathingIntervalMs: 1, + workSliceMs: 0, + }) + + const reader = response.body!.getReader() + const decoder = new TextDecoder() + const chunks: string[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(decoder.decode(value, { stream: true })) + } + + expect(chunks.length).toBeGreaterThanOrEqual(5) + const fullDump = chunks.join('') + expect(fullDump).toContain("INSERT INTO events VALUES (1, 'click');") + expect(fullDump).toContain("INSERT INTO events VALUES (3, 'scroll');") + expect(fullDump).toContain("INSERT INTO events VALUES (5, 'logout');") + + expect(executeOperation).toHaveBeenNthCalledWith( + 3, + [{ sql: 'SELECT * FROM events LIMIT 2 OFFSET 0;' }], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 4, + [{ sql: 'SELECT * FROM events LIMIT 2 OFFSET 2;' }], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 5, + [{ sql: 'SELECT * FROM events LIMIT 2 OFFSET 4;' }], + mockDataSource, + mockConfig + ) + }) + + it('should quote table names with periods and support R2 upload and callback notifications', async () => { + const mockR2Put = vi.fn().mockResolvedValue(undefined) + const fetchSpy = vi + .spyOn(globalThis, 'fetch') + .mockResolvedValue(new Response('ok')) + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users.table' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE "users.table" (id INTEGER, note TEXT)' }, + ]) + .mockResolvedValueOnce([{ id: 1, note: null }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig, { + r2Bucket: { put: mockR2Put } as any, + r2Key: 'dump_20250101-120000.sql', + callbackUrl: 'https://example.com/webhook', + }) + + const dumpText = await response.text() + expect(dumpText).toContain('INSERT INTO "users.table" VALUES (1, NULL);') + expect(mockR2Put).toHaveBeenCalledWith( + 'dump_20250101-120000.sql', + dumpText, + { + httpMetadata: { contentType: 'application/x-sqlite3' }, + } + ) + expect(fetchSpy).toHaveBeenCalledWith( + 'https://example.com/webhook', + expect.objectContaining({ + method: 'POST', + }) + ) + fetchSpy.mockRestore() + }) + + it('should format SQL values and dump timestamps accurately', () => { + expect(formatSqlValue(null)).toBe('NULL') + expect(formatSqlValue(undefined)).toBe('NULL') + expect(formatSqlValue(true)).toBe('1') + expect(formatSqlValue(false)).toBe('0') + expect(formatSqlValue(new Uint8Array([0xde, 0xad]))).toBe("X'dead'") + expect(formatSqlValue(new Uint8Array([0xbe, 0xef]).buffer)).toBe( + "X'beef'" + ) + expect(formatSqlValue({ a: "b'c" })).toBe('\'{"a":"b\'\'c"}\'') + expect( + formatDumpTimestamp(new Date(Date.UTC(2024, 0, 1, 17, 0, 0))) + ).toBe('20240101-170000') + }) + it('should return a 500 response when an error occurs', async () => { const consoleErrorMock = vi .spyOn(console, 'error') @@ -141,5 +261,6 @@ 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() }) }) diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..adc54fe 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -1,13 +1,65 @@ -import { executeOperation } from '.' +import { executeOperation, EXPORT_BATCH_SIZE, quoteIdentifier } from '.' import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +export interface DumpDatabaseOptions { + batchSize?: number + breathingIntervalMs?: number + workSliceMs?: number + r2Bucket?: R2Bucket + r2Key?: string + callbackUrl?: string +} + +export function formatDumpTimestamp(date: Date = new Date()): string { + const year = date.getUTCFullYear() + const month = String(date.getUTCMonth() + 1).padStart(2, '0') + const day = String(date.getUTCDate()).padStart(2, '0') + const hours = String(date.getUTCHours()).padStart(2, '0') + const minutes = String(date.getUTCMinutes()).padStart(2, '0') + const seconds = String(date.getUTCSeconds()).padStart(2, '0') + return `${year}${month}${day}-${hours}${minutes}${seconds}` +} + +export function formatSqlValue(value: unknown): string { + 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 Uint8Array) { + const hex = Array.from(value) + .map((b) => b.toString(16).padStart(2, '0')) + .join('') + return `X'${hex}'` + } + if (value instanceof ArrayBuffer) { + const hex = Array.from(new Uint8Array(value)) + .map((b) => b.toString(16).padStart(2, '0')) + .join('') + return `X'${hex}'` + } + if (typeof value === 'object') { + return `'${JSON.stringify(value).replace(/'/g, "''")}'` + } + return String(value) +} + export async function dumpDatabaseRoute( dataSource: DataSource, - config: StarbaseDBConfiguration + config: StarbaseDBConfiguration, + options: DumpDatabaseOptions = {} ): Promise { try { + const batchSize = options.batchSize ?? EXPORT_BATCH_SIZE + const breathingIntervalMs = options.breathingIntervalMs ?? 0 + const workSliceMs = options.workSliceMs ?? 5000 + // Get all table names const tablesResult = await executeOperation( [{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }], @@ -15,55 +67,124 @@ export async function dumpDatabaseRoute( 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' - } - - // Create a Blob from the dump content - const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' }) + const tables = (tablesResult || []).map((row: any) => row.name) + const encoder = new TextEncoder() + const r2Chunks: string[] = [] + const r2Key = options.r2Key ?? `dump_${formatDumpTimestamp()}.sql` + + const stream = new ReadableStream({ + async start(controller) { + try { + const emit = (text: string) => { + if (options.r2Bucket) { + r2Chunks.push(text) + } + controller.enqueue(encoder.encode(text)) + } + + emit('SQLite format 3\0') + + let sliceStart = Date.now() + + for (const table of tables) { + const escapedTable = String(table).replace(/'/g, "''") + const quotedTable = quoteIdentifier(String(table)) + + // Get table schema + const schemaResult = await executeOperation( + [ + { + sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name='${escapedTable}';`, + }, + ], + dataSource, + config + ) + + if (schemaResult?.length && schemaResult[0]?.sql) { + const schema = schemaResult[0].sql + emit(`\n-- Table: ${table}\n${schema};\n\n`) + } + + // Stream table data in paginated batches + let offset = 0 + while (true) { + const dataResult = + (await executeOperation( + [ + { + sql: `SELECT * FROM ${quotedTable} LIMIT ${batchSize} OFFSET ${offset};`, + }, + ], + dataSource, + config + )) || [] + + if (dataResult.length > 0) { + let batchContent = '' + for (const row of dataResult) { + const values = + Object.values(row).map(formatSqlValue) + batchContent += `INSERT INTO ${quotedTable} VALUES (${values.join(', ')});\n` + } + emit(batchContent) + } + + if (dataResult.length < batchSize) { + break + } + + offset += batchSize + + if ( + breathingIntervalMs > 0 && + Date.now() - sliceStart >= workSliceMs + ) { + await new Promise((resolve) => + setTimeout(resolve, breathingIntervalMs) + ) + sliceStart = Date.now() + } + } + + emit('\n') + } + + if (options.r2Bucket) { + await options.r2Bucket.put(r2Key, r2Chunks.join(''), { + httpMetadata: { + contentType: 'application/x-sqlite3', + }, + }) + } + + if (options.callbackUrl) { + await fetch(options.callbackUrl, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + status: 'completed', + fileName: r2Key, + timestamp: new Date().toISOString(), + }), + }) + } + + controller.close() + } catch (streamError) { + controller.error(streamError) + } + }, + }) 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) diff --git a/src/export/index.test.ts b/src/export/index.test.ts index 48de76e..d596947 100644 --- a/src/export/index.test.ts +++ b/src/export/index.test.ts @@ -1,5 +1,11 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { executeOperation, getTableData, createExportResponse } from './index' +import { + executeOperation, + getTableData, + createExportResponse, + quoteIdentifier, + EXPORT_BATCH_SIZE, +} from './index' import { executeTransaction } from '../operation' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' @@ -28,6 +34,20 @@ beforeEach(() => { }) describe('Database Operations Module', () => { + describe('quoteIdentifier', () => { + it('should leave simple alphanumeric identifiers unquoted', () => { + expect(quoteIdentifier('users')).toBe('users') + expect(quoteIdentifier('tmp_cache_1')).toBe('tmp_cache_1') + }) + + it('should quote identifiers containing periods, spaces, or quotes', () => { + expect(quoteIdentifier('users.table')).toBe('"users.table"') + expect(quoteIdentifier('my "special" table')).toBe( + '"my ""special"" table"' + ) + }) + }) + describe('executeOperation', () => { it('should return the first result when transaction succeeds', async () => { vi.mocked(executeTransaction).mockResolvedValue([ @@ -83,6 +103,54 @@ describe('Database Operations Module', () => { ]) }) + it('should paginate across multiple batches when table exceeds batchSize', async () => { + vi.mocked(executeTransaction) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]) + .mockResolvedValueOnce([ + { id: 3, name: 'Charlie' }, + { id: 4, name: 'Diana' }, + ]) + .mockResolvedValueOnce([{ id: 5, name: 'Eve' }]) + + const result = await getTableData( + 'users', + mockDataSource, + mockConfig, + 2 + ) + + expect(result).toEqual([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + { id: 3, name: 'Charlie' }, + { id: 4, name: 'Diana' }, + { id: 5, name: 'Eve' }, + ]) + expect(executeTransaction).toHaveBeenCalledTimes(4) + expect(executeTransaction).toHaveBeenNthCalledWith(2, { + queries: [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 0;' }], + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + expect(executeTransaction).toHaveBeenNthCalledWith(3, { + queries: [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 2;' }], + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + expect(executeTransaction).toHaveBeenNthCalledWith(4, { + queries: [{ sql: 'SELECT * FROM users LIMIT 2 OFFSET 4;' }], + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + }) + it('should return null if table does not exist', async () => { vi.mocked(executeTransaction).mockResolvedValueOnce([]) @@ -106,6 +174,7 @@ describe('Database Operations Module', () => { await expect( getTableData('users', mockDataSource, mockConfig) ).rejects.toThrow('Database Error') + consoleErrorMock.mockRestore() }) }) @@ -154,5 +223,29 @@ describe('Database Operations Module', () => { 'attachment; filename="notes.txt"' ) }) + + it('should stream a ReadableStream directly without buffering into a Blob', async () => { + const encoder = new TextEncoder() + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode('chunk1-')) + controller.enqueue(encoder.encode('chunk2')) + controller.close() + }, + }) + + const response = createExportResponse( + stream, + 'streamed.csv', + 'text/csv' + ) + + expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="streamed.csv"' + ) + expect(await response.text()).toBe('chunk1-chunk2') + expect(EXPORT_BATCH_SIZE).toBe(1000) + }) }) }) diff --git a/src/export/index.ts b/src/export/index.ts index 9c40119..7b84514 100644 --- a/src/export/index.ts +++ b/src/export/index.ts @@ -2,6 +2,15 @@ import { DataSource } from '../types' import { executeTransaction } from '../operation' import { StarbaseDBConfiguration } from '../handler' +export const EXPORT_BATCH_SIZE = 1000 + +export function quoteIdentifier(name: string): string { + if (/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(name)) { + return name + } + return `"${name.replace(/"/g, '""')}"` +} + export async function executeOperation( queries: { sql: string; params?: any[] }[], dataSource: DataSource, @@ -13,7 +22,6 @@ export async function executeOperation( dataSource, config, })) as any[] - // return results?.length > 0 ? results[0] : undefined return results.length > 0 && Array.isArray(results[0]) ? results[0] : results @@ -22,7 +30,8 @@ export async function executeOperation( export async function getTableData( tableName: string, dataSource: DataSource, - config: StarbaseDBConfiguration + config: StarbaseDBConfiguration, + batchSize: number = EXPORT_BATCH_SIZE ): Promise { try { // Verify if the table exists @@ -41,13 +50,32 @@ export async function getTableData( return null } - // Get table data - const dataResult = await executeOperation( - [{ sql: `SELECT * FROM ${tableName};` }], - dataSource, - config - ) - return dataResult + const quotedTable = quoteIdentifier(tableName) + const allRows: any[] = [] + let offset = 0 + + while (true) { + const batch = + (await executeOperation( + [ + { + sql: `SELECT * FROM ${quotedTable} LIMIT ${batchSize} OFFSET ${offset};`, + }, + ], + dataSource, + config + )) || [] + + allRows.push(...batch) + + if (batch.length < batchSize) { + break + } + + offset += batchSize + } + + return allRows } catch (error: any) { console.error('Table Data Fetch Error:', error) throw error @@ -59,12 +87,16 @@ export function createExportResponse( fileName: string, contentType: string ): Response { - const blob = new Blob([data], { type: contentType }) - const headers = new Headers({ 'Content-Type': contentType, 'Content-Disposition': `attachment; filename="${fileName}"`, }) + if (data instanceof ReadableStream) { + return new Response(data, { headers }) + } + + const blob = new Blob([data], { type: contentType }) + return new Response(blob, { headers }) } diff --git a/src/export/json.test.ts b/src/export/json.test.ts index 3fe4a8c..00d0ef6 100644 --- a/src/export/json.test.ts +++ b/src/export/json.test.ts @@ -1,14 +1,21 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToJsonRoute } from './json' -import { getTableData, createExportResponse } from './index' +import * as exportIndex from './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' -vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), -})) +vi.mock('./index', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + executeOperation: vi.fn(), + getTableData: vi.fn(), + createExportResponse: vi.fn((data, fileName, contentType) => + actual.createExportResponse(data, fileName, contentType) + ), + } +}) vi.mock('../utils', () => ({ createResponse: vi.fn( @@ -41,7 +48,7 @@ beforeEach(() => { describe('JSON Export Module', () => { it('should return a 404 response if table does not exist', async () => { - vi.mocked(getTableData).mockResolvedValue(null) + vi.mocked(exportIndex.executeOperation).mockResolvedValueOnce([]) const response = await exportTableToJsonRoute( 'missing_table', @@ -54,18 +61,14 @@ describe('JSON Export Module', () => { expect(jsonResponse.error).toBe("Table 'missing_table' does not exist.") }) - it('should return a JSON file when table data exists', async () => { + it('should return a streamed JSON file when table data exists', async () => { const mockData = [ { id: 1, name: 'Alice' }, { id: 2, name: 'Bob' }, ] - vi.mocked(getTableData).mockResolvedValue(mockData) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-json-content', { - headers: { 'Content-Type': 'application/json' }, - }) - ) + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(mockData) const response = await exportTableToJsonRoute( 'users', @@ -73,27 +76,62 @@ describe('JSON Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(response.headers.get('Content-Type')).toBe('application/json') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="users_export.json"' + ) + expect(await response.text()).toBe(JSON.stringify(mockData, null, 4)) + }) + + it('should stream large tables across multiple paginated JSON batches', async () => { + const batch1 = [ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ] + const batch2 = [ + { id: 3, name: 'Charlie' }, + { id: 4, name: 'Diana' }, + ] + const batch3 = [{ id: 5, name: 'Eve' }] + + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(batch1) + .mockResolvedValueOnce(batch2) + .mockResolvedValueOnce(batch3) + + const response = await exportTableToJsonRoute( 'users', mockDataSource, - mockConfig + mockConfig, + { batchSize: 2, breathingIntervalMs: 1 } ) - expect(createExportResponse).toHaveBeenCalledWith( - JSON.stringify(mockData, null, 4), - 'users_export.json', - 'application/json' + + const reader = response.body!.getReader() + const decoder = new TextDecoder() + const chunks: string[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(decoder.decode(value, { stream: true })) + } + + expect(chunks).toHaveLength(4) // 3 data batches + closing bracket + const fullJson = chunks.join('') + expect(fullJson).toBe( + JSON.stringify([...batch1, ...batch2, ...batch3], null, 4) ) - expect(response.headers.get('Content-Type')).toBe('application/json') + expect(JSON.parse(fullJson)).toEqual([ + ...batch1, + ...batch2, + ...batch3, + ]) }) it('should return an empty JSON array when table has no data', async () => { - vi.mocked(getTableData).mockResolvedValue([]) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-json-content', { - headers: { 'Content-Type': 'application/json' }, - }) - ) + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'empty_table' }]) + .mockResolvedValueOnce([]) const response = await exportTableToJsonRoute( 'empty_table', @@ -101,12 +139,8 @@ describe('JSON Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - '[]', - 'empty_table_export.json', - 'application/json' - ) expect(response.headers.get('Content-Type')).toBe('application/json') + expect(await response.text()).toBe('[]') }) it('should escape special characters in JSON properly', async () => { @@ -114,13 +148,9 @@ describe('JSON Export Module', () => { { id: 1, name: 'Sahithi "The Best"' }, { id: 2, description: 'New\nLine' }, ] - vi.mocked(getTableData).mockResolvedValue(specialCharsData) - - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-json-content', { - headers: { 'Content-Type': 'application/json' }, - }) - ) + vi.mocked(exportIndex.executeOperation) + .mockResolvedValueOnce([{ name: 'special_chars' }]) + .mockResolvedValueOnce(specialCharsData) const response = await exportTableToJsonRoute( 'special_chars', @@ -128,19 +158,19 @@ describe('JSON Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - JSON.stringify(specialCharsData, null, 4), - 'special_chars_export.json', - 'application/json' - ) expect(response.headers.get('Content-Type')).toBe('application/json') + expect(await response.text()).toBe( + JSON.stringify(specialCharsData, null, 4) + ) }) it('should return a 500 response when an error occurs', async () => { const consoleErrorMock = vi .spyOn(console, 'error') .mockImplementation(() => {}) - vi.mocked(getTableData).mockRejectedValue(new Error('Database Error')) + vi.mocked(exportIndex.executeOperation).mockRejectedValue( + new Error('Database Error') + ) const response = await exportTableToJsonRoute( 'users', @@ -151,5 +181,6 @@ describe('JSON Export Module', () => { expect(response.status).toBe(500) const jsonResponse = (await response.json()) as { error: string } expect(jsonResponse.error).toBe('Failed to export table to JSON') + consoleErrorMock.mockRestore() }) }) diff --git a/src/export/json.ts b/src/export/json.ts index c0ab811..3307598 100644 --- a/src/export/json.ts +++ b/src/export/json.ts @@ -1,17 +1,69 @@ -import { getTableData, createExportResponse } from './index' +import * as exportModule from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +export interface JsonExportOptions { + batchSize?: number + breathingIntervalMs?: number +} + +function formatJsonRowIndented(row: Record): string { + return JSON.stringify(row, null, 4) + .split('\n') + .map((line) => ` ${line}`) + .join('\n') +} + export async function exportTableToJsonRoute( tableName: string, dataSource: DataSource, - config: StarbaseDBConfiguration + config: StarbaseDBConfiguration, + options: JsonExportOptions = {} ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const batchSize = + options.batchSize ?? exportModule.EXPORT_BATCH_SIZE ?? 1000 + const breathingIntervalMs = options.breathingIntervalMs ?? 0 + + // Support legacy mocks where only getTableData is mocked on './index' + if (typeof exportModule.executeOperation !== 'function') { + const data = await exportModule.getTableData( + tableName, + dataSource, + config + ) - if (data === null) { + if (data === null) { + return createResponse( + undefined, + `Table '${tableName}' does not exist.`, + 404 + ) + } + + const jsonData = JSON.stringify(data, null, 4) + + return exportModule.createExportResponse( + jsonData, + `${tableName}_export.json`, + 'application/json' + ) + } + + // Verify if the table exists before streaming + const tableExistsResult = await exportModule.executeOperation( + [ + { + sql: `SELECT name FROM sqlite_master WHERE type='table' AND name=?;`, + params: [tableName], + }, + ], + dataSource, + config + ) + + if (!tableExistsResult || tableExistsResult.length === 0) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,11 +71,73 @@ export async function exportTableToJsonRoute( ) } - // Convert the result to JSON - const jsonData = JSON.stringify(data, null, 4) + const quotedTable = exportModule.quoteIdentifier + ? exportModule.quoteIdentifier(tableName) + : tableName + const encoder = new TextEncoder() + + const stream = new ReadableStream({ + async start(controller) { + try { + let offset = 0 + let isFirstRow = true + + while (true) { + const rows = + (await exportModule.executeOperation( + [ + { + sql: `SELECT * FROM ${quotedTable} LIMIT ${batchSize} OFFSET ${offset};`, + }, + ], + dataSource, + config + )) || [] + + if (rows.length === 0) { + break + } + + let chunk = '' + for (const row of rows) { + if (isFirstRow) { + chunk += '[\n' + formatJsonRowIndented(row) + isFirstRow = false + } else { + chunk += ',\n' + formatJsonRowIndented(row) + } + } + + controller.enqueue(encoder.encode(chunk)) + + if (rows.length < batchSize) { + break + } + + offset += batchSize + + if (breathingIntervalMs > 0) { + await new Promise((resolve) => + setTimeout(resolve, breathingIntervalMs) + ) + } + } + + if (isFirstRow) { + controller.enqueue(encoder.encode('[]')) + } else { + controller.enqueue(encoder.encode('\n]')) + } + + controller.close() + } catch (streamError) { + controller.error(streamError) + } + }, + }) - return createExportResponse( - jsonData, + return exportModule.createExportResponse( + stream, `${tableName}_export.json`, 'application/json' )