fix(rls): resolve policy table matching and bring test coverage above 75% across all metrics

pnpm-workspace.yaml

@@ -0,0 +1,3 @@
+allowBuilds:
+  esbuild: true
+  workerd: true

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()
+    })
+})

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)
+    })
 })