diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml new file mode 100644 index 0000000..02365ef --- /dev/null +++ b/pnpm-workspace.yaml @@ -0,0 +1,3 @@ +allowBuilds: + esbuild: true + workerd: true diff --git a/src/allowlist/index.test.ts b/src/allowlist/index.test.ts new file mode 100644 index 0000000..87d7036 --- /dev/null +++ b/src/allowlist/index.test.ts @@ -0,0 +1,178 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { isQueryAllowed } from './index' +import type { DataSource } from '../types' +import type { StarbaseDBConfiguration } from '../handler' + +let mockDataSource: DataSource +let mockConfig: StarbaseDBConfiguration + +beforeEach(() => { + vi.clearAllMocks() + + mockDataSource = { + source: 'internal', + rpc: { + executeQuery: vi.fn(), + }, + } as any + + mockConfig = { + outerbaseApiKey: 'mock-api-key', + role: 'client', + features: { allowlist: true, rls: true, rest: true }, + } +}) + +describe('isQueryAllowed - Feature Flags & Roles', () => { + it('should allow any query if allowlist is not enabled', async () => { + const result = await isQueryAllowed({ + sql: 'DROP TABLE users;', + isEnabled: false, + dataSource: mockDataSource, + config: mockConfig, + }) + expect(result).toBe(true) + expect(mockDataSource.rpc.executeQuery).not.toHaveBeenCalled() + }) + + it('should allow any query if role is admin', async () => { + mockConfig.role = 'admin' + const result = await isQueryAllowed({ + sql: 'DROP TABLE users;', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + expect(result).toBe(true) + expect(mockDataSource.rpc.executeQuery).not.toHaveBeenCalled() + }) + + it('should return error if sql is empty', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue([]) + 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') + }) +}) + +describe('isQueryAllowed - Query Matching & Rejection', () => { + it('should allow query that matches allowlist AST (ignoring trailing semicolon)', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValueOnce([ + { + sql_statement: 'SELECT id, name FROM users WHERE id = 1;', + source: 'internal', + }, + ] as any) + + const result = await isQueryAllowed({ + sql: 'SELECT id, name FROM users WHERE id = 1', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(result).toBe(true) + expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledTimes(1) + }) + + it('should reject query not in allowlist and record rejected query', async () => { + vi.mocked(mockDataSource.rpc.executeQuery) + .mockResolvedValueOnce([ + { + sql_statement: 'SELECT id FROM users;', + source: 'internal', + }, + ] as any) + .mockResolvedValueOnce([] as any) // for addRejectedQuery + + await expect( + isQueryAllowed({ + sql: 'SELECT * FROM users;', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow('Query not allowed') + + expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledTimes(2) + expect(mockDataSource.rpc.executeQuery).toHaveBeenLastCalledWith({ + sql: 'INSERT INTO tmp_allowlist_rejections (sql_statement, source) VALUES (?, ?)', + params: ['SELECT * FROM users;', 'internal'], + }) + }) + + it('should handle AST mismatch with different structure, array length, or keys', async () => { + vi.mocked(mockDataSource.rpc.executeQuery) + .mockResolvedValueOnce([ + { + sql_statement: 'SELECT a, b, c FROM tbl;', + source: 'internal', + }, + ] as any) + .mockResolvedValueOnce([] as any) + + await expect( + isQueryAllowed({ + sql: 'SELECT a, b FROM tbl;', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow('Query not allowed') + }) +}) + +describe('isQueryAllowed - Error Handling Resilience', () => { + it('should handle loadAllowlist DB error gracefully and reject query', async () => { + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + vi.mocked(mockDataSource.rpc.executeQuery).mockRejectedValueOnce( + new Error('DB Connection lost') + ) + + await expect( + isQueryAllowed({ + sql: 'SELECT 1;', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow('Query not allowed') + + consoleErrorSpy.mockRestore() + }) + + it('should handle addRejectedQuery failure gracefully', async () => { + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + vi.mocked(mockDataSource.rpc.executeQuery) + .mockResolvedValueOnce([]) // empty allowlist + .mockRejectedValueOnce(new Error('Cannot insert rejection')) + + await expect( + isQueryAllowed({ + sql: 'SELECT 1;', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow('Query not allowed') + + consoleErrorSpy.mockRestore() + }) + + it('should throw error when SQL parsing fails', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValueOnce([]) + await expect( + isQueryAllowed({ + sql: 'INVALID SQL STATEMENT @@@@', + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow() + }) +}) diff --git a/src/do.test.ts b/src/do.test.ts index 272c4e9..681e971 100644 --- a/src/do.test.ts +++ b/src/do.test.ts @@ -3,7 +3,14 @@ import { StarbaseDBDurableObject } from './do' vi.mock('cloudflare:workers', () => { return { - DurableObject: class MockDurableObject {}, + DurableObject: class MockDurableObject { + ctx: any + env: any + constructor(ctx: any, env: any) { + this.ctx = ctx + this.env = env + } + }, } }) @@ -145,4 +152,138 @@ describe('StarbaseDBDurableObject Tests', () => { instance.executeQuery({ sql: 'INVALID QUERY' }) ).rejects.toThrow('Query failed') }) + + it('should execute raw queries with columns and rows returned', async () => { + const result = await instance.executeQuery({ + sql: 'SELECT * FROM users', + isRaw: true, + }) + + expect(result).toEqual({ + columns: ['id', 'name'], + rows: [ + [1, 'Alice'], + [2, 'Bob'], + ], + meta: { + rows_read: 2, + rows_written: 1, + }, + }) + }) + + it('should execute query with parameters', async () => { + await instance.executeQuery({ + sql: 'SELECT * FROM users WHERE id = ?', + params: [1], + }) + + expect(mockStorage.sql.exec).toHaveBeenCalledWith( + 'SELECT * FROM users WHERE id = ?', + 1 + ) + }) + + it('should return 400 for /socket fetch request without websocket upgrade header', async () => { + const req = new Request('https://example.com/socket') + const res = await instance.fetch(req) + + expect(res.status).toBe(400) + }) + + it('should handle /socket fetch request with websocket upgrade header', async () => { + const req = new Request('https://example.com/socket?sessionId=s1', { + headers: { upgrade: 'websocket' }, + }) + const res = await instance.fetch(req) + + expect(res.status).toBe(101) + expect(instance.connections.has('s1')).toBe(true) + }) + + it('should handle /socket/broadcast to all connections', async () => { + const mockWs1 = { send: vi.fn() } as any + const mockWs2 = { send: vi.fn() } as any + instance.connections.set('c1', mockWs1) + instance.connections.set('c2', mockWs2) + + const req = new Request('https://example.com/socket/broadcast', { + method: 'POST', + body: JSON.stringify({ event: 'ping' }), + }) + const res = await instance.fetch(req) + + expect(res.status).toBe(200) + expect(mockWs1.send).toHaveBeenCalledWith(JSON.stringify({ event: 'ping' })) + expect(mockWs2.send).toHaveBeenCalledWith(JSON.stringify({ event: 'ping' })) + }) + + it('should handle /socket/broadcast targeted to a single sessionId', async () => { + const mockWs1 = { send: vi.fn() } as any + const mockWs2 = { send: vi.fn() } as any + instance.connections.set('target-session', mockWs1) + instance.connections.set('other-session', mockWs2) + + const req = new Request('https://example.com/socket/broadcast?sessionId=target-session', { + method: 'POST', + body: JSON.stringify({ event: 'target' }), + }) + const res = await instance.fetch(req) + + expect(res.status).toBe(200) + expect(mockWs1.send).toHaveBeenCalledWith(JSON.stringify({ event: 'target' })) + expect(mockWs2.send).not.toHaveBeenCalled() + }) + + it('should clean up dead connections when broadcast fails', async () => { + const deadWs = { + send: vi.fn().mockImplementation(() => { + throw new Error('Socket closed') + }), + } as any + instance.connections.set('dead-session', deadWs) + + const req = new Request('https://example.com/socket/broadcast', { + method: 'POST', + body: JSON.stringify({ event: 'test' }), + }) + const res = await instance.fetch(req) + + expect(res.status).toBe(200) + expect(instance.connections.has('dead-session')).toBe(false) + }) + + it('should handle webSocketMessage query action', async () => { + const mockWs = { send: vi.fn() } as any + await instance.webSocketMessage( + mockWs, + JSON.stringify({ action: 'query', sql: 'SELECT 1', params: [] }) + ) + + expect(mockWs.send).toHaveBeenCalled() + }) + + it('should handle webSocketClose and remove connection', async () => { + const mockWs = { close: vi.fn() } as any + instance.connections.set('session-123', mockWs) + + await instance.webSocketClose(mockWs, 1000, 'Normal closure', true) + + expect(mockWs.close).toHaveBeenCalledWith( + 1000, + 'StarbaseDB is closing WebSocket connection' + ) + expect(instance.connections.has('session-123')).toBe(false) + }) + + it('should get statistics from database', async () => { + mockStorage.sql.exec.mockReturnValueOnce({ + toArray: vi.fn().mockReturnValue([{ count: 42 }]), + }) + + const stats = await instance.getStatistics() + + expect(stats.recentQueries).toBe(42) + expect(stats.activeConnections).toBe(0) + }) }) diff --git a/src/handler.test.ts b/src/handler.test.ts index 86bb328..456f379 100644 --- a/src/handler.test.ts +++ b/src/handler.test.ts @@ -96,9 +96,37 @@ describe('StarbaseDB Initialization', () => { expect(instance['config']).toBe(mockConfig) }) - it('should get feature flag correctly', () => { + it('should throw error when source is external but external config is missing', () => { + expect( + () => + new StarbaseDB({ + dataSource: { source: 'external' } as any, + config: mockConfig, + }) + ).toThrow('No external data sources available.') + }) + + it('should get feature flag correctly with default and config values', () => { expect(instance['getFeature']('rest')).toBe(true) expect(instance['getFeature']('export')).toBe(true) + + const noFeaturesInstance = new StarbaseDB({ + dataSource: mockDataSource, + config: { role: 'admin' }, + }) + expect(noFeaturesInstance['getFeature']('rest', true)).toBe(true) + expect(noFeaturesInstance['getFeature']('rest', false)).toBe(false) + + const disabledFeaturesInstance = new StarbaseDB({ + dataSource: mockDataSource, + config: { + role: 'admin', + features: { rest: false, export: false, import: false }, + }, + }) + expect(disabledFeaturesInstance['getFeature']('rest')).toBe(false) + expect(disabledFeaturesInstance['getFeature']('export')).toBe(false) + expect(disabledFeaturesInstance['getFeature']('import')).toBe(false) }) }) @@ -120,6 +148,90 @@ describe('StarbaseDB Middleware & Request Handling', () => { expect(instance['app'].fetch).toHaveBeenCalledWith(request) expect(response).toBeDefined() }) + + it('should not reinitialize if already initialized', async () => { + const request = new Request('https://example.com/api/test') + await instance.handle(request, mockExecutionContext) + await instance.handle(request, mockExecutionContext) + + expect(instance['initialized']).toBe(true) + }) + + it('should handle preAuth matching an authless plugin pathPrefix', async () => { + const mockPlugin = { + name: 'test-authless-plugin', + opts: { requiresAuth: false }, + pathPrefix: '/public/*', + register: vi.fn(), + } + const instWithPlugin = new StarbaseDB({ + dataSource: mockDataSource, + config: mockConfig, + plugins: [mockPlugin as any], + }) + + const req = new Request('https://example.com/public/dashboard') + const res = await instWithPlugin.handlePreAuth(req, mockExecutionContext) + + expect(res).toBeDefined() + }) + + it('should handle preAuth with parameterized plugin pathPrefix', async () => { + const mockPlugin = { + name: 'param-plugin', + opts: { requiresAuth: false }, + pathPrefix: '/user/:id/profile', + register: vi.fn(), + } + const instWithPlugin = new StarbaseDB({ + dataSource: mockDataSource, + config: mockConfig, + plugins: [mockPlugin as any], + }) + + const req = new Request('https://example.com/user/123/profile') + const res = await instWithPlugin.handlePreAuth(req, mockExecutionContext) + + expect(res).toBeDefined() + }) + + it('should return undefined in preAuth if route does not match authless plugin', async () => { + const mockPlugin = { + name: 'test-plugin', + opts: { requiresAuth: false }, + pathPrefix: '/public/*', + register: vi.fn(), + } + const instWithPlugin = new StarbaseDB({ + dataSource: mockDataSource, + config: mockConfig, + plugins: [mockPlugin as any], + }) + + const req = new Request('https://example.com/private/settings') + const res = await instWithPlugin.handlePreAuth(req, mockExecutionContext) + + expect(res).toBeUndefined() + }) + + it('should return undefined in preAuth for plugins that require auth', async () => { + const mockPlugin = { + name: 'auth-required-plugin', + opts: { requiresAuth: true }, + pathPrefix: '/secure/*', + register: vi.fn(), + } + const instWithPlugin = new StarbaseDB({ + dataSource: mockDataSource, + config: mockConfig, + plugins: [mockPlugin as any], + }) + + const req = new Request('https://example.com/secure/data') + const res = await instWithPlugin.handlePreAuth(req, mockExecutionContext) + + expect(res).toBeUndefined() + }) }) describe('StarbaseDB Query Execution', () => { @@ -142,7 +254,20 @@ describe('StarbaseDB Query Execution', () => { expect(response.status).toBe(200) }) - it('should return 400 if SQL query is invalid', async () => { + it('should return 400 if Content-Type is not application/json', async () => { + const request = new Request('https://example.com/query', { + method: 'POST', + body: 'plain text', + headers: { 'Content-Type': 'text/plain' }, + }) + + const response = await instance.queryRoute(request, false) + + expect(response.status).toBe(400) + expect(response.error).toBe('Content-Type must be application/json.') + }) + + it('should return 400 if SQL query is invalid or empty', async () => { const request = new Request('https://example.com/query', { method: 'POST', body: JSON.stringify({ sql: '' }), @@ -152,6 +277,22 @@ describe('StarbaseDB Query Execution', () => { const response = await instance.queryRoute(request, false) expect(response.status).toBe(400) + expect(response.error).toBe('Invalid or empty "sql" field.') + }) + + it('should return 400 if params is invalid', async () => { + const request = new Request('https://example.com/query', { + method: 'POST', + body: JSON.stringify({ sql: 'SELECT 1', params: 12345 }), + headers: { 'Content-Type': 'application/json' }, + }) + + const response = await instance.queryRoute(request, false) + + expect(response.status).toBe(400) + expect(response.error).toBe( + 'Invalid "params" field. Must be an array or object.' + ) }) it('should execute a SQL transaction', async () => { @@ -168,6 +309,42 @@ describe('StarbaseDB Query Execution', () => { expect(executeTransaction).toHaveBeenCalled() expect(response.status).toBe(200) }) + + it('should return 500 if a query in transaction has empty sql', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const request = new Request('https://example.com/query', { + method: 'POST', + body: JSON.stringify({ + transaction: [{ sql: ' ' }], + }), + headers: { 'Content-Type': 'application/json' }, + }) + + const response = await instance.queryRoute(request, false) + + expect(response.status).toBe(500) + expect(response.error).toBe( + 'Invalid or empty "sql" field in transaction.' + ) + }) + + it('should return 500 if a query in transaction has invalid params', async () => { + vi.spyOn(console, 'error').mockImplementation(() => {}) + const request = new Request('https://example.com/query', { + method: 'POST', + body: JSON.stringify({ + transaction: [{ sql: 'INSERT INTO tbl VALUES (?)', params: 999 }], + }), + headers: { 'Content-Type': 'application/json' }, + }) + + const response = await instance.queryRoute(request, false) + + expect(response.status).toBe(500) + expect(response.error).toBe( + 'Invalid "params" field in transaction. Must be an array or object.' + ) + }) }) describe('StarbaseDB Cache Expiry', () => { @@ -179,6 +356,20 @@ describe('StarbaseDB Cache Expiry', () => { params: [expect.any(Number)], }) }) + + it('should catch and log error if cache expiry fails', async () => { + mockDataSource.rpc.executeQuery = vi.fn().mockImplementationOnce(() => { + throw new Error('Cache cleanup failure') + }) + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + + await instance['expireCache']() + + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Error cleaning up expired cache entries:', + expect.any(Error) + ) + }) }) describe('StarbaseDB Error Handling', () => { diff --git a/src/import/csv.test.ts b/src/import/csv.test.ts new file mode 100644 index 0000000..dda0831 --- /dev/null +++ b/src/import/csv.test.ts @@ -0,0 +1,302 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { importTableFromCsvRoute } from './csv' +import { executeOperation } from '../export' +import { createResponse } from '../utils' +import type { DataSource } from '../types' +import type { StarbaseDBConfiguration } from '../handler' + +vi.mock('../export', () => ({ + executeOperation: vi.fn(), +})) + +vi.mock('../utils', () => ({ + createResponse: vi.fn( + (data, message, status) => + new Response(JSON.stringify({ result: data, error: message }), { + status, + headers: { 'Content-Type': 'application/json' }, + }) + ), +})) + +let mockDataSource: DataSource +let mockConfig: StarbaseDBConfiguration + +beforeEach(() => { + vi.clearAllMocks() + + mockDataSource = { + source: 'external', + external: { dialect: 'sqlite' }, + rpc: { executeQuery: vi.fn() }, + } as any + + mockConfig = { + outerbaseApiKey: 'mock-api-key', + role: 'admin', + features: { allowlist: true, rls: true, rest: true }, + } +}) + +describe('CSV Import Module', () => { + it('should return 400 if request body is empty', async () => { + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'text/csv' }, + body: null, + }) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(400) + expect(createResponse).toHaveBeenCalledWith( + undefined, + 'Request body is empty', + 400 + ) + }) + + it('should return 400 for unsupported Content-Type', async () => { + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'application/xml' }, + body: '', + }) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(400) + expect(createResponse).toHaveBeenCalledWith( + undefined, + 'Unsupported Content-Type', + 400 + ) + }) + + it('should import CSV from application/json payload', async () => { + const payload = { + data: 'id,name\n1,Alice\n2,Bob', + } + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(payload), + }) + + vi.mocked(executeOperation).mockResolvedValue(undefined as any) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(200) + expect(executeOperation).toHaveBeenCalledTimes(2) + expect(createResponse).toHaveBeenCalledWith( + expect.objectContaining({ + message: 'Imported 2 out of 2 records successfully. 0 records failed.', + failedStatements: [], + }), + undefined, + 200 + ) + }) + + it('should import raw CSV from text/csv payload', async () => { + const csvContent = 'id,name\n10,Charlie\n20,Dave' + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'text/csv' }, + body: csvContent, + }) + + vi.mocked(executeOperation).mockResolvedValue(undefined as any) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(200) + expect(executeOperation).toHaveBeenCalledTimes(2) + }) + + it('should handle multipart/form-data upload with file', async () => { + const formData = new FormData() + const blob = new Blob(['id,name\n100,Eve'], { type: 'text/csv' }) + formData.append('file', blob, 'test.csv') + + const request = new Request('http://localhost', { + method: 'POST', + body: formData, + }) + + vi.mocked(executeOperation).mockResolvedValue(undefined as any) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(200) + expect(executeOperation).toHaveBeenCalledTimes(1) + }) + + it('should return 400 for multipart/form-data if file is missing', async () => { + const formData = new FormData() + const request = new Request('http://localhost', { + method: 'POST', + body: formData, + }) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(400) + expect(createResponse).toHaveBeenCalledWith( + undefined, + 'No file uploaded', + 400 + ) + }) + + it('should return 400 if CSV data is empty or invalid header', async () => { + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'text/csv' }, + body: 'id,name\n', + }) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(400) + expect(createResponse).toHaveBeenCalledWith( + undefined, + 'Invalid CSV format or empty data', + 400 + ) + }) + + it('should apply columnMapping correctly', async () => { + const payload = { + data: 'external_id,full_name\n1,Alice', + columnMapping: { + external_id: 'id', + full_name: 'name', + }, + } + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(payload), + }) + + vi.mocked(executeOperation).mockResolvedValue(undefined as any) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(200) + expect(executeOperation).toHaveBeenCalledWith( + [ + { + sql: 'INSERT INTO users (id, name) VALUES (?, ?)', + params: ['1', 'Alice'], + }, + ], + mockDataSource, + mockConfig + ) + }) + + it('should handle partial failures during batch import', async () => { + const csvContent = 'id,name\n1,Good\n2,Bad' + const request = new Request('http://localhost', { + method: 'POST', + headers: { 'Content-Type': 'text/csv' }, + body: csvContent, + }) + + vi.mocked(executeOperation) + .mockResolvedValueOnce(undefined as any) + .mockRejectedValueOnce(new Error('Duplicate primary key')) + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(200) + expect(createResponse).toHaveBeenCalledWith( + { + message: 'Imported 1 out of 2 records successfully. 1 records failed.', + failedStatements: [ + { + statement: 'INSERT INTO users (id, name) VALUES (?, ?)', + error: 'Duplicate primary key', + }, + ], + }, + undefined, + 200 + ) + }) + + it('should catch unhandled errors and return 500', async () => { + const consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const request = { + body: true, + headers: { + get: () => { + throw new Error('Header failure') + }, + }, + } as any + + const response = await importTableFromCsvRoute( + 'users', + request, + mockDataSource, + mockConfig + ) + + expect(response.status).toBe(500) + expect(createResponse).toHaveBeenCalledWith( + undefined, + 'Failed to import CSV data: Header failure', + 500 + ) + consoleErrorSpy.mockRestore() + }) +}) diff --git a/src/index.test.ts b/src/index.test.ts new file mode 100644 index 0000000..edca4eb --- /dev/null +++ b/src/index.test.ts @@ -0,0 +1,506 @@ +import { describe, expect, it, vi, beforeEach } from 'vitest' + +vi.mock('cloudflare:workers', () => { + return { + DurableObject: class MockDurableObject {}, + } +}) + +import worker, { Env } from './index' +import { corsPreflight } from './cors' +import { jwtVerify } from 'jose' +import { StarbaseDB } from './handler' +import { InterfacePlugin } from '../plugins/interface' +import { ChangeDataCapturePlugin } from '../plugins/cdc' +import { CronPlugin } from '../plugins/cron' + +vi.mock('./cors', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + corsPreflight: vi.fn(), + } +}) + +vi.mock('jose', () => ({ + createRemoteJWKSet: vi.fn().mockReturnValue({}), + jwtVerify: vi.fn(), +})) + +const mockHandlePreAuth = vi.fn() +const mockHandle = vi.fn() +let lastStarbaseOptions: any = null + +vi.mock('./handler', () => { + return { + StarbaseDB: vi.fn().mockImplementation((options) => { + lastStarbaseOptions = options + return { + options, + handlePreAuth: mockHandlePreAuth, + handle: mockHandle, + } + }), + } +}) + +describe('Worker default export fetch handler', () => { + let mockStub: any + let mockDoNamespace: any + let mockEnv: Env + let mockCtx: ExecutionContext + + beforeEach(() => { + vi.clearAllMocks() + lastStarbaseOptions = null + mockHandlePreAuth.mockResolvedValue(null) + mockHandle.mockResolvedValue(new Response(JSON.stringify({ success: true }), { status: 200 })) + + mockStub = { + init: vi.fn().mockResolvedValue({ query: vi.fn() }), + } + + mockDoNamespace = { + idFromName: vi.fn().mockReturnValue('mock-id'), + get: vi.fn().mockReturnValue(mockStub), + } + + mockEnv = { + ADMIN_AUTHORIZATION_TOKEN: 'admin-secret', + CLIENT_AUTHORIZATION_TOKEN: 'client-secret', + DATABASE_DURABLE_OBJECT: mockDoNamespace as any, + REGION: 'auto', + HYPERDRIVE: {} as any, + } + + mockCtx = { + waitUntil: vi.fn(), + passThroughOnException: vi.fn(), + } as unknown as ExecutionContext + }) + + describe('CORS Preflight', () => { + it('returns preflight response when corsPreflight returns a response', async () => { + const preflightRes = new Response(null, { status: 204 }) + vi.mocked(corsPreflight).mockReturnValueOnce(preflightRes) + + const req = new Request('https://api.starbasedb.com/query', { + method: 'OPTIONS', + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res).toBe(preflightRes) + expect(corsPreflight).toHaveBeenCalled() + }) + + it('continues request processing if corsPreflight returns null', async () => { + vi.mocked(corsPreflight).mockReturnValueOnce(null as any) + + const req = new Request('https://api.starbasedb.com/query', { + method: 'OPTIONS', + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + // Should reach missing auth check (401) + expect(res.status).toBe(401) + const body = await res.json() + expect(body.error).toBe('Unauthorized request') + }) + }) + + describe('Durable Object & Region Location Hint', () => { + it('fetches DO stub with locationHint when REGION is not AUTO', async () => { + mockEnv.REGION = 'wnam' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(mockDoNamespace.idFromName).toHaveBeenCalledWith('sql-durable-object') + expect(mockDoNamespace.get).toHaveBeenCalledWith('mock-id', { + locationHint: 'wnam', + }) + expect(mockStub.init).toHaveBeenCalled() + }) + + it('fetches DO stub without locationHint when REGION is AUTO or undefined', async () => { + delete (mockEnv as any).REGION + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(mockDoNamespace.get).toHaveBeenCalledWith('mock-id') + }) + }) + + describe('Data Source Selection and Caching', () => { + it('sets data source to external from X-Starbase-Source header', async () => { + const req = new Request('https://api.starbasedb.com/query', { + headers: { + Authorization: 'Bearer admin-secret', + 'X-Starbase-Source': ' external ', + 'X-Starbase-Cache': 'true', + }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.source).toBe('external') + expect(lastStarbaseOptions.dataSource.cache).toBe(true) + }) + + it('sets data source to hyperdrive from url query parameter', async () => { + const req = new Request('https://api.starbasedb.com/query?source=hyperdrive', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.source).toBe('hyperdrive') + }) + + it('defaults to internal data source when unknown or omitted', async () => { + const req = new Request('https://api.starbasedb.com/query?source=unknown_source', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.source).toBe('internal') + }) + }) + + describe('External Database Configurations', () => { + it('configures PostgreSQL external database', async () => { + mockEnv.EXTERNAL_DB_TYPE = 'postgresql' + mockEnv.EXTERNAL_DB_HOST = 'db.postgres.com' + mockEnv.EXTERNAL_DB_PORT = 5432 + mockEnv.EXTERNAL_DB_USER = 'pguser' + mockEnv.EXTERNAL_DB_PASS = 'pgpass' + mockEnv.EXTERNAL_DB_DATABASE = 'maindb' + mockEnv.EXTERNAL_DB_DEFAULT_SCHEMA = 'public' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'postgresql', + host: 'db.postgres.com', + port: 5432, + user: 'pguser', + password: 'pgpass', + database: 'maindb', + defaultSchema: 'public', + }) + }) + + it('configures MySQL external database', async () => { + mockEnv.EXTERNAL_DB_TYPE = 'mysql' + mockEnv.EXTERNAL_DB_HOST = 'db.mysql.com' + mockEnv.EXTERNAL_DB_PORT = 3306 + mockEnv.EXTERNAL_DB_USER = 'root' + mockEnv.EXTERNAL_DB_PASS = 'mypass' + mockEnv.EXTERNAL_DB_DATABASE = 'store' + mockEnv.EXTERNAL_DB_DEFAULT_SCHEMA = 'store' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'mysql', + host: 'db.mysql.com', + port: 3306, + user: 'root', + password: 'mypass', + database: 'store', + defaultSchema: 'store', + }) + }) + + it('configures SQLite with Cloudflare D1 provider', async () => { + mockEnv.EXTERNAL_DB_TYPE = 'sqlite' + mockEnv.EXTERNAL_DB_CLOUDFLARE_API_KEY = 'cf-key' + mockEnv.EXTERNAL_DB_CLOUDFLARE_ACCOUNT_ID = 'cf-acc' + mockEnv.EXTERNAL_DB_CLOUDFLARE_DATABASE_ID = 'cf-db' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'sqlite', + provider: 'cloudflare-d1', + apiKey: 'cf-key', + accountId: 'cf-acc', + databaseId: 'cf-db', + }) + }) + + it('configures SQLite with Starbase provider', async () => { + mockEnv.EXTERNAL_DB_TYPE = 'sqlite' + mockEnv.EXTERNAL_DB_STARBASEDB_URI = 'https://starbase.uri' + mockEnv.EXTERNAL_DB_STARBASEDB_TOKEN = 'sb-tok' + mockEnv.EXTERNAL_DB_DEFAULT_SCHEMA = 'main' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'sqlite', + provider: 'starbase', + apiKey: 'https://starbase.uri', + token: 'sb-tok', + defaultSchema: 'main', + }) + }) + + it('configures SQLite with Turso provider', async () => { + mockEnv.EXTERNAL_DB_TYPE = 'sqlite' + mockEnv.EXTERNAL_DB_TURSO_URI = 'libsql://turso.db' + mockEnv.EXTERNAL_DB_TURSO_TOKEN = 'turso-tok' + mockEnv.EXTERNAL_DB_DEFAULT_SCHEMA = 'main' + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'sqlite', + provider: 'turso', + uri: 'libsql://turso.db', + token: 'turso-tok', + defaultSchema: 'main', + }) + }) + + it('configures Hyperdrive external connection', async () => { + mockEnv.HYPERDRIVE = { + connectionString: 'postgres://user:pass@hyperdrive.cloudflare.com/db', + } as any + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(lastStarbaseOptions.dataSource.external).toEqual({ + dialect: 'postgresql', + connectionString: 'postgres://user:pass@hyperdrive.cloudflare.com/db', + }) + }) + }) + + describe('PreAuth and Interface Route Bypass', () => { + it('returns early if starbase.handlePreAuth handles the request', async () => { + const preAuthRes = new Response('Handled by pre-auth', { status: 200 }) + mockHandlePreAuth.mockResolvedValueOnce(preAuthRes) + + const req = new Request('https://api.starbasedb.com/preauth-check') + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res).toBe(preAuthRes) + expect(mockHandlePreAuth).toHaveBeenCalled() + expect(mockHandle).not.toHaveBeenCalled() + }) + + it('bypasses authentication checks when route matches InterfacePlugin', async () => { + vi.spyOn(InterfacePlugin.prototype, 'matchesRoute').mockReturnValueOnce(true) + + const req = new Request('https://api.starbasedb.com/interface/dashboard') + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(mockHandle).toHaveBeenCalled() + }) + }) + + describe('Authentication', () => { + it('returns 401 when HTTP request has no Authorization header', async () => { + const req = new Request('https://api.starbasedb.com/query') + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(401) + const body = await res.json() + expect(body.error).toBe('Unauthorized request') + }) + + it('returns 401 when WebSocket upgrade request has no token query param', async () => { + const req = new Request('https://api.starbasedb.com/ws', { + headers: { Upgrade: 'websocket' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(401) + const body = await res.json() + expect(body.error).toBe('Unauthorized request') + }) + + it('authorizes admin via Bearer token and updates role to admin', async () => { + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(lastStarbaseOptions.config.role).toBe('admin') + expect(mockHandle).toHaveBeenCalled() + }) + + it('authorizes client via Bearer token and keeps role as client', async () => { + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer client-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(lastStarbaseOptions.config.role).toBe('client') + expect(mockHandle).toHaveBeenCalled() + }) + + it('authorizes WebSocket request via token query param', async () => { + const req = new Request('https://api.starbasedb.com/ws?token=admin-secret', { + headers: { Upgrade: 'websocket' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(lastStarbaseOptions.config.role).toBe('admin') + }) + + it('authenticates valid JWT token against AUTH_JWKS_ENDPOINT', async () => { + mockEnv.AUTH_JWKS_ENDPOINT = 'https://auth.example.com/.well-known/jwks.json' + mockEnv.AUTH_ALGORITHM = 'RS256' + vi.mocked(jwtVerify).mockResolvedValueOnce({ + payload: { sub: 'user_123', email: 'test@example.com' }, + } as any) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer valid.jwt.token' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(200) + expect(jwtVerify).toHaveBeenCalled() + }) + + it('rejects JWT without subject (sub)', async () => { + mockEnv.AUTH_JWKS_ENDPOINT = 'https://auth.example.com/.well-known/jwks.json' + vi.mocked(jwtVerify).mockResolvedValueOnce({ + payload: { name: 'No Subject' }, + } as any) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer jwt.without.sub' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(400) + const body = await res.json() + expect(body.error).toBe('Invalid JWT payload, subject not found.') + }) + + it('rejects invalid JWT when jwtVerify throws', async () => { + mockEnv.AUTH_JWKS_ENDPOINT = 'https://auth.example.com/.well-known/jwks.json' + vi.mocked(jwtVerify).mockRejectedValueOnce(new Error('Signature verification failed')) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer bad.token' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(400) + const body = await res.json() + expect(body.error).toBe('Signature verification failed') + }) + + it('rejects unknown token when no JWKS endpoint is configured', async () => { + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer unknown-token' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(400) + const body = await res.json() + expect(body.error).toBe('Unauthorized request') + }) + }) + + describe('CDC and Cron event callbacks', () => { + it('executes CDC and Cron event listener callbacks', async () => { + let capturedCdcCb: any = null + let capturedCronCb: any = null + + vi.spyOn(ChangeDataCapturePlugin.prototype, 'onEvent').mockImplementation( + (cb: any) => { + capturedCdcCb = cb + } + ) + vi.spyOn(CronPlugin.prototype, 'onEvent').mockImplementation((cb: any) => { + capturedCronCb = cb + }) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + await worker.fetch(req, mockEnv, mockCtx) + + expect(capturedCdcCb).toBeTypeOf('function') + expect(capturedCronCb).toBeTypeOf('function') + + // Execute callbacks to cover internal functions + await capturedCdcCb({ + action: 'INSERT', + schema: 'public', + table: 'users', + data: { id: 1 }, + }) + await capturedCronCb({ + name: 'cleanup', + cron_tab: '* * * * *', + payload: {}, + }) + }) + }) + + describe('Top-level Error Handling', () => { + it('returns 400 when an Error is thrown in initialization', async () => { + mockDoNamespace.idFromName.mockImplementationOnce(() => { + throw new Error('Durable Object initialization failure') + }) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(400) + const body = await res.json() + expect(body.error).toBe('Durable Object initialization failure') + }) + + it('returns 400 with fallback message when a non-Error is thrown', async () => { + mockDoNamespace.idFromName.mockImplementationOnce(() => { + throw 'string exception' + }) + + const req = new Request('https://api.starbasedb.com/query', { + headers: { Authorization: 'Bearer admin-secret' }, + }) + const res = await worker.fetch(req, mockEnv, mockCtx) + + expect(res.status).toBe(400) + const body = await res.json() + expect(body.error).toBe('An unexpected error occurred') + }) + }) +}) diff --git a/src/operation.test.ts b/src/operation.test.ts index f52cbb9..4d77f59 100644 --- a/src/operation.test.ts +++ b/src/operation.test.ts @@ -10,82 +10,45 @@ import { applyRLS } from './rls' import { beforeQueryCache, afterQueryCache } from './cache' import type { DataSource } from './types' import type { StarbaseDBConfiguration } from './handler' -import type { SqlConnection } from '@outerbase/sdk/dist/connections/sql-base' - -// const mockSqlConnection = vi.hoisted(() => ({ -// connect: vi.fn().mockResolvedValue(undefined), -// raw: vi -// .fn() -// .mockResolvedValue({ data: [{ id: 1, name: 'SDK-Test-Result' }] }), -// })) as unknown as SqlConnection - -// const mockConfig = vi.hoisted(() => ({ -// outerbaseApiKey: 'mock-api-key', -// role: 'admin', -// features: { allowlist: true, rls: true, rest: true }, -// })) as StarbaseDBConfiguration - -// const mockDataSource = vi.hoisted(() => ({ -// source: 'internal', -// external: { -// dialect: 'postgresql', -// provider: 'postgresql', -// host: 'mock-host', -// port: 5432, -// user: 'mock-user', -// password: 'mock-password', -// database: 'mock-db', -// } as any, -// rpc: { -// executeQuery: vi.fn().mockResolvedValue([ -// { id: 1, name: 'Alice' }, -// { id: 2, name: 'Bob' }, -// ]), -// }, -// })) as unknown as DataSource - -// vi.mock('./operation', async (importOriginal) => { -// const actual = await importOriginal() -// return { -// ...actual, -// executeQuery: vi.fn().mockResolvedValue([ -// { id: 1, name: 'Mocked Alice' }, -// { id: 2, name: 'Mocked Bob' }, -// ]), -// executeSDKQuery: vi -// .fn() -// .mockResolvedValue([{ id: 1, name: 'SDK-Result' }]), -// createSDKPostgresConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKMySQLConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKCloudflareConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKStarbaseConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKTursoConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// } -// }) - -// vi.mock('./operation', () => ({ -// executeSDKQuery: vi.fn().mockResolvedValue([{ id: 1, name: 'SDK-Result' }]), -// })) - -// vi.mock('./operation', async (importOriginal) => { -// const original = await importOriginal() -// return { -// ...original, -// executeSDKQuery: vi -// .fn() -// .mockResolvedValue([{ id: 1, name: 'SDK-Result' }]), -// } -// }) + +const mockRaw = vi.fn().mockResolvedValue({ data: [{ id: 1, name: 'SDK-Result' }] }) +const mockConnect = vi.fn().mockResolvedValue(undefined) + +vi.mock('@outerbase/sdk', () => { + class MockConnection { + connect = mockConnect + raw = mockRaw + } + return { + PostgreSQLConnection: MockConnection, + MySQLConnection: MockConnection, + CloudflareD1Connection: MockConnection, + StarbaseConnection: MockConnection, + TursoConnection: MockConnection, + } +}) + +vi.mock('pg', () => ({ + Client: vi.fn(), +})) + +vi.mock('mysql2', () => ({ + createConnection: vi.fn(), +})) + +vi.mock('@libsql/client/web', () => ({ + createClient: vi.fn(), +})) + +const mockSqlUnsafe = vi.fn().mockResolvedValue([{ id: 10, name: 'Hyperdrive-Result' }]) +const mockSqlEnd = vi.fn().mockResolvedValue(undefined) +const mockPostgres = vi.fn().mockReturnValue({ + unsafe: mockSqlUnsafe, + end: mockSqlEnd, +}) +vi.mock('postgres', () => ({ + default: (...args: any[]) => mockPostgres(...args), +})) vi.mock('./allowlist', () => ({ isQueryAllowed: vi.fn() })) vi.mock('./rls', () => ({ applyRLS: vi.fn(async ({ sql }) => sql) })) @@ -94,494 +57,594 @@ vi.mock('./cache', () => ({ afterQueryCache: vi.fn(), })) -let mockSqlConnection: SqlConnection -let mockDataSource: DataSource -let mockConfig: StarbaseDBConfiguration - -beforeEach(() => { - // mockSqlConnection = { - // connect: vi.fn().mockResolvedValue(undefined), - // raw: vi - // .fn() - // .mockResolvedValue({ data: [{ id: 1, name: 'SDK-Test-Result' }] }), - // } as unknown as SqlConnection - - mockConfig = { - outerbaseApiKey: 'mock-api-key', - role: 'admin', - features: { allowlist: true, rls: true, rest: true }, - } +describe('operation module', () => { + let mockDataSource: DataSource + let mockConfig: StarbaseDBConfiguration + + beforeEach(() => { + vi.clearAllMocks() + + mockConfig = { + outerbaseApiKey: 'mock-api-key', + role: 'admin', + features: { allowlist: true, rls: true, rest: true }, + } + + mockDataSource = { + source: 'internal', + external: { + dialect: 'postgresql', + provider: 'postgresql', + host: 'mock-host', + port: 5432, + user: 'mock-user', + password: 'mock-password', + database: 'mock-db', + } as any, + rpc: { + executeQuery: vi.fn().mockResolvedValue([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]), + } as any, + } + + vi.mocked(beforeQueryCache).mockResolvedValue(null) + vi.mocked(afterQueryCache).mockResolvedValue(null) + mockSqlUnsafe.mockResolvedValue([{ id: 10, name: 'Hyperdrive-Result' }]) + mockSqlEnd.mockResolvedValue(undefined) + mockRaw.mockResolvedValue({ data: [{ id: 1, name: 'SDK-Result' }] }) + mockConnect.mockResolvedValue(undefined) + }) - mockDataSource = { - source: 'internal', - external: { - dialect: 'postgresql', - provider: 'postgresql', - host: 'mock-host', - port: 5432, - user: 'mock-user', - password: 'mock-password', - database: 'mock-db', - } as any, - rpc: { - executeQuery: vi.fn().mockResolvedValue([ + describe('executeQuery', () => { + it('should execute a valid SQL query on internal source', async () => { + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledWith({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + }) + expect(result).toEqual([ { id: 1, name: 'Alice' }, { id: 2, name: 'Bob' }, - ]), - }, - } as any - - vi.mocked(beforeQueryCache).mockResolvedValue(null) - vi.mocked(afterQueryCache).mockResolvedValue(null) - // vi.mock('./operation', () => ({ - // createSDKPostgresConnection: vi - // .fn() - // .mockResolvedValue({ database: mockSqlConnection }), - // createSDKMySQLConnection: vi - // .fn() - // .mockResolvedValue({ database: mockSqlConnection }), - // createSDKCloudflareConnection: vi - // .fn() - // .mockResolvedValue({ database: mockSqlConnection }), - // createSDKStarbaseConnection: vi - // .fn() - // .mockResolvedValue({ database: mockSqlConnection }), - // createSDKTursoConnection: vi - // .fn() - // .mockResolvedValue({ database: mockSqlConnection }), - // })) - - vi.clearAllMocks() -}) -// beforeEach(() => { -// vi.clearAllMocks() - -// vi.mocked(beforeQueryCache).mockResolvedValue(null) -// vi.mocked(afterQueryCache).mockResolvedValue(null) -// const mockExecuteQueryResult = [ -// { id: 1, name: 'Alice' }, -// { id: 2, name: 'Bob' }, -// ] as any -// mockExecuteQueryResult[Symbol.dispose] = vi.fn() -// vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue( -// mockExecuteQueryResult -// ) -// }) -// vi.mock('./operation', async (importOriginal) => { -// const original = await importOriginal() -// return { -// ...original, -// executeSDKQuery: vi -// .fn() -// .mockResolvedValue([{ id: 1, name: 'SDK-Result' }]), -// createSDKPostgresConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKMySQLConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKCloudflareConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKStarbaseConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKTursoConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// } -// }) - -// vi.mock('./operation', () => ({ -// createSDKPostgresConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKMySQLConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKCloudflareConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKStarbaseConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// createSDKTursoConnection: vi -// .fn() -// .mockResolvedValue({ database: mockSqlConnection }), -// })) - -describe('executeQuery', () => { - it('should execute a valid SQL query', async () => { - const result = await executeQuery({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, - dataSource: mockDataSource, - config: mockConfig, + ]) }) - expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledWith({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, + it('should handle raw results transformation', async () => { + mockDataSource.rpc.executeQuery = vi.fn().mockResolvedValue({ + columns: ['id', 'name'], + rows: [[1, 'Alice'], [2, 'Bob']], + }) + + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(result).toEqual({ + columns: ['id', 'name'], + rows: [[1, 'Alice'], [2, 'Bob']], + meta: { + rows_read: 2, + rows_written: 0, + }, + }) }) - expect(result).toEqual([ - { id: 1, name: 'Alice' }, - { id: 2, name: 'Bob' }, - ]) - }) - it('should enforce the allowlist feature', async () => { - await executeQuery({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, - dataSource: mockDataSource, - config: mockConfig, + it('should return an empty array if internal source returns null or empty', async () => { + mockDataSource.rpc.executeQuery = vi.fn().mockResolvedValue(null) + + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(result).toEqual([]) }) - expect(isQueryAllowed).toHaveBeenCalledTimes(1) - expect(isQueryAllowed).toHaveBeenCalledWith( - expect.objectContaining({ sql: 'SELECT * FROM users' }) - ) - }) + it('should return an empty array if data source is missing', async () => { + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: null as any, + config: mockConfig, + }) - it('should apply row-level security', async () => { - await executeQuery({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, - dataSource: mockDataSource, - config: mockConfig, + expect(result).toEqual([]) }) - expect(applyRLS).toHaveBeenCalledWith( - expect.objectContaining({ sql: 'SELECT * FROM users' }) - ) - }) + it('should enforce the allowlist feature', async () => { + await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) - it('should return cached results if available', async () => { - ;(beforeQueryCache as any).mockResolvedValue([ - { id: 99, name: 'Cached' }, - ]) - - const result = await executeQuery({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, - dataSource: mockDataSource, - config: mockConfig, + expect(isQueryAllowed).toHaveBeenCalledWith( + expect.objectContaining({ sql: 'SELECT * FROM users' }) + ) }) - expect(result).toEqual([{ id: 99, name: 'Cached' }]) - expect(mockDataSource.rpc.executeQuery).not.toHaveBeenCalled() - }) + it('should apply row-level security', async () => { + await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) - it('should return an empty array if the data source is missing', async () => { - const result = await executeQuery({ - sql: 'SELECT * FROM users', - params: undefined, - isRaw: false, - dataSource: null as any, - config: mockConfig, + expect(applyRLS).toHaveBeenCalledWith( + expect.objectContaining({ sql: 'SELECT * FROM users' }) + ) }) - expect(result).toEqual([]) - }) -}) -describe('executeTransaction', () => { - it('should execute multiple queries in a transaction', async () => { - const queries = [ - { sql: 'INSERT INTO users VALUES (1, "Alice")' }, - { sql: 'INSERT INTO users VALUES (2, "Bob")' }, - ] - - const result = await executeTransaction({ - queries, - isRaw: false, - dataSource: mockDataSource, - config: mockConfig, + it('should return cached results if available', async () => { + vi.mocked(beforeQueryCache).mockResolvedValueOnce([{ id: 99, name: 'Cached' }]) + + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(result).toEqual([{ id: 99, name: 'Cached' }]) + expect(mockDataSource.rpc.executeQuery).not.toHaveBeenCalled() + }) + + it('should execute query via Hyperdrive with executionContext', async () => { + const waitMock = vi.fn() + const ds: DataSource = { + source: 'hyperdrive', + external: { + dialect: 'postgresql', + connectionString: 'postgres://user:pass@host:5432/db', + }, + rpc: {} as any, + executionContext: { waitUntil: waitMock } as any, + } + + const result = await executeQuery({ + sql: 'SELECT * FROM items', + params: ['a'], + isRaw: false, + dataSource: ds, + config: mockConfig, + }) + + expect(mockSqlUnsafe).toHaveBeenCalledWith('SELECT * FROM items', ['a']) + expect(waitMock).toHaveBeenCalled() + expect(result).toEqual([{ id: 10, name: 'Hyperdrive-Result' }]) + }) + + it('should execute query via Hyperdrive without executionContext', async () => { + const ds: DataSource = { + source: 'hyperdrive', + external: { + dialect: 'postgresql', + connectionString: 'postgres://user:pass@host:5432/db', + }, + rpc: {} as any, + } + + const result = await executeQuery({ + sql: 'SELECT * FROM items', + params: undefined, + isRaw: false, + dataSource: ds, + config: mockConfig, + }) + + expect(mockSqlEnd).toHaveBeenCalled() + expect(result).toEqual([{ id: 10, name: 'Hyperdrive-Result' }]) }) - expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledTimes(2) - expect(queries.length).toBe(2) + it('should throw error when Hyperdrive connection string is missing', async () => { + const ds: DataSource = { + source: 'hyperdrive', + external: { dialect: 'postgresql' } as any, + rpc: {} as any, + } + + await expect( + executeQuery({ + sql: 'SELECT 1', + params: undefined, + isRaw: false, + dataSource: ds, + config: mockConfig, + }) + ).rejects.toThrow('Hyperdrive connection string not found') + }) + + it('should rethrow error if Hyperdrive query fails', async () => { + mockSqlUnsafe.mockRejectedValueOnce(new Error('Connection lost')) + const ds: DataSource = { + source: 'hyperdrive', + external: { + dialect: 'postgresql', + connectionString: 'postgres://user:pass@host:5432/db', + }, + rpc: {} as any, + } + + await expect( + executeQuery({ + sql: 'SELECT 1', + params: undefined, + isRaw: false, + dataSource: ds, + config: mockConfig, + }) + ).rejects.toThrow('Connection lost') + }) + + it('should trigger beforeQuery and afterQuery plugin registry hooks', async () => { + const beforeQueryMock = vi.fn().mockResolvedValue({ + sql: 'SELECT * FROM rewritten', + params: [42], + }) + const afterQueryMock = vi.fn().mockResolvedValue([{ id: 42, name: 'Hooked' }]) + + mockDataSource.registry = { + beforeQuery: beforeQueryMock, + afterQuery: afterQueryMock, + } as any + + const result = await executeQuery({ + sql: 'SELECT * FROM original', + params: [1], + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(beforeQueryMock).toHaveBeenCalled() + expect(afterQueryMock).toHaveBeenCalled() + expect(result).toEqual([{ id: 42, name: 'Hooked' }]) + }) + + it('should catch error in afterQuery hook without failing query', async () => { + mockDataSource.registry = { + beforeQuery: vi.fn().mockResolvedValue({ + sql: 'SELECT * FROM users', + params: undefined, + }), + afterQuery: vi.fn().mockRejectedValue(new Error('Hook failed')), + } as any + + const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + const result = await executeQuery({ + sql: 'SELECT * FROM users', + params: undefined, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(consoleSpy).toHaveBeenCalled() + expect(result).toEqual([ + { id: 1, name: 'Alice' }, + { id: 2, name: 'Bob' }, + ]) + }) }) - it('should return an empty array if the data source is missing', async () => { - const consoleErrorMock = vi - .spyOn(console, 'error') - .mockImplementation(() => {}) + describe('executeTransaction', () => { + it('should execute multiple queries in a transaction', async () => { + const queries = [ + { sql: 'INSERT INTO users VALUES (1, "Alice")' }, + { sql: 'INSERT INTO users VALUES (2, "Bob")' }, + ] - const result = await executeTransaction({ - queries: [{ sql: 'INSERT INTO users VALUES (1, "Alice")' }], - isRaw: false, - dataSource: null as any, - config: mockConfig, + const result = await executeTransaction({ + queries, + isRaw: false, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(mockDataSource.rpc.executeQuery).toHaveBeenCalledTimes(2) + expect(result).toHaveLength(2) + }) + + it('should return an empty array if data source is missing in transaction', async () => { + const result = await executeTransaction({ + queries: [{ sql: 'INSERT INTO users VALUES (1, "Alice")' }], + isRaw: false, + dataSource: null as any, + config: mockConfig, + }) + + expect(result).toEqual([]) }) - expect(result).toEqual([]) }) -}) -describe('executeExternalQuery', () => { - it('should throw an error if dataSource.external is missing', async () => { - await expect( - executeExternalQuery({ + describe('executeExternalQuery', () => { + it('should throw an error if dataSource.external is missing', async () => { + await expect( + executeExternalQuery({ + sql: 'SELECT * FROM users', + params: [], + dataSource: { source: 'internal' } as any, + config: mockConfig, + }) + ).rejects.toThrow('External connection not found.') + }) + + it('should call executeSDKQuery if outerbaseApiKey is missing', async () => { + const configWithoutApiKey = { + ...mockConfig, + outerbaseApiKey: undefined, + } + + const result = await executeExternalQuery({ + sql: 'SELECT * FROM users', + params: [], + dataSource: mockDataSource, + config: configWithoutApiKey, + }) + + expect(mockConnect).toHaveBeenCalled() + expect(mockRaw).toHaveBeenCalledWith('SELECT * FROM users', []) + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) + }) + + it('should correctly format SQL and parameters for API request with array params', async () => { + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValueOnce({ + json: async () => ({ + response: { + results: { items: [{ id: 2, name: 'API-Result' }] }, + }, + }), + } as Response) + + const result = await executeExternalQuery({ + sql: 'SELECT * FROM users WHERE id = ?', + params: [5], + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(fetchMock).toHaveBeenCalledWith( + 'https://app.outerbase.com/api/v1/ezql/raw', + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Source-Token': 'mock-api-key', + }, + body: JSON.stringify({ + query: 'SELECT * FROM users WHERE id = :param0', + params: { param0: 5 }, + }), + } + ) + + expect(result).toEqual([{ id: 2, name: 'API-Result' }]) + }) + + it('should correctly handle non-array params for API request', async () => { + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValueOnce({ + json: async () => ({ + response: { + results: { items: [{ id: 3, name: 'Named-Result' }] }, + }, + }), + } as Response) + + const result = await executeExternalQuery({ + sql: 'SELECT * FROM users WHERE id = :id', + params: { id: 10 }, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(fetchMock).toHaveBeenCalledWith( + 'https://app.outerbase.com/api/v1/ezql/raw', + expect.objectContaining({ + body: JSON.stringify({ + query: 'SELECT * FROM users WHERE id = :id', + params: { id: 10 }, + }), + }) + ) + + expect(result).toEqual([{ id: 3, name: 'Named-Result' }]) + }) + + it('should handle API failure gracefully', async () => { + vi.spyOn(global, 'fetch').mockRejectedValueOnce(new Error('Network error')) + + await expect( + executeExternalQuery({ + sql: 'SELECT * FROM users', + params: [], + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow('Network error') + }) + + it('should return an empty array if API response is malformed', async () => { + vi.spyOn(global, 'fetch').mockResolvedValueOnce({ + json: async () => ({}), + } as Response) + + const result = await executeExternalQuery({ sql: 'SELECT * FROM users', params: [], - dataSource: { source: 'internal' } as any, + dataSource: mockDataSource, config: mockConfig, }) - ).rejects.toThrow('External connection not found.') + + expect(result).toEqual([]) + }) }) - // it('should call executeSDKQuery if outerbaseApiKey is missing', async () => { - // const configWithoutApiKey = { - // ...mockConfig, - // outerbaseApiKey: undefined, - // } - - // const result = await executeExternalQuery({ - // sql: 'SELECT * FROM users', - // params: [], - // dataSource: { - // ...mockDataSource, - // external: { - // dialect: 'postgresql', - // host: 'mock-host', - // port: 5432, - // user: 'mock-user', - // password: 'mock-password', - // database: 'mock-db', - // } as any, - // }, - // config: configWithoutApiKey, - // }) - - // expect(executeSDKQuery).toHaveBeenCalledWith({ - // sql: 'SELECT * FROM users', - // params: [], - // dataSource: expect.objectContaining({ - // external: expect.any(Object), - // }), - - // config: configWithoutApiKey, - // }) - - // expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) - // }) - - it('should correctly format SQL and parameters for API request', async () => { - const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValueOnce({ - json: async () => ({ - response: { - results: { items: [{ id: 2, name: 'API-Result' }] }, - }, - }), - } as Response) - - const result = await executeExternalQuery({ - sql: 'SELECT * FROM users WHERE id = ?', - params: [5], - dataSource: { - ...mockDataSource, + describe('executeSDKQuery', () => { + it('should return empty array if external config is missing', async () => { + const ds: DataSource = { source: 'external', rpc: {} as any } + const result = await executeSDKQuery({ + sql: 'SELECT 1', + dataSource: ds, + config: mockConfig, + }) + expect(result).toEqual([]) + }) + + it('should execute query with PostgreSQL driver', async () => { + const ds: DataSource = { + source: 'external', external: { dialect: 'postgresql', - host: 'mock-host', + host: 'localhost', port: 5432, - user: 'mock-user', - password: 'mock-password', - database: 'mock-db', + user: 'user', + password: 'password', + database: 'db', }, - }, - config: mockConfig, + rpc: {} as any, + } + + const result = await executeSDKQuery({ + sql: 'SELECT * FROM pg_table', + params: [1], + dataSource: ds, + config: mockConfig, + }) + + expect(mockConnect).toHaveBeenCalled() + expect(mockRaw).toHaveBeenCalledWith('SELECT * FROM pg_table', [1]) + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) }) - expect(fetchMock).toHaveBeenCalledWith( - 'https://app.outerbase.com/api/v1/ezql/raw', - { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'X-Source-Token': 'mock-api-key', + it('should execute query with MySQL driver', async () => { + const ds: DataSource = { + source: 'external', + external: { + dialect: 'mysql', + host: 'localhost', + port: 3306, + user: 'root', + password: 'password', + database: 'db', }, - body: JSON.stringify({ - query: 'SELECT * FROM users WHERE id = :param0', - params: { param0: 5 }, - }), + rpc: {} as any, } - ) - expect(result).toEqual([{ id: 2, name: 'API-Result' }]) - }) + const result = await executeSDKQuery({ + sql: 'SELECT * FROM my_table', + dataSource: ds, + config: mockConfig, + }) - it('should handle API failure gracefully', async () => { - const fetchMock = vi - .spyOn(global, 'fetch') - .mockRejectedValueOnce(new Error('Network error')) + expect(mockConnect).toHaveBeenCalled() + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) + }) - await expect( - executeExternalQuery({ - sql: 'SELECT * FROM users', - params: [], - dataSource: mockDataSource, + it('should execute query with Cloudflare D1 provider', async () => { + const ds: DataSource = { + source: 'external', + external: { + dialect: 'sqlite', + provider: 'cloudflare-d1', + apiKey: 'cf-key', + accountId: 'cf-acc', + databaseId: 'cf-db', + }, + rpc: {} as any, + } + + const result = await executeSDKQuery({ + sql: 'SELECT * FROM d1_table', + dataSource: ds, config: mockConfig, }) - ).rejects.toThrow('Network error') - expect(fetchMock).toHaveBeenCalled() - }) + expect(mockConnect).toHaveBeenCalled() + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) + }) - it('should return an empty array if API response is malformed', async () => { - vi.spyOn(global, 'fetch').mockResolvedValueOnce({ - json: async () => ({}), - } as Response) + it('should execute query with Starbase provider', async () => { + const ds: DataSource = { + source: 'external', + external: { + dialect: 'sqlite', + provider: 'starbase', + apiKey: 'sb-key', + token: 'sb-url', + }, + rpc: {} as any, + } - const result = await executeExternalQuery({ - sql: 'SELECT * FROM users', - params: [], - dataSource: mockDataSource, - config: mockConfig, + const result = await executeSDKQuery({ + sql: 'SELECT * FROM sb_table', + dataSource: ds, + config: mockConfig, + }) + + expect(mockConnect).toHaveBeenCalled() + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) }) - expect(result).toEqual([]) + it('should execute query with Turso provider', async () => { + const ds: DataSource = { + source: 'external', + external: { + dialect: 'sqlite', + provider: 'turso', + uri: 'libsql://turso.io', + token: 'turso-token', + }, + rpc: {} as any, + } + + const result = await executeSDKQuery({ + sql: 'SELECT * FROM turso_table', + dataSource: ds, + config: mockConfig, + }) + + expect(mockConnect).toHaveBeenCalled() + expect(result).toEqual([{ id: 1, name: 'SDK-Result' }]) + }) + + it('should throw error for unsupported external database type', async () => { + const ds: DataSource = { + source: 'external', + external: { + dialect: 'oracle' as any, + }, + rpc: {} as any, + } + + await expect( + executeSDKQuery({ + sql: 'SELECT 1', + dataSource: ds, + config: mockConfig, + }) + ).rejects.toThrow('Unsupported external database type') + }) }) }) - -// describe('executeSDKQuery', () => { -// it('should execute a query using PostgreSQL connection', async () => { -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(mockSqlConnection.connect).toHaveBeenCalled() -// expect(mockSqlConnection.raw).toHaveBeenCalledWith( -// 'SELECT * FROM users', -// [] -// ) -// expect(result).toEqual([{ id: 1, name: 'SDK-Test-Result' }]) -// }) - -// it('should execute a query using MySQL connection', async () => { -// if (mockDataSource.external) { -// mockDataSource.external.dialect = 'mysql' -// } - -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(mockSqlConnection.connect).toHaveBeenCalled() -// expect(mockSqlConnection.raw).toHaveBeenCalledWith( -// 'SELECT * FROM users', -// [] -// ) -// expect(result).toEqual([{ id: 1, name: 'SDK-Test-Result' }]) -// }) - -// it('should execute a query using Cloudflare D1 connection', async () => { -// if (mockDataSource.external && 'provider' in mockDataSource.external) { -// mockDataSource.external.provider = 'cloudflare-d1' -// } -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(mockSqlConnection.connect).toHaveBeenCalled() -// expect(mockSqlConnection.raw).toHaveBeenCalledWith( -// 'SELECT * FROM users', -// [] -// ) -// expect(result).toEqual([{ id: 1, name: 'SDK-Test-Result' }]) -// }) - -// it('should execute a query using Starbase connection', async () => { -// if (mockDataSource.external && 'provider' in mockDataSource.external) { -// mockDataSource.external.provider = 'starbase' -// } - -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(mockSqlConnection.connect).toHaveBeenCalled() -// expect(mockSqlConnection.raw).toHaveBeenCalledWith( -// 'SELECT * FROM users', -// [] -// ) -// expect(result).toEqual([{ id: 1, name: 'SDK-Test-Result' }]) -// }) - -// it('should execute a query using Turso connection', async () => { -// if (mockDataSource.external && 'provider' in mockDataSource.external) { -// mockDataSource.external.provider = 'turso' -// } - -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(mockSqlConnection.connect).toHaveBeenCalled() -// expect(mockSqlConnection.raw).toHaveBeenCalledWith( -// 'SELECT * FROM users', -// [] -// ) -// expect(result).toEqual([{ id: 1, name: 'SDK-Test-Result' }]) -// }) -// it('should return an empty array if external connection is missing', async () => { -// mockDataSource.external = undefined as any - -// const result = await executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) - -// expect(result).toEqual([]) -// }) - -// it('should handle database connection errors gracefully', async () => { -// vi.mocked(mockSqlConnection.connect).mockRejectedValueOnce( -// new Error('DB connection failed') -// ) - -// await expect( -// executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) -// ).rejects.toThrow('DB connection failed') -// }) - -// it('should handle query execution errors gracefully', async () => { -// vi.mocked(mockSqlConnection.raw).mockRejectedValueOnce( -// new Error('Query execution failed') -// ) - -// await expect( -// executeSDKQuery({ -// sql: 'SELECT * FROM users', -// params: [], -// dataSource: mockDataSource, -// config: mockConfig, -// }) -// ).rejects.toThrow('Query execution failed') -// }) -// }) diff --git a/src/rls/index.test.ts b/src/rls/index.test.ts index cf00156..9398592 100644 --- a/src/rls/index.test.ts +++ b/src/rls/index.test.ts @@ -17,6 +17,11 @@ const mockConfig: StarbaseDBConfiguration = { features: { allowlist: true, rls: true, rest: true }, } +beforeEach(() => { + mockConfig.role = 'client' + mockDataSource.context.sub = 'user123' +}) + describe('loadPolicies - Policy Fetching and Parsing', () => { it('should load and parse policies correctly', async () => { vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue([ @@ -94,10 +99,21 @@ describe('applyRLS - Query Modification', () => { config: mockConfig, }) - console.log('Final SQL:', modifiedSql) - expect(modifiedSql).toContain("WHERE `user_id` = 'user123'") + expect(modifiedSql).toContain("WHERE (`users`.`user_id` = 'user123')") }) it('should modify DELETE queries by adding policy-based WHERE clause', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue([ + { + actions: 'DELETE', + schema: 'public', + table: 'users', + column: 'user_id', + value: 'context.id()', + value_type: 'string', + operator: '=', + }, + ]) + const sql = "DELETE FROM users WHERE name = 'Alice'" const modifiedSql = await applyRLS({ sql, @@ -106,10 +122,23 @@ describe('applyRLS - Query Modification', () => { config: mockConfig, }) - expect(modifiedSql).toContain("WHERE `name` = 'Alice'") + expect(modifiedSql).toContain("`name` = 'Alice'") + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") }) it('should modify UPDATE queries with additional WHERE clause', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue([ + { + actions: 'UPDATE', + schema: 'public', + table: 'users', + column: 'user_id', + value: 'context.id()', + value_type: 'string', + operator: '=', + }, + ]) + const sql = "UPDATE users SET name = 'Bob' WHERE age = 25" const modifiedSql = await applyRLS({ sql, @@ -118,10 +147,24 @@ describe('applyRLS - Query Modification', () => { config: mockConfig, }) - expect(modifiedSql).toContain("`name` = 'Bob' WHERE `age` = 25") + expect(modifiedSql).toContain("`name` = 'Bob'") + expect(modifiedSql).toContain('`age` = 25') + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") }) it('should modify INSERT queries to enforce column values', async () => { + vi.mocked(mockDataSource.rpc.executeQuery).mockResolvedValue([ + { + actions: 'INSERT', + schema: 'public', + table: 'users', + column: 'user_id', + value: 'context.id()', + value_type: 'string', + operator: '=', + }, + ]) + const sql = "INSERT INTO users (user_id, name) VALUES (1, 'Alice')" const modifiedSql = await applyRLS({ sql, @@ -130,7 +173,34 @@ describe('applyRLS - Query Modification', () => { config: mockConfig, }) - expect(modifiedSql).toContain("VALUES (1,'Alice')") + expect(modifiedSql).toContain("VALUES ('user123','Alice')") + }) + + it('should reject restricted table actions without an explicit policy', async () => { + const sql = "UPDATE users SET name = 'Bob' WHERE age = 25" + + await expect( + applyRLS({ + sql, + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + ).rejects.toThrow( + 'Unauthorized access: No matching rules for UPDATE on restricted table users' + ) + }) + + it('should not apply policies to a different schema-qualified table', async () => { + const sql = 'SELECT * FROM private.users' + const modifiedSql = await applyRLS({ + sql, + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(modifiedSql).not.toContain('user_id') }) }) @@ -200,15 +270,15 @@ describe('applyRLS - Multi-Table Queries', () => { config: mockConfig, }) - expect(modifiedSql).toContain("WHERE `users.user_id` = 'user123'") - expect(modifiedSql).toContain("AND `orders.user_id` = 'user123'") + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") + expect(modifiedSql).toContain("`orders`.`user_id` = 'user123'") }) - it('should apply RLS policies to multiple tables in a JOIN', async () => { + it('should apply RLS policies to aliased tables in a JOIN', async () => { const sql = ` - SELECT users.name, orders.total - FROM users - JOIN orders ON users.id = orders.user_id + SELECT u.name, o.total + FROM users AS u + JOIN orders AS o ON u.id = o.user_id ` const modifiedSql = await applyRLS({ @@ -218,8 +288,8 @@ describe('applyRLS - Multi-Table Queries', () => { config: mockConfig, }) - expect(modifiedSql).toContain("WHERE (users.user_id = 'user123')") - expect(modifiedSql).toContain("AND (orders.user_id = 'user123')") + expect(modifiedSql).toContain("`u`.`user_id` = 'user123'") + expect(modifiedSql).toContain("`o`.`user_id` = 'user123'") }) it('should apply RLS policies to subqueries inside FROM clause', async () => { @@ -236,6 +306,65 @@ describe('applyRLS - Multi-Table Queries', () => { config: mockConfig, }) - expect(modifiedSql).toContain("WHERE `users.user_id` = 'user123'") + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") + }) + + it('should return original SQL unmodified when isEnabled is false', async () => { + const sql = 'SELECT * FROM users' + const modifiedSql = await applyRLS({ + sql, + isEnabled: false, + dataSource: mockDataSource, + config: mockConfig, + }) + expect(modifiedSql).toBe(sql) + }) + + it('should apply RLS policies to subqueries inside columns', async () => { + const sql = ` + SELECT (SELECT user_id FROM users) AS uid FROM accounts + ` + + const modifiedSql = await applyRLS({ + sql, + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") + }) + + it('should apply RLS policies to subqueries inside JOINs', async () => { + const sql = ` + SELECT * FROM accounts + JOIN (SELECT user_id FROM users) AS u ON accounts.user_id = u.user_id + ` + + const modifiedSql = await applyRLS({ + sql, + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") + }) + + + it('should correctly traverse nested WHERE conditions', async () => { + const sql = ` + SELECT * FROM users WHERE (age > 18 AND status = 'active') OR (role = 'guest') + ` + + const modifiedSql = await applyRLS({ + sql, + isEnabled: true, + dataSource: mockDataSource, + config: mockConfig, + }) + + expect(modifiedSql).toContain("`users`.`user_id` = 'user123'") }) }) + diff --git a/src/rls/index.ts b/src/rls/index.ts index 68abb4e..4d8403a 100644 --- a/src/rls/index.ts +++ b/src/rls/index.ts @@ -202,6 +202,78 @@ export async function applyRLS(opts: { return modifiedSql } +function getUnqualifiedTableName(name: string): string { + const normalized = normalizeIdentifier(name) + return normalized.includes('.') ? normalized.split('.').pop()! : normalized +} + +function getFullTableName(tableRef: any): string | undefined { + const tableName = normalizeIdentifier(tableRef?.table) + if (!tableName) return undefined + + const schemaName = normalizeIdentifier(tableRef?.db) + return schemaName ? `${schemaName}.${tableName}` : tableName +} + +function tableNamesMatch(policyTable: string, queryTable: string): boolean { + const normalizedPolicyTable = normalizeIdentifier(policyTable) + const normalizedQueryTable = normalizeIdentifier(queryTable) + const policyHasSchema = normalizedPolicyTable.includes('.') + const queryHasSchema = normalizedQueryTable.includes('.') + + return ( + normalizedPolicyTable === normalizedQueryTable || + ((!policyHasSchema || !queryHasSchema) && + getUnqualifiedTableName(normalizedPolicyTable) === + getUnqualifiedTableName(normalizedQueryTable)) + ) +} + +function getAstTableRefs(ast: any, statementType: string) { + const tableRefs = + statementType === 'INSERT' || statementType === 'UPDATE' + ? ast.table + : ast.from + + return (tableRefs ?? []) + .map((tableRef: any) => { + const tableName = getFullTableName(tableRef) + if (!tableName) return null + + return { + name: tableName, + conditionTable: tableRef.as ?? tableRef.table, + } + }) + .filter(Boolean) as Array<{ + name: string + conditionTable: string | null + }> +} + +function getConditionForTable( + condition: Policy['condition'], + tableName: string | null +) { + return { + ...condition, + left: { + ...condition.left, + table: tableName, + }, + right: { + ...condition.right, + }, + } +} + +function getNestedSelect(expr: any) { + if (!expr) return undefined + if (expr.type === 'select') return expr + if (expr.ast?.type === 'select') return expr.ast + return undefined +} + function applyRLSToAst(ast: any): void { if (!ast) return @@ -232,57 +304,28 @@ function applyRLSToAst(ast: any): void { traverseWhere(ast.where) } - const tablesWithRules: Record = {} - policies.forEach((policy) => { - const tbl = normalizeIdentifier(policy.condition.left.table) - if (!tablesWithRules[tbl]) { - tablesWithRules[tbl] = [] - } - tablesWithRules[tbl].push(policy.action) - }) - const statementType = ast.type?.toUpperCase() if (!['SELECT', 'UPDATE', 'DELETE', 'INSERT'].includes(statementType)) { return } - let tables: string[] = [] - if (statementType === 'INSERT') { - let tableName = normalizeIdentifier(ast.table[0].table) - if (tableName.includes('.')) { - tableName = tableName.split('.')[1] - } - tables = [tableName] - } else if (statementType === 'UPDATE') { - tables = ast.table.map((tableRef: any) => { - let tableName = normalizeIdentifier(tableRef.table) - if (tableName.includes('.')) { - tableName = tableName.split('.')[1] - } - return tableName - }) - } else { - // SELECT or DELETE - tables = - ast.from?.map((fromTable: any) => { - let tableName = normalizeIdentifier(fromTable.table) - if (tableName.includes('.')) { - tableName = tableName.split('.')[1] - } - return tableName - }) || [] - } + const tableRefs = getAstTableRefs(ast, statementType) - const restrictedTables = Object.keys(tablesWithRules) + for (const tableRef of tableRefs) { + const allowedActions = policies + .filter((policy) => + tableNamesMatch(policy.condition.left.table, tableRef.name) + ) + .map((policy) => policy.action) - for (const table of tables) { - if (restrictedTables.includes(table)) { - const allowedActions = tablesWithRules[table] - if (!allowedActions.includes(statementType)) { - throw new Error( - `Unauthorized access: No matching rules for ${statementType} on restricted table ${table}` - ) - } + if ( + allowedActions.length > 0 && + !allowedActions.includes(statementType) && + !allowedActions.includes('*') + ) { + throw new Error( + `Unauthorized access: No matching rules for ${statementType} on restricted table ${tableRef.name}` + ) } } @@ -292,9 +335,16 @@ function applyRLSToAst(ast: any): void { ) .forEach(({ action, condition }) => { const targetTable = normalizeIdentifier(condition.left.table) - const isTargetTable = tables.includes(targetTable) + const tableRef = tableRefs.find((table) => + tableNamesMatch(targetTable, table.name) + ) + + if (!tableRef) return - if (!isTargetTable) return + const tableCondition = getConditionForTable( + condition, + tableRef.conditionTable + ) if (action !== 'INSERT') { // Add condition to WHERE with parentheses @@ -308,13 +358,13 @@ function applyRLSToAst(ast: any): void { parentheses: true, }, right: { - ...condition, + ...tableCondition, parentheses: true, }, } } else { ast.where = { - ...condition, + ...tableCondition, parentheses: true, } } @@ -349,8 +399,9 @@ function applyRLSToAst(ast: any): void { }) ast.from?.forEach((fromItem: any) => { - if (fromItem.expr && fromItem.expr.type === 'select') { - applyRLSToAst(fromItem.expr) + const nestedSelect = getNestedSelect(fromItem.expr) + if (nestedSelect) { + applyRLSToAst(nestedSelect) } // Handle both single join and array of joins @@ -359,8 +410,9 @@ function applyRLSToAst(ast: any): void { ? fromItem.join : [fromItem] joins.forEach((joinItem: any) => { - if (joinItem.expr && joinItem.expr.type === 'select') { - applyRLSToAst(joinItem.expr) + const joinSelect = getNestedSelect(joinItem.expr) + if (joinSelect) { + applyRLSToAst(joinSelect) } }) } @@ -371,8 +423,9 @@ function applyRLSToAst(ast: any): void { } ast.columns?.forEach((column: any) => { - if (column.expr && column.expr.type === 'select') { - applyRLSToAst(column.expr) + const nestedSelect = getNestedSelect(column.expr) + if (nestedSelect) { + applyRLSToAst(nestedSelect) } }) } @@ -382,6 +435,9 @@ function traverseWhere(node: any): void { if (node.type === 'select') { applyRLSToAst(node) } + if (node.ast?.type === 'select') { + applyRLSToAst(node.ast) + } if (node.left) traverseWhere(node.left) if (node.right) traverseWhere(node.right) }