From 20eba6e1ad9ae6d119cf3a7a3fcf1742bc289e63 Mon Sep 17 00:00:00 2001 From: mdp28 <212812676+mdp28@users.noreply.github.com> Date: Fri, 15 May 2026 16:44:46 +1000 Subject: [PATCH] Stream large database exports --- src/export/csv.test.ts | 150 +++++++++++++++----------- src/export/csv.ts | 56 +++++----- src/export/dump.test.ts | 68 +++++++++++- src/export/dump.ts | 93 ++++++++-------- src/export/json.test.ts | 111 ++++++++++--------- src/export/json.ts | 34 ++++-- src/export/streaming.test.ts | 155 +++++++++++++++++++++++++++ src/export/streaming.ts | 200 +++++++++++++++++++++++++++++++++++ 8 files changed, 666 insertions(+), 201 deletions(-) create mode 100644 src/export/streaming.test.ts create mode 100644 src/export/streaming.ts diff --git a/src/export/csv.test.ts b/src/export/csv.test.ts index b186aeb..428c1df 100644 --- a/src/export/csv.test.ts +++ b/src/export/csv.test.ts @@ -1,13 +1,12 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToCsvRoute } from './csv' -import { getTableData, createExportResponse } from './index' +import { executeOperation } from '.' 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('.', () => ({ + executeOperation: vi.fn(), })) vi.mock('../utils', () => ({ @@ -42,17 +41,18 @@ 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 stream a CSV file when table data exists', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { name: 'id' }, + { name: 'name' }, + { name: 'age' }, + ]) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + ]) const response = await exportTableToCsvRoute( 'users', @@ -60,21 +60,17 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( - 'users', - mockDataSource, - mockConfig + expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="users_export.csv"' ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,age\n1,Alice,30\n2,Bob,25\n', - 'users_export.csv', - 'text/csv' + await expect(response.text()).resolves.toBe( + 'id,name,age\n1,Alice,30\n2,Bob,25\n' ) - 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(executeOperation).mockResolvedValueOnce([]) const response = await exportTableToCsvRoute( 'non_existent_table', @@ -82,11 +78,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 +86,11 @@ 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 include headers for empty tables', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'empty_table' }]) + .mockResolvedValueOnce([{ name: 'id' }, { name: 'name' }]) + .mockResolvedValueOnce([]) const response = await exportTableToCsvRoute( 'empty_table', @@ -110,49 +98,86 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( - 'empty_table', + await expect(response.text()).resolves.toBe('id,name\n') + }) + + it('should escape commas, quotes, and newlines in CSV values', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'special_chars' }]) + .mockResolvedValueOnce([ + { name: 'id' }, + { name: 'name' }, + { name: 'bio' }, + ]) + .mockResolvedValueOnce([ + { + id: 1, + name: 'Sahithi, is', + bio: 'my forever "penguin"\nline', + }, + ]) + + const response = await exportTableToCsvRoute( + 'special_chars', mockDataSource, mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - '', - 'empty_table_export.csv', - 'text/csv' + + await expect(response.text()).resolves.toBe( + 'id,name,bio\n1,"Sahithi, is","my forever ""penguin""\nline"\n' ) - expect(response.headers.get('Content-Type')).toBe('text/csv') }) - it('should escape commas and quotes in CSV values', async () => { - vi.mocked(getTableData).mockResolvedValue([ - { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' }, - ]) + it('should page table data instead of loading the full table', async () => { + const firstPage = Array.from({ length: 1000 }, (_, index) => ({ + id: index + 1, + name: `User ${index + 1}`, + })) - vi.mocked(createExportResponse).mockReturnValue( - new Response('mocked-csv-content', { - headers: { 'Content-Type': 'text/csv' }, - }) - ) + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([{ name: 'id' }, { name: 'name' }]) + .mockResolvedValueOnce(firstPage) + .mockResolvedValueOnce([{ id: 1001, name: 'Last User' }]) const response = await exportTableToCsvRoute( - 'special_chars', + 'users', mockDataSource, mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n', - 'special_chars_export.csv', - 'text/csv' + const csv = await response.text() + + expect(csv).toContain('1001,Last User\n') + expect(executeOperation).toHaveBeenNthCalledWith( + 3, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 0], + }, + ], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 4, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 1000], + }, + ], + mockDataSource, + mockConfig ) - expect(response.headers.get('Content-Type')).toBe('text/csv') }) - it('should return 500 on an unexpected error', async () => { + it('should return 500 on an unexpected error before streaming starts', async () => { const consoleErrorMock = vi .spyOn(console, 'error') .mockImplementation(() => {}) - vi.mocked(getTableData).mockRejectedValue(new Error('Database Error')) + vi.mocked(executeOperation).mockRejectedValue(new Error('Database Error')) const response = await exportTableToCsvRoute( 'users', @@ -163,5 +188,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..dd343d6 100644 --- a/src/export/csv.ts +++ b/src/export/csv.ts @@ -1,7 +1,28 @@ -import { getTableData, createExportResponse } from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +import { + createStreamingExportResponse, + formatCsvValue, + getTableColumns, + iterateTableRows, + tableExists, +} from './streaming' + +async function* csvTableChunks( + tableName: string, + columns: string[], + dataSource: DataSource, + config: StarbaseDBConfiguration +): AsyncGenerator { + if (columns.length) { + yield `${columns.map(formatCsvValue).join(',')}\n` + } + + for await (const row of iterateTableRows(tableName, dataSource, config)) { + yield `${columns.map((column) => formatCsvValue(row[column])).join(',')}\n` + } +} export async function exportTableToCsvRoute( tableName: string, @@ -9,9 +30,7 @@ export async function exportTableToCsvRoute( config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) - - if (data === null) { + if (!(await tableExists(tableName, dataSource, config))) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,33 +38,10 @@ 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 columns = await getTableColumns(tableName, dataSource, config) - return createExportResponse( - csvContent, + return createStreamingExportResponse( + csvTableChunks(tableName, columns, dataSource, config), `${tableName}_export.csv`, 'text/csv' ) diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..66a16e1 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 () => { @@ -124,10 +124,68 @@ 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 page table data instead of loading a full table at once', async () => { + const firstPage = Array.from({ length: 1000 }, (_, index) => ({ + id: index + 1, + name: `User ${index + 1}`, + })) + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + ]) + .mockResolvedValueOnce(firstPage) + .mockResolvedValueOnce([{ id: 1001, name: 'Last User' }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await response.text() + + expect(dumpText).toContain( + 'INSERT INTO "users" VALUES (1001, \'Last User\');' + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 3, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 0], + }, + ], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 4, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 1000], + }, + ], + mockDataSource, + mockConfig + ) + }) + + it('should emit SQL NULL for nullish values', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, bio TEXT);' }, + ]) + .mockResolvedValueOnce([{ id: 1, bio: null }]) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await response.text() + + expect(dumpText).toContain('INSERT INTO "users" VALUES (1, NULL);') + }) + it('should return a 500 response when an error occurs', async () => { const consoleErrorMock = vi .spyOn(console, 'error') diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..a9a3299 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -2,6 +2,47 @@ import { executeOperation } from '.' import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +import { + createStreamingExportResponse, + formatSqlValue, + iterateTableRows, + quoteSqlIdentifier, +} from './streaming' + +async function* dumpDatabaseChunks( + tables: string[], + dataSource: DataSource, + config: StarbaseDBConfiguration +): AsyncGenerator { + yield 'SQLite format 3\0' + + for (const table of tables) { + const schemaResult = await executeOperation( + [ + { + sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name=?;`, + params: [table], + }, + ], + dataSource, + config + ) + + if (schemaResult.length) { + const schema = schemaResult[0].sql + yield `\n-- Table: ${table}\n${schema};\n\n` + } + + const quotedTableName = quoteSqlIdentifier(table) + + for await (const row of iterateTableRows(table, dataSource, config)) { + const values = Object.values(row).map(formatSqlValue) + yield `INSERT INTO ${quotedTableName} VALUES (${values.join(', ')});\n` + } + + yield '\n' + } +} export async function dumpDatabaseRoute( dataSource: DataSource, @@ -16,54 +57,12 @@ export async function dumpDatabaseRoute( ) 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 headers = new Headers({ - 'Content-Type': 'application/x-sqlite3', - 'Content-Disposition': 'attachment; filename="database_dump.sql"', - }) - - return new Response(blob, { headers }) + return createStreamingExportResponse( + dumpDatabaseChunks(tables, dataSource, config), + 'database_dump.sql', + 'application/x-sqlite3' + ) } catch (error: any) { console.error('Database Dump Error:', error) return createResponse(undefined, 'Failed to create database dump', 500) diff --git a/src/export/json.test.ts b/src/export/json.test.ts index 3fe4a8c..f12114f 100644 --- a/src/export/json.test.ts +++ b/src/export/json.test.ts @@ -1,13 +1,12 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToJsonRoute } from './json' -import { getTableData, createExportResponse } from './index' +import { executeOperation } from '.' 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('.', () => ({ + executeOperation: vi.fn(), })) vi.mock('../utils', () => ({ @@ -41,7 +40,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(executeOperation).mockResolvedValueOnce([]) const response = await exportTableToJsonRoute( 'missing_table', @@ -54,18 +53,15 @@ 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 stream a 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(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(mockData) const response = await exportTableToJsonRoute( 'users', @@ -73,27 +69,17 @@ describe('JSON Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( - 'users', - mockDataSource, - mockConfig - ) - expect(createExportResponse).toHaveBeenCalledWith( - JSON.stringify(mockData, null, 4), - 'users_export.json', - 'application/json' - ) expect(response.headers.get('Content-Type')).toBe('application/json') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="users_export.json"' + ) + await expect(response.json()).resolves.toEqual(mockData) }) 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(executeOperation) + .mockResolvedValueOnce([{ name: 'empty_table' }]) + .mockResolvedValueOnce([]) const response = await exportTableToJsonRoute( 'empty_table', @@ -101,12 +87,7 @@ describe('JSON Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - '[]', - 'empty_table_export.json', - 'application/json' - ) - expect(response.headers.get('Content-Type')).toBe('application/json') + await expect(response.json()).resolves.toEqual([]) }) it('should escape special characters in JSON properly', async () => { @@ -114,13 +95,10 @@ 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(executeOperation) + .mockResolvedValueOnce([{ name: 'special_chars' }]) + .mockResolvedValueOnce(specialCharsData) const response = await exportTableToJsonRoute( 'special_chars', @@ -128,19 +106,55 @@ describe('JSON Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - JSON.stringify(specialCharsData, null, 4), - 'special_chars_export.json', - 'application/json' + await expect(response.json()).resolves.toEqual(specialCharsData) + }) + + it('should page table data instead of loading the full table', async () => { + const firstPage = Array.from({ length: 1000 }, (_, index) => ({ + id: index + 1, + })) + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(firstPage) + .mockResolvedValueOnce([{ id: 1001 }]) + + const response = await exportTableToJsonRoute( + 'users', + mockDataSource, + mockConfig + ) + + await expect(response.json()).resolves.toHaveLength(1001) + expect(executeOperation).toHaveBeenNthCalledWith( + 2, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 0], + }, + ], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 3, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1000, 1000], + }, + ], + mockDataSource, + mockConfig ) - expect(response.headers.get('Content-Type')).toBe('application/json') }) - it('should return a 500 response when an error occurs', async () => { + it('should return a 500 response when an error occurs before streaming starts', async () => { const consoleErrorMock = vi .spyOn(console, 'error') .mockImplementation(() => {}) - vi.mocked(getTableData).mockRejectedValue(new Error('Database Error')) + vi.mocked(executeOperation).mockRejectedValue(new Error('Database Error')) const response = await exportTableToJsonRoute( 'users', @@ -151,5 +165,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..b5a831c 100644 --- a/src/export/json.ts +++ b/src/export/json.ts @@ -1,7 +1,28 @@ -import { getTableData, createExportResponse } from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +import { + createStreamingExportResponse, + iterateTableRows, + tableExists, +} from './streaming' + +async function* jsonTableChunks( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): AsyncGenerator { + let isFirstRow = true + + yield '[' + + for await (const row of iterateTableRows(tableName, dataSource, config)) { + yield `${isFirstRow ? '' : ','}\n${JSON.stringify(row, null, 4)}` + isFirstRow = false + } + + yield isFirstRow ? ']' : '\n]' +} export async function exportTableToJsonRoute( tableName: string, @@ -9,9 +30,7 @@ export async function exportTableToJsonRoute( config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) - - if (data === null) { + if (!(await tableExists(tableName, dataSource, config))) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,11 +38,8 @@ export async function exportTableToJsonRoute( ) } - // Convert the result to JSON - const jsonData = JSON.stringify(data, null, 4) - - return createExportResponse( - jsonData, + return createStreamingExportResponse( + jsonTableChunks(tableName, dataSource, config), `${tableName}_export.json`, 'application/json' ) diff --git a/src/export/streaming.test.ts b/src/export/streaming.test.ts new file mode 100644 index 0000000..07b77e2 --- /dev/null +++ b/src/export/streaming.test.ts @@ -0,0 +1,155 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import { executeOperation } from '.' +import type { DataSource } from '../types' +import type { StarbaseDBConfiguration } from '../handler' +import { + createStreamingExportResponse, + formatCsvValue, + formatSqlValue, + getTableColumns, + iterateTableRows, + quoteSqlIdentifier, + tableExists, +} from './streaming' + +vi.mock('.', () => ({ + executeOperation: vi.fn(), +})) + +let mockDataSource: DataSource +let mockConfig: StarbaseDBConfiguration + +beforeEach(() => { + vi.clearAllMocks() + + mockDataSource = { + source: 'external', + external: { dialect: 'sqlite' }, + rpc: { executeQuery: vi.fn() }, + } as any + + mockConfig = { + outerbaseApiKey: 'mock-api-key', + role: 'admin', + features: { allowlist: true, rls: true, rest: true }, + } +}) + +afterEach(() => { + vi.unstubAllGlobals() +}) + +describe('Export Streaming Helpers', () => { + it('should quote SQL identifiers safely', () => { + expect(quoteSqlIdentifier('users')).toBe('"users"') + expect(quoteSqlIdentifier('weird"name')).toBe('"weird""name"') + }) + + it('should format SQL values without producing invalid undefined literals', () => { + expect(formatSqlValue(null)).toBe('NULL') + expect(formatSqlValue(undefined)).toBe('NULL') + expect(formatSqlValue("Alice's")).toBe("'Alice''s'") + expect(formatSqlValue(true)).toBe('1') + expect(formatSqlValue(new Uint8Array([0, 15, 255]))).toBe("X'000fff'") + }) + + it('should escape CSV values only when needed', () => { + expect(formatCsvValue('Alice')).toBe('Alice') + expect(formatCsvValue('Alice, Bob')).toBe('"Alice, Bob"') + expect(formatCsvValue('Alice "A"')).toBe('"Alice ""A"""') + expect(formatCsvValue(null)).toBe('') + }) + + it('should check table existence with a bound table name', async () => { + vi.mocked(executeOperation).mockResolvedValueOnce([{ name: 'users' }]) + + await expect( + tableExists('users', mockDataSource, mockConfig) + ).resolves.toBe(true) + + expect(executeOperation).toHaveBeenCalledWith( + [ + { + sql: "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", + params: ['users'], + }, + ], + mockDataSource, + mockConfig + ) + }) + + it('should read column names from table metadata', async () => { + vi.mocked(executeOperation).mockResolvedValueOnce([ + { name: 'id' }, + { name: 'name' }, + ]) + + await expect( + getTableColumns('users', mockDataSource, mockConfig) + ).resolves.toEqual(['id', 'name']) + }) + + it('should iterate table rows in bounded pages and yield between full pages', async () => { + const wait = vi.fn().mockResolvedValue(undefined) + vi.stubGlobal('scheduler', { wait }) + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ id: 1 }, { id: 2 }]) + .mockResolvedValueOnce([{ id: 3 }]) + + const rows = [] + + for await (const row of iterateTableRows( + 'users', + mockDataSource, + mockConfig, + 2 + )) { + rows.push(row) + } + + expect(rows).toEqual([{ id: 1 }, { id: 2 }, { id: 3 }]) + expect(wait).toHaveBeenCalledWith(0) + expect(executeOperation).toHaveBeenNthCalledWith( + 1, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [2, 0], + }, + ], + mockDataSource, + mockConfig + ) + expect(executeOperation).toHaveBeenNthCalledWith( + 2, + [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [2, 2], + }, + ], + mockDataSource, + mockConfig + ) + }) + + it('should create a readable streaming response', async () => { + async function* chunks() { + yield 'hello' + yield ' ' + yield 'world' + } + + const response = createStreamingExportResponse( + chunks(), + 'hello.txt', + 'text/plain' + ) + + expect(response.headers.get('Content-Type')).toBe('text/plain') + expect(response.headers.get('Cache-Control')).toBe('no-store') + await expect(response.text()).resolves.toBe('hello world') + }) +}) diff --git a/src/export/streaming.ts b/src/export/streaming.ts new file mode 100644 index 0000000..a8024de --- /dev/null +++ b/src/export/streaming.ts @@ -0,0 +1,200 @@ +import { StarbaseDBConfiguration } from '../handler' +import { DataSource } from '../types' +import { executeOperation } from '.' + +export const DEFAULT_EXPORT_PAGE_SIZE = 1000 + +type SchedulerWithWait = { + wait?: (duration: number) => Promise +} + +export function quoteSqlIdentifier(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"` +} + +function bytesToHex(value: Uint8Array): string { + return Array.from(value) + .map((byte) => byte.toString(16).padStart(2, '0')) + .join('') +} + +function valueToText(value: unknown): string { + if (value === null || value === undefined) { + return '' + } + + if (value instanceof Uint8Array) { + return bytesToHex(value) + } + + if (value instanceof ArrayBuffer) { + return bytesToHex(new Uint8Array(value)) + } + + return String(value) +} + +export function formatSqlValue(value: unknown): string { + if (value === null || value === undefined) { + return 'NULL' + } + + if (typeof value === 'number' || typeof value === 'bigint') { + return String(value) + } + + if (typeof value === 'boolean') { + return value ? '1' : '0' + } + + if (value instanceof Uint8Array) { + return `X'${bytesToHex(value)}'` + } + + if (value instanceof ArrayBuffer) { + return `X'${bytesToHex(new Uint8Array(value))}'` + } + + return `'${String(value).replace(/'/g, "''")}'` +} + +export function formatCsvValue(value: unknown): string { + const text = valueToText(value) + + if (/[",\r\n]/.test(text)) { + return `"${text.replace(/"/g, '""')}"` + } + + return text +} + +export function sanitizeExportFileName(fileName: string): string { + return fileName.replace(/[\x00-\x1f\x7f"\\/:*?<>|]+/g, '_') +} + +export async function yieldToRuntime(): Promise { + const scheduler = (globalThis as { scheduler?: SchedulerWithWait }) + .scheduler + + if (scheduler?.wait) { + await scheduler.wait(0) + return + } + + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +export async function tableExists( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + const result = await executeOperation( + [ + { + sql: "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", + params: [tableName], + }, + ], + dataSource, + config + ) + + return result.length > 0 +} + +export async function getTableColumns( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + const result = await executeOperation( + [ + { + sql: `PRAGMA table_info(${quoteSqlIdentifier(tableName)});`, + }, + ], + dataSource, + config + ) + + return result + .map((column: any) => column.name) + .filter((name: unknown): name is string => typeof name === 'string') +} + +export async function* iterateTableRows( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration, + pageSize = DEFAULT_EXPORT_PAGE_SIZE +): AsyncGenerator> { + const quotedTableName = quoteSqlIdentifier(tableName) + let offset = 0 + + while (true) { + const rows = await executeOperation( + [ + { + sql: `SELECT * FROM ${quotedTableName} LIMIT ? OFFSET ?;`, + params: [pageSize, offset], + }, + ], + dataSource, + config + ) + + if (!rows.length) { + return + } + + for (const row of rows) { + yield row + } + + offset += rows.length + + if (rows.length < pageSize) { + return + } + + await yieldToRuntime() + } +} + +export function createStreamingExportResponse( + chunks: AsyncIterable, + fileName: string, + contentType: string +): Response { + const encoder = new TextEncoder() + const iterator = chunks[Symbol.asyncIterator]() + + const body = new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await iterator.next() + + if (done) { + controller.close() + return + } + + controller.enqueue(encoder.encode(value)) + } catch (error) { + controller.error(error) + } + }, + async cancel() { + await iterator.return?.() + }, + }) + + return new Response(body, { + headers: { + 'Cache-Control': 'no-store', + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${sanitizeExportFileName(fileName)}"`, + }, + }) +}