Skip to content

Commit 3131b54

Browse files
committed
Consolidate multiple DB function tools into a single unified query function
This commit simplifies the API by replacing multiple specialized tools (read-query, write-query, list-tables, describe-table, create-table) with a single unified 'query' function. Benefits: - Simpler, more intuitive interface for users - Consistent access to all DuckDB features - Relies on DuckDB's native readonly mode for permission control - Removes unnecessary validation checks - Improves maintainability of the codebase
1 parent 62967f9 commit 3131b54

File tree

1 file changed

+4
-83
lines changed

1 file changed

+4
-83
lines changed

src/mcp_server_duckdb/server.py

+4-83
Original file line numberDiff line numberDiff line change
@@ -88,77 +88,21 @@ async def handle_list_tools() -> list[types.Tool]:
8888
"""List available tools"""
8989
tools = [
9090
types.Tool(
91-
name="read-query",
92-
description="Execute a SELECT query on the DuckDB database",
91+
name="query",
92+
description="Execute a query on the DuckDB database",
9393
inputSchema={
9494
"type": "object",
9595
"properties": {
9696
"query": {
9797
"type": "string",
98-
"description": "SELECT SQL query to execute",
98+
"description": "SQL query to execute",
9999
},
100100
},
101101
"required": ["query"],
102102
},
103103
),
104-
types.Tool(
105-
name="list-tables",
106-
description="List all tables in the DuckDB database",
107-
inputSchema={
108-
"type": "object",
109-
"properties": {},
110-
},
111-
),
112-
types.Tool(
113-
name="describe-table",
114-
description="Get the schema information for a specific table",
115-
inputSchema={
116-
"type": "object",
117-
"properties": {
118-
"table_name": {
119-
"type": "string",
120-
"description": "Name of the table to describe",
121-
},
122-
},
123-
"required": ["table_name"],
124-
},
125-
),
126104
]
127105

128-
if not config.readonly:
129-
tools.extend(
130-
[
131-
types.Tool(
132-
name="write-query",
133-
description="Execute an INSERT, UPDATE, or DELETE query on the DuckDB database",
134-
inputSchema={
135-
"type": "object",
136-
"properties": {
137-
"query": {
138-
"type": "string",
139-
"description": "SQL query to execute",
140-
},
141-
},
142-
"required": ["query"],
143-
},
144-
),
145-
types.Tool(
146-
name="create-table",
147-
description="Create a new table in the DuckDB database",
148-
inputSchema={
149-
"type": "object",
150-
"properties": {
151-
"query": {
152-
"type": "string",
153-
"description": "CREATE TABLE SQL statement",
154-
},
155-
},
156-
"required": ["query"],
157-
},
158-
),
159-
]
160-
)
161-
162106
return tools
163107

164108
@server.call_tool()
@@ -167,35 +111,12 @@ async def handle_call_tool(
167111
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
168112
"""Handle tool execution requests"""
169113
try:
170-
if name == "list-tables":
171-
results = db.execute_query("SELECT * FROM information_schema.tables;")
172-
return [types.TextContent(type="text", text=str(results))]
173-
174-
elif name == "describe-table":
175-
if not arguments or "table_name" not in arguments:
176-
raise ValueError("Missing table_name argument")
177-
results = db.execute_query("PRAGMA table_info(?)", [arguments["table_name"]])
178-
return [types.TextContent(type="text", text=str(results))]
179-
180114
if not arguments:
181115
raise ValueError("Missing arguments")
182116

183-
if name == "read-query":
184-
results = db.execute_query(arguments["query"])
185-
return [types.TextContent(type="text", text=str(results))]
186-
187-
elif name == "write-query":
188-
if arguments["query"].strip().upper().startswith("SELECT"):
189-
raise ValueError("SELECT queries are not allowed for write-query")
117+
if name == "query":
190118
results = db.execute_query(arguments["query"])
191119
return [types.TextContent(type="text", text=str(results))]
192-
193-
elif name == "create-table":
194-
if not arguments["query"].strip().upper().startswith("CREATE TABLE"):
195-
raise ValueError("Only CREATE TABLE statements are allowed")
196-
db.execute_query(arguments["query"])
197-
return [types.TextContent(type="text", text="Table created successfully")]
198-
199120
else:
200121
raise ValueError(f"Unknown tool: {name}")
201122

0 commit comments

Comments
 (0)