Application layout: Refactor core tool functions to separate module

Author: amotl

cratedb_mcp/__main__.py

@@ -1,73 +1,66 @@
-import httpx
 from fastmcp import FastMCP
+from fastmcp.tools import Tool
 
 from . import __appname__
-from .knowledge import DocumentationIndex, Queries
-from .settings import HTTP_URL, Settings
-from .util.sql import sql_is_permitted
-
-# Load CrateDB documentation outline.
-documentation_index = DocumentationIndex()
+from .tool import (
+    fetch_cratedb_docs,
+    get_cluster_health,
+    get_cratedb_documentation_index,
+    get_table_metadata,
+    query_sql,
+)
 
 # Create FastMCP application object.
 mcp: FastMCP = FastMCP(__appname__)
 
 
 # ------------------------------------------
-#              Text-to-SQL
+#            Health / Status
 # ------------------------------------------
-
-
-def query_cratedb(query: str) -> list[dict]:
-    """Sends a `query` to the set `$CRATEDB_CLUSTER_URL`"""
-    url = HTTP_URL
-    if url.endswith("/"):
-        url = url.removesuffix("/")
-
-    return httpx.post(f"{url}/_sql", json={"stmt": query}, timeout=Settings.http_timeout()).json()
-
-
-@mcp.tool(description="Send an SQL query to CrateDB. Only 'SELECT' queries are allowed.")
-def query_sql(query: str):
-    if not sql_is_permitted(query):
-        raise PermissionError("Only queries that have a SELECT statement are allowed.")
-    return query_cratedb(query)
-
-
-@mcp.tool(description="Return an aggregation of all CrateDB's schema, tables and their metadata.")
-def get_table_metadata() -> list[dict]:
-    """
-    Return an aggregation of schema:tables, e.g.: {'doc': [{name:'mytable', ...}, ...]}
-
-    The tables have metadata datapoints like replicas, shards,
-    name, version, total_shards, total_records.
-    """
-    return query_cratedb(Queries.TABLES_METADATA)
-
-
-@mcp.tool(description="Return the health of the CrateDB cluster.")
-def get_cluster_health() -> list[dict]:
-    """Query sys.health ordered by severity."""
-    return query_cratedb(Queries.HEALTH)
+mcp.add_tool(
+    Tool.from_function(
+        fn=get_cluster_health,
+        description="Return the health of the CrateDB cluster.",
+        tags={"health", "monitoring", "status"},
+    )
+)
 
 
 # ------------------------------------------
-#          Documentation Inquiry
+#              Text-to-SQL
 # ------------------------------------------
-
-
-@mcp.tool(
-    description="Get an index of CrateDB documentation links for fetching. "
-    "Should download docs before answering questions. "
-    "Has documentation title, description, and link."
+mcp.add_tool(
+    Tool.from_function(
+        fn=query_sql,
+        description="Send an SQL query to CrateDB. Only 'SELECT' queries are allowed.",
+        tags={"text-to-sql"},
+    )
+)
+mcp.add_tool(
+    Tool.from_function(
+        fn=get_table_metadata,
+        description="Return an aggregation of all CrateDB's schema, tables and their metadata.",
+        tags={"text-to-sql"},
+    )
 )
-def get_cratedb_documentation_index():
-    return documentation_index.items()
 
 
-@mcp.tool(description="Download individual CrateDB documentation pages by link.")
-def fetch_cratedb_docs(link: str):
-    """Fetch a CrateDB documentation link."""
-    if not documentation_index.url_permitted(link):
-        raise ValueError(f"Link is not permitted: {link}")
-    return documentation_index.client.get(link, timeout=Settings.http_timeout()).text
+# ------------------------------------------
+#          Documentation inquiry
+# ------------------------------------------
+mcp.add_tool(
+    Tool.from_function(
+        fn=get_cratedb_documentation_index,
+        description="Get an index of CrateDB documentation links for fetching. "
+        "Should download docs before answering questions. "
+        "Has documentation title, description, and link.",
+        tags={"documentation"},
+    )
+)
+mcp.add_tool(
+    Tool.from_function(
+        fn=fetch_cratedb_docs,
+        description="Download individual CrateDB documentation pages by link.",
+        tags={"documentation"},
+    )
+)

cratedb_mcp/tool.py

@@ -0,0 +1,60 @@
+import httpx
+
+from cratedb_mcp.knowledge import DocumentationIndex, Queries
+from cratedb_mcp.settings import HTTP_URL, Settings
+from cratedb_mcp.util.sql import sql_is_permitted
+
+
+# ------------------------------------------
+#            Health / Status
+# ------------------------------------------
+def get_cluster_health() -> list[dict]:
+    """Query sys.health ordered by severity."""
+    return query_cratedb(Queries.HEALTH)
+
+
+# ------------------------------------------
+#              Text-to-SQL
+# ------------------------------------------
+def query_cratedb(query: str) -> list[dict]:
+    """Sends a `query` to the set `$CRATEDB_CLUSTER_URL`"""
+    url = HTTP_URL
+    if url.endswith("/"):
+        url = url.removesuffix("/")
+
+    return httpx.post(f"{url}/_sql", json={"stmt": query}, timeout=Settings.http_timeout()).json()
+
+
+def query_sql(query: str):
+    if not sql_is_permitted(query):
+        raise PermissionError("Only queries that have a SELECT statement are allowed.")
+    return query_cratedb(query)
+
+
+def get_table_metadata() -> list[dict]:
+    """
+    Return an aggregation of schema:tables, e.g.: {'doc': [{name:'mytable', ...}, ...]}
+
+    The tables have metadata datapoints like replicas, shards,
+    name, version, total_shards, total_records.
+    """
+    return query_cratedb(Queries.TABLES_METADATA)
+
+
+# ------------------------------------------
+#          Documentation inquiry
+# ------------------------------------------
+# Load CrateDB documentation outline.
+documentation_index = DocumentationIndex()
+
+
+def get_cratedb_documentation_index():
+    """Get curated CrateDB documentation index."""
+    return documentation_index.items()
+
+
+def fetch_cratedb_docs(link: str):
+    """Fetch a CrateDB documentation link."""
+    if not documentation_index.url_permitted(link):
+        raise ValueError(f"Link is not permitted: {link}")
+    return documentation_index.client.get(link, timeout=Settings.http_timeout()).text

pyproject.toml

@@ -74,7 +74,7 @@ dependencies = [
   "cachetools<7",
   "click<9",
   "cratedb-about==0.0.5",
-  "fastmcp<2.7",
+  "fastmcp>=2.7,<2.10",
   "hishel<0.2",
   "pueblo==0.0.11",
   "sqlparse<0.6",

tests/test_mcp.py

@@ -1,6 +1,6 @@
 import pytest
 
-from cratedb_mcp.__main__ import (
+from cratedb_mcp.tool import (
     fetch_cratedb_docs,
     get_cluster_health,
     get_cratedb_documentation_index,
@@ -39,7 +39,7 @@ def test_query_sql_permitted():
 
 def test_query_sql_trailing_slash(mocker):
     """Verify that query_sql works correctly when HTTP_URL has a trailing slash."""
-    mocker.patch("cratedb_mcp.__main__.HTTP_URL", "http://localhost:4200/")
+    mocker.patch("cratedb_mcp.tool.HTTP_URL", "http://localhost:4200/")
     assert query_sql("SELECT 42")["rows"] == [[42]]