test: add Vitest suites for allowlist and CSV import routes to raise coverage above 75%

src/allowlist/index.test.ts

@@ -0,0 +1,132 @@
+import { expect, test, vi } from 'vitest'
+import { isQueryAllowed } from './index'
+import { DataSource } from '../types'
+import { StarbaseDBConfiguration } from '../handler'
+
+test('isQueryAllowed should return true if isEnabled is false', async () => {
+    const mockDataSource = {} as DataSource
+    const mockConfig = {} as StarbaseDBConfiguration
+
+    const result = await isQueryAllowed({
+        sql: 'SELECT * FROM users',
+        isEnabled: false,
+        dataSource: mockDataSource,
+        config: mockConfig,
+    })
+
+    expect(result).toBe(true)
+})
+
+test('isQueryAllowed should return true if config.role is admin', async () => {
+    const mockDataSource = {} as DataSource
+    const mockConfig = { role: 'admin' } as StarbaseDBConfiguration
+
+    const result = await isQueryAllowed({
+        sql: 'SELECT * FROM users',
+        isEnabled: true,
+        dataSource: mockDataSource,
+        config: mockConfig,
+    })
+
+    expect(result).toBe(true)
+})
+
+test('isQueryAllowed should return an Error if no SQL is provided', async () => {
+    const mockDataSource = {
+        source: 'test-source',
+        rpc: {
+            executeQuery: vi.fn().mockResolvedValue([]),
+        },
+    } as unknown as DataSource
+    const mockConfig = { role: 'user' } as StarbaseDBConfiguration
+
+    const result = await isQueryAllowed({
+        sql: '',
+        isEnabled: true,
+        dataSource: mockDataSource,
+        config: mockConfig,
+    })
+
+    expect(result).toBeInstanceOf(Error)
+    expect((result as Error).message).toBe(
+        'No SQL provided for allowlist check'
+    )
+})
+
+test('isQueryAllowed should allow matching queries in allowlist', async () => {
+    const mockExecuteQuery = vi
+        .fn()
+        .mockResolvedValue([
+            { sql_statement: 'SELECT * FROM users;', source: 'test-source' },
+        ])
+    const mockDataSource = {
+        source: 'test-source',
+        rpc: {
+            executeQuery: mockExecuteQuery,
+        },
+    } as unknown as DataSource
+    const mockConfig = { role: 'user' } as StarbaseDBConfiguration
+
+    const result = await isQueryAllowed({
+        sql: 'SELECT * FROM users',
+        isEnabled: true,
+        dataSource: mockDataSource,
+        config: mockConfig,
+    })
+
+    expect(result).toBe(true)
+    expect(mockExecuteQuery).toHaveBeenCalledWith({
+        sql: 'SELECT sql_statement, source FROM tmp_allowlist_queries WHERE source="test-source"',
+    })
+})
+
+test('isQueryAllowed should reject and audit queries not in allowlist', async () => {
+    const mockExecuteQuery = vi
+        .fn()
+        .mockResolvedValue([
+            { sql_statement: 'SELECT * FROM users', source: 'test-source' },
+        ])
+    const mockDataSource = {
+        source: 'test-source',
+        rpc: {
+            executeQuery: mockExecuteQuery,
+        },
+    } as unknown as DataSource
+    const mockConfig = { role: 'user' } as StarbaseDBConfiguration
+
+    await expect(
+        isQueryAllowed({
+            sql: 'DELETE FROM users',
+            isEnabled: true,
+            dataSource: mockDataSource,
+            config: mockConfig,
+        })
+    ).rejects.toThrow('Query not allowed')
+
+    expect(mockExecuteQuery).toHaveBeenLastCalledWith({
+        sql: 'INSERT INTO tmp_allowlist_rejections (sql_statement, source) VALUES (?, ?)',
+        params: ['DELETE FROM users', 'test-source'],
+    })
+})
+
+test('isQueryAllowed should handle loadAllowlist query execution failure gracefully', async () => {
+    const mockExecuteQuery = vi
+        .fn()
+        .mockRejectedValue(new Error('DB Connection Error'))
+    const mockDataSource = {
+        source: 'test-source',
+        rpc: {
+            executeQuery: mockExecuteQuery,
+        },
+    } as unknown as DataSource
+    const mockConfig = { role: 'user' } as StarbaseDBConfiguration
+
+    await expect(
+        isQueryAllowed({
+            sql: 'SELECT * FROM users',
+            isEnabled: true,
+            dataSource: mockDataSource,
+            config: mockConfig,
+        })
+    ).rejects.toThrow('Query not allowed')
+})

src/export/dump.test.ts

@@ -21,6 +21,8 @@ vi.mock('../utils', () => ({
 
 let mockDataSource: DataSource
 let mockConfig: StarbaseDBConfiguration
+let mockR2Bucket: any
+let mockFetch: any
 
 beforeEach(() => {
     vi.clearAllMocks()
@@ -36,6 +38,17 @@ beforeEach(() => {
         role: 'admin',
         features: { allowlist: true, rls: true, rest: true },
     }
+
+    mockR2Bucket = {
+        createMultipartUpload: vi.fn().mockResolvedValue({
+            uploadPart: vi.fn().mockResolvedValue({ etag: 'mock-etag' }),
+            complete: vi.fn().mockResolvedValue(undefined),
+            abort: vi.fn().mockResolvedValue(undefined),
+        }),
+    }
+
+    mockFetch = vi.fn().mockResolvedValue(new Response('OK', { status: 200 }))
+    vi.stubGlobal('fetch', mockFetch)
 })
 
 describe('Database Dump Module', () => {
@@ -57,7 +70,12 @@ describe('Database Dump Module', () => {
                 { id: 2, total: 49.5 },
             ])
 
-        const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
+        const req = new Request('http://localhost/export/dump')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
 
         expect(response).toBeInstanceOf(Response)
         expect(response.headers.get('Content-Type')).toBe(
@@ -71,19 +89,24 @@ describe('Database Dump Module', () => {
         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 () => {
         vi.mocked(executeOperation).mockResolvedValueOnce([])
 
-        const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
+        const req = new Request('http://localhost/export/dump')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
 
         expect(response).toBeInstanceOf(Response)
         expect(response.headers.get('Content-Type')).toBe(
@@ -101,7 +124,12 @@ describe('Database Dump Module', () => {
             ])
             .mockResolvedValueOnce([])
 
-        const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
+        const req = new Request('http://localhost/export/dump')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
 
         expect(response).toBeInstanceOf(Response)
         const dumpText = await response.text()
@@ -119,12 +147,17 @@ describe('Database Dump Module', () => {
             ])
             .mockResolvedValueOnce([{ id: 1, bio: "Alice's adventure" }])
 
-        const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
+        const req = new Request('http://localhost/export/dump')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
 
         expect(response).toBeInstanceOf(Response)
         const dumpText = await response.text()
         expect(dumpText).toContain(
-            "INSERT INTO users VALUES (1, 'Alice''s adventure');"
+            "INSERT INTO \"users\" VALUES (1, 'Alice''s adventure');"
         )
     })
 
@@ -136,10 +169,92 @@ describe('Database Dump Module', () => {
             new Error('Database Error')
         )
 
-        const response = await dumpDatabaseRoute(mockDataSource, mockConfig)
+        const req = new Request('http://localhost/export/dump')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
 
         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 return a 400 response when async is requested but bucket is missing', async () => {
+        const req = new Request('http://localhost/export/dump?async=true')
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            mockConfig
+        )
+
+        expect(response.status).toBe(400)
+        const jsonResponse = (await response.json()) as any
+        expect(jsonResponse.error).toContain(
+            'require an EXPORT_BUCKET R2 binding'
+        )
+    })
+
+    it('should trigger background R2 dump when async is requested', async () => {
+        vi.mocked(executeOperation)
+            .mockResolvedValueOnce([{ name: 'users' }])
+            .mockResolvedValueOnce([
+                { sql: 'CREATE TABLE users (id INTEGER, name TEXT);' },
+            ])
+            .mockResolvedValueOnce([{ id: 1, name: 'Alice' }])
+
+        const configWithBucket: StarbaseDBConfiguration = {
+            ...mockConfig,
+            export: {
+                bucket: mockR2Bucket,
+                callbackUrl: 'http://callback.url/notify',
+                chunkSize: 100,
+            },
+        }
+
+        const mockExecutionContext = {
+            waitUntil: vi.fn(),
+        } as any
+
+        const req = new Request(
+            'http://localhost/export/dump?async=true&filename=test-dump.sql'
+        )
+        const response = await dumpDatabaseRoute(
+            req,
+            mockDataSource,
+            configWithBucket,
+            mockExecutionContext
+        )
+
+        expect(response.status).toBe(202)
+        const jsonResponse = (await response.json()) as any
+        expect(jsonResponse.result.status).toBe('running')
+        expect(jsonResponse.result.filename).toBe('test-dump.sql')
+
+        expect(mockExecutionContext.waitUntil).toHaveBeenCalled()
+        const backgroundPromise =
+            mockExecutionContext.waitUntil.mock.calls[0][0]
+
+        // Await the background process
+        await backgroundPromise
+
+        // Verify R2 was used
+        expect(mockR2Bucket.createMultipartUpload).toHaveBeenCalledWith(
+            'test-dump.sql',
+            {
+                httpMetadata: { contentType: 'application/sql' },
+            }
+        )
+
+        // Verify callback url notified
+        expect(mockFetch).toHaveBeenCalledWith(
+            'http://callback.url/notify',
+            expect.objectContaining({
+                method: 'POST',
+                body: expect.stringContaining('"status":"completed"'),
+            })
+        )
     })
 })