Skip to content

Commit 2bb62b7

Browse files
committed
fix: fix new add_tool interface
1 parent fdc5fba commit 2bb62b7

File tree

2 files changed

+11
-10
lines changed

2 files changed

+11
-10
lines changed

mcp_clickhouse/mcp_server.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from clickhouse_connect.driver.binding import format_query_value
1111
from dotenv import load_dotenv
1212
from fastmcp import FastMCP
13+
from fastmcp.tools import Tool
1314
from fastmcp.prompts import Prompt
1415
from dataclasses import dataclass, field, asdict, is_dataclass
1516

@@ -321,9 +322,9 @@ def _init_chdb_client():
321322

322323
# Register tools based on configuration
323324
if os.getenv("CLICKHOUSE_ENABLED", "true").lower() == "true":
324-
mcp.add_tool(list_databases)
325-
mcp.add_tool(list_tables)
326-
mcp.add_tool(run_select_query)
325+
mcp.add_tool(Tool.from_function(list_databases))
326+
mcp.add_tool(Tool.from_function(list_tables))
327+
mcp.add_tool(Tool.from_function(run_select_query))
327328
logger.info("ClickHouse tools registered")
328329

329330

@@ -332,7 +333,7 @@ def _init_chdb_client():
332333
if _chdb_client:
333334
atexit.register(lambda: _chdb_client.close())
334335

335-
mcp.add_tool(run_chdb_select_query)
336+
mcp.add_tool(Tool.from_function(run_chdb_select_query))
336337
chdb_prompt = Prompt.from_function(
337338
chdb_initial_prompt,
338339
name="chdb_initial_prompt",

tests/test_tool.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,29 +42,29 @@ def tearDownClass(cls):
4242

4343
def test_list_databases(self):
4444
"""Test listing databases."""
45-
result = list_databases.fn()
45+
result = list_databases()
4646
# Parse JSON response
4747
databases = json.loads(result)
4848
self.assertIn(self.test_db, databases)
4949

5050
def test_list_tables_without_like(self):
5151
"""Test listing tables without a 'LIKE' filter."""
52-
result = list_tables.fn(self.test_db)
52+
result = list_tables(self.test_db)
5353
self.assertIsInstance(result, list)
5454
self.assertEqual(len(result), 1)
5555
self.assertEqual(result[0]["name"], self.test_table)
5656

5757
def test_list_tables_with_like(self):
5858
"""Test listing tables with a 'LIKE' filter."""
59-
result = list_tables.fn(self.test_db, like=f"{self.test_table}%")
59+
result = list_tables(self.test_db, like=f"{self.test_table}%")
6060
self.assertIsInstance(result, list)
6161
self.assertEqual(len(result), 1)
6262
self.assertEqual(result[0]["name"], self.test_table)
6363

6464
def test_run_select_query_success(self):
6565
"""Test running a SELECT query successfully."""
6666
query = f"SELECT * FROM {self.test_db}.{self.test_table}"
67-
result = run_select_query.fn(query)
67+
result = run_select_query(query)
6868
self.assertIsInstance(result, dict)
6969
self.assertEqual(len(result["rows"]), 2)
7070
self.assertEqual(result["rows"][0][0], 1)
@@ -73,14 +73,14 @@ def test_run_select_query_success(self):
7373
def test_run_select_query_failure(self):
7474
"""Test running a SELECT query with an error."""
7575
query = f"SELECT * FROM {self.test_db}.non_existent_table"
76-
result = run_select_query.fn(query)
76+
result = run_select_query(query)
7777
self.assertIsInstance(result, dict)
7878
self.assertEqual(result["status"], "error")
7979
self.assertIn("Query failed", result["message"])
8080

8181
def test_table_and_column_comments(self):
8282
"""Test that table and column comments are correctly retrieved."""
83-
result = list_tables.fn(self.test_db)
83+
result = list_tables(self.test_db)
8484
self.assertIsInstance(result, list)
8585
self.assertEqual(len(result), 1)
8686

0 commit comments

Comments
 (0)