diff --git a/cratedb_mcp/__main__.py b/cratedb_mcp/__main__.py index 02de1f6..53432ee 100644 --- a/cratedb_mcp/__main__.py +++ b/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"}, + ) +) diff --git a/cratedb_mcp/tool.py b/cratedb_mcp/tool.py new file mode 100644 index 0000000..33d7aaa --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index e13dfe8..c17d410 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 933b8b2..15859a9 100644 --- a/tests/test_mcp.py +++ b/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]]