Stream large database exports

src/export/csv.test.ts

@@ -1,23 +1,11 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest'
 import { exportTableToCsvRoute } from './csv'
-import { getTableData, createExportResponse } from './index'
-import { createResponse } from '../utils'
+import { executeTransaction } from '../operation'
 import type { DataSource } from '../types'
 import type { StarbaseDBConfiguration } from '../handler'
 
-vi.mock('./index', () => ({
-    getTableData: vi.fn(),
-    createExportResponse: vi.fn(),
-}))
-
-vi.mock('../utils', () => ({
-    createResponse: vi.fn(
-        (data, message, status) =>
-            new Response(JSON.stringify({ result: data, error: message }), {
-                status,
-                headers: { 'Content-Type': 'application/json' },
-            })
-    ),
+vi.mock('../operation', () => ({
+    executeTransaction: vi.fn(),
 }))
 
 let mockDataSource: DataSource
@@ -27,8 +15,7 @@ beforeEach(() => {
     vi.clearAllMocks()
 
     mockDataSource = {
-        source: 'external',
-        external: { dialect: 'sqlite' },
+        source: 'internal',
         rpc: {
             executeQuery: vi.fn(),
         },
@@ -42,51 +29,43 @@ 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(executeTransaction)
+            .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',
             mockDataSource,
             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'
+        expect(await response.text()).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(executeTransaction).mockResolvedValueOnce([[]])
 
         const response = await exportTableToCsvRoute(
             'non_existent_table',
             mockDataSource,
             mockConfig
         )
 
-        expect(getTableData).toHaveBeenCalledWith(
-            'non_existent_table',
-            mockDataSource,
-            mockConfig
-        )
         expect(response.status).toBe(404)
 
         const jsonResponse: { error: string } = await response.json()
@@ -95,64 +74,81 @@ 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 when an existing table has no rows', async () => {
+        vi.mocked(executeTransaction)
+            .mockResolvedValueOnce([[{ name: 'empty_table' }]])
+            .mockResolvedValueOnce([[{ name: 'id' }, { name: 'name' }]])
+            .mockResolvedValueOnce([[]])
 
         const response = await exportTableToCsvRoute(
             'empty_table',
             mockDataSource,
             mockConfig
         )
 
-        expect(getTableData).toHaveBeenCalledWith(
-            'empty_table',
+        expect(await response.text()).toBe('id,name\n')
+    })
+
+    it('should escape commas, quotes, and new lines in CSV values', async () => {
+        vi.mocked(executeTransaction)
+            .mockResolvedValueOnce([[{ name: 'special_chars' }]])
+            .mockResolvedValueOnce([
+                [{ name: 'id' }, { name: 'name' }, { name: 'bio' }],
+            ])
+            .mockResolvedValueOnce([
+                [
+                    {
+                        id: 1,
+                        name: 'Sahithi, is',
+                        bio: 'line one\nwith a "quote"',
+                    },
+                ],
+            ])
+
+        const response = await exportTableToCsvRoute(
+            'special_chars',
             mockDataSource,
             mockConfig
         )
-        expect(createExportResponse).toHaveBeenCalledWith(
-            '',
-            'empty_table_export.csv',
-            'text/csv'
+
+        expect(await response.text()).toBe(
+            'id,name,bio\n1,"Sahithi, is","line one\nwith a ""quote"""\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"' },
-        ])
-
-        vi.mocked(createExportResponse).mockReturnValue(
-            new Response('mocked-csv-content', {
-                headers: { 'Content-Type': 'text/csv' },
-            })
-        )
+    it('should page additional CSV rows only as the stream is read', async () => {
+        vi.mocked(executeTransaction)
+            .mockResolvedValueOnce([[{ name: 'users' }]])
+            .mockResolvedValueOnce([[{ name: 'id' }, { name: 'name' }]])
+            .mockResolvedValueOnce([[{ id: 1, name: 'Alice' }]])
+            .mockResolvedValueOnce([[{ id: 2, name: 'Bob' }]])
+            .mockResolvedValueOnce([[]])
 
         const response = await exportTableToCsvRoute(
-            'special_chars',
+            'users',
             mockDataSource,
-            mockConfig
+            mockConfig,
+            { batchSize: 1 }
         )
 
-        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(executeTransaction).toHaveBeenCalledTimes(3)
+        expect(await response.text()).toBe('id,name\n1,Alice\n2,Bob\n')
+        expect(executeTransaction).toHaveBeenCalledTimes(5)
+        expect(vi.mocked(executeTransaction).mock.calls[3][0].queries).toEqual([
+            {
+                sql: 'SELECT * FROM "users" LIMIT ? OFFSET ?;',
+                params: [1, 1],
+            },
+        ])
     })
 
-    it('should return 500 on an unexpected error', async () => {
+    it('should return 500 on an unexpected error before streaming', async () => {
         const consoleErrorMock = vi
             .spyOn(console, 'error')
             .mockImplementation(() => {})
-        vi.mocked(getTableData).mockRejectedValue(new Error('Database Error'))
+        vi.mocked(executeTransaction).mockRejectedValue(
+            new Error('Database Error')
+        )
 
         const response = await exportTableToCsvRoute(
             'users',
@@ -163,5 +159,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()
     })
 })

src/export/csv.ts

@@ -1,51 +1,76 @@
-import { getTableData, createExportResponse } from './index'
+import {
+    createExportResponse,
+    createTextStream,
+    formatCsvValue,
+    getTableColumns,
+    getTableDataPage,
+    iterateTableRows,
+    resolveExportBatchSize,
+    tableExists,
+    type ExportOptions,
+    type ExportRow,
+} from './index'
 import { createResponse } from '../utils'
-import { DataSource } from '../types'
-import { StarbaseDBConfiguration } from '../handler'
+import type { DataSource } from '../types'
+import type { StarbaseDBConfiguration } from '../handler'
+
+async function* createCsvExportIterator(opts: {
+    tableName: string
+    dataSource: DataSource
+    config: StarbaseDBConfiguration
+    batchSize: number
+    columns: string[]
+    firstPage: ExportRow[]
+}): AsyncGenerator<string> {
+    const { columns } = opts
+
+    if (columns.length) {
+        yield columns.map(formatCsvValue).join(',') + '\n'
+    }
+
+    for await (const row of iterateTableRows(opts)) {
+        yield columns.map((column) => formatCsvValue(row[column])).join(',') +
+            '\n'
+    }
+}
 
 export async function exportTableToCsvRoute(
     tableName: string,
     dataSource: DataSource,
-    config: StarbaseDBConfiguration
+    config: StarbaseDBConfiguration,
+    options: ExportOptions = {}
 ): Promise<Response> {
     try {
-        const data = await getTableData(tableName, dataSource, config)
+        const batchSize = resolveExportBatchSize(options.batchSize)
+        const exists = await tableExists(tableName, dataSource, config)
 
-        if (data === null) {
+        if (!exists) {
             return createResponse(
                 undefined,
                 `Table '${tableName}' does not exist.`,
                 404
             )
         }
 
-        // 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 [columnsResult, firstPage] = await Promise.all([
+            getTableColumns(tableName, dataSource, config),
+            getTableDataPage(tableName, dataSource, config, batchSize, 0),
+        ])
+        const columns = columnsResult.length
+            ? columnsResult
+            : Object.keys(firstPage[0] ?? {})
 
         return createExportResponse(
-            csvContent,
+            createTextStream(
+                createCsvExportIterator({
+                    tableName,
+                    dataSource,
+                    config,
+                    batchSize,
+                    columns,
+                    firstPage,
+                })
+            ),
             `${tableName}_export.csv`,
             'text/csv'
         )