fix: stream-based export for large databases (SQL, CSV, JSON)

src/export/csv.test.ts

@@ -1,37 +1,93 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest'
-import { exportTableToCsvRoute } from './csv'
-import { getTableData, createExportResponse } from './index'
-import { createResponse } from '../utils'
 import type { DataSource } from '../types'
 import type { StarbaseDBConfiguration } from '../handler'
 
+const mockExecuteOperation = vi.fn()
+
 vi.mock('./index', () => ({
+    executeOperation: (...args: any[]) => mockExecuteOperation(...args),
+    getTableDataChunked: async function* (
+        tableName: string,
+        dataSource: any,
+        config: any,
+        chunkSize: number = 1000
+    ) {
+        let offset = 0
+        while (true) {
+            const chunk = await mockExecuteOperation(
+                [{ sql: `SELECT * FROM ${tableName} LIMIT ? OFFSET ?;`, params: [chunkSize, offset] }],
+                dataSource,
+                config
+            )
+            if (!chunk || chunk.length === 0) break
+            yield chunk
+            if (chunk.length < chunkSize) break
+            offset += chunkSize
+        }
+    },
+    createStreamingExportResponse: (
+        producer: any,
+        fileName: string,
+        contentType: string
+    ) => {
+        const { readable, writable } = new TransformStream()
+        const writer = writable.getWriter()
+
+        const done = (async () => {
+            try {
+                await producer(writer)
+            } finally {
+                await writer.close()
+            }
+        })()
+
+        const response = new Response(readable, {
+            headers: {
+                'Content-Type': contentType,
+                'Content-Disposition': `attachment; filename="${fileName}"`,
+            },
+        })
+        ;(response as any).__producerDone = done
+        return response
+    },
+    writeChunk: async (writer: WritableStreamDefaultWriter, content: string) => {
+        await writer.write(new TextEncoder().encode(content))
+    },
+    createExportResponse: (data: any, fileName: string, contentType: string) => {
+        const blob = new Blob([data], { type: contentType })
+        return new Response(blob, {
+            headers: {
+                'Content-Type': contentType,
+                'Content-Disposition': `attachment; filename="${fileName}"`,
+            },
+        })
+    },
     getTableData: vi.fn(),
-    createExportResponse: vi.fn(),
 }))
 
 vi.mock('../utils', () => ({
     createResponse: vi.fn(
-        (data, message, status) =>
+        (data: any, message: any, status: any) =>
             new Response(JSON.stringify({ result: data, error: message }), {
                 status,
                 headers: { 'Content-Type': 'application/json' },
             })
     ),
 }))
 
+import { exportTableToCsvRoute } from './csv'
+
 let mockDataSource: DataSource
 let mockConfig: StarbaseDBConfiguration
 
 beforeEach(() => {
     vi.clearAllMocks()
+    mockExecuteOperation.mockReset()
 
     mockDataSource = {
         source: 'external',
         external: { dialect: 'sqlite' },
-        rpc: {
-            executeQuery: vi.fn(),
-        },
+        rpc: { executeQuery: vi.fn() },
     } as any
 
     mockConfig = {
@@ -43,116 +99,88 @@ beforeEach(() => {
 
 describe('CSV Export Module', () => {
     it('should return a CSV file when table data exists', async () => {
-        vi.mocked(getTableData).mockResolvedValue([
+        // Table exists check
+        mockExecuteOperation.mockResolvedValueOnce([{ name: 'users' }])
+        // Data chunk
+        mockExecuteOperation.mockResolvedValueOnce([
             { 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' },
-            })
-        )
-
         const response = await exportTableToCsvRoute(
             'users',
             mockDataSource,
             mockConfig
         )
 
-        expect(getTableData).toHaveBeenCalledWith(
-            'users',
-            mockDataSource,
-            mockConfig
-        )
-        expect(createExportResponse).toHaveBeenCalledWith(
-            'id,name,age\n1,Alice,30\n2,Bob,25\n',
-            'users_export.csv',
-            'text/csv'
-        )
+        expect(response).toBeInstanceOf(Response)
         expect(response.headers.get('Content-Type')).toBe('text/csv')
+
+        const csvText = await response.text()
+        expect(csvText).toContain('id,name,age')
+        expect(csvText).toContain('1,Alice,30')
+        expect(csvText).toContain('2,Bob,25')
     })
 
     it('should return 404 if table does not exist', async () => {
-        vi.mocked(getTableData).mockResolvedValue(null)
+        mockExecuteOperation.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()
         expect(jsonResponse.error).toBe(
             "Table 'non_existent_table' does not exist."
         )
     })
 
-    it('should handle empty table (return only headers)', async () => {
-        vi.mocked(getTableData).mockResolvedValue([])
-
-        vi.mocked(createExportResponse).mockReturnValue(
-            new Response('mocked-csv-content', {
-                headers: { 'Content-Type': 'text/csv' },
-            })
-        )
+    it('should handle empty table (no output)', async () => {
+        // Table exists
+        mockExecuteOperation.mockResolvedValueOnce([{ name: 'empty_table' }])
+        // Empty data
+        mockExecuteOperation.mockResolvedValueOnce([])
 
         const response = await exportTableToCsvRoute(
             'empty_table',
             mockDataSource,
             mockConfig
         )
 
-        expect(getTableData).toHaveBeenCalledWith(
-            'empty_table',
-            mockDataSource,
-            mockConfig
-        )
-        expect(createExportResponse).toHaveBeenCalledWith(
-            '',
-            'empty_table_export.csv',
-            'text/csv'
-        )
+        expect(response).toBeInstanceOf(Response)
         expect(response.headers.get('Content-Type')).toBe('text/csv')
+
+        const csvText = await response.text()
+        expect(csvText).toBe('') // No headers, no data
     })
 
     it('should escape commas and quotes in CSV values', async () => {
-        vi.mocked(getTableData).mockResolvedValue([
+        // Table exists
+        mockExecuteOperation.mockResolvedValueOnce([{ name: 'special_chars' }])
+        // Data with special chars
+        mockExecuteOperation.mockResolvedValueOnce([
             { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' },
         ])
 
-        vi.mocked(createExportResponse).mockReturnValue(
-            new Response('mocked-csv-content', {
-                headers: { 'Content-Type': 'text/csv' },
-            })
-        )
-
         const response = await exportTableToCsvRoute(
             'special_chars',
             mockDataSource,
             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')
+        const csvText = await response.text()
+        expect(csvText).toContain('id,name,bio')
+        expect(csvText).toContain('1,"Sahithi, is","my forever ""penguin"""')
     })
 
     it('should return 500 on an unexpected error', async () => {
         const consoleErrorMock = vi
             .spyOn(console, 'error')
             .mockImplementation(() => {})
-        vi.mocked(getTableData).mockRejectedValue(new Error('Database Error'))
+        mockExecuteOperation.mockRejectedValue(new Error('Database Error'))
 
         const response = await exportTableToCsvRoute(
             'users',
@@ -163,5 +191,6 @@ describe('CSV Export Module', () => {
         expect(response.status).toBe(500)
         const jsonResponse: { error: string } = await response.json()
         expect(jsonResponse.error).toBe('Failed to export table to CSV')
+        consoleErrorMock.mockRestore()
     })
 })

src/export/csv.ts

@@ -1,51 +1,92 @@
-import { getTableData, createExportResponse } from './index'
+import {
+    getTableDataChunked,
+    executeOperation,
+    createStreamingExportResponse,
+    writeChunk,
+    createExportResponse,
+} from './index'
 import { createResponse } from '../utils'
 import { DataSource } from '../types'
 import { StarbaseDBConfiguration } from '../handler'
 
+const BREATHE_MS = 10
+
+function formatCsvRow(row: any): string {
+    return Object.values(row)
+        .map((value) => {
+            if (
+                typeof value === 'string' &&
+                (value.includes(',') ||
+                    value.includes('"') ||
+                    value.includes('\n'))
+            ) {
+                return `"${value.replace(/"/g, '""')}"`
+            }
+            return value === null ? '' : value
+        })
+        .join(',')
+}
+
 export async function exportTableToCsvRoute(
     tableName: string,
     dataSource: DataSource,
     config: StarbaseDBConfiguration
 ): Promise<Response> {
     try {
-        const data = await getTableData(tableName, dataSource, config)
+        // Verify table exists
+        const tableExistsResult = await executeOperation(
+            [
+                {
+                    sql: `SELECT name FROM sqlite_master WHERE type='table' AND name=?;`,
+                    params: [tableName],
+                },
+            ],
+            dataSource,
+            config
+        )
 
-        if (data === null) {
+        if (!tableExistsResult || tableExistsResult.length === 0) {
             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'
-            })
-        }
+        return createStreamingExportResponse(
+            async (writer) => {
+                let headersWritten = false
+
+                for await (const chunk of getTableDataChunked(
+                    tableName,
+                    dataSource,
+                    config,
+                    1000
+                )) {
+                    if (chunk.length === 0) continue
+
+                    // Write CSV headers from first row of first chunk
+                    if (!headersWritten) {
+                        await writeChunk(
+                            writer,
+                            Object.keys(chunk[0]).join(',') + '\n'
+                        )
+                        headersWritten = true
+                    }
+
+                    // Write rows
+                    let batchContent = ''
+                    for (const row of chunk) {
+                        batchContent += formatCsvRow(row) + '\n'
+                    }
+                    await writeChunk(writer, batchContent)
 
-        return createExportResponse(
-            csvContent,
+                    // Breathing interval
+                    if (BREATHE_MS > 0) {
+                        await new Promise((r) => setTimeout(r, BREATHE_MS))
+                    }
+                }
+            },
             `${tableName}_export.csv`,
             'text/csv'
         )