fix(export): stream chunks incrementally for SQL, CSV, and JSON dumps to prevent OOM

src/export/csv.test.ts

@@ -1,14 +1,17 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest'
 import { exportTableToCsvRoute } from './csv'
-import { getTableData, createExportResponse } from './index'
-import { createResponse } from '../utils'
+import { executeOperation } from './index'
+import type * as IndexModule from './index'
 import type { DataSource } from '../types'
 import type { StarbaseDBConfiguration } from '../handler'
 
-vi.mock('./index', () => ({
-    getTableData: vi.fn(),
-    createExportResponse: vi.fn(),
-}))
+vi.mock('./index', async () => {
+    const actual = await vi.importActual<typeof IndexModule>('./index')
+    return {
+        ...actual,
+        executeOperation: vi.fn(),
+    }
+})
 
 vi.mock('../utils', () => ({
     createResponse: vi.fn(
@@ -43,116 +46,161 @@ 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([
+                { 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(createExportResponse).toHaveBeenCalledWith(
-            'id,name,age\n1,Alice,30\n2,Bob,25\n',
-            'users_export.csv',
-            'text/csv'
-        )
         expect(response.headers.get('Content-Type')).toBe('text/csv')
+        expect(response.headers.get('Content-Disposition')).toBe(
+            'attachment; filename="users_export.csv"'
+        )
+
+        const csvText = await response.text()
+        expect(csvText).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(executeOperation).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 (return empty content)', async () => {
+        vi.mocked(executeOperation)
+            .mockResolvedValueOnce([{ name: 'empty_table' }])
+            .mockResolvedValueOnce([])
 
         const response = await exportTableToCsvRoute(
             'empty_table',
             mockDataSource,
             mockConfig
         )
 
-        expect(getTableData).toHaveBeenCalledWith(
-            'empty_table',
+        expect(response.status).toBe(200)
+        expect(response.headers.get('Content-Type')).toBe('text/csv')
+        expect(response.headers.get('Content-Disposition')).toBe(
+            'attachment; filename="empty_table_export.csv"'
+        )
+
+        const csvText = await response.text()
+        expect(csvText).toBe('')
+    })
+
+    it('should escape commas and quotes in CSV values', async () => {
+        vi.mocked(executeOperation)
+            .mockResolvedValueOnce([{ name: 'special_chars' }])
+            .mockResolvedValueOnce([
+                { id: 1, name: 'Sahithi, is', bio: 'my forever "penguin"' },
+            ])
+
+        const response = await exportTableToCsvRoute(
+            'special_chars',
             mockDataSource,
             mockConfig
         )
-        expect(createExportResponse).toHaveBeenCalledWith(
-            '',
-            'empty_table_export.csv',
-            'text/csv'
-        )
+
         expect(response.headers.get('Content-Type')).toBe('text/csv')
+        const csvText = await response.text()
+        expect(csvText).toBe(
+            'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n'
+        )
     })
 
-    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 paginate large tables across multiple queries', async () => {
+        const firstPage = Array.from({ length: 1000 }, (_, i) => ({
+            id: i + 1,
+            name: `User${i + 1}`,
+        }))
+        const secondPage = [{ id: 1001, name: 'User1001' }]
 
-        vi.mocked(createExportResponse).mockReturnValue(
-            new Response('mocked-csv-content', {
-                headers: { 'Content-Type': 'text/csv' },
-            })
-        )
+        vi.mocked(executeOperation)
+            .mockResolvedValueOnce([{ name: 'users' }])
+            .mockResolvedValueOnce(firstPage)
+            .mockResolvedValueOnce(secondPage)
 
         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 csvText = await response.text()
+        expect(csvText).toContain('id,name\n')
+        expect(csvText).toContain('1,User1\n')
+        expect(csvText).toContain('1000,User1000\n')
+        expect(csvText).toContain('1001,User1001\n')
+
+        const issuedSql = vi
+            .mocked(executeOperation)
+            .mock.calls.map((c) => (c[0] as Array<{ sql: string }>)[0].sql)
+        expect(issuedSql).toContain('SELECT * FROM users LIMIT 1000 OFFSET 0;')
+        expect(issuedSql).toContain(
+            'SELECT * FROM users LIMIT 1000 OFFSET 1000;'
         )
-        expect(response.headers.get('Content-Type')).toBe('text/csv')
+    })
+
+    it('should stream chunks incrementally without buffering entire dataset', async () => {
+        const page1 = Array.from({ length: 1000 }, (_, i) => ({
+            id: i + 1,
+            name: `User${i + 1}`,
+        }))
+        const page2 = [{ id: 1001, name: 'User1001' }]
+
+        vi.mocked(executeOperation)
+            .mockResolvedValueOnce([{ name: 'users' }])
+            .mockResolvedValueOnce(page1)
+            .mockResolvedValueOnce(page2)
+
+        const response = await exportTableToCsvRoute(
+            'users',
+            mockDataSource,
+            mockConfig
+        )
+
+        expect(response.body).toBeInstanceOf(ReadableStream)
+        const reader = response.body!.getReader()
+        const decoder = new TextDecoder()
+        let receivedChunks = 0
+        let combined = ''
+
+        while (true) {
+            const { done, value } = await reader.read()
+            if (done) break
+            receivedChunks++
+            combined += decoder.decode(value, { stream: true })
+        }
+
+        expect(receivedChunks).toBeGreaterThan(0)
+        expect(combined).toContain('id,name\n1,User1\n')
+        expect(combined).toContain('1001,User1001\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(executeOperation).mockRejectedValue(
+            new Error('Database Error')
+        )
 
         const response = await exportTableToCsvRoute(
             'users',

src/export/csv.ts

@@ -1,4 +1,8 @@
-import { getTableData, createExportResponse } from './index'
+import {
+    executeOperation,
+    createExportResponse,
+    EXPORT_PAGE_SIZE,
+} from './index'
 import { createResponse } from '../utils'
 import { DataSource } from '../types'
 import { StarbaseDBConfiguration } from '../handler'
@@ -9,47 +13,96 @@ export async function exportTableToCsvRoute(
     config: StarbaseDBConfiguration
 ): Promise<Response> {
     try {
-        const data = await getTableData(tableName, dataSource, config)
+        // Verify if the table exists up front
+        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'
-            })
-        }
+        const encoder = new TextEncoder()
+        const stream = new ReadableStream<Uint8Array>({
+            async start(controller) {
+                try {
+                    let offset = 0
+                    let headerWritten = false
+
+                    while (true) {
+                        const rows = await executeOperation(
+                            [
+                                {
+                                    sql: `SELECT * FROM ${tableName} LIMIT ${EXPORT_PAGE_SIZE} OFFSET ${offset};`,
+                                },
+                            ],
+                            dataSource,
+                            config
+                        )
+
+                        if (!rows.length) {
+                            break
+                        }
+
+                        let chunk = ''
+                        if (!headerWritten) {
+                            chunk += Object.keys(rows[0]).join(',') + '\n'
+                            headerWritten = true
+                        }
+
+                        for (const row of rows) {
+                            chunk +=
+                                Object.values(row)
+                                    .map((value) => {
+                                        if (
+                                            typeof value === 'string' &&
+                                            (value.includes(',') ||
+                                                value.includes('"') ||
+                                                value.includes('\n') ||
+                                                value.includes('\r'))
+                                        ) {
+                                            return `"${value.replace(/"/g, '""')}"`
+                                        }
+                                        return value === null ||
+                                            value === undefined
+                                            ? ''
+                                            : value
+                                    })
+                                    .join(',') + '\n'
+                        }
+
+                        controller.enqueue(encoder.encode(chunk))
+
+                        if (rows.length < EXPORT_PAGE_SIZE) {
+                            break
+                        }
+                        offset += EXPORT_PAGE_SIZE
+                    }
+
+                    controller.close()
+                } catch (error) {
+                    controller.error(error)
+                }
+            },
+        })
 
         return createExportResponse(
-            csvContent,
+            stream,
             `${tableName}_export.csv`,
             'text/csv'
         )
-    } catch (error: any) {
+    } catch (error: unknown) {
         console.error('CSV Export Error:', error)
         return createResponse(undefined, 'Failed to export table to CSV', 500)
     }