import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; import { ERROR_CODES } from '@insforge/shared-schemas'; const { connectMock } = vi.hoisted(() => ({ connectMock: vi.fn(), })); vi.mock('../../src/infra/database/database.manager', () => ({ DatabaseManager: { getInstance: vi.fn(() => ({ getPool: vi.fn(() => ({ connect: connectMock, })), })), clearColumnTypeCache: vi.fn(), }, })); import { DatabaseAdvanceService } from '../../src/services/database/database-advance.service'; import { initSqlParser } from '../../src/utils/sql-parser'; describe('DatabaseAdvanceService - admin SQL execution', () => { beforeAll(async () => { await initSqlParser(); }); beforeEach(() => { vi.clearAllMocks(); }); it('sanitizes raw SQL before opening a database connection', async () => { const service = DatabaseAdvanceService.getInstance(); await expect(service.executeRawSQL('RESET ROLE')).rejects.toMatchObject({ statusCode: 403, code: ERROR_CODES.FORBIDDEN, }); expect(connectMock).not.toHaveBeenCalled(); }); it('executes raw SQL under project_admin and resets the pooled session', async () => { const queryMock = vi .fn() .mockResolvedValueOnce({}) // SET statement_timeout .mockResolvedValueOnce({}) // SET ROLE project_admin .mockResolvedValueOnce({}) // set request.jwt.claims .mockResolvedValueOnce({ rows: [{ id: 1 }], rowCount: 1, fields: [{ name: 'id', dataTypeID: 23 }], }) // execute user SQL .mockResolvedValueOnce({}) // RESET ROLE .mockResolvedValueOnce({}) // reset request.jwt.claims .mockResolvedValueOnce({}) // NOTIFY pgrst .mockResolvedValueOnce({}); // reset statement_timeout connectMock.mockResolvedValue({ query: queryMock, release: vi.fn(), }); const service = DatabaseAdvanceService.getInstance(); const result = await service.executeRawSQL('CREATE TABLE public.products (id integer)', []); expect(result).toEqual({ rows: [{ id: 1 }], rowCount: 1, fields: [{ name: 'id', dataTypeID: 23 }], }); expect(queryMock).toHaveBeenNthCalledWith(1, 'SET statement_timeout = 30000'); expect(queryMock).toHaveBeenNthCalledWith(2, 'SET ROLE project_admin'); expect(queryMock).toHaveBeenNthCalledWith(3, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', JSON.stringify({ role: 'project_admin' }), false, ]); expect(queryMock).toHaveBeenNthCalledWith(4, 'CREATE TABLE public.products (id integer)', []); expect(queryMock).toHaveBeenNthCalledWith(5, 'RESET ROLE'); expect(queryMock).toHaveBeenNthCalledWith(6, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', '{}', false, ]); expect(queryMock).toHaveBeenNthCalledWith(7, `NOTIFY pgrst, 'reload schema';`); expect(queryMock).toHaveBeenNthCalledWith(8, 'SET statement_timeout = 0'); }); it('keeps unrestricted raw SQL on the root session', async () => { const queryMock = vi .fn() .mockResolvedValueOnce({}) // SET statement_timeout .mockResolvedValueOnce({ rows: [{ current_user: 'postgres' }], rowCount: 1, fields: [{ name: 'current_user', dataTypeID: 19 }], }) // execute user SQL .mockResolvedValueOnce({}); // reset statement_timeout connectMock.mockResolvedValue({ query: queryMock, release: vi.fn(), }); const service = DatabaseAdvanceService.getInstance(); const result = await service.executeRawSQL('SELECT current_user', [], true); expect(result.rows).toEqual([{ current_user: 'postgres' }]); expect(queryMock).toHaveBeenNthCalledWith(1, 'SET statement_timeout = 30000'); expect(queryMock).toHaveBeenNthCalledWith(2, 'SELECT current_user', []); expect(queryMock).toHaveBeenNthCalledWith(3, 'SET statement_timeout = 0'); expect(queryMock).not.toHaveBeenCalledWith('SET ROLE project_admin'); expect(queryMock).not.toHaveBeenCalledWith('RESET ROLE'); }); it('returns healthy clients to the pool after ordinary SQL errors', async () => { const releaseMock = vi.fn(); const queryError = new Error('syntax error'); const queryMock = vi .fn() .mockResolvedValueOnce({}) // SET statement_timeout .mockResolvedValueOnce({}) // SET ROLE project_admin .mockResolvedValueOnce({}) // set request.jwt.claims .mockRejectedValueOnce(queryError) // execute user SQL .mockResolvedValueOnce({}) // RESET ROLE .mockResolvedValueOnce({}) // reset request.jwt.claims .mockResolvedValueOnce({}); // reset statement_timeout connectMock.mockResolvedValue({ query: queryMock, release: releaseMock, }); const service = DatabaseAdvanceService.getInstance(); await expect(service.executeRawSQL('SELECT broken')).rejects.toBe(queryError); expect(releaseMock).toHaveBeenCalledWith(undefined); }); it('discards the pooled client when admin context cleanup fails', async () => { const releaseMock = vi.fn(); const resetError = new Error('reset failed'); const queryMock = vi .fn() .mockResolvedValueOnce({}) // SET statement_timeout .mockResolvedValueOnce({}) // SET ROLE project_admin .mockResolvedValueOnce({}) // set request.jwt.claims .mockResolvedValueOnce({ rows: [], rowCount: 0, fields: [] }) // execute user SQL .mockRejectedValueOnce(resetError) // RESET ROLE .mockResolvedValueOnce({}); // reset statement_timeout connectMock.mockResolvedValue({ query: queryMock, release: releaseMock, }); const service = DatabaseAdvanceService.getInstance(); await expect(service.executeRawSQL('SELECT 1')).rejects.toBe(resetError); expect(releaseMock).toHaveBeenCalledWith(resetError); }); it('imports SQL files under project_admin', async () => { const releaseMock = vi.fn(); const queryMock = vi .fn() .mockResolvedValueOnce({}) // BEGIN .mockResolvedValueOnce({}) // SET LOCAL ROLE project_admin .mockResolvedValueOnce({}) // set local request.jwt.claims .mockResolvedValueOnce({ rowCount: 1, rows: [] }) // execute imported SQL .mockResolvedValueOnce({}) // RESET ROLE .mockResolvedValueOnce({}) // reset request.jwt.claims .mockResolvedValueOnce({}) // NOTIFY pgrst .mockResolvedValueOnce({}); // COMMIT connectMock.mockResolvedValue({ query: queryMock, release: releaseMock, }); const service = DatabaseAdvanceService.getInstance(); const result = await service.importDatabase( Buffer.from('INSERT INTO products (id) VALUES (1);'), 'seed.sql', 36 ); expect(result.rowsImported).toBe(1); expect(queryMock).toHaveBeenNthCalledWith(1, 'BEGIN'); expect(queryMock).toHaveBeenNthCalledWith(2, 'SET LOCAL ROLE project_admin'); expect(queryMock).toHaveBeenNthCalledWith(3, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', JSON.stringify({ role: 'project_admin' }), true, ]); expect(String(queryMock.mock.calls[3][0])).toContain('INSERT INTO products'); expect(queryMock).toHaveBeenNthCalledWith(5, 'RESET ROLE'); expect(queryMock).toHaveBeenNthCalledWith(6, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', '{}', true, ]); expect(queryMock).toHaveBeenNthCalledWith(7, `NOTIFY pgrst, 'reload schema';`); expect(queryMock).toHaveBeenNthCalledWith(8, 'COMMIT'); expect(releaseMock).toHaveBeenCalled(); }); it('handles table truncate failures using savepoints in importDatabase', async () => { const releaseMock = vi.fn(); const queryMock = vi .fn() .mockResolvedValueOnce({}) // BEGIN .mockResolvedValueOnce({}) // SET LOCAL ROLE project_admin .mockResolvedValueOnce({}) // set local request.jwt.claims // For the SELECT tablename query: .mockResolvedValueOnce({ rows: [{ tablename: 'table1' }, { tablename: 'table2' }], }) // Truncate table1: .mockResolvedValueOnce({}) // SAVEPOINT truncate_attempt .mockRejectedValueOnce(new Error('Truncate table1 failed (mock locked)')) // TRUNCATE TABLE table1 CASCADE (FAILS) .mockResolvedValueOnce({}) // ROLLBACK TO SAVEPOINT truncate_attempt .mockResolvedValueOnce({}) // RELEASE SAVEPOINT truncate_attempt (after rollback) // Truncate table2: .mockResolvedValueOnce({}) // SAVEPOINT truncate_attempt .mockResolvedValueOnce({}) // TRUNCATE TABLE table2 CASCADE (SUCCEEDS) .mockResolvedValueOnce({}) // RELEASE SAVEPOINT truncate_attempt // SQL statements: .mockResolvedValueOnce({ rowCount: 1, rows: [] }) // execute imported SQL // Reset context: .mockResolvedValueOnce({}) // RESET ROLE .mockResolvedValueOnce({}) // reset request.jwt.claims .mockResolvedValueOnce({}) // NOTIFY pgrst .mockResolvedValueOnce({}); // COMMIT connectMock.mockResolvedValue({ query: queryMock, release: releaseMock, }); const service = DatabaseAdvanceService.getInstance(); const result = await service.importDatabase( Buffer.from('INSERT INTO products (id) VALUES (1);'), 'seed.sql', 36, true // truncate = true ); expect(result.rowsImported).toBe(1); expect(queryMock).toHaveBeenCalledWith('SAVEPOINT truncate_attempt'); expect(queryMock).toHaveBeenCalledWith('ROLLBACK TO SAVEPOINT truncate_attempt'); expect(queryMock).toHaveBeenCalledWith('RELEASE SAVEPOINT truncate_attempt'); expect(queryMock).toHaveBeenLastCalledWith('COMMIT'); expect(releaseMock).toHaveBeenCalled(); }); it('bulk upserts under project_admin', async () => { const releaseMock = vi.fn(); const queryMock = vi .fn() .mockResolvedValueOnce({}) // 1. SET ROLE project_admin .mockResolvedValueOnce({}) // 2. set request.jwt.claims .mockResolvedValueOnce({}) // 3. BEGIN .mockResolvedValueOnce({ rowCount: 1, rows: [] }) // 4. execute bulk upsert .mockResolvedValueOnce({}) // 5. NOTIFY pgrst .mockResolvedValueOnce({}) // 6. COMMIT .mockResolvedValueOnce({}) // 7. RESET ROLE .mockResolvedValueOnce({}); // 8. reset request.jwt.claims connectMock.mockResolvedValue({ query: queryMock, release: releaseMock, }); const service = DatabaseAdvanceService.getInstance(); const result = await service.bulkUpsertFromFile( 'public', 'profiles', Buffer.from('id,name\n1,Alice\n'), 'profiles.csv', 'id' ); expect(result.rowsAffected).toBe(1); expect(queryMock).toHaveBeenNthCalledWith(1, 'SET ROLE project_admin'); expect(queryMock).toHaveBeenNthCalledWith(2, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', JSON.stringify({ role: 'project_admin' }), false, ]); expect(queryMock).toHaveBeenNthCalledWith(3, 'BEGIN'); expect(String(queryMock.mock.calls[3][0])).toContain('INSERT INTO public.profiles'); expect(String(queryMock.mock.calls[3][0])).toContain('ON CONFLICT (id) DO UPDATE'); expect(queryMock).toHaveBeenNthCalledWith(5, `NOTIFY pgrst, 'reload schema';`); expect(queryMock).toHaveBeenNthCalledWith(6, 'COMMIT'); expect(queryMock).toHaveBeenNthCalledWith(7, 'RESET ROLE'); expect(queryMock).toHaveBeenNthCalledWith(8, 'SELECT set_config($1, $2, $3)', [ 'request.jwt.claims', '{}', false, ]); expect(releaseMock).toHaveBeenCalled(); }); });