diff --git a/src/export/csv.test.ts b/src/export/csv.test.ts index b186aeb..8b81639 100644 --- a/src/export/csv.test.ts +++ b/src/export/csv.test.ts @@ -1,13 +1,44 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToCsvRoute } from './csv' -import { getTableData, createExportResponse } from './index' +import { getTableDataBatches, tableExists } from './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' +function createTestTextStream( + write: (enqueue: (chunk: string) => void) => Promise +): ReadableStream { + const encoder = new TextEncoder() + + return new ReadableStream({ + async start(controller) { + await write((chunk) => controller.enqueue(encoder.encode(chunk))) + controller.close() + }, + }) +} + +async function* batches(...chunks: any[][]) { + for (const chunk of chunks) { + yield chunk + } +} + vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), + tableExists: vi.fn(), + getTableDataBatches: vi.fn(), + createTextStream: createTestTextStream, + createExportStreamResponse: ( + stream: ReadableStream, + fileName: string, + contentType: string + ) => + new Response(stream, { + headers: { + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${fileName}"`, + }, + }), })) vi.mock('../utils', () => ({ @@ -43,15 +74,12 @@ 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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + ]) ) const response = await exportTableToCsvRoute( @@ -60,21 +88,24 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(tableExists).toHaveBeenCalledWith( 'users', mockDataSource, mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - 'id,name,age\n1,Alice,30\n2,Bob,25\n', - 'users_export.csv', - 'text/csv' + expect(getTableDataBatches).toHaveBeenCalledWith( + 'users', + mockDataSource, + mockConfig ) expect(response.headers.get('Content-Type')).toBe('text/csv') + expect(await response.text()).toBe( + 'id,name,age\n1,Alice,30\n2,Bob,25\n' + ) }) it('should return 404 if table does not exist', async () => { - vi.mocked(getTableData).mockResolvedValue(null) + vi.mocked(tableExists).mockResolvedValue(false) const response = await exportTableToCsvRoute( 'non_existent_table', @@ -82,7 +113,7 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(tableExists).toHaveBeenCalledWith( 'non_existent_table', mockDataSource, mockConfig @@ -96,13 +127,8 @@ 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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue(batches()) const response = await exportTableToCsvRoute( 'empty_table', @@ -110,28 +136,21 @@ describe('CSV Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(tableExists).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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([ + { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' }, + ]) ) const response = await exportTableToCsvRoute( @@ -140,19 +159,32 @@ 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 stream multiple batches into one CSV file', async () => { + vi.mocked(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([{ id: 1, name: 'Alice' }], [{ id: 2, name: 'Bob' }]) + ) + + const response = await exportTableToCsvRoute( + 'users', + mockDataSource, + mockConfig + ) + + expect(await response.text()).toBe('id,name\n1,Alice\n2,Bob\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(tableExists).mockRejectedValue(new Error('Database Error')) const response = await exportTableToCsvRoute( 'users', diff --git a/src/export/csv.ts b/src/export/csv.ts index 22a4591..3fae0ee 100644 --- a/src/export/csv.ts +++ b/src/export/csv.ts @@ -1,17 +1,37 @@ -import { getTableData, createExportResponse } from './index' +import { + createExportStreamResponse, + createTextStream, + getTableDataBatches, + tableExists, +} from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' +function escapeCsvValue(value: unknown): string { + const stringValue = + value === null || value === undefined ? '' : String(value) + + if ( + stringValue.includes(',') || + stringValue.includes('"') || + stringValue.includes('\n') + ) { + return `"${stringValue.replace(/"/g, '""')}"` + } + + return stringValue +} + export async function exportTableToCsvRoute( tableName: string, dataSource: DataSource, config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const exists = await tableExists(tableName, dataSource, config) - if (data === null) { + if (!exists) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,33 +39,32 @@ 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 stream = createTextStream(async (enqueue) => { + let hasHeader = false + + for await (const rows of getTableDataBatches( + tableName, + dataSource, + config + )) { + if (!hasHeader && rows.length > 0) { + enqueue( + Object.keys(rows[0]).map(escapeCsvValue).join(',') + + '\n' + ) + hasHeader = true + } + + for (const row of rows) { + enqueue( + Object.values(row).map(escapeCsvValue).join(',') + '\n' + ) + } + } + }) - return createExportResponse( - csvContent, + return createExportStreamResponse( + stream, `${tableName}_export.csv`, 'text/csv' ) diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..7919217 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -1,12 +1,46 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { dumpDatabaseRoute } from './dump' -import { executeOperation } from '.' +import { executeOperation, getTableDataBatches } from '.' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' +function createTestTextStream( + write: (enqueue: (chunk: string) => void) => Promise +): ReadableStream { + const encoder = new TextEncoder() + + return new ReadableStream({ + async start(controller) { + await write((chunk) => controller.enqueue(encoder.encode(chunk))) + controller.close() + }, + }) +} + +async function* batches(...chunks: any[][]) { + for (const chunk of chunks) { + yield chunk + } +} + vi.mock('.', () => ({ executeOperation: vi.fn(), + getTableDataBatches: vi.fn(), + quoteSqlIdentifier: (identifier: string) => + `"${identifier.replace(/"/g, '""')}"`, + createTextStream: createTestTextStream, + createExportStreamResponse: ( + stream: ReadableStream, + fileName: string, + contentType: string + ) => + new Response(stream, { + headers: { + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${fileName}"`, + }, + }), })) vi.mock('../utils', () => ({ @@ -45,17 +79,22 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { 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);' }, ]) - .mockResolvedValueOnce([ + vi.mocked(getTableDataBatches).mockImplementation((tableName) => { + if (tableName === 'users') { + return batches([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]) + } + + return batches([ { id: 1, total: 99.99 }, { id: 2, total: 49.5 }, ]) + }) const response = await dumpDatabaseRoute(mockDataSource, mockConfig) @@ -71,13 +110,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,7 +138,7 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, ]) - .mockResolvedValueOnce([]) + vi.mocked(getTableDataBatches).mockReturnValue(batches()) const response = await dumpDatabaseRoute(mockDataSource, mockConfig) @@ -117,15 +156,50 @@ describe('Database Dump Module', () => { .mockResolvedValueOnce([ { sql: 'CREATE TABLE users (id INTEGER, bio TEXT);' }, ]) - .mockResolvedValueOnce([{ id: 1, bio: "Alice's adventure" }]) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([{ id: 1, bio: "Alice's adventure" }]) + ) const response = await dumpDatabaseRoute(mockDataSource, mockConfig) 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 stream rows from multiple batches', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + ]) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([{ id: 1, name: 'Alice' }], [{ id: 2, name: 'Bob' }]) ) + + const response = await dumpDatabaseRoute(mockDataSource, mockConfig) + const dumpText = await response.text() + + expect(dumpText).toContain('INSERT INTO "users" VALUES (1, \'Alice\');') + expect(dumpText).toContain('INSERT INTO "users" VALUES (2, \'Bob\');') + }) + + it('should render null values as SQL NULL', async () => { + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'users' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' }, + ]) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([{ id: 1, name: 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 () => { diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..37e350a 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -1,8 +1,26 @@ -import { executeOperation } from '.' +import { + createExportStreamResponse, + createTextStream, + executeOperation, + getTableDataBatches, + quoteSqlIdentifier, +} from '.' import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +function escapeSqlValue(value: unknown): string { + if (value === null || value === undefined) { + return 'NULL' + } + + if (typeof value === 'string') { + return `'${value.replace(/'/g, "''")}'` + } + + return String(value) +} + export async function dumpDatabaseRoute( dataSource: DataSource, config: StarbaseDBConfiguration @@ -16,54 +34,53 @@ 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 - ) + const stream = createTextStream(async (enqueue) => { + enqueue('SQLite format 3\0') // SQLite file header - 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 - ) + // Iterate through all tables without building the full dump in memory. + for (const table of tables) { + const tableIdentifier = quoteSqlIdentifier(table) - for (const row of dataResult) { - const values = Object.values(row).map((value) => - typeof value === 'string' - ? `'${value.replace(/'/g, "''")}'` - : value + // Get table schema + const schemaResult = await executeOperation( + [ + { + sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name=?;`, + params: [table], + }, + ], + dataSource, + config ) - dumpContent += `INSERT INTO ${table} VALUES (${values.join(', ')});\n` - } - dumpContent += '\n' - } + if (schemaResult.length) { + const schema = schemaResult[0].sql + enqueue(`\n-- Table: ${table}\n${schema};\n\n`) + } - // Create a Blob from the dump content - const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' }) + for await (const rows of getTableDataBatches( + table, + dataSource, + config + )) { + for (const row of rows) { + const values = Object.values(row).map(escapeSqlValue) + enqueue( + `INSERT INTO ${tableIdentifier} VALUES (${values.join(', ')});\n` + ) + } + } - const headers = new Headers({ - 'Content-Type': 'application/x-sqlite3', - 'Content-Disposition': 'attachment; filename="database_dump.sql"', + enqueue('\n') + } }) - return new Response(blob, { headers }) + return createExportStreamResponse( + stream, + '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/index.test.ts b/src/export/index.test.ts index 48de76e..58fa6bb 100644 --- a/src/export/index.test.ts +++ b/src/export/index.test.ts @@ -1,5 +1,14 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { executeOperation, getTableData, createExportResponse } from './index' +import { + createExportResponse, + createExportStreamResponse, + createTextStream, + executeOperation, + getTableData, + getTableDataBatches, + quoteSqlIdentifier, + tableExists, +} from './index' import { executeTransaction } from '../operation' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' @@ -109,6 +118,114 @@ describe('Database Operations Module', () => { }) }) + describe('tableExists', () => { + it('should return true if the table exists', async () => { + vi.mocked(executeTransaction).mockResolvedValueOnce([ + { name: 'users' }, + ]) + + await expect( + tableExists('users', mockDataSource, mockConfig) + ).resolves.toBe(true) + }) + + it('should return false if the table does not exist', async () => { + vi.mocked(executeTransaction).mockResolvedValueOnce([]) + + await expect( + tableExists('missing_table', mockDataSource, mockConfig) + ).resolves.toBe(false) + }) + }) + + describe('getTableDataBatches', () => { + it('should fetch rows in batches until no rows remain', async () => { + vi.mocked(executeTransaction) + .mockResolvedValueOnce([{ id: 1 }]) + .mockResolvedValueOnce([{ id: 2 }]) + .mockResolvedValueOnce([]) + + const rows: any[] = [] + + for await (const batch of getTableDataBatches( + 'users', + mockDataSource, + mockConfig, + 1 + )) { + rows.push(...batch) + } + + expect(rows).toEqual([{ id: 1 }, { id: 2 }]) + expect(executeTransaction).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + queries: [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1, 0], + }, + ], + }) + ) + expect(executeTransaction).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + queries: [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1, 1], + }, + ], + }) + ) + expect(executeTransaction).toHaveBeenNthCalledWith( + 3, + expect.objectContaining({ + queries: [ + { + sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;', + params: [1, 2], + }, + ], + }) + ) + }) + + it('should quote table identifiers used in batch queries', async () => { + vi.mocked(executeTransaction).mockResolvedValueOnce([]) + + const rows: any[] = [] + + for await (const batch of getTableDataBatches( + 'user"events', + mockDataSource, + mockConfig, + 10 + )) { + rows.push(...batch) + } + + expect(rows).toEqual([]) + expect(executeTransaction).toHaveBeenCalledWith( + expect.objectContaining({ + queries: [ + { + sql: 'SELECT * FROM "user""events" LIMIT ? OFFSET ?;', + params: [10, 0], + }, + ], + }) + ) + }) + }) + + describe('quoteSqlIdentifier', () => { + it('should quote double quotes inside identifiers', () => { + expect(quoteSqlIdentifier('user"events')).toBe('"user""events"') + }) + }) + describe('createExportResponse', () => { it('should create a valid response for a CSV file', () => { const response = createExportResponse( @@ -155,4 +272,24 @@ describe('Database Operations Module', () => { ) }) }) + + describe('streaming export responses', () => { + it('should create a readable text stream response', async () => { + const stream = createTextStream(async (enqueue) => { + enqueue('hello') + enqueue(' world') + }) + const response = createExportStreamResponse( + stream, + 'notes.txt', + 'text/plain' + ) + + expect(response.headers.get('Content-Type')).toBe('text/plain') + expect(response.headers.get('Content-Disposition')).toBe( + 'attachment; filename="notes.txt"' + ) + await expect(response.text()).resolves.toBe('hello world') + }) + }) }) diff --git a/src/export/index.ts b/src/export/index.ts index 9c40119..e112e4c 100644 --- a/src/export/index.ts +++ b/src/export/index.ts @@ -2,6 +2,12 @@ import { DataSource } from '../types' import { executeTransaction } from '../operation' import { StarbaseDBConfiguration } from '../handler' +export const EXPORT_BATCH_SIZE = 1000 + +export function quoteSqlIdentifier(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"` +} + export async function executeOperation( queries: { sql: string; params?: any[] }[], dataSource: DataSource, @@ -42,8 +48,9 @@ export async function getTableData( } // Get table data + const tableIdentifier = quoteSqlIdentifier(tableName) const dataResult = await executeOperation( - [{ sql: `SELECT * FROM ${tableName};` }], + [{ sql: `SELECT * FROM ${tableIdentifier};` }], dataSource, config ) @@ -54,6 +61,75 @@ export async function getTableData( } } +export async function tableExists( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration +): Promise { + const tableExistsResult = await executeOperation( + [ + { + sql: `SELECT name FROM sqlite_master WHERE type='table' AND name=?;`, + params: [tableName], + }, + ], + dataSource, + config + ) + + return !!tableExistsResult && tableExistsResult.length > 0 +} + +export async function* getTableDataBatches( + tableName: string, + dataSource: DataSource, + config: StarbaseDBConfiguration, + batchSize = EXPORT_BATCH_SIZE +): AsyncGenerator { + let offset = 0 + const tableIdentifier = quoteSqlIdentifier(tableName) + const normalizedBatchSize = Math.max(1, Math.floor(batchSize)) + + while (true) { + const rows = await executeOperation( + [ + { + sql: `SELECT * FROM ${tableIdentifier} LIMIT ? OFFSET ?;`, + params: [normalizedBatchSize, offset], + }, + ], + dataSource, + config + ) + + if (rows.length === 0) { + break + } + + yield rows + offset += rows.length + } +} + +export function createTextStream( + write: (enqueue: (chunk: string) => void) => Promise +): ReadableStream { + const encoder = new TextEncoder() + + return new ReadableStream({ + async start(controller) { + try { + await write((chunk) => + controller.enqueue(encoder.encode(chunk)) + ) + controller.close() + } catch (error) { + controller.error(error) + } + }, + }) +} + export function createExportResponse( data: any, fileName: string, @@ -68,3 +144,16 @@ export function createExportResponse( return new Response(blob, { headers }) } + +export function createExportStreamResponse( + stream: ReadableStream, + fileName: string, + contentType: string +): Response { + const headers = new Headers({ + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${fileName}"`, + }) + + return new Response(stream, { headers }) +} diff --git a/src/export/json.test.ts b/src/export/json.test.ts index 3fe4a8c..45f32c4 100644 --- a/src/export/json.test.ts +++ b/src/export/json.test.ts @@ -1,13 +1,44 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' import { exportTableToJsonRoute } from './json' -import { getTableData, createExportResponse } from './index' +import { getTableDataBatches, tableExists } from './index' import { createResponse } from '../utils' import type { DataSource } from '../types' import type { StarbaseDBConfiguration } from '../handler' +function createTestTextStream( + write: (enqueue: (chunk: string) => void) => Promise +): ReadableStream { + const encoder = new TextEncoder() + + return new ReadableStream({ + async start(controller) { + await write((chunk) => controller.enqueue(encoder.encode(chunk))) + controller.close() + }, + }) +} + +async function* batches(...chunks: any[][]) { + for (const chunk of chunks) { + yield chunk + } +} + vi.mock('./index', () => ({ - getTableData: vi.fn(), - createExportResponse: vi.fn(), + tableExists: vi.fn(), + getTableDataBatches: vi.fn(), + createTextStream: createTestTextStream, + createExportStreamResponse: ( + stream: ReadableStream, + fileName: string, + contentType: string + ) => + new Response(stream, { + headers: { + 'Content-Type': contentType, + 'Content-Disposition': `attachment; filename="${fileName}"`, + }, + }), })) vi.mock('../utils', () => ({ @@ -41,7 +72,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(tableExists).mockResolvedValue(false) const response = await exportTableToJsonRoute( 'missing_table', @@ -59,13 +90,8 @@ 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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue(batches(mockData)) const response = await exportTableToJsonRoute( 'users', @@ -73,27 +99,23 @@ describe('JSON Export Module', () => { mockConfig ) - expect(getTableData).toHaveBeenCalledWith( + expect(tableExists).toHaveBeenCalledWith( 'users', mockDataSource, mockConfig ) - expect(createExportResponse).toHaveBeenCalledWith( - JSON.stringify(mockData, null, 4), - 'users_export.json', - 'application/json' + expect(getTableDataBatches).toHaveBeenCalledWith( + 'users', + mockDataSource, + mockConfig ) expect(response.headers.get('Content-Type')).toBe('application/json') + expect(await response.json()).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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue(batches()) const response = await exportTableToJsonRoute( 'empty_table', @@ -101,12 +123,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.json()).toEqual([]) }) it('should escape special characters in JSON properly', async () => { @@ -114,12 +132,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(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue( + batches(specialCharsData) ) const response = await exportTableToJsonRoute( @@ -128,19 +143,33 @@ 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.json()).toEqual(specialCharsData) + }) + + it('should stream multiple batches into one JSON array', async () => { + vi.mocked(tableExists).mockResolvedValue(true) + vi.mocked(getTableDataBatches).mockReturnValue( + batches([{ id: 1, name: 'Alice' }], [{ id: 2, name: 'Bob' }]) + ) + + const response = await exportTableToJsonRoute( + 'users', + mockDataSource, + mockConfig + ) + + expect(await response.json()).toEqual([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]) }) 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(tableExists).mockRejectedValue(new Error('Database Error')) const response = await exportTableToJsonRoute( 'users', diff --git a/src/export/json.ts b/src/export/json.ts index c0ab811..9630404 100644 --- a/src/export/json.ts +++ b/src/export/json.ts @@ -1,4 +1,9 @@ -import { getTableData, createExportResponse } from './index' +import { + createExportStreamResponse, + createTextStream, + getTableDataBatches, + tableExists, +} from './index' import { createResponse } from '../utils' import { DataSource } from '../types' import { StarbaseDBConfiguration } from '../handler' @@ -9,9 +14,9 @@ export async function exportTableToJsonRoute( config: StarbaseDBConfiguration ): Promise { try { - const data = await getTableData(tableName, dataSource, config) + const exists = await tableExists(tableName, dataSource, config) - if (data === null) { + if (!exists) { return createResponse( undefined, `Table '${tableName}' does not exist.`, @@ -19,11 +24,27 @@ export async function exportTableToJsonRoute( ) } - // Convert the result to JSON - const jsonData = JSON.stringify(data, null, 4) + const stream = createTextStream(async (enqueue) => { + let isFirstRow = true - return createExportResponse( - jsonData, + enqueue('[') + + for await (const rows of getTableDataBatches( + tableName, + dataSource, + config + )) { + for (const row of rows) { + enqueue(`${isFirstRow ? '' : ','}${JSON.stringify(row)}`) + isFirstRow = false + } + } + + enqueue(']') + }) + + return createExportStreamResponse( + stream, `${tableName}_export.json`, 'application/json' )