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

dbmcp_postgres/
handler.rs
