256 lines
8.0 KiB
TypeScript
256 lines
8.0 KiB
TypeScript
|
import { NetaClawDataSourceEntity } from '../src/modules/netaclaw/entity/data_source.js';
|
||
|
|
import { MysqlQueryService } from '../src/modules/netaclaw/service/mysql_query.js';
|
||
|
|
|
||
|
|
function makeSource(overrides: Partial<NetaClawDataSourceEntity> = {}) {
|
||
|
|
return Object.assign(new NetaClawDataSourceEntity(), {
|
||
|
|
id: 42,
|
||
|
|
name: 'orders',
|
||
|
|
type: 'mysql',
|
||
|
|
database: 'shop',
|
||
|
|
extra: {
|
||
|
|
allowedTables: ['customers', 'orders'],
|
||
|
|
blockedTables: [],
|
||
|
|
maxRows: 2,
|
||
|
|
maxJoinTables: 4,
|
||
|
|
maskedColumns: { 'customers.phone': 'partial' },
|
||
|
|
queryTimeoutMs: 5000,
|
||
|
|
},
|
||
|
|
...overrides,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
function makeService(pool: any, auditRepo: any) {
|
||
|
|
const service = new MysqlQueryService();
|
||
|
|
service.mysqlPoolManager = {
|
||
|
|
getPool: jest.fn().mockResolvedValue(pool),
|
||
|
|
} as any;
|
||
|
|
service.auditRepo = auditRepo;
|
||
|
|
return service;
|
||
|
|
}
|
||
|
|
|
||
|
|
describe('MysqlQueryService', () => {
|
||
|
|
it('executes guarded query and writes success audit', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1, name: 'Ada' }],
|
||
|
|
[{ name: 'id' }, { name: 'name' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
const result = await service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT id, name FROM customers LIMIT 1',
|
||
|
|
agentId: 7,
|
||
|
|
userId: 8,
|
||
|
|
toolCallId: 'tool-1',
|
||
|
|
});
|
||
|
|
|
||
|
|
expect(execute).toHaveBeenCalledWith(
|
||
|
|
{ sql: 'SELECT id, name FROM customers LIMIT 1', timeout: 5000 },
|
||
|
|
[]
|
||
|
|
);
|
||
|
|
expect(result).toMatchObject({
|
||
|
|
columns: ['id', 'name'],
|
||
|
|
rows: [{ id: 1, name: 'Ada' }],
|
||
|
|
rowCount: 1,
|
||
|
|
truncated: false,
|
||
|
|
sql: 'SELECT id, name FROM customers LIMIT 1',
|
||
|
|
});
|
||
|
|
expect(result.elapsedMs).toEqual(expect.any(Number));
|
||
|
|
expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
|
||
|
|
dataSourceId: 42,
|
||
|
|
agentId: 7,
|
||
|
|
userId: 8,
|
||
|
|
toolCallId: 'tool-1',
|
||
|
|
status: 'success',
|
||
|
|
rejectReason: null,
|
||
|
|
rowCount: 1,
|
||
|
|
errorCode: null,
|
||
|
|
sqlPreview: 'SELECT id, name FROM customers LIMIT 1',
|
||
|
|
}));
|
||
|
|
expect(auditRepo.save.mock.calls[0][0].sqlHash).toMatch(/^[a-f0-9]{64}$/);
|
||
|
|
});
|
||
|
|
|
||
|
|
it('rejects masked column before querying MySQL and audits rejected', async () => {
|
||
|
|
const execute = jest.fn();
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await expect(service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT phone FROM customers LIMIT 1',
|
||
|
|
})).rejects.toThrow('mysql_sql_rejected: masked_column_denied');
|
||
|
|
|
||
|
|
expect(execute).not.toHaveBeenCalled();
|
||
|
|
expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
|
||
|
|
status: 'rejected',
|
||
|
|
rejectReason: 'masked_column_denied',
|
||
|
|
rowCount: 0,
|
||
|
|
errorCode: null,
|
||
|
|
}));
|
||
|
|
});
|
||
|
|
|
||
|
|
it('audits rejected SQL', async () => {
|
||
|
|
const execute = jest.fn();
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await expect(service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'DELETE FROM customers',
|
||
|
|
agentId: null,
|
||
|
|
userId: null,
|
||
|
|
toolCallId: null,
|
||
|
|
})).rejects.toThrow('mysql_sql_rejected: dml_sql_denied');
|
||
|
|
|
||
|
|
expect(execute).not.toHaveBeenCalled();
|
||
|
|
expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
|
||
|
|
dataSourceId: 42,
|
||
|
|
agentId: null,
|
||
|
|
userId: null,
|
||
|
|
toolCallId: null,
|
||
|
|
status: 'rejected',
|
||
|
|
rejectReason: 'dml_sql_denied',
|
||
|
|
rowCount: 0,
|
||
|
|
}));
|
||
|
|
});
|
||
|
|
|
||
|
|
it('uses no-limit wrapper params, truncates node-side, and audits fetched row count', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1 }, { id: 2 }, { id: 3 }],
|
||
|
|
[{ name: 'id' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
const result = await service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT id FROM customers',
|
||
|
|
});
|
||
|
|
|
||
|
|
expect(execute).toHaveBeenCalledWith(
|
||
|
|
{
|
||
|
|
sql: 'SELECT * FROM (SELECT id FROM customers) AS neta_limited_query LIMIT ?',
|
||
|
|
timeout: 5000,
|
||
|
|
},
|
||
|
|
[3]
|
||
|
|
);
|
||
|
|
expect(result.rows).toEqual([{ id: 1 }, { id: 2 }]);
|
||
|
|
expect(result.rowCount).toBe(2);
|
||
|
|
expect(result.truncated).toBe(true);
|
||
|
|
expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
|
||
|
|
status: 'success',
|
||
|
|
rowCount: 2,
|
||
|
|
}));
|
||
|
|
});
|
||
|
|
|
||
|
|
it('audits mysql failures with sanitized error codes', async () => {
|
||
|
|
const execute = jest.fn().mockRejectedValue({ code: 'ER_PARSE_ERROR: detail' });
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await expect(service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT id FROM customers LIMIT 1',
|
||
|
|
})).rejects.toThrow('mysql_query_failed:ER_PARSE_ERROR_detail');
|
||
|
|
|
||
|
|
expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
|
||
|
|
status: 'failed',
|
||
|
|
rejectReason: null,
|
||
|
|
rowCount: 0,
|
||
|
|
errorCode: 'ER_PARSE_ERROR_detail',
|
||
|
|
}));
|
||
|
|
});
|
||
|
|
|
||
|
|
it('returns successful query results when success audit write fails', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1 }],
|
||
|
|
[{ name: 'id' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockRejectedValue(new Error('audit_down')) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await expect(service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT id FROM customers LIMIT 1',
|
||
|
|
})).resolves.toMatchObject({
|
||
|
|
rows: [{ id: 1 }],
|
||
|
|
rowCount: 1,
|
||
|
|
truncated: false,
|
||
|
|
});
|
||
|
|
});
|
||
|
|
|
||
|
|
it('preserves guard rejection reason when rejected audit write fails', async () => {
|
||
|
|
const execute = jest.fn();
|
||
|
|
const auditRepo = { save: jest.fn().mockRejectedValue(new Error('audit_down')) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await expect(service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'DELETE FROM customers',
|
||
|
|
})).rejects.toThrow('mysql_sql_rejected: dml_sql_denied');
|
||
|
|
|
||
|
|
expect(execute).not.toHaveBeenCalled();
|
||
|
|
});
|
||
|
|
|
||
|
|
it('does not pass wrapper limit params for explicit LIMIT SQL containing a question mark literal', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1 }],
|
||
|
|
[{ name: 'id' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: 'SELECT id FROM customers WHERE name = "?" LIMIT 1',
|
||
|
|
});
|
||
|
|
|
||
|
|
expect(execute).toHaveBeenCalledWith(
|
||
|
|
{ sql: 'SELECT id FROM customers WHERE name = "?" LIMIT 1', timeout: 5000 },
|
||
|
|
[]
|
||
|
|
);
|
||
|
|
});
|
||
|
|
|
||
|
|
it('does not pass wrapper limit params when explicit LIMIT SQL has surrounding whitespace', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1 }],
|
||
|
|
[{ name: 'id' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: ' SELECT id FROM customers LIMIT 1 ',
|
||
|
|
});
|
||
|
|
|
||
|
|
expect(execute).toHaveBeenCalledWith(
|
||
|
|
{ sql: 'SELECT id FROM customers LIMIT 1', timeout: 5000 },
|
||
|
|
[]
|
||
|
|
);
|
||
|
|
});
|
||
|
|
|
||
|
|
it('passes wrapper limit params when LIMIT appears only in a string literal', async () => {
|
||
|
|
const execute = jest.fn().mockResolvedValue([
|
||
|
|
[{ id: 1 }, { id: 2 }, { id: 3 }],
|
||
|
|
[{ name: 'id' }],
|
||
|
|
]);
|
||
|
|
const auditRepo = { save: jest.fn().mockResolvedValue({}) };
|
||
|
|
const service = makeService({ execute }, auditRepo);
|
||
|
|
|
||
|
|
await service.executeReadOnly({
|
||
|
|
source: makeSource(),
|
||
|
|
sql: "SELECT id FROM customers WHERE note = 'limit 1'",
|
||
|
|
});
|
||
|
|
|
||
|
|
expect(execute).toHaveBeenCalledWith(
|
||
|
|
{
|
||
|
|
sql: "SELECT * FROM (SELECT id FROM customers WHERE note = 'limit 1') AS neta_limited_query LIMIT ?",
|
||
|
|
timeout: 5000,
|
||
|
|
},
|
||
|
|
[3]
|
||
|
|
);
|
||
|
|
});
|
||
|
|
});
|