Skip to content

Commit 6591cc0

Browse files
tianzhouCopilot
andauthored
fix: read-only flag (#155)
* fix: read-only flag * Update src/tools/execute-sql.ts Co-authored-by: Copilot <[email protected]> * Update src/tools/execute-sql.ts Co-authored-by: Copilot <[email protected]> --------- Co-authored-by: Copilot <[email protected]>
1 parent a052a0a commit 6591cc0

File tree

4 files changed

+92
-123
lines changed

4 files changed

+92
-123
lines changed

docs/config/server-options.mdx

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,8 +330,9 @@ DBHub follows this priority order for configuration:
330330
**`--config` (TOML) is mutually exclusive with:**
331331
- `--dsn` - TOML configuration defines database connections directly
332332
- `--id` - TOML configuration defines source IDs directly in the file
333+
- `--readonly` - TOML configuration defines readonly mode per-source using `readonly = true`
333334

334-
If using TOML config, remove these flags. If you need `--id` or `--dsn`, use command-line/environment configuration instead.
335+
If using TOML config, remove these flags. If you need `--id`, `--dsn`, or `--readonly`, use command-line/environment configuration instead.
335336
</Note>
336337

337338
## Quick Reference

src/config/env.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,14 @@ export async function resolveSourceConfigs(): Promise<{ sources: SourceConfig[];
487487
"Either remove the --id flag or use command-line DSN configuration instead."
488488
);
489489
}
490+
// Validate that --readonly flag is not used with TOML config
491+
if (isReadOnlyMode()) {
492+
throw new Error(
493+
"The --readonly flag cannot be used with TOML configuration. " +
494+
"TOML config defines readonly mode per-source using 'readonly = true'. " +
495+
"Either remove the --readonly flag or use command-line DSN configuration instead."
496+
);
497+
}
490498
return tomlConfig;
491499
}
492500
}
Lines changed: 78 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,10 @@
11
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
22
import { createExecuteSqlToolHandler } from '../execute-sql.js';
33
import { ConnectorManager } from '../../connectors/manager.js';
4-
import { isReadOnlyMode } from '../../config/env.js';
54
import type { Connector, ConnectorType, SQLResult } from '../../connectors/interface.js';
65

76
// Mock dependencies
87
vi.mock('../../connectors/manager.js');
9-
vi.mock('../../config/env.js');
108

119
// Mock connector for testing
1210
const createMockConnector = (id: ConnectorType = 'sqlite'): Connector => ({
@@ -34,21 +32,19 @@ describe('execute-sql tool', () => {
3432
let mockConnector: Connector;
3533
const mockGetCurrentConnector = vi.mocked(ConnectorManager.getCurrentConnector);
3634
const mockGetCurrentExecuteOptions = vi.mocked(ConnectorManager.getCurrentExecuteOptions);
37-
const mockIsReadOnlyMode = vi.mocked(isReadOnlyMode);
3835

3936
beforeEach(() => {
4037
mockConnector = createMockConnector('sqlite');
4138
mockGetCurrentConnector.mockReturnValue(mockConnector);
4239
mockGetCurrentExecuteOptions.mockReturnValue({});
43-
mockIsReadOnlyMode.mockReturnValue(false);
4440
});
4541

4642
afterEach(() => {
4743
vi.clearAllMocks();
4844
});
4945

50-
describe('single statement execution', () => {
51-
it('should execute a single SELECT statement successfully', async () => {
46+
describe('basic execution', () => {
47+
it('should execute SELECT and return rows', async () => {
5248
const mockResult: SQLResult = { rows: [{ id: 1, name: 'test' }] };
5349
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
5450

@@ -62,21 +58,6 @@ describe('execute-sql tool', () => {
6258
expect(mockConnector.executeSQL).toHaveBeenCalledWith('SELECT * FROM users', {});
6359
});
6460

65-
it('should handle execution errors', async () => {
66-
vi.mocked(mockConnector.executeSQL).mockRejectedValue(new Error('Database error'));
67-
68-
const handler = createExecuteSqlToolHandler('test_source');
69-
const result = await handler({ sql: 'SELECT * FROM invalid_table' }, null);
70-
71-
expect(result.isError).toBe(true);
72-
const parsedResult = parseToolResponse(result);
73-
expect(parsedResult.success).toBe(false);
74-
expect(parsedResult.error).toBe('Database error');
75-
expect(parsedResult.code).toBe('EXECUTION_ERROR');
76-
});
77-
});
78-
79-
describe('multi-statement execution', () => {
8061
it('should pass multi-statement SQL directly to connector', async () => {
8162
const mockResult: SQLResult = { rows: [{ id: 1 }] };
8263
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
@@ -89,14 +70,28 @@ describe('execute-sql tool', () => {
8970
expect(parsedResult.success).toBe(true);
9071
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
9172
});
73+
74+
it('should handle execution errors', async () => {
75+
vi.mocked(mockConnector.executeSQL).mockRejectedValue(new Error('Database error'));
76+
77+
const handler = createExecuteSqlToolHandler('test_source');
78+
const result = await handler({ sql: 'SELECT * FROM invalid_table' }, null);
79+
80+
expect(result.isError).toBe(true);
81+
const parsedResult = parseToolResponse(result);
82+
expect(parsedResult.success).toBe(false);
83+
expect(parsedResult.error).toBe('Database error');
84+
expect(parsedResult.code).toBe('EXECUTION_ERROR');
85+
});
9286
});
9387

94-
describe('read-only mode validation', () => {
88+
describe('read-only mode enforcement', () => {
9589
beforeEach(() => {
96-
mockIsReadOnlyMode.mockReturnValue(true);
90+
// Set per-source readonly mode via executeOptions (simulates TOML config)
91+
mockGetCurrentExecuteOptions.mockReturnValue({ readonly: true });
9792
});
9893

99-
it('should allow single SELECT statement in read-only mode', async () => {
94+
it('should allow SELECT statements', async () => {
10095
const mockResult: SQLResult = { rows: [{ id: 1 }] };
10196
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
10297

@@ -105,159 +100,124 @@ describe('execute-sql tool', () => {
105100
const parsedResult = parseToolResponse(result);
106101

107102
expect(parsedResult.success).toBe(true);
108-
expect(mockConnector.executeSQL).toHaveBeenCalled();
103+
expect(mockConnector.executeSQL).toHaveBeenCalledWith('SELECT * FROM users', { readonly: true });
109104
});
110105

111-
it('should allow multiple read-only statements in read-only mode', async () => {
106+
it('should allow multiple read-only statements', async () => {
112107
const mockResult: SQLResult = { rows: [] };
113108
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
114109

115110
const sql = 'SELECT * FROM users; SELECT * FROM roles;';
116111
const handler = createExecuteSqlToolHandler('test_source');
117112
const result = await handler({ sql }, null);
118-
const parsedResult = parseToolResponse(result);
119113

120-
expect(parsedResult.success).toBe(true);
121-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
114+
expect(parseToolResponse(result).success).toBe(true);
122115
});
123116

124-
it('should reject single INSERT statement in read-only mode', async () => {
117+
it.each([
118+
['INSERT', "INSERT INTO users (name) VALUES ('test')"],
119+
['UPDATE', "UPDATE users SET name = 'x' WHERE id = 1"],
120+
['DELETE', "DELETE FROM users WHERE id = 1"],
121+
['DROP', "DROP TABLE users"],
122+
['CREATE', "CREATE TABLE test (id INT)"],
123+
['ALTER', "ALTER TABLE users ADD COLUMN email VARCHAR(255)"],
124+
['TRUNCATE', "TRUNCATE TABLE users"],
125+
])('should reject %s statement', async (_, sql) => {
125126
const handler = createExecuteSqlToolHandler('test_source');
126-
const result = await handler({ sql: "INSERT INTO users (name) VALUES ('test')" }, null);
127+
const result = await handler({ sql }, null);
127128

128129
expect(result.isError).toBe(true);
129130
const parsedResult = parseToolResponse(result);
130-
expect(parsedResult.success).toBe(false);
131-
expect(parsedResult.error).toContain('Read-only mode is enabled');
132131
expect(parsedResult.code).toBe('READONLY_VIOLATION');
133132
expect(mockConnector.executeSQL).not.toHaveBeenCalled();
134133
});
135134

136-
it('should reject multi-statement with any write operation in read-only mode', async () => {
137-
const sql = "SELECT * FROM users; INSERT INTO users (name) VALUES ('test'); SELECT COUNT(*) FROM users;";
135+
it('should reject multi-statement with any write operation', async () => {
136+
const sql = "SELECT * FROM users; INSERT INTO users (name) VALUES ('test');";
138137
const handler = createExecuteSqlToolHandler('test_source');
139138
const result = await handler({ sql }, null);
140139

141140
expect(result.isError).toBe(true);
142-
const parsedResult = parseToolResponse(result);
143-
expect(parsedResult.success).toBe(false);
144-
expect(parsedResult.error).toContain('Read-only mode is enabled');
145-
expect(parsedResult.code).toBe('READONLY_VIOLATION');
146-
expect(mockConnector.executeSQL).not.toHaveBeenCalled();
141+
expect(parseToolResponse(result).code).toBe('READONLY_VIOLATION');
147142
});
148143

149-
});
150-
151-
describe('SQL comments handling', () => {
152-
it('should allow SELECT with single-line comment in read-only mode', async () => {
153-
mockIsReadOnlyMode.mockReturnValue(true);
154-
const mockResult: SQLResult = { rows: [{ id: 1 }] };
155-
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
156-
157-
const sql = '-- Fetch active users\nSELECT * FROM users WHERE active = TRUE';
158-
const handler = createExecuteSqlToolHandler('test_source');
159-
const result = await handler({ sql }, null);
160-
const parsedResult = parseToolResponse(result);
144+
it('should include source_id in error message', async () => {
145+
const handler = createExecuteSqlToolHandler('prod_db');
146+
const result = await handler({ sql: "DROP TABLE users" }, null);
161147

162-
expect(parsedResult.success).toBe(true);
163-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
148+
expect(parseToolResponse(result).error).toContain('prod_db');
164149
});
150+
});
151+
152+
describe('readonly per-source isolation', () => {
153+
// Verifies readonly is enforced per-source from executeOptions, not globally
165154

166-
it('should allow SELECT with multi-line comment in read-only mode', async () => {
167-
mockIsReadOnlyMode.mockReturnValue(true);
155+
it.each([
156+
['readonly: false', { readonly: false }],
157+
['readonly: undefined', {}],
158+
])('should allow writes when %s', async (_, options) => {
159+
mockGetCurrentExecuteOptions.mockReturnValue(options);
168160
const mockResult: SQLResult = { rows: [] };
169161
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
170162

171-
const sql = '/* This query fetches\n all products */\nSELECT * FROM products';
172-
const handler = createExecuteSqlToolHandler('test_source');
173-
const result = await handler({ sql }, null);
174-
const parsedResult = parseToolResponse(result);
163+
const handler = createExecuteSqlToolHandler('writable_source');
164+
const result = await handler({ sql: "INSERT INTO users (name) VALUES ('test')" }, null);
175165

176-
expect(parsedResult.success).toBe(true);
177-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
166+
expect(parseToolResponse(result).success).toBe(true);
167+
expect(mockConnector.executeSQL).toHaveBeenCalled();
178168
});
179169

180-
it('should handle multiple statements with comments in read-only mode', async () => {
181-
mockIsReadOnlyMode.mockReturnValue(true);
182-
const mockResult: SQLResult = { rows: [] };
183-
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
170+
it('should enforce readonly even with other options set', async () => {
171+
mockGetCurrentExecuteOptions.mockReturnValue({ readonly: true, maxRows: 100 });
184172

185-
const sql = '-- First query\nSELECT * FROM users;\n/* Second query */\nSELECT * FROM roles;';
186-
const handler = createExecuteSqlToolHandler('test_source');
187-
const result = await handler({ sql }, null);
188-
const parsedResult = parseToolResponse(result);
173+
const handler = createExecuteSqlToolHandler('limited_source');
174+
const result = await handler({ sql: "DELETE FROM users" }, null);
189175

190-
expect(parsedResult.success).toBe(true);
191-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
176+
expect(parseToolResponse(result).code).toBe('READONLY_VIOLATION');
192177
});
178+
});
193179

194-
it('should reject INSERT with comment in read-only mode', async () => {
195-
mockIsReadOnlyMode.mockReturnValue(true);
196-
197-
const sql = '-- Insert new user\nINSERT INTO users (name) VALUES (\'test\')';
198-
const handler = createExecuteSqlToolHandler('test_source');
199-
const result = await handler({ sql }, null);
200-
201-
expect(result.isError).toBe(true);
202-
const parsedResult = parseToolResponse(result);
203-
expect(parsedResult.success).toBe(false);
204-
expect(parsedResult.code).toBe('READONLY_VIOLATION');
205-
expect(mockConnector.executeSQL).not.toHaveBeenCalled();
180+
describe('SQL comments handling in readonly mode', () => {
181+
beforeEach(() => {
182+
mockGetCurrentExecuteOptions.mockReturnValue({ readonly: true });
206183
});
207184

208-
it('should handle query that is only comments as read-only', async () => {
209-
mockIsReadOnlyMode.mockReturnValue(true);
185+
it.each([
186+
['single-line comment', '-- Fetch users\nSELECT * FROM users'],
187+
['multi-line comment', '/* Fetch all */\nSELECT * FROM products'],
188+
['inline comments', 'SELECT id, -- user id\n name FROM users'],
189+
['only comments', '-- Just a comment\n/* Another */'],
190+
])('should allow SELECT with %s', async (_, sql) => {
210191
const mockResult: SQLResult = { rows: [] };
211192
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
212193

213-
const sql = '-- Just a comment\n/* Another comment */';
214194
const handler = createExecuteSqlToolHandler('test_source');
215195
const result = await handler({ sql }, null);
216-
const parsedResult = parseToolResponse(result);
217196

218-
expect(parsedResult.success).toBe(true);
219-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
197+
expect(parseToolResponse(result).success).toBe(true);
220198
});
221199

222-
it('should handle inline comments correctly', async () => {
223-
mockIsReadOnlyMode.mockReturnValue(true);
224-
const mockResult: SQLResult = { rows: [] };
225-
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
226-
227-
const sql = 'SELECT id, -- user id\n name -- user name\nFROM users';
200+
it('should reject write statement hidden after comment', async () => {
201+
const sql = '-- Insert new user\nINSERT INTO users (name) VALUES (\'test\')';
228202
const handler = createExecuteSqlToolHandler('test_source');
229203
const result = await handler({ sql }, null);
230-
const parsedResult = parseToolResponse(result);
231204

232-
expect(parsedResult.success).toBe(true);
233-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(sql, {});
205+
expect(parseToolResponse(result).code).toBe('READONLY_VIOLATION');
234206
});
235207
});
236208

237-
238209
describe('edge cases', () => {
239-
it('should handle empty SQL string', async () => {
210+
it.each([
211+
['empty string', ''],
212+
['only semicolons and whitespace', ' ; ; ; '],
213+
])('should handle %s', async (_, sql) => {
240214
const mockResult: SQLResult = { rows: [] };
241215
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
242216

243217
const handler = createExecuteSqlToolHandler('test_source');
244-
const result = await handler({ sql: '' }, null);
245-
const parsedResult = parseToolResponse(result);
246-
247-
expect(parsedResult.success).toBe(true);
248-
expect(mockConnector.executeSQL).toHaveBeenCalledWith('', {});
249-
});
250-
251-
it('should handle SQL with only semicolons and whitespace', async () => {
252-
const mockResult: SQLResult = { rows: [] };
253-
vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
254-
255-
const handler = createExecuteSqlToolHandler('test_source');
256-
const result = await handler({ sql: ' ; ; ; ' }, null);
257-
const parsedResult = parseToolResponse(result);
218+
const result = await handler({ sql }, null);
258219

259-
expect(parsedResult.success).toBe(true);
260-
expect(mockConnector.executeSQL).toHaveBeenCalledWith(' ; ; ; ', {});
220+
expect(parseToolResponse(result).success).toBe(true);
261221
});
262222
});
263-
});
223+
});

src/tools/execute-sql.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import { z } from "zod";
22
import { ConnectorManager } from "../connectors/manager.js";
33
import { createToolSuccessResponse, createToolErrorResponse } from "../utils/response-formatter.js";
4-
import { isReadOnlyMode } from "../config/env.js";
54
import { allowedKeywords } from "../utils/allowed-keywords.js";
65
import { ConnectorType } from "../connectors/interface.js";
76
import { requestStore } from "../requests/index.js";
@@ -108,9 +107,10 @@ export function createExecuteSqlToolHandler(sourceId?: string) {
108107
const connector = ConnectorManager.getCurrentConnector(sourceId);
109108
const executeOptions = ConnectorManager.getCurrentExecuteOptions(sourceId);
110109

111-
// Check if SQL is allowed based on readonly mode
112-
if (isReadOnlyMode() && !areAllStatementsReadOnly(sql, connector.id)) {
113-
errorMessage = `Read-only mode is enabled. Only the following SQL operations are allowed: ${allowedKeywords[connector.id]?.join(", ") || "none"}`;
110+
// Check if SQL is allowed based on readonly mode (per-source)
111+
const isReadonly = executeOptions.readonly === true;
112+
if (isReadonly && !areAllStatementsReadOnly(sql, connector.id)) {
113+
errorMessage = `Read-only mode is enabled for source '${effectiveSourceId}'. Only the following SQL operations are allowed: ${allowedKeywords[connector.id]?.join(", ") || "none"}`;
114114
success = false;
115115
return createToolErrorResponse(errorMessage, "READONLY_VIOLATION");
116116
}

0 commit comments

Comments
 (0)