Skip to content

fix: stream database dumps to support large databases #128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
avasis-ai wants to merge 1 commit into
base: main
Choose a base branch
from
+248 −64
Open

fix: stream database dumps to support large databases #128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
117 changes: 99 additions & 18 deletions src/export/dump.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,12 @@ beforeEach(() => {
}
})

async function collectStream(response: Response): Promise<string> {
return await response.text()
}

describe('Database Dump Module', () => {
it('should return a database dump when tables exist', async () => {
it('should return a streaming database dump when tables exist', async () => {
vi.mocked(executeOperation)
.mockResolvedValueOnce([{ name: 'users' }, { name: 'orders' }])
.mockResolvedValueOnce([
Expand All @@ -60,24 +64,24 @@ describe('Database Dump Module', () => {
const response = await dumpDatabaseRoute(mockDataSource, mockConfig)

expect(response).toBeInstanceOf(Response)
expect(response.headers.get('Content-Type')).toBe(
'application/x-sqlite3'
)
expect(response.headers.get('Content-Type')).toBe('text/sql')
expect(response.headers.get('Content-Disposition')).toBe(
'attachment; filename="database_dump.sql"'
)

const dumpText = await response.text()
const dumpText = await collectStream(response)
expect(dumpText).toContain('BEGIN TRANSACTION;')
expect(dumpText).toContain('COMMIT;')
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 () => {
Expand All @@ -86,11 +90,11 @@ describe('Database Dump Module', () => {
const response = await dumpDatabaseRoute(mockDataSource, mockConfig)

expect(response).toBeInstanceOf(Response)
expect(response.headers.get('Content-Type')).toBe(
'application/x-sqlite3'
)
const dumpText = await response.text()
expect(dumpText).toBe('SQLite format 3\0')
expect(response.headers.get('Content-Type')).toBe('text/sql')
const dumpText = await collectStream(response)
expect(dumpText).toContain('BEGIN TRANSACTION;')
expect(dumpText).toContain('COMMIT;')
expect(dumpText).not.toContain('INSERT INTO')
})

it('should handle databases with tables but no data', async () => {
Expand All @@ -104,11 +108,11 @@ describe('Database Dump Module', () => {
const response = await dumpDatabaseRoute(mockDataSource, mockConfig)

expect(response).toBeInstanceOf(Response)
const dumpText = await response.text()
const dumpText = await collectStream(response)
expect(dumpText).toContain(
'CREATE TABLE users (id INTEGER, name TEXT);'
)
expect(dumpText).not.toContain('INSERT INTO users VALUES')
expect(dumpText).not.toContain('INSERT INTO "users" VALUES')
})

it('should escape single quotes properly in string values', async () => {
Expand All @@ -122,12 +126,55 @@ describe('Database Dump Module', () => {
const response = await dumpDatabaseRoute(mockDataSource, mockConfig)

expect(response).toBeInstanceOf(Response)
const dumpText = await response.text()
const dumpText = await collectStream(response)
expect(dumpText).toContain(
"INSERT INTO users VALUES (1, 'Alice''s adventure');"
"INSERT INTO \"users\" VALUES (1, 'Alice''s adventure');"
)
})

it('should handle NULL values', async () => {
vi.mocked(executeOperation)
.mockResolvedValueOnce([{ name: 'users' }])
.mockResolvedValueOnce([
{ sql: 'CREATE TABLE users (id INTEGER, name TEXT);' },
])
.mockResolvedValueOnce([{ id: 1, name: null }])

const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
const dumpText = await collectStream(response)
expect(dumpText).toContain('INSERT INTO "users" VALUES (1, NULL);')
})

it('should handle boolean values', async () => {
vi.mocked(executeOperation)
.mockResolvedValueOnce([{ name: 'flags' }])
.mockResolvedValueOnce([
{ sql: 'CREATE TABLE flags (id INTEGER, active INTEGER);' },
])
.mockResolvedValueOnce([
{ id: 1, active: true },
{ id: 2, active: false },
])

const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
const dumpText = await collectStream(response)
expect(dumpText).toContain('INSERT INTO "flags" VALUES (1, 1);')
expect(dumpText).toContain('INSERT INTO "flags" VALUES (2, 0);')
})

it('should escape table names with special characters', async () => {
vi.mocked(executeOperation)
.mockResolvedValueOnce([{ name: 'my-table' }])
.mockResolvedValueOnce([
{ sql: 'CREATE TABLE "my-table" (id INTEGER);' },
])
.mockResolvedValueOnce([{ id: 1 }])

const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
const dumpText = await collectStream(response)
expect(dumpText).toContain('INSERT INTO "my-table" VALUES (1);')
})

it('should return a 500 response when an error occurs', async () => {
const consoleErrorMock = vi
.spyOn(console, 'error')
Expand All @@ -141,5 +188,39 @@ describe('Database Dump Module', () => {
expect(response.status).toBe(500)
const jsonResponse: { error: string } = await response.json()
expect(jsonResponse.error).toBe('Failed to create database dump')
consoleErrorMock.mockRestore()
})

it('should fetch data in batches for large tables', async () => {
const firstBatch = Array.from({ length: 500 }, (_, i) => ({
id: i + 1,
name: `user${i + 1}`,
}))
const secondBatch = Array.from({ length: 100 }, (_, i) => ({
id: i + 501,
name: `user${i + 501}`,
}))

vi.mocked(executeOperation)
.mockResolvedValueOnce([{ name: 'users' }])
.mockResolvedValueOnce([
{ sql: 'CREATE TABLE users (id INTEGER, name TEXT);' },
])
.mockResolvedValueOnce(firstBatch)
.mockResolvedValueOnce(secondBatch)

const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
const dumpText = await collectStream(response)

expect(dumpText).toContain('INSERT INTO "users" VALUES (1, \'user1\');')
expect(dumpText).toContain(
'INSERT INTO "users" VALUES (500, \'user500\');'
)
expect(dumpText).toContain(
'INSERT INTO "users" VALUES (501, \'user501\');'
)
expect(dumpText).toContain(
'INSERT INTO "users" VALUES (600, \'user600\');'
)
})
})
195 changes: 149 additions & 46 deletions src/export/dump.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,67 +3,170 @@ import { StarbaseDBConfiguration } from '../handler'
import { DataSource } from '../types'
import { createResponse } from '../utils'

const BATCH_SIZE = 500

function escapeSqlIdentifier(name: string): string {
return `"${name.replace(/"/g, '""')}"`
}

function escapeSqlValue(value: unknown): string {
if (value === null || value === undefined) {
return 'NULL'
}
if (typeof value === 'number') {
return String(value)
}
if (typeof value === 'boolean') {
return value ? '1' : '0'
}
if (value instanceof ArrayBuffer || ArrayBuffer.isView(value)) {
const bytes = new Uint8Array(
value instanceof ArrayBuffer ? value : value.buffer
)
const hex = Array.from(bytes)
.map((b) => b.toString(16).padStart(2, '0'))
.join('')
return `X'${hex}'`
}
return `'${String(value).replace(/'/g, "''")}'`
}

async function* streamTableData(
tableName: string,
dataSource: DataSource,
config: StarbaseDBConfiguration
): AsyncGenerator<string> {
const escapedName = escapeSqlIdentifier(tableName)

const schemaResult = await executeOperation(
[
{
sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name=?;`,
params: [tableName],
},
],
dataSource,
config
)

if (schemaResult.length) {
const schema = schemaResult[0].sql
yield `\n-- Table: ${tableName}\n${schema};\n\n`
}

let offset = 0
let hasMore = true

while (hasMore) {
const dataResult = await executeOperation(
[
{
sql: `SELECT * FROM ${escapedName} LIMIT ? OFFSET ?;`,
params: [BATCH_SIZE, offset],
},
],
dataSource,
config
)

if (!dataResult || dataResult.length === 0) {
hasMore = false
continue
}

for (const row of dataResult) {
const values = Object.values(row).map((v) => escapeSqlValue(v))
yield `INSERT INTO ${escapedName} VALUES (${values.join(', ')});\n`
}

if (dataResult.length < BATCH_SIZE) {
hasMore = false
} else {
offset += BATCH_SIZE
}
}

yield '\n'
}

function createSqlDumpStream(
tableNames: string[],
dataSource: DataSource,
config: StarbaseDBConfiguration
): ReadableStream<Uint8Array> {
const encoder = new TextEncoder()
let tableIterator: AsyncGenerator<string> | null = null
let tableIndex = 0
let chunkBuffer = 'BEGIN TRANSACTION;\n'

return new ReadableStream<Uint8Array>({
async pull(controller) {
try {
while (tableIndex < tableNames.length) {
if (!tableIterator) {
tableIterator = streamTableData(
tableNames[tableIndex],
dataSource,
config
)
}

const result = await tableIterator.next()
if (result.done) {
tableIterator = null
tableIndex++
continue
}

chunkBuffer += result.value

if (chunkBuffer.length >= 8192) {
controller.enqueue(encoder.encode(chunkBuffer))
chunkBuffer = ''
return
}
}

chunkBuffer += 'COMMIT;\n'
if (chunkBuffer.length > 0) {
controller.enqueue(encoder.encode(chunkBuffer))
chunkBuffer = ''
}
controller.close()
} catch (error: any) {
console.error('Database Dump Stream Error:', error)
controller.error(error)
}
},
})
}

export async function dumpDatabaseRoute(
dataSource: DataSource,
config: StarbaseDBConfiguration
): Promise<Response> {
try {
// Get all table names
const tablesResult = await executeOperation(
[{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }],
[
{
sql: "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'tmp_%' ORDER BY name;",
},
],
dataSource,
config
)

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
)

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
)

for (const row of dataResult) {
const values = Object.values(row).map((value) =>
typeof value === 'string'
? `'${value.replace(/'/g, "''")}'`
: value
)
dumpContent += `INSERT INTO ${table} VALUES (${values.join(', ')});\n`
}

dumpContent += '\n'
}
const tableNames = tablesResult.map((row: any) => row.name)

// Create a Blob from the dump content
const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' })
const stream = createSqlDumpStream(tableNames, dataSource, config)

const headers = new Headers({
'Content-Type': 'application/x-sqlite3',
'Content-Disposition': 'attachment; filename="database_dump.sql"',
'Content-Type': 'text/sql',
'Content-Disposition': `attachment; filename="database_dump.sql"`,
'Transfer-Encoding': 'chunked',
})

return new Response(blob, { headers })
return new Response(stream, { headers })
} catch (error: any) {
console.error('Database Dump Error:', error)
return createResponse(undefined, 'Failed to create database dump', 500)
Expand Down