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::Redactor;
10use dbmcp_server::{Server, Shutdown, ShutdownFuture, ToolRouterExt, ToolSpec, server_info};
11use rmcp::RoleServer;
12use rmcp::handler::server::router::tool::ToolRouter;
13use rmcp::handler::server::tool::ToolCallContext;
14use rmcp::model::{
15    CallToolRequestParams, CallToolResponse, ListToolsResult, PaginatedRequestParams, ServerConfig, Tool,
16};
17use rmcp::service::RequestContext;
18use rmcp::{ErrorData, ServerHandler};
19
20use crate::connection::PostgresConnection;
21
22use crate::tools::{
23    CreateDatabaseTool, DropDatabaseTool, ListDatabasesTool, PinnedDropTableTool, PinnedExplainQueryTool,
24    PinnedListFunctionsTool, PinnedListMaterializedViewsTool, PinnedListProceduresTool, PinnedListTablesTool,
25    PinnedListTriggersTool, PinnedListViewsTool, PinnedReadQueryTool, PinnedWriteQueryTool, UnpinnedDropTableTool,
26    UnpinnedExplainQueryTool, UnpinnedListFunctionsTool, UnpinnedListMaterializedViewsTool, UnpinnedListProceduresTool,
27    UnpinnedListTablesTool, UnpinnedListTriggersTool, UnpinnedListViewsTool, UnpinnedReadQueryTool,
28    UnpinnedWriteQueryTool,
29};
30
31/// Backend-specific description for `PostgreSQL`.
32const DESCRIPTION: &str = "Database MCP Server for PostgreSQL";
33
34/// Backend-specific instructions for `PostgreSQL` in read-write mode.
35const INSTRUCTIONS: &str = include_str!("../assets/instructions/default.md");
36
37/// Backend-specific instructions for `PostgreSQL` in read-only mode.
38const INSTRUCTIONS_READ_ONLY: &str = include_str!("../assets/instructions/read-only.md");
39
40/// Backend-specific instructions when a database name is pinned.
41const INSTRUCTIONS_PINNED: &str = include_str!("../assets/instructions/default.pinned.md");
42
43/// Backend-specific instructions for read-only mode with a pinned database.
44const INSTRUCTIONS_READ_ONLY_PINNED: &str = include_str!("../assets/instructions/read-only.pinned.md");
45
46/// Declarative tool table: `(tool, pinned, read_only)`.
47///
48/// Per-database tools expose a `Pinned*Tool` variant (no `database` field)
49/// when the config pins a db name, and an `Unpinned*Tool` variant
50/// (carries a `database` field) otherwise. Cross-database tools
51/// (`listDatabases`, `createDatabase`, `dropDatabase`) are hidden in pinned
52/// mode altogether.
53const TOOLS: &[ToolSpec<PostgresHandler>] = &[
54    ToolSpec::async_tool::<ListDatabasesTool>(false, false),
55    ToolSpec::async_tool::<PinnedListTablesTool>(true, false),
56    ToolSpec::async_tool::<UnpinnedListTablesTool>(false, false),
57    ToolSpec::async_tool::<PinnedListViewsTool>(true, false),
58    ToolSpec::async_tool::<UnpinnedListViewsTool>(false, false),
59    ToolSpec::async_tool::<PinnedListTriggersTool>(true, false),
60    ToolSpec::async_tool::<UnpinnedListTriggersTool>(false, false),
61    ToolSpec::async_tool::<PinnedListFunctionsTool>(true, false),
62    ToolSpec::async_tool::<UnpinnedListFunctionsTool>(false, false),
63    ToolSpec::async_tool::<PinnedListProceduresTool>(true, false),
64    ToolSpec::async_tool::<UnpinnedListProceduresTool>(false, false),
65    ToolSpec::async_tool::<PinnedListMaterializedViewsTool>(true, false),
66    ToolSpec::async_tool::<UnpinnedListMaterializedViewsTool>(false, false),
67    ToolSpec::async_tool::<PinnedReadQueryTool>(true, false),
68    ToolSpec::async_tool::<UnpinnedReadQueryTool>(false, false),
69    ToolSpec::async_tool::<PinnedExplainQueryTool>(true, false),
70    ToolSpec::async_tool::<UnpinnedExplainQueryTool>(false, false),
71    ToolSpec::async_tool::<CreateDatabaseTool>(false, true),
72    ToolSpec::async_tool::<DropDatabaseTool>(false, true),
73    ToolSpec::async_tool::<PinnedDropTableTool>(true, true),
74    ToolSpec::async_tool::<UnpinnedDropTableTool>(false, true),
75    ToolSpec::async_tool::<PinnedWriteQueryTool>(true, true),
76    ToolSpec::async_tool::<UnpinnedWriteQueryTool>(false, true),
77];
78
79/// `PostgreSQL` database handler.
80///
81/// Composes one [`PostgresConnection`] (which owns every pool
82/// and the pool initialization logic) with the per-backend MCP tool
83/// router.
84#[derive(Clone)]
85pub struct PostgresHandler {
86    pub(crate) config: DatabaseConfig,
87    pub(crate) connection: PostgresConnection,
88    pub(crate) redactor: Option<Redactor>,
89    tool_router: ToolRouter<Self>,
90}
91
92impl std::fmt::Debug for PostgresHandler {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        f.debug_struct("PostgresHandler")
95            .field("read_only", &self.config.read_only)
96            .field("redact_pii", &self.redactor.is_some())
97            .field("connection", &self.connection)
98            .finish_non_exhaustive()
99    }
100}
101
102impl PostgresHandler {
103    /// Creates a new `PostgreSQL` handler.
104    ///
105    /// Constructs the [`PostgresConnection`] (which builds the
106    /// lazy default pool and the per-database pool cache) and the MCP
107    /// tool router. No network I/O happens here.
108    /// # Errors
109    ///
110    /// Returns [`dbmcp_pii::RedactorInitError`] when PII redaction is enabled
111    /// with a NER model that fails to load (fail-closed startup).
112    pub fn new(config: &Config) -> Result<Self, dbmcp_pii::RedactorInitError> {
113        Ok(Self {
114            config: config.database.clone(),
115            connection: PostgresConnection::new(&config.database),
116            redactor: Redactor::from_config(&config.pii)?,
117            tool_router: ToolRouter::from_specs(TOOLS, config.database.read_only, config.database.name.is_some()),
118        })
119    }
120}
121
122impl From<PostgresHandler> for Server {
123    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
124    fn from(handler: PostgresHandler) -> Self {
125        Self::new(handler)
126    }
127}
128
129impl Shutdown for PostgresHandler {
130    fn shutdown(&self) -> ShutdownFuture<'_> {
131        Box::pin(async move { self.connection.shutdown().await })
132    }
133}
134
135impl ServerHandler for PostgresHandler {
136    fn get_info(&self) -> ServerConfig {
137        let mut info = server_info();
138        info.server_info.description = Some(DESCRIPTION.into());
139        info.instructions = Some(
140            match (self.config.read_only, self.config.name.is_some()) {
141                (false, false) => INSTRUCTIONS,
142                (true, false) => INSTRUCTIONS_READ_ONLY,
143                (false, true) => INSTRUCTIONS_PINNED,
144                (true, true) => INSTRUCTIONS_READ_ONLY_PINNED,
145            }
146            .into(),
147        );
148        info
149    }
150
151    async fn call_tool(
152        &self,
153        request: CallToolRequestParams,
154        context: RequestContext<RoleServer>,
155    ) -> Result<CallToolResponse, ErrorData> {
156        let tcc = ToolCallContext::new(self, request, context);
157        self.tool_router.call(tcc).await
158    }
159
160    async fn list_tools(
161        &self,
162        _request: Option<PaginatedRequestParams>,
163        _context: RequestContext<RoleServer>,
164    ) -> Result<ListToolsResult, ErrorData> {
165        Ok(ListToolsResult::with_all_items(self.tool_router.list_all()))
166    }
167
168    fn get_tool(&self, name: &str) -> Option<Tool> {
169        self.tool_router.get(name).cloned()
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use dbmcp_config::DatabaseBackend;
177
178    fn base_config() -> DatabaseConfig {
179        DatabaseConfig {
180            backend: DatabaseBackend::Postgres,
181            host: "pg.example.com".into(),
182            port: 5433,
183            user: "pgadmin".into(),
184            password: Some("pgpass".into()),
185            name: None,
186            ..DatabaseConfig::default()
187        }
188    }
189
190    fn handler(read_only: bool) -> PostgresHandler {
191        PostgresHandler::new(&Config {
192            database: DatabaseConfig {
193                read_only,
194                ..base_config()
195            },
196            http: None,
197            pii: dbmcp_config::PiiConfig::default(),
198        })
199        .expect("handler builds in test")
200    }
201
202    /// Handler whose config pins a specific database name.
203    fn pinned_handler(read_only: bool) -> PostgresHandler {
204        PostgresHandler::new(&Config {
205            database: DatabaseConfig {
206                read_only,
207                name: Some("mydb".into()),
208                ..base_config()
209            },
210            http: None,
211            pii: dbmcp_config::PiiConfig::default(),
212        })
213        .expect("handler builds in test")
214    }
215
216    #[tokio::test]
217    async fn handler_exposes_connection_default_db() {
218        let handler = PostgresHandler::new(&Config {
219            database: DatabaseConfig {
220                name: Some("mydb".into()),
221                ..base_config()
222            },
223            http: None,
224            pii: dbmcp_config::PiiConfig::default(),
225        })
226        .expect("handler builds in test");
227        assert_eq!(handler.connection.default_database_name(), "mydb");
228    }
229
230    #[tokio::test]
231    async fn router_exposes_all_thirteen_tools_in_read_write_mode() {
232        let router = handler(false).tool_router;
233        for name in [
234            "listDatabases",
235            "listTables",
236            "listViews",
237            "listTriggers",
238            "listFunctions",
239            "listProcedures",
240            "listMaterializedViews",
241            "readQuery",
242            "explainQuery",
243            "createDatabase",
244            "dropDatabase",
245            "dropTable",
246            "writeQuery",
247        ] {
248            assert!(router.has_route(name), "missing tool: {name}");
249        }
250    }
251
252    #[tokio::test]
253    async fn router_hides_write_tools_in_read_only_mode() {
254        let router = handler(true).tool_router;
255        assert!(router.has_route("listDatabases"));
256        assert!(router.has_route("listTables"));
257        assert!(router.has_route("listViews"));
258        assert!(router.has_route("listTriggers"));
259        assert!(router.has_route("listFunctions"));
260        assert!(router.has_route("listProcedures"));
261        assert!(router.has_route("listMaterializedViews"));
262        assert!(router.has_route("readQuery"));
263        assert!(router.has_route("explainQuery"));
264        assert!(!router.has_route("writeQuery"));
265        assert!(!router.has_route("createDatabase"));
266        assert!(!router.has_route("dropDatabase"));
267        assert!(!router.has_route("dropTable"));
268    }
269
270    #[tokio::test]
271    async fn instructions_match_read_only_mode() {
272        let read_write = handler(false).get_info().instructions.expect("instructions present");
273        assert!(
274            read_write.contains("writeQuery"),
275            "read-write instructions mention writeQuery"
276        );
277
278        let read_only = handler(true).get_info().instructions.expect("instructions present");
279        for tool in ["writeQuery", "createDatabase", "dropDatabase", "dropTable"] {
280            assert!(
281                !read_only.contains(tool),
282                "read-only instructions must not mention {tool}"
283            );
284        }
285    }
286
287    #[tokio::test]
288    async fn router_does_not_expose_get_table_schema() {
289        // Postgres drops `getTableSchema` in favour of `listTables({detailed: true})`.
290        // MySQL and SQLite still expose it; their routers are tested separately.
291        let rw = handler(false).tool_router;
292        let ro = handler(true).tool_router;
293        assert!(
294            !rw.has_route("getTableSchema"),
295            "read-write router must not expose getTableSchema"
296        );
297        assert!(
298            !ro.has_route("getTableSchema"),
299            "read-only router must not expose getTableSchema"
300        );
301    }
302
303    #[tokio::test]
304    async fn router_hides_cross_database_tools_when_name_pinned() {
305        let router = pinned_handler(false).tool_router;
306        for present in ["listTables", "readQuery", "explainQuery", "dropTable", "writeQuery"] {
307            assert!(router.has_route(present), "missing tool: {present}");
308        }
309        for absent in ["listDatabases", "createDatabase", "dropDatabase"] {
310            assert!(!router.has_route(absent), "pinned router must not expose {absent}");
311        }
312    }
313
314    #[tokio::test]
315    async fn router_hides_list_databases_when_name_pinned_read_only() {
316        let router = pinned_handler(true).tool_router;
317        assert!(!router.has_route("listDatabases"));
318        assert!(!router.has_route("createDatabase"));
319        assert!(!router.has_route("dropDatabase"));
320        assert!(router.has_route("listTables"));
321        assert!(router.has_route("readQuery"));
322    }
323
324    #[tokio::test]
325    async fn instructions_match_pinned_mode() {
326        for read_only in [false, true] {
327            let instructions = pinned_handler(read_only)
328                .get_info()
329                .instructions
330                .expect("instructions present");
331            for tool in ["listDatabases", "createDatabase", "dropDatabase"] {
332                assert!(
333                    !instructions.contains(tool),
334                    "pinned instructions must not mention {tool} (read_only={read_only})"
335                );
336            }
337        }
338    }
339}
dbmcp_postgres/handler.rs

dbmcp_postgres/
handler.rs
