MCP: Generalize workhorse code. Evaluate available prompts.

Author: amotl

framework/mcp/.gitignore

@@ -0,0 +1 @@
+uv.lock

framework/mcp/README.md

@@ -93,7 +93,12 @@ uv run example_jdbc.py
 
 Run integration tests.
 ```bash
-pytest
+uv run pytest
+```
+
+Run tests selectively.
+```bash
+uv run pytest -k dbhub
 ```
 
 ## Development

framework/mcp/backlog.md

@@ -4,12 +4,14 @@
 - Builtin: Submit a CrateDB adapter, getting rid of the
   `Could not roll back transaction: error`
 - JDBC: Tool `list_tables` currently blocks.
+- JDBC: Tool `describe_table` does not work.
+  Problem with `SELECT current_database()`.
+  https://github.com/crate/crate/issues/17393
 - DBHub: Reading resource `tables` does not work,
   because `WHERE table_schema = 'public'`
 
 ## Iteration +2
 - General: Evaluate all connectors per `stdio` and `sse`, where possible
-- General: Evaluate available prompts, there are no test cases yet
 - General: Improve test cases to not just look at stdout/stderr streams,
   do regular method-based testing instead
 
@@ -27,6 +29,7 @@
 ## Done
 - DBHub: https://github.com/bytebase/dbhub/issues/5
   Resolved: https://github.com/bytebase/dbhub/pull/6
+- General: Evaluate available prompts, there are no test cases yet
 
 
 [SQLite Explorer]: https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#sqlite-explorer

framework/mcp/example_builtin.py

@@ -3,10 +3,11 @@
 #
 # Derived from:
 # https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#writing-mcp-clients
+from cratedb_toolkit.util import DatabaseAdapter
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
-from pprint import pprint
 
+from mcp_utils import McpDatabaseConversation
 
 # Create server parameters for stdio connection.
 server_params = StdioServerParameters(
@@ -28,27 +29,25 @@ async def run():
             # Initialize the connection
             await session.initialize()
 
-            # List available prompts
-            # TODO: Not available on this server.
-            #print("Prompts:")
-            #pprint(await session.list_prompts())
-            #print()
+            client = McpDatabaseConversation(session)
+            await client.inquire()
 
-            # List available resources
-            print("Resources:")
-            pprint(await session.list_resources())
+            print("## MCP server conversations")
             print()
 
-            # List available tools
-            print("Tools:")
-            pprint(await session.list_tools())
-            print()
+            # Call a few tools.
+            await client.call_tool("query", arguments={"sql": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
 
-            print("Calling tool: read_query")
-            result = await session.call_tool("query", arguments={"sql": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
-            pprint(result)
-            print()
+            # Validate database content.
+            db = DatabaseAdapter("crate://crate@localhost:4200/")
+            db.run_sql("CREATE TABLE IF NOT EXISTS public.testdrive (id INT, data TEXT)")
+            db.run_sql("INSERT INTO public.testdrive (id, data) VALUES (42, 'Hotzenplotz')")
+            db.refresh_table("public.testdrive")
 
+            # Read a few resources.
+            # FIXME: Only works on schema=public, because the PostgreSQL adapter hard-codes `WHERE table_schema = 'public'`.
+            # https://github.com/bytebase/dbhub/blob/09424c8513c8c7bef7f66377b46a2b93a69a57d2/src/connectors/postgres/index.ts#L89-L107
+            await client.read_resource("postgres://crate@localhost:5432/testdrive/schema")
 
 
 if __name__ == "__main__":

framework/mcp/example_dbhub.py

@@ -3,12 +3,11 @@
 #
 # Derived from:
 # https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#writing-mcp-clients
-import shlex
-
+from cratedb_toolkit.util import DatabaseAdapter
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
-from pprint import pprint
 
+from mcp_utils import McpDatabaseConversation
 
 # Create server parameters for stdio connection.
 server_params_npx = StdioServerParameters(
@@ -24,22 +23,6 @@
     env=None,
 )
 
-docker_command = """
-docker run --rm --init \
-   --name dbhub \
-   --publish 8080:8080 \
-   bytebase/dbhub:latest \
-   --transport sse \
-   --port 8080 \
-   --dsn="postgres://[email protected]:5432/doc?sslmode=disable"
-"""
-server_params_docker = StdioServerParameters(
-    command="npx",
-    args=shlex.split(docker_command),
-    env=None,
-)
-
-
 async def run():
     async with stdio_client(server_params_npx) as (read, write):
         async with ClientSession(
@@ -48,40 +31,30 @@ async def run():
             # Initialize the connection
             await session.initialize()
 
-            # List available prompts
-            print("Prompts:")
-            pprint(await session.list_prompts())
-            print()
-
-            # List available resources
-            print("Resources:")
-            pprint(await session.list_resources())
-            print()
+            client = McpDatabaseConversation(session)
+            await client.inquire()
 
-            # List available tools
-            print("Tools:")
-            pprint(await session.list_tools())
+            print("## MCP server conversations")
             print()
 
-            print("Calling tool: run_query")
-            result = await session.call_tool("run_query", arguments={"query": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
-            pprint(result)
-            print()
+            # Call a few tools.
+            await client.call_tool("run_query", arguments={"query": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
+            await client.call_tool("list_connectors", arguments={})
 
-            print("Calling tool: list_connectors")
-            result = await session.call_tool("list_connectors", arguments={})
-            pprint(result)
-            print()
+            # Validate database content.
+            db = DatabaseAdapter("crate://crate@localhost:4200/")
+            db.run_sql("CREATE TABLE IF NOT EXISTS public.testdrive (id INT, data TEXT)")
+            db.run_sql("INSERT INTO public.testdrive (id, data) VALUES (42, 'Hotzenplotz')")
+            db.refresh_table("public.testdrive")
 
-            # Read a resource
-            # FIXME: Does not work, because the PostgreSQL adapters hard-codes `WHERE table_schema = 'public'`.
+            # Read a few resources.
+            # FIXME: Only works on schema=public, because the PostgreSQL adapter hard-codes `WHERE table_schema = 'public'`.
             # https://github.com/bytebase/dbhub/blob/09424c8513c8c7bef7f66377b46a2b93a69a57d2/src/connectors/postgres/index.ts#L89-L107
-            """
-            print("Reading resource: tables")
-            content, mime_type = await session.read_resource("db://tables")
-            print("MIME type:", mime_type)
-            pprint(content)
-            """
+            await client.read_resource("db://tables")
+
+            # Get a few prompts.
+            await client.get_prompt("generate_sql", arguments={"description": "Please enumerate the highest five mountains.", "dialect": "postgres"})
+            await client.get_prompt("explain_db", arguments={"target": "testdrive"})
 
 
 if __name__ == "__main__":

framework/mcp/example_jdbc.py

@@ -5,8 +5,8 @@
 # https://github.com/modelcontextprotocol/python-sdk?tab=readme-ov-file#writing-mcp-clients
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
-from pprint import pprint
 
+from mcp_utils import McpDatabaseConversation
 
 # Create server parameters for stdio connection.
 server_params = StdioServerParameters(
@@ -30,61 +30,24 @@ async def run():
             # Initialize the connection
             await session.initialize()
 
-            # List available prompts
-            print("Prompts:")
-            pprint(await session.list_prompts())
-            print()
-
-            # List available resources
-            print("Resources:")
-            pprint(await session.list_resources())
-            print()
-
-            # List available tools
-            print("Tools:")
-            pprint(await session.list_tools())
-            print()
+            client = McpDatabaseConversation(session)
+            await client.inquire()
 
-            # Call a tool
-            print("Calling tool: database_info")
-            result = await session.call_tool("database_info")
-            pprint(result)
+            print("## MCP server conversations")
             print()
 
+            # Call a few tools.
+            await client.call_tool("database_info")
             # FIXME: This operation currently blocks.
-            #print("Calling tool: list_tables")
-            #result = await session.call_tool("list_tables")
-            #pprint(result)
-            #print()
-
-            print("Calling tool: describe_table")
-            result = await session.call_tool("describe_table", arguments={"schema": "sys", "table": "summits"})
-            pprint(result)
-            print()
-
-            print("Calling tool: read_query")
-            result = await session.call_tool("read_query", arguments={"query": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
-            pprint(result)
-            print()
-
-            print("Calling tool: create_table")
-            result = await session.call_tool("create_table", arguments={"query": "CREATE TABLE IF NOT EXISTS testdrive (id INT, data TEXT)"})
-            pprint(result)
-            print()
-
-            print("Calling tool: write_query")
-            result = await session.call_tool("write_query", arguments={"query": "INSERT INTO testdrive (id, data) VALUES (42, 'foobar')"})
-            pprint(result)
-            print()
-
-            # Read a resource
-            #content, mime_type = await session.read_resource("file://some/path")
-
-            # Get a prompt
-            #prompt = await session.get_prompt(
-            #    "example-prompt", arguments={"arg1": "value"}
-            #)
-
+            # await client.call_tool("list_tables", arguments={})
+            await client.call_tool("describe_table", arguments={"schema": "sys", "table": "summits"})
+            await client.call_tool("read_query", arguments={"query": "SELECT * FROM sys.summits ORDER BY height DESC LIMIT 3"})
+            await client.call_tool("create_table", arguments={"query": "CREATE TABLE IF NOT EXISTS testdrive (id INT, data TEXT)"})
+            await client.call_tool("write_query", arguments={"query": "INSERT INTO testdrive (id, data) VALUES (42, 'foobar')"})
+
+            # Get a few prompts.
+            await client.get_prompt("er_diagram")
+            await client.get_prompt("sample_data", arguments={"topic": "wind energy"})
 
 
 if __name__ == "__main__":

framework/mcp/mcp_utils.py

@@ -0,0 +1,84 @@
+import io
+import json
+import mcp.types as types
+from typing import Any
+
+import pydantic_core
+import yaml
+from mcp import ClientSession, McpError
+from pydantic import AnyUrl
+
+
+class McpDatabaseConversation:
+    """
+    Wrap database conversations through MCP servers.
+    """
+    def __init__(self, session: ClientSession):
+        self.session = session
+
+    def decode_items(self, items):
+        return list(map(self.decode_item, json.loads(pydantic_core.to_json(items))))
+
+    @staticmethod
+    def decode_item(item):
+        try:
+            item["text"] = json.loads(item["text"])
+        except Exception:
+            pass
+        return item
+
+    def list_items(self, items):
+        buffer = io.StringIO()
+        if items:
+            data = self.decode_items(items)
+            buffer.write("```yaml\n")
+            buffer.write(yaml.dump(data, sort_keys=False, width=100))
+            buffer.write("```\n")
+        return buffer.getvalue()
+
+    async def inquire(self):
+        print("# MCP server inquiry")
+        print()
+
+        # List available prompts
+        print("## Prompts")
+        try:
+            print(self.list_items((await self.session.list_prompts()).prompts))
+        except McpError as e:
+            print(f"Not implemented on this server: {e}")
+        print()
+
+        # List available resources
+        print("## Resources")
+        print(self.list_items((await self.session.list_resources()).resources))
+        print()
+
+        # List available tools
+        print("## Tools")
+        print(self.list_items((await self.session.list_tools()).tools))
+        print()
+
+    async def call_tool(
+        self, name: str, arguments: dict[str, Any] | None = None
+    ) -> types.CallToolResult:
+        print(f"Calling tool: {name} with arguments: {arguments}")
+        result = await self.session.call_tool(name, arguments)
+        print(self.list_items(result.content))
+        print()
+        return result
+
+    async def get_prompt(
+        self, name: str, arguments: dict[str, str] | None = None
+    ) -> types.GetPromptResult:
+        print(f"Getting prompt: {name} with arguments: {arguments}")
+        result = await self.session.get_prompt(name, arguments)
+        print(self.list_items(result.messages))
+        print()
+        return result
+
+    async def read_resource(self, uri: AnyUrl) -> types.ReadResourceResult:
+        print(f"Reading resource: {uri}")
+        result = await self.session.read_resource(uri)
+        print(self.list_items(result.contents))
+        print()
+        return result

framework/mcp/requirements-test.txt

@@ -1,2 +1 @@
-cratedb-toolkit
 pytest<9

framework/mcp/requirements.txt

@@ -1 +1,2 @@
+cratedb-toolkit
 mcp<1.5