database_mcp_postgres/
handler.rs

1//! `PostgreSQL` handler: composes a [`PostgresConnection`] with the MCP tool router.
2//!
3//! All pool ownership and pool initialization logic lives in the
4//! [`PostgresConnection`]. This module wires the connection into
5//! the MCP `ServerHandler` surface and exposes a small set of thin
6//! delegator methods that the per-tool implementations call.
7
8use database_mcp_config::DatabaseConfig;
9use database_mcp_server::{Server, server_info};
10use rmcp::RoleServer;
11use rmcp::handler::server::router::tool::ToolRouter;
12use rmcp::handler::server::tool::ToolCallContext;
13use rmcp::model::{CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, ServerInfo, Tool};
14use rmcp::service::RequestContext;
15use rmcp::{ErrorData, ServerHandler};
16
17use crate::connection::PostgresConnection;
18
19use crate::tools::{
20    CreateDatabaseTool, DropDatabaseTool, DropTableTool, ExplainQueryTool, GetTableSchemaTool, ListDatabasesTool,
21    ListTablesTool, ReadQueryTool, WriteQueryTool,
22};
23
24/// Backend-specific description for `PostgreSQL`.
25const DESCRIPTION: &str = "Database MCP Server for PostgreSQL";
26
27/// Backend-specific instructions for `PostgreSQL`.
28const INSTRUCTIONS: &str = r"## Workflow
29
301. Call `list_databases` to discover available databases.
312. Call `list_tables` with a `database_name` to see its tables.
323. Call `get_table_schema` with `database_name` and `table_name` to inspect columns, types, and foreign keys before writing queries.
334. Use `read_query` for read-only SQL (SELECT, SHOW, EXPLAIN).
345. Use `write_query` for data changes (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
356. Use `explain_query` to analyze query execution plans and diagnose slow queries.
367. Use `create_database` to create a new database.
378. Use `drop_database` to drop an existing database.
389. Use `drop_table` to remove a table from a database (supports `cascade` for foreign key dependencies).
39
40Tools accept an optional `database_name` parameter to query across databases without reconnecting.
41
42## Constraints
43
44- The `write_query`, `create_database`, `drop_database`, and `drop_table` tools are hidden when read-only mode is active.
45- Multi-statement queries are not supported. Send one statement per request.";
46
47/// `PostgreSQL` database handler.
48///
49/// Composes one [`PostgresConnection`] (which owns every pool
50/// and the pool initialization logic) with the per-backend MCP tool
51/// router.
52#[derive(Clone)]
53pub struct PostgresHandler {
54    pub(crate) config: DatabaseConfig,
55    pub(crate) connection: PostgresConnection,
56    tool_router: ToolRouter<Self>,
57}
58
59impl std::fmt::Debug for PostgresHandler {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        f.debug_struct("PostgresHandler")
62            .field("read_only", &self.config.read_only)
63            .field("connection", &self.connection)
64            .finish_non_exhaustive()
65    }
66}
67
68impl PostgresHandler {
69    /// Creates a new `PostgreSQL` handler.
70    ///
71    /// Constructs the [`PostgresConnection`] (which builds the
72    /// lazy default pool and the per-database pool cache) and the MCP
73    /// tool router. No network I/O happens here.
74    #[must_use]
75    pub fn new(config: &DatabaseConfig) -> Self {
76        Self {
77            config: config.clone(),
78            connection: PostgresConnection::new(config),
79            tool_router: build_tool_router(config.read_only),
80        }
81    }
82}
83
84impl From<PostgresHandler> for Server {
85    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
86    fn from(handler: PostgresHandler) -> Self {
87        Self::new(handler)
88    }
89}
90
91/// Builds the tool router, including write tools only when not in read-only mode.
92fn build_tool_router(read_only: bool) -> ToolRouter<PostgresHandler> {
93    let mut router = ToolRouter::new()
94        .with_async_tool::<ListDatabasesTool>()
95        .with_async_tool::<ListTablesTool>()
96        .with_async_tool::<GetTableSchemaTool>()
97        .with_async_tool::<ReadQueryTool>()
98        .with_async_tool::<ExplainQueryTool>();
99
100    if !read_only {
101        router = router
102            .with_async_tool::<CreateDatabaseTool>()
103            .with_async_tool::<DropDatabaseTool>()
104            .with_async_tool::<DropTableTool>()
105            .with_async_tool::<WriteQueryTool>();
106    }
107    router
108}
109
110impl ServerHandler for PostgresHandler {
111    fn get_info(&self) -> ServerInfo {
112        let mut info = server_info();
113        info.server_info.description = Some(DESCRIPTION.into());
114        info.instructions = Some(INSTRUCTIONS.into());
115        info
116    }
117
118    async fn call_tool(
119        &self,
120        request: CallToolRequestParams,
121        context: RequestContext<RoleServer>,
122    ) -> Result<CallToolResult, ErrorData> {
123        let tcc = ToolCallContext::new(self, request, context);
124        self.tool_router.call(tcc).await
125    }
126
127    async fn list_tools(
128        &self,
129        _request: Option<PaginatedRequestParams>,
130        _context: RequestContext<RoleServer>,
131    ) -> Result<ListToolsResult, ErrorData> {
132        Ok(ListToolsResult {
133            tools: self.tool_router.list_all(),
134            next_cursor: None,
135            meta: None,
136        })
137    }
138
139    fn get_tool(&self, name: &str) -> Option<Tool> {
140        self.tool_router.get(name).cloned()
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use database_mcp_config::DatabaseBackend;
148
149    fn base_config() -> DatabaseConfig {
150        DatabaseConfig {
151            backend: DatabaseBackend::Postgres,
152            host: "pg.example.com".into(),
153            port: 5433,
154            user: "pgadmin".into(),
155            password: Some("pgpass".into()),
156            name: Some("mydb".into()),
157            ..DatabaseConfig::default()
158        }
159    }
160
161    fn handler(read_only: bool) -> PostgresHandler {
162        PostgresHandler::new(&DatabaseConfig {
163            read_only,
164            ..base_config()
165        })
166    }
167
168    #[tokio::test]
169    async fn handler_exposes_connection_default_db() {
170        let handler = PostgresHandler::new(&base_config());
171        assert_eq!(handler.connection.default_database_name(), "mydb");
172    }
173
174    #[tokio::test]
175    async fn router_exposes_all_nine_tools_in_read_write_mode() {
176        let router = handler(false).tool_router;
177        for name in [
178            "list_databases",
179            "list_tables",
180            "get_table_schema",
181            "read_query",
182            "explain_query",
183            "create_database",
184            "drop_database",
185            "drop_table",
186            "write_query",
187        ] {
188            assert!(router.has_route(name), "missing tool: {name}");
189        }
190    }
191
192    #[tokio::test]
193    async fn router_hides_write_tools_in_read_only_mode() {
194        let router = handler(true).tool_router;
195        assert!(router.has_route("list_databases"));
196        assert!(router.has_route("list_tables"));
197        assert!(router.has_route("get_table_schema"));
198        assert!(router.has_route("read_query"));
199        assert!(router.has_route("explain_query"));
200        assert!(!router.has_route("write_query"));
201        assert!(!router.has_route("create_database"));
202        assert!(!router.has_route("drop_database"));
203        assert!(!router.has_route("drop_table"));
204    }
205}
database_mcp_postgres/handler.rs

database_mcp_postgres/
handler.rs
