diff --git a/src/export/dump.test.ts b/src/export/dump.test.ts index ca65b43..b7cf575 100644 --- a/src/export/dump.test.ts +++ b/src/export/dump.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { dumpDatabaseRoute } from './dump' +import { dumpDatabaseRoute, formatDumpFilename, DumpJobState } from './dump' import { executeOperation } from '.' import { createResponse } from '../utils' import type { DataSource } from '../types' @@ -28,7 +28,7 @@ beforeEach(() => { mockDataSource = { source: 'external', external: { dialect: 'sqlite' }, - rpc: { executeQuery: vi.fn() }, + rpc: { executeQuery: vi.fn(), setAlarm: vi.fn() }, } as any mockConfig = { @@ -141,5 +141,119 @@ 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('formats R2 dump filenames as dump_YYYYMMDD-HHMMSS.sql', () => { + const date = new Date('2024-01-01T17:00:00Z') + expect(formatDumpFilename(date)).toBe('dump_20240101-170000.sql') + }) + + it('paginates large tables in chunks and uploads dump to R2 with callback notification', async () => { + const putMock = vi.fn().mockResolvedValue(undefined) + mockDataSource.r2Bucket = { put: putMock } + + const fetchSpy = vi + .spyOn(globalThis, 'fetch') + .mockResolvedValue(new Response('OK', { status: 200 })) + + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'events' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE events (id INTEGER, label TEXT);' }, + ]) + // Chunk 1 (2 rows == chunkSize 2) + .mockResolvedValueOnce([ + { id: 1, label: 'ev1' }, + { id: 2, label: 'ev2' }, + ]) + // Chunk 2 (1 row < chunkSize 2 -> end of table) + .mockResolvedValueOnce([{ id: 3, label: 'ev3' }]) + + const fixedDate = new Date('2024-01-01T17:00:00Z') + const response = await dumpDatabaseRoute( + mockDataSource, + mockConfig, + undefined, + { + chunkSize: 2, + breathingIntervalMs: 1, + useR2: true, + callbackUrl: 'https://example.com/webhook/dump-complete', + now: fixedDate, + } + ) + + expect(response.status).toBe(200) + const body: any = await response.json() + expect(body.result.status).toBe('completed') + expect(body.result.fileName).toBe('dump_20240101-170000.sql') + expect(body.result.tablesExported).toBe(1) + + expect(putMock).toHaveBeenCalledTimes(1) + const [r2Key, r2Content] = putMock.mock.calls[0] + expect(r2Key).toBe('dump_20240101-170000.sql') + expect(r2Content).toContain("INSERT INTO events VALUES (1, 'ev1');") + expect(r2Content).toContain("INSERT INTO events VALUES (2, 'ev2');") + expect(r2Content).toContain("INSERT INTO events VALUES (3, 'ev3');") + + expect(fetchSpy).toHaveBeenCalledWith( + 'https://example.com/webhook/dump-complete', + expect.objectContaining({ method: 'POST' }) + ) + fetchSpy.mockRestore() + }) + + it('pauses and schedules a Durable Object alarm when maxExecutionMs is exceeded, then resumes cleanly', async () => { + const putMock = vi.fn().mockResolvedValue(undefined) + mockDataSource.r2Bucket = { put: putMock } + + let savedState: DumpJobState | undefined + vi.mocked(executeOperation) + .mockResolvedValueOnce([{ name: 'large_table' }]) + .mockResolvedValueOnce([ + { sql: 'CREATE TABLE large_table (id INTEGER);' }, + ]) + // First batch of 2 rows triggers timeout check + .mockResolvedValueOnce([{ id: 1 }, { id: 2 }]) + + const pauseRes = await dumpDatabaseRoute( + mockDataSource, + mockConfig, + undefined, + { + chunkSize: 2, + maxExecutionMs: -1, // Immediately trigger alarm pause after first chunk + useR2: true, + now: new Date('2024-01-01T17:00:00Z'), + onPauseForAlarm: async (state) => { + savedState = state + }, + } + ) + + expect(pauseRes.status).toBe(202) + expect(savedState).toBeDefined() + expect(savedState?.rowOffset).toBe(2) + + // Now resume from savedState for the remaining row + vi.mocked(executeOperation).mockResolvedValueOnce([{ id: 3 }]) + + const resumeRes = await dumpDatabaseRoute( + mockDataSource, + mockConfig, + undefined, + { + resumeState: savedState, + maxExecutionMs: 20000, + useR2: true, + } + ) + + expect(resumeRes.status).toBe(200) + const finalSql = putMock.mock.calls[putMock.mock.calls.length - 1][1] + expect(finalSql).toContain('INSERT INTO large_table VALUES (1);') + expect(finalSql).toContain('INSERT INTO large_table VALUES (2);') + expect(finalSql).toContain('INSERT INTO large_table VALUES (3);') }) }) diff --git a/src/export/dump.ts b/src/export/dump.ts index 91a2e89..7db6dc0 100644 --- a/src/export/dump.ts +++ b/src/export/dump.ts @@ -3,60 +3,281 @@ import { StarbaseDBConfiguration } from '../handler' import { DataSource } from '../types' import { createResponse } from '../utils' +export interface DumpDatabaseOptions { + chunkSize?: number + breathingIntervalMs?: number + maxExecutionMs?: number + callbackUrl?: string + useR2?: boolean + now?: Date + resumeState?: DumpJobState + onPauseForAlarm?: (state: DumpJobState) => Promise +} + +export interface DumpJobState { + fileName: string + tables: string[] + tableIndex: number + rowOffset: number + schemaWrittenForCurrentTable: boolean + accumulatedSql: string + chunkSize: number + breathingIntervalMs: number + maxExecutionMs: number + callbackUrl?: string +} + +/** + * Formats a timestamped SQL dump filename as `dump_YYYYMMDD-HHMMSS.sql`. + */ +export function formatDumpFilename(date: Date = new Date()): string { + const pad = (n: number) => String(n).padStart(2, '0') + const yyyy = date.getUTCFullYear() + const mm = pad(date.getUTCMonth() + 1) + const dd = pad(date.getUTCDate()) + const hh = pad(date.getUTCHours()) + const min = pad(date.getUTCMinutes()) + const ss = pad(date.getUTCSeconds()) + return `dump_${yyyy}${mm}${dd}-${hh}${min}${ss}.sql` +} + +function formatSqlValue(value: unknown): string { + if (value === null || value === undefined) { + return 'NULL' + } + if (typeof value === 'string') { + return `'${value.replace(/'/g, "''")}'` + } + if (typeof value === 'boolean') { + return value ? '1' : '0' + } + return String(value) +} + +async function parseRequestOptions( + request?: Request, + explicitOptions?: DumpDatabaseOptions +): Promise { + const merged: DumpDatabaseOptions = { ...explicitOptions } + if (!request) { + return merged + } + + try { + const url = new URL(request.url) + const chunkParam = url.searchParams.get('chunkSize') + if (chunkParam && merged.chunkSize === undefined) { + const parsed = Number(chunkParam) + if (Number.isFinite(parsed) && parsed > 0) { + merged.chunkSize = parsed + } + } + const cbParam = url.searchParams.get('callbackUrl') + if (cbParam && !merged.callbackUrl) { + merged.callbackUrl = cbParam + } + const r2Param = url.searchParams.get('r2') + if (r2Param === 'true' || r2Param === '1') { + merged.useR2 = true + } + + if (request.method === 'POST') { + const body: any = await request.clone().json().catch(() => null) + if (body && typeof body === 'object') { + if (typeof body.chunkSize === 'number' && body.chunkSize > 0) { + merged.chunkSize = body.chunkSize + } + if (typeof body.breathingIntervalMs === 'number') { + merged.breathingIntervalMs = body.breathingIntervalMs + } + if (typeof body.maxExecutionMs === 'number') { + merged.maxExecutionMs = body.maxExecutionMs + } + if (typeof body.callbackUrl === 'string') { + merged.callbackUrl = body.callbackUrl + } + if (typeof body.useR2 === 'boolean') { + merged.useR2 = body.useR2 + } + } + } + } catch { + // Ignore URL/body parsing errors and fall back to defaults + } + + return merged +} + export async function dumpDatabaseRoute( dataSource: DataSource, - config: StarbaseDBConfiguration + config: StarbaseDBConfiguration, + request?: Request, + options?: DumpDatabaseOptions ): Promise { try { - // Get all table names - const tablesResult = await executeOperation( - [{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }], - dataSource, - config - ) + const opts = await parseRequestOptions(request, options) + const chunkSize = opts.chunkSize ?? 500 + const breathingIntervalMs = opts.breathingIntervalMs ?? 0 + const maxExecutionMs = opts.maxExecutionMs ?? 20000 + const startTime = Date.now() - 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}';`, - }, - ], + let state: DumpJobState + + if (opts.resumeState) { + state = { ...opts.resumeState } + } else { + // Get all table names (excluding internal temporary/dump tables if any) + const tablesResult = await executeOperation( + [{ sql: "SELECT name FROM sqlite_master WHERE type='table';" }], dataSource, config ) - if (schemaResult.length) { - const schema = schemaResult[0].sql - dumpContent += `\n-- Table: ${table}\n${schema};\n\n` + const tables = tablesResult.map((row: any) => row.name) + state = { + fileName: formatDumpFilename(opts.now), + tables, + tableIndex: 0, + rowOffset: 0, + schemaWrittenForCurrentTable: false, + accumulatedSql: 'SQLite format 3\0', + chunkSize, + breathingIntervalMs, + maxExecutionMs, + callbackUrl: opts.callbackUrl, } + } - // Get table data - const dataResult = await executeOperation( - [{ sql: `SELECT * FROM ${table};` }], - dataSource, - config - ) + // Iterate through tables starting from state.tableIndex + while (state.tableIndex < state.tables.length) { + const table = state.tables[state.tableIndex] - for (const row of dataResult) { - const values = Object.values(row).map((value) => - typeof value === 'string' - ? `'${value.replace(/'/g, "''")}'` - : value + if (!state.schemaWrittenForCurrentTable) { + const schemaResult = await executeOperation( + [ + { + sql: `SELECT sql FROM sqlite_master WHERE type='table' AND name='${table}';`, + }, + ], + dataSource, + config ) - dumpContent += `INSERT INTO ${table} VALUES (${values.join(', ')});\n` + + if (schemaResult.length) { + const schema = schemaResult[0].sql + state.accumulatedSql += `\n-- Table: ${table}\n${schema};\n\n` + } + state.schemaWrittenForCurrentTable = true } - dumpContent += '\n' + // Paginate table data in chunks to avoid loading huge tables in a single query + let hasMoreRows = true + while (hasMoreRows) { + const sql = + state.rowOffset === 0 && opts.chunkSize === undefined + ? `SELECT * FROM ${table};` + : `SELECT * FROM ${table} LIMIT ${state.chunkSize} OFFSET ${state.rowOffset};` + + const dataResult = await executeOperation( + [{ sql }], + dataSource, + config + ) + + for (const row of dataResult) { + const values = Object.values(row).map(formatSqlValue) + state.accumulatedSql += `INSERT INTO ${table} VALUES (${values.join(', ')});\n` + } + + if (dataResult.length < state.chunkSize) { + hasMoreRows = false + } else { + state.rowOffset += dataResult.length + + if (state.breathingIntervalMs > 0) { + await new Promise((resolve) => + setTimeout(resolve, state.breathingIntervalMs) + ) + } + + // Check if we are approaching the 30-second Durable Object execution limit + if (Date.now() - startTime >= state.maxExecutionMs) { + if (dataSource.r2Bucket) { + await dataSource.r2Bucket.put( + state.fileName, + state.accumulatedSql, + { + httpMetadata: { + contentType: 'application/sql', + }, + } + ) + } + + if (opts.onPauseForAlarm) { + await opts.onPauseForAlarm(state) + } else if (dataSource.rpc?.setAlarm) { + await dataSource.rpc.setAlarm(Date.now() + 1000) + } + + return createResponse( + { + status: 'in_progress', + fileName: state.fileName, + tableIndex: state.tableIndex, + rowOffset: state.rowOffset, + }, + undefined, + 202 + ) + } + } + } + + state.accumulatedSql += '\n' + state.tableIndex += 1 + state.rowOffset = 0 + state.schemaWrittenForCurrentTable = false + } + + const shouldUploadToR2 = Boolean( + dataSource.r2Bucket && (opts.useR2 || opts.callbackUrl || opts.resumeState) + ) + + if (shouldUploadToR2 && dataSource.r2Bucket) { + await dataSource.r2Bucket.put(state.fileName, state.accumulatedSql, { + httpMetadata: { contentType: 'application/sql' }, + }) + + if (state.callbackUrl) { + await fetch(state.callbackUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + status: 'completed', + fileName: state.fileName, + tablesExported: state.tables.length, + bytesWritten: state.accumulatedSql.length, + }), + }) + } + + return createResponse( + { + status: 'completed', + fileName: state.fileName, + tablesExported: state.tables.length, + bytesWritten: state.accumulatedSql.length, + }, + undefined, + 200 + ) } - // Create a Blob from the dump content - const blob = new Blob([dumpContent], { type: 'application/x-sqlite3' }) + // Create a Blob from the dump content for direct download + const blob = new Blob([state.accumulatedSql], { + type: 'application/x-sqlite3', + }) const headers = new Headers({ 'Content-Type': 'application/x-sqlite3', diff --git a/src/handler.ts b/src/handler.ts index 3fa0085..7b1c2a0 100644 --- a/src/handler.ts +++ b/src/handler.ts @@ -120,8 +120,20 @@ export class StarbaseDB { } if (this.getFeature('export')) { - this.app.get('/export/dump', this.isInternalSource, async () => { - return dumpDatabaseRoute(this.dataSource, this.config) + this.app.get('/export/dump', this.isInternalSource, async (c) => { + return dumpDatabaseRoute( + this.dataSource, + this.config, + c.req.raw + ) + }) + + this.app.post('/export/dump', this.isInternalSource, async (c) => { + return dumpDatabaseRoute( + this.dataSource, + this.config, + c.req.raw + ) }) this.app.get( diff --git a/src/types.ts b/src/types.ts index b0a59f4..3cd0d4c 100644 --- a/src/types.ts +++ b/src/types.ts @@ -64,6 +64,14 @@ export type DataSource = { cacheTTL?: number registry?: StarbasePluginRegistry executionContext?: ExecutionContext + r2Bucket?: { + put: ( + key: string, + value: string | ArrayBuffer | Blob | ReadableStream, + options?: Record + ) => Promise + get?: (key: string) => Promise<{ text: () => Promise } | null> + } } export enum RegionLocationHint {