diff --git a/src/export/csv.test.ts b/src/export/csv.test.ts index b186aeb..dd9efc5 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 './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), + executeOperation: vi.fn(), })) vi.mock('../utils', () => ({ @@ -23,6 +22,16 @@ vi.mock('../utils', () => ({ let mockDataSource: DataSource let mockConfig: StarbaseDBConfiguration +const tableColumns = (names: string[]) => + names.map((name, index) => ({ + cid: index, + name, + type: '', + notnull: 0, + dflt_value: null, + pk: name === 'id' ? 1 : 0, + })) + beforeEach(() => { vi.clearAllMocks() @@ -43,16 +52,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' }, - }) - ) + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(tableColumns(['id', 'name', 'age'])) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + ]) const response = await exportTableToCsvRoute( 'users', @@ -60,21 +66,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 +84,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() @@ -96,13 +93,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' }, - }) - ) + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'empty_table' }]) + .mockResolvedValueOnce(tableColumns(['id', 'name'])) + .mockResolvedValueOnce([]) const response = await exportTableToCsvRoute( 'empty_table', @@ -110,29 +104,17 @@ 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') + await expect(response.text()).resolves.toBe('id,name\n') }) 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(executeOperation) + .mockResolvedValueOnce([{ name: 'special_chars' }]) + .mockResolvedValueOnce(tableColumns(['id', 'name', 'bio'])) + .mockResolvedValueOnce([ + { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' }, + ]) const response = await exportTableToCsvRoute( 'special_chars', @@ -140,10 +122,8 @@ describe('CSV Export Module', () => { mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n', - 'special_chars_export.csv', - 'text/csv' + await expect(response.text()).resolves.toBe( + 'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n' ) expect(response.headers.get('Content-Type')).toBe('text/csv') }) @@ -152,7 +132,9 @@ describe('CSV Export Module', () => { 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', diff --git a/src/export/csv.ts b/src/export/csv.ts index 22a4591..b2bc556 100644 --- a/src/export/csv.ts +++ b/src/export/csv.ts @@ -1,7 +1,11 @@ -import { getTableData, createExportResponse } from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +import { + createStreamingExportResponse, + csvTableChunks, + getTablePagePlan, +} from './streaming' export async function exportTableToCsvRoute( tableName: string, @@ -9,9 +13,9 @@ export async function exportTableToCsvRoute( config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const pagePlan = await getTablePagePlan(tableName, dataSource, config) - if (data === null) { + if (!pagePlan) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,33 +23,8 @@ 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' - }) - } - - return createExportResponse( - csvContent, + return createStreamingExportResponse( + csvTableChunks(tableName, dataSource, config, pagePlan), `${tableName}_export.csv`, 'text/csv' ) diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..56c03a5 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -22,6 +22,16 @@ vi.mock('../utils', () => ({ let mockDataSource: DataSource let mockConfig: StarbaseDBConfiguration +const tableColumns = (names: string[]) => + names.map((name, index) => ({ + cid: index, + name, + type: '', + notnull: 0, + dflt_value: null, + pk: name === 'id' ? 1 : 0, + })) + beforeEach(() => { vi.clearAllMocks() @@ -45,6 +55,8 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, ]) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(tableColumns(['id', 'name'])) .mockResolvedValueOnce([ { id: 1, name: 'Alice' }, { id: 2, name: 'Bob' }, @@ -52,6 +64,8 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE orders (id INTEGER, total REAL);' }, ]) + .mockResolvedValueOnce([{ name: 'orders' }]) + .mockResolvedValueOnce(tableColumns(['id', 'total'])) .mockResolvedValueOnce([ { id: 1, total: 99.99 }, { id: 2, total: 49.5 }, @@ -71,13 +85,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 () => { @@ -99,6 +113,8 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, ]) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(tableColumns(['id', 'name'])) .mockResolvedValueOnce([]) const response = await dumpDatabaseRoute(mockDataSource, mockConfig) @@ -117,6 +133,8 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE users (id INTEGER, bio TEXT);' }, ]) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce(tableColumns(['id', 'bio'])) .mockResolvedValueOnce([{ id: 1, bio: "Alice's adventure" }]) const response = await dumpDatabaseRoute(mockDataSource, mockConfig) @@ -124,7 +142,7 @@ 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');" ) }) diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..f038b6a 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -1,69 +1,24 @@ -import { executeOperation } from '.' import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +import { + createStreamingExportResponse, + listExportableTables, + sqlDumpChunks, +} from './streaming' 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';" }], - 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 - ) + const tables = await listExportableTables(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( + sqlDumpChunks(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..b4aeca9 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 './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), + executeOperation: vi.fn(), })) vi.mock('../utils', () => ({ @@ -23,6 +22,16 @@ vi.mock('../utils', () => ({ let mockDataSource: DataSource let mockConfig: StarbaseDBConfiguration +const tableColumns = (names: string[]) => + names.map((name, index) => ({ + cid: index, + name, + type: '', + notnull: 0, + dflt_value: null, + pk: name === 'id' ? 1 : 0, + })) + beforeEach(() => { vi.clearAllMocks() @@ -41,7 +50,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', @@ -59,13 +68,10 @@ describe('JSON Export Module', () => { { 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(tableColumns(['id', 'name'])) + .mockResolvedValueOnce(mockData) const response = await exportTableToJsonRoute( 'users', @@ -73,27 +79,15 @@ 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') + 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(tableColumns(['id', 'name'])) + .mockResolvedValueOnce([]) const response = await exportTableToJsonRoute( 'empty_table', @@ -101,12 +95,8 @@ 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 +104,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(tableColumns(['id', 'name', 'description'])) + .mockResolvedValueOnce(specialCharsData) const response = await exportTableToJsonRoute( 'special_chars', @@ -128,19 +115,17 @@ 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') + await expect(response.json()).resolves.toEqual(specialCharsData) }) 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(executeOperation).mockRejectedValue( + new Error('Database Error') + ) const response = await exportTableToJsonRoute( 'users', diff --git a/src/export/json.ts b/src/export/json.ts index c0ab811..e6736a6 100644 --- a/src/export/json.ts +++ b/src/export/json.ts @@ -1,7 +1,11 @@ -import { getTableData, createExportResponse } from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +import { + createStreamingExportResponse, + getTablePagePlan, + jsonTableChunks, +} from './streaming' export async function exportTableToJsonRoute( tableName: string, @@ -9,9 +13,9 @@ export async function exportTableToJsonRoute( config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const pagePlan = await getTablePagePlan(tableName, dataSource, config) - if (data === null) { + if (!pagePlan) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,11 +23,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, pagePlan), `${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..cc373de --- /dev/null +++ b/src/export/streaming.test.ts @@ -0,0 +1,143 @@ +import { describe, expect, it, vi, beforeEach } from 'vitest' +import { executeOperation } from './index' +import { + createStreamingExportResponse, + formatCsvValue, + formatSqlValue, + iterateTableRows, + quoteSqlIdentifier, + type TablePagePlan, +} from './streaming' +import type { DataSource } from '../types' +import type { StarbaseDBConfiguration } from '../handler' + +vi.mock('./index', () => ({ + 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 }, + } +}) + +describe('Streaming export helpers', () => { + it('quotes SQL identifiers and formats SQL values safely', () => { + expect(quoteSqlIdentifier('weird"name')).toBe('"weird""name"') + expect(formatSqlValue(null)).toBe('NULL') + expect(formatSqlValue(Number.NaN)).toBe('NULL') + expect(formatSqlValue(true)).toBe('1') + expect(formatSqlValue("Alice's adventure")).toBe("'Alice''s adventure'") + expect(formatSqlValue(new Uint8Array([0, 15, 255]))).toBe("X'000fff'") + }) + + it('formats CSV values without buffering a full export', () => { + expect(formatCsvValue(null)).toBe('') + expect(formatCsvValue('plain')).toBe('plain') + expect(formatCsvValue('a,b')).toBe('"a,b"') + expect(formatCsvValue('say "hi"')).toBe('"say ""hi"""') + }) + + it('streams iterable chunks with sanitized attachment names', async () => { + async function* chunks() { + yield 'a' + yield 'b' + } + + const response = createStreamingExportResponse( + chunks(), + '../bad:name.csv', + 'text/csv' + ) + + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename=".._bad_name.csv"' + ) + await expect(response.text()).resolves.toBe('ab') + }) + + it('uses keyset pagination after a full page', async () => { + const pagePlan: TablePagePlan = { + columns: ['id', 'name'], + selectList: '*', + orderBy: ['"id"'], + cursorColumns: [{ expression: '"id"', resultKey: 'id' }], + } + + vi.mocked(executeOperation) + .mockResolvedValueOnce([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]) + .mockResolvedValueOnce([{ id: 3, name: 'Cleo' }]) + + const rows = [] + for await (const row of iterateTableRows( + 'users', + mockDataSource, + mockConfig, + { pagePlan, pageSize: 2 } + )) { + rows.push(row) + } + + expect(rows).toEqual([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + { id: 3, name: 'Cleo' }, + ]) + + expect(vi.mocked(executeOperation).mock.calls[0][0][0]).toMatchObject({ + sql: 'SELECT * FROM "users" ORDER BY "id" ASC LIMIT ?;', + params: [2], + }) + expect(vi.mocked(executeOperation).mock.calls[1][0][0]).toMatchObject({ + sql: 'SELECT * FROM "users" WHERE ("id" > ?) ORDER BY "id" ASC LIMIT ?;', + params: [2, 2], + }) + }) + + it('keeps internal rowid cursors out of exported rows', async () => { + const pagePlan: TablePagePlan = { + columns: ['name'], + selectList: '*, rowid AS "__starbasedb_export_rowid"', + orderBy: ['rowid'], + cursorColumns: [ + { + expression: 'rowid', + resultKey: '__starbasedb_export_rowid', + }, + ], + rowidAlias: '__starbasedb_export_rowid', + } + + vi.mocked(executeOperation).mockResolvedValueOnce([ + { name: 'Alice', __starbasedb_export_rowid: 1 }, + ]) + + const rows = [] + for await (const row of iterateTableRows( + 'users', + mockDataSource, + mockConfig, + { pagePlan, pageSize: 2 } + )) { + rows.push(row) + } + + expect(rows).toEqual([{ name: 'Alice' }]) + }) +}) diff --git a/src/export/streaming.ts b/src/export/streaming.ts new file mode 100644 index 0000000..1e5a530 --- /dev/null +++ b/src/export/streaming.ts @@ -0,0 +1,458 @@ +import { executeOperation } from '.' +import { StarbaseDBConfiguration } from '../handler' +import { DataSource } from '../types' + +export const DEFAULT_EXPORT_PAGE_SIZE = 500 + +type TableColumn = { + name: string + pk?: number +} + +type CursorColumn = { + expression: string + resultKey: string +} + +export type TablePagePlan = { + columns: string[] + selectList: string + orderBy: string[] + cursorColumns: CursorColumn[] + rowidAlias?: string +} + +export function quoteSqlIdentifier(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"` +} + +function arrayBufferToHex(buffer: ArrayBuffer): string { + return Array.from(new Uint8Array(buffer)) + .map((byte) => byte.toString(16).padStart(2, '0')) + .join('') +} + +function isArrayBufferView(value: unknown): value is ArrayBufferView { + return ArrayBuffer.isView(value) +} + +function formatBinaryValue(value: ArrayBuffer | ArrayBufferView): string { + const buffer = isArrayBufferView(value) + ? value.buffer.slice( + value.byteOffset, + value.byteOffset + value.byteLength + ) + : value + + return `X'${arrayBufferToHex(buffer)}'` +} + +export function formatSqlValue(value: unknown): string { + if (value === null || value === undefined) { + return 'NULL' + } + + if (typeof value === 'number') { + return Number.isFinite(value) ? String(value) : 'NULL' + } + + if (typeof value === 'bigint') { + return value.toString() + } + + if (typeof value === 'boolean') { + return value ? '1' : '0' + } + + if (value instanceof ArrayBuffer || isArrayBufferView(value)) { + return formatBinaryValue(value) + } + + return `'${String(value).replace(/'/g, "''")}'` +} + +export function formatCsvValue(value: unknown): string { + if (value === null || value === undefined) { + return '' + } + + const stringValue = + value instanceof ArrayBuffer || isArrayBufferView(value) + ? formatSqlValue(value) + : String(value) + + if (/["\r\n,]/.test(stringValue)) { + return `"${stringValue.replace(/"/g, '""')}"` + } + + return stringValue +} + +export function sanitizeExportFileName(fileName: string): string { + return fileName.replace(/[\x00-\x1f"\\/:*?<>|]+/g, '_') +} + +export async function yieldToRuntime(): Promise { + const scheduler = ( + globalThis as typeof globalThis & { + scheduler?: { wait?: (delay: number) => Promise } + } + ).scheduler + + if (scheduler?.wait) { + await scheduler.wait(0) + return + } + + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +export function createStreamingExportResponse( + chunks: AsyncIterable, + fileName: string, + contentType: string +): Response { + const iterator = chunks[Symbol.asyncIterator]() + const encoder = new TextEncoder() + + const stream = 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(reason) { + if (iterator.return) { + await iterator.return(reason) + } + }, + }) + + return new Response(stream, { + headers: { + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${sanitizeExportFileName( + fileName + )}"`, + }, + }) +} + +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 listExportableTables( + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + const result = await executeOperation( + [ + { + sql: "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name;", + }, + ], + dataSource, + config + ) + + return result.map((row: { name: string }) => row.name) +} + +export async function getTableSchema( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + const result = await executeOperation( + [ + { + sql: "SELECT sql FROM sqlite_master WHERE type='table' AND name=?;", + params: [tableName], + }, + ], + dataSource, + config + ) + + const schema = result[0]?.sql + return typeof schema === 'string' ? schema : undefined +} + +export async function getTableColumns( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + return executeOperation( + [{ sql: `PRAGMA table_info(${quoteSqlIdentifier(tableName)});` }], + dataSource, + config + ) +} + +function createRowidAlias(columns: TableColumn[]): string { + const columnNames = new Set(columns.map((column) => column.name)) + let alias = '__starbasedb_export_rowid' + + while (columnNames.has(alias)) { + alias = `${alias}_` + } + + return alias +} + +export async function getTablePagePlan( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + if (!(await tableExists(tableName, dataSource, config))) { + return null + } + + const columns = await getTableColumns(tableName, dataSource, config) + const columnNames = columns.map((column) => column.name) + const primaryKeyColumns = columns + .filter((column) => Number(column.pk ?? 0) > 0) + .sort((left, right) => Number(left.pk ?? 0) - Number(right.pk ?? 0)) + + if (primaryKeyColumns.length > 0) { + const cursorColumns = primaryKeyColumns.map((column) => ({ + expression: quoteSqlIdentifier(column.name), + resultKey: column.name, + })) + + return { + columns: columnNames, + selectList: '*', + orderBy: cursorColumns.map((column) => column.expression), + cursorColumns, + } + } + + const rowidAlias = createRowidAlias(columns) + + return { + columns: columnNames, + selectList: `*, rowid AS ${quoteSqlIdentifier(rowidAlias)}`, + orderBy: ['rowid'], + cursorColumns: [{ expression: 'rowid', resultKey: rowidAlias }], + rowidAlias, + } +} + +function buildKeysetCondition( + cursorColumns: CursorColumn[], + cursorValues: unknown[] +): { sql: string; params: unknown[] } { + const clauses: string[] = [] + const params: unknown[] = [] + + for (let index = 0; index < cursorColumns.length; index++) { + const equalityColumns = cursorColumns.slice(0, index) + const equality = equalityColumns + .map((column) => `${column.expression} = ?`) + .join(' AND ') + const comparison = `${cursorColumns[index].expression} > ?` + const clause = equality ? `(${equality} AND ${comparison})` : comparison + + clauses.push(`(${clause})`) + params.push(...cursorValues.slice(0, index), cursorValues[index]) + } + + return { sql: clauses.join(' OR '), params } +} + +function removeInternalCursorColumns( + row: Record, + pagePlan: TablePagePlan +): Record { + if (!pagePlan.rowidAlias) { + return row + } + + const { [pagePlan.rowidAlias]: _rowid, ...exportedRow } = row + return exportedRow +} + +export async function* iterateTableRows( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration, + options: { + pagePlan?: TablePagePlan + pageSize?: number + } = {} +): AsyncGenerator> { + const pagePlan = + options.pagePlan ?? + (await getTablePagePlan(tableName, dataSource, config)) + + if (!pagePlan) { + return + } + + const pageSize = options.pageSize ?? DEFAULT_EXPORT_PAGE_SIZE + const quotedTableName = quoteSqlIdentifier(tableName) + const orderBy = pagePlan.orderBy + .map((expression) => `${expression} ASC`) + .join(', ') + let cursorValues: unknown[] | undefined + let offset = 0 + + while (true) { + let params: unknown[] = [pageSize] + let whereClause = '' + let offsetClause = '' + + if ( + cursorValues?.length === pagePlan.cursorColumns.length && + cursorValues.every((value) => value !== null && value !== undefined) + ) { + const keyset = buildKeysetCondition( + pagePlan.cursorColumns, + cursorValues + ) + whereClause = ` WHERE ${keyset.sql}` + params = [...keyset.params, pageSize] + } else if (offset > 0) { + offsetClause = ' OFFSET ?' + params = [pageSize, offset] + } + + const pageRows = await executeOperation( + [ + { + sql: `SELECT ${pagePlan.selectList} FROM ${quotedTableName}${whereClause} ORDER BY ${orderBy} LIMIT ?${offsetClause};`, + params: params as any[], + }, + ], + dataSource, + config + ) + + if (pageRows.length === 0) { + return + } + + for (const row of pageRows as Record[]) { + yield removeInternalCursorColumns(row, pagePlan) + } + + if (pageRows.length < pageSize) { + return + } + + const lastRow = pageRows[pageRows.length - 1] as Record + cursorValues = pagePlan.cursorColumns.map( + (column) => lastRow[column.resultKey] + ) + offset += pageSize + await yieldToRuntime() + } +} + +export async function* sqlDumpChunks( + tables: string[], + dataSource: DataSource, + config: StarbaseDBConfiguration +): AsyncGenerator { + yield 'SQLite format 3\0' + + for (const tableName of tables) { + const schema = await getTableSchema(tableName, dataSource, config) + const pagePlan = await getTablePagePlan(tableName, dataSource, config) + + if (!pagePlan) { + continue + } + + if (schema) { + const normalizedSchema = schema.trim().replace(/;+\s*$/, '') + yield `\n-- Table: ${tableName}\n${normalizedSchema};\n\n` + } + + const quotedTableName = quoteSqlIdentifier(tableName) + + for await (const row of iterateTableRows( + tableName, + dataSource, + config, + { + pagePlan, + } + )) { + const values = pagePlan.columns.map((columnName) => + formatSqlValue(row[columnName]) + ) + yield `INSERT INTO ${quotedTableName} VALUES (${values.join(', ')});\n` + } + + yield '\n' + } +} + +export async function* csvTableChunks( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration, + pagePlan: TablePagePlan +): AsyncGenerator { + if (pagePlan.columns.length === 0) { + return + } + + yield `${pagePlan.columns.map(formatCsvValue).join(',')}\n` + + for await (const row of iterateTableRows(tableName, dataSource, config, { + pagePlan, + })) { + yield `${pagePlan.columns + .map((columnName) => formatCsvValue(row[columnName])) + .join(',')}\n` + } +} + +export async function* jsonTableChunks( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration, + pagePlan: TablePagePlan +): AsyncGenerator { + let isFirstRow = true + yield '[' + + for await (const row of iterateTableRows(tableName, dataSource, config, { + pagePlan, + })) { + yield `${isFirstRow ? '' : ','}${JSON.stringify(row, null, 4)}` + isFirstRow = false + } + + yield ']' +}