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

dbmcp_postgres/
handler.rs
