dbmcp_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 dbmcp_config::DatabaseConfig;
9use dbmcp_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, ListDatabasesTool, ListFunctionsTool,
21    ListMaterializedViewsTool, ListProceduresTool, ListTablesTool, ListTriggersTool, ListViewsTool, ReadQueryTool,
22    WriteQueryTool,
23};
24
25/// Backend-specific description for `PostgreSQL`.
26const DESCRIPTION: &str = "Database MCP Server for PostgreSQL";
27
28/// Backend-specific instructions for `PostgreSQL`.
29const INSTRUCTIONS: &str = r"## Workflow
30
311. Call `listDatabases` to discover available databases.
322. Call `listTables` to see tables. Pass `search` to filter by name (case-insensitive substring). Pass `detailed: true` to get columns, constraints, indexes, and triggers in the same call — this supersedes the legacy `getTableSchema` workflow.
333. Call `listViews` to see views in the `public` schema.
344. Call `listTriggers` to see user-defined triggers in the `public` schema.
355. Call `listFunctions` to see user-defined functions in the `public` schema.
366. Call `listProcedures` to see user-defined procedures in the `public` schema.
377. Call `listMaterializedViews` to see materialized views in the `public` schema.
388. Use `readQuery` for read-only SQL (SELECT, SHOW, EXPLAIN).
399. Use `writeQuery` for data changes (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
4010. Use `explainQuery` to analyze query execution plans and diagnose slow queries.
4111. Use `createDatabase` to create a new database.
4212. Use `dropDatabase` to drop an existing database.
4313. Use `dropTable` to remove a table from a database (supports `cascade` for foreign key dependencies).
44
45Per-database tools default to the active database; pass `database` to target another.
46
47## Constraints
48
49- The `writeQuery`, `createDatabase`, `dropDatabase`, and `dropTable` tools are hidden when read-only mode is active.
50- Multi-statement queries are not supported. Send one statement per request.";
51
52/// `PostgreSQL` database handler.
53///
54/// Composes one [`PostgresConnection`] (which owns every pool
55/// and the pool initialization logic) with the per-backend MCP tool
56/// router.
57#[derive(Clone)]
58pub struct PostgresHandler {
59    pub(crate) config: DatabaseConfig,
60    pub(crate) connection: PostgresConnection,
61    tool_router: ToolRouter<Self>,
62}
63
64impl std::fmt::Debug for PostgresHandler {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        f.debug_struct("PostgresHandler")
67            .field("read_only", &self.config.read_only)
68            .field("connection", &self.connection)
69            .finish_non_exhaustive()
70    }
71}
72
73impl PostgresHandler {
74    /// Creates a new `PostgreSQL` handler.
75    ///
76    /// Constructs the [`PostgresConnection`] (which builds the
77    /// lazy default pool and the per-database pool cache) and the MCP
78    /// tool router. No network I/O happens here.
79    #[must_use]
80    pub fn new(config: &DatabaseConfig) -> Self {
81        Self {
82            config: config.clone(),
83            connection: PostgresConnection::new(config),
84            tool_router: build_tool_router(config.read_only),
85        }
86    }
87}
88
89impl From<PostgresHandler> for Server {
90    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
91    fn from(handler: PostgresHandler) -> Self {
92        Self::new(handler)
93    }
94}
95
96/// Builds the tool router, including write tools only when not in read-only mode.
97fn build_tool_router(read_only: bool) -> ToolRouter<PostgresHandler> {
98    let mut router = ToolRouter::new()
99        .with_async_tool::<ListDatabasesTool>()
100        .with_async_tool::<ListTablesTool>()
101        .with_async_tool::<ListViewsTool>()
102        .with_async_tool::<ListTriggersTool>()
103        .with_async_tool::<ListFunctionsTool>()
104        .with_async_tool::<ListProceduresTool>()
105        .with_async_tool::<ListMaterializedViewsTool>()
106        .with_async_tool::<ReadQueryTool>()
107        .with_async_tool::<ExplainQueryTool>();
108
109    if !read_only {
110        router = router
111            .with_async_tool::<CreateDatabaseTool>()
112            .with_async_tool::<DropDatabaseTool>()
113            .with_async_tool::<DropTableTool>()
114            .with_async_tool::<WriteQueryTool>();
115    }
116    router
117}
118
119impl ServerHandler for PostgresHandler {
120    fn get_info(&self) -> ServerInfo {
121        let mut info = server_info();
122        info.server_info.description = Some(DESCRIPTION.into());
123        info.instructions = Some(INSTRUCTIONS.into());
124        info
125    }
126
127    async fn call_tool(
128        &self,
129        request: CallToolRequestParams,
130        context: RequestContext<RoleServer>,
131    ) -> Result<CallToolResult, ErrorData> {
132        let tcc = ToolCallContext::new(self, request, context);
133        self.tool_router.call(tcc).await
134    }
135
136    async fn list_tools(
137        &self,
138        _request: Option<PaginatedRequestParams>,
139        _context: RequestContext<RoleServer>,
140    ) -> Result<ListToolsResult, ErrorData> {
141        Ok(ListToolsResult {
142            tools: self.tool_router.list_all(),
143            next_cursor: None,
144            meta: None,
145        })
146    }
147
148    fn get_tool(&self, name: &str) -> Option<Tool> {
149        self.tool_router.get(name).cloned()
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use dbmcp_config::DatabaseBackend;
157
158    fn base_config() -> DatabaseConfig {
159        DatabaseConfig {
160            backend: DatabaseBackend::Postgres,
161            host: "pg.example.com".into(),
162            port: 5433,
163            user: "pgadmin".into(),
164            password: Some("pgpass".into()),
165            name: Some("mydb".into()),
166            ..DatabaseConfig::default()
167        }
168    }
169
170    fn handler(read_only: bool) -> PostgresHandler {
171        PostgresHandler::new(&DatabaseConfig {
172            read_only,
173            ..base_config()
174        })
175    }
176
177    #[tokio::test]
178    async fn handler_exposes_connection_default_db() {
179        let handler = PostgresHandler::new(&base_config());
180        assert_eq!(handler.connection.default_database_name(), "mydb");
181    }
182
183    #[tokio::test]
184    async fn router_exposes_all_thirteen_tools_in_read_write_mode() {
185        let router = handler(false).tool_router;
186        for name in [
187            "listDatabases",
188            "listTables",
189            "listViews",
190            "listTriggers",
191            "listFunctions",
192            "listProcedures",
193            "listMaterializedViews",
194            "readQuery",
195            "explainQuery",
196            "createDatabase",
197            "dropDatabase",
198            "dropTable",
199            "writeQuery",
200        ] {
201            assert!(router.has_route(name), "missing tool: {name}");
202        }
203    }
204
205    #[tokio::test]
206    async fn router_hides_write_tools_in_read_only_mode() {
207        let router = handler(true).tool_router;
208        assert!(router.has_route("listDatabases"));
209        assert!(router.has_route("listTables"));
210        assert!(router.has_route("listViews"));
211        assert!(router.has_route("listTriggers"));
212        assert!(router.has_route("listFunctions"));
213        assert!(router.has_route("listProcedures"));
214        assert!(router.has_route("listMaterializedViews"));
215        assert!(router.has_route("readQuery"));
216        assert!(router.has_route("explainQuery"));
217        assert!(!router.has_route("writeQuery"));
218        assert!(!router.has_route("createDatabase"));
219        assert!(!router.has_route("dropDatabase"));
220        assert!(!router.has_route("dropTable"));
221    }
222
223    #[tokio::test]
224    async fn router_does_not_expose_get_table_schema() {
225        // Postgres drops `getTableSchema` in favour of `listTables({detailed: true})`.
226        // MySQL and SQLite still expose it; their routers are tested separately.
227        let rw = handler(false).tool_router;
228        let ro = handler(true).tool_router;
229        assert!(
230            !rw.has_route("getTableSchema"),
231            "read-write router must not expose getTableSchema"
232        );
233        assert!(
234            !ro.has_route("getTableSchema"),
235            "read-only router must not expose getTableSchema"
236        );
237    }
238}
dbmcp_postgres/handler.rs

dbmcp_postgres/
handler.rs
