304 lines
11 KiB
TypeScript
304 lines
11 KiB
TypeScript
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();
|
|
});
|
|
});
|