Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 36 additions & 45 deletions src/export/csv.test.ts
Original file line number Diff line number Diff line change
@@ -1,13 +1,22 @@
import { describe, it, expect, vi, beforeEach } from 'vitest'
import { exportTableToCsvRoute } from './csv'
import { getTableData, createExportResponse } from './index'
import { tableExists, getTableDataChunked } 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(),
tableExists: vi.fn(),
getTableDataChunked: vi.fn(),
createStreamingExportResponse: vi.fn((stream, fileName, contentType) => {
const headers = new Headers({
'Content-Type': contentType,
'Content-Disposition': `attachment; filename="${fileName}"`,
'Transfer-Encoding': 'chunked',
})
return new Response(stream, { headers })
}),
CHUNK_SIZE: 5000,
}))

vi.mock('../utils', () => ({
Expand Down Expand Up @@ -43,46 +52,41 @@ beforeEach(() => {

describe('CSV Export Module', () => {
it('should return a CSV file when table data exists', async () => {
vi.mocked(getTableData).mockResolvedValue([
vi.mocked(tableExists).mockResolvedValue(true)

// Fewer rows than CHUNK_SIZE -> loop breaks after first call
vi.mocked(getTableDataChunked).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(
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(response.headers.get('Content-Type')).toBe('text/csv')

const text = await response.text()
expect(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',
mockDataSource,
mockConfig
)

expect(getTableData).toHaveBeenCalledWith(
expect(tableExists).toHaveBeenCalledWith(
'non_existent_table',
mockDataSource,
mockConfig
Expand All @@ -95,64 +99,51 @@ 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 handle empty table (return empty stream)', async () => {
vi.mocked(tableExists).mockResolvedValue(true)
vi.mocked(getTableDataChunked).mockResolvedValueOnce([])

const response = await exportTableToCsvRoute(
'empty_table',
mockDataSource,
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')

const text = await response.text()
expect(text).toBe('')
})

it('should escape commas and quotes in CSV values', async () => {
vi.mocked(getTableData).mockResolvedValue([
vi.mocked(tableExists).mockResolvedValue(true)

vi.mocked(getTableDataChunked).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'
const text = await response.text()
expect(text).toBe(
'id,name,bio\n1,"Sahithi, is","my forever ""penguin"""\n'
)
expect(response.headers.get('Content-Type')).toBe('text/csv')
})

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',
Expand Down
108 changes: 79 additions & 29 deletions src/export/csv.ts
Original file line number Diff line number Diff line change
@@ -1,51 +1,101 @@
import { getTableData, createExportResponse } from './index'
import {
tableExists,
getTableDataChunked,
createStreamingExportResponse,
CHUNK_SIZE,
} from './index'
import { createResponse } from '../utils'
import { DataSource } from '../types'
import { StarbaseDBConfiguration } from '../handler'

/**
* Format a single value for CSV output.
* Wraps in double quotes and escapes inner quotes when the value contains
* commas, double quotes, or newlines.
*/
function formatCsvValue(value: unknown): string {
if (value === null || value === undefined) {
return ''
}

const str = String(value)
if (str.includes(',') || str.includes('"') || str.includes('\n')) {
return `"${str.replace(/"/g, '""')}"`
}
return str
}

export async function exportTableToCsvRoute(
tableName: string,
dataSource: DataSource,
config: StarbaseDBConfiguration
): Promise<Response> {
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.`,
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()
let headerWritten = false

const stream = new ReadableStream({
async start(controller) {
try {
let offset = 0

while (true) {
const rows = await getTableDataChunked(
tableName,
offset,
CHUNK_SIZE,
dataSource,
config
)

if (!rows || rows.length === 0) {
break
}

let chunk = ''

// Write header row from first batch of results
if (!headerWritten) {
chunk += Object.keys(rows[0]).join(',') + '\n'
headerWritten = true
}

// Write data rows
for (const row of rows) {
chunk +=
Object.values(row)
.map(formatCsvValue)
.join(',') + '\n'
}

controller.enqueue(encoder.encode(chunk))

if (rows.length < CHUNK_SIZE) {
break
}

offset += CHUNK_SIZE
}

controller.close()
} catch (error) {
controller.error(error)
}
},
})

return createExportResponse(
csvContent,
return createStreamingExportResponse(
stream,
`${tableName}_export.csv`,
'text/csv'
)
Expand Down
Loading