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    #[must_use]
107    pub fn new(config: &Config) -> Self {
108        Self {
109            config: config.database.clone(),
110            connection: PostgresConnection::new(&config.database),
111            redactor: Redactor::from_config(&config.pii),
112            tool_router: ToolRouter::from_specs(TOOLS, config.database.read_only, config.database.name.is_some()),
113        }
114    }
115}
116
117impl From<PostgresHandler> for Server {
118    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
119    fn from(handler: PostgresHandler) -> Self {
120        Self::new(handler)
121    }
122}
123
124impl ServerHandler for PostgresHandler {
125    fn get_info(&self) -> ServerInfo {
126        let mut info = server_info();
127        info.server_info.description = Some(DESCRIPTION.into());
128        info.instructions = Some(
129            match (self.config.read_only, self.config.name.is_some()) {
130                (false, false) => INSTRUCTIONS,
131                (true, false) => INSTRUCTIONS_READ_ONLY,
132                (false, true) => INSTRUCTIONS_PINNED,
133                (true, true) => INSTRUCTIONS_READ_ONLY_PINNED,
134            }
135            .into(),
136        );
137        info
138    }
139
140    async fn call_tool(
141        &self,
142        request: CallToolRequestParams,
143        context: RequestContext<RoleServer>,
144    ) -> Result<CallToolResult, ErrorData> {
145        let tcc = ToolCallContext::new(self, request, context);
146        self.tool_router.call(tcc).await
147    }
148
149    async fn list_tools(
150        &self,
151        _request: Option<PaginatedRequestParams>,
152        _context: RequestContext<RoleServer>,
153    ) -> Result<ListToolsResult, ErrorData> {
154        Ok(ListToolsResult {
155            tools: self.tool_router.list_all(),
156            next_cursor: None,
157            meta: None,
158        })
159    }
160
161    fn get_tool(&self, name: &str) -> Option<Tool> {
162        self.tool_router.get(name).cloned()
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use dbmcp_config::DatabaseBackend;
170
171    fn base_config() -> DatabaseConfig {
172        DatabaseConfig {
173            backend: DatabaseBackend::Postgres,
174            host: "pg.example.com".into(),
175            port: 5433,
176            user: "pgadmin".into(),
177            password: Some("pgpass".into()),
178            name: None,
179            ..DatabaseConfig::default()
180        }
181    }
182
183    fn handler(read_only: bool) -> PostgresHandler {
184        PostgresHandler::new(&Config {
185            database: DatabaseConfig {
186                read_only,
187                ..base_config()
188            },
189            http: None,
190            pii: dbmcp_config::PiiConfig::default(),
191        })
192    }
193
194    /// Handler whose config pins a specific database name.
195    fn pinned_handler(read_only: bool) -> PostgresHandler {
196        PostgresHandler::new(&Config {
197            database: DatabaseConfig {
198                read_only,
199                name: Some("mydb".into()),
200                ..base_config()
201            },
202            http: None,
203            pii: dbmcp_config::PiiConfig::default(),
204        })
205    }
206
207    #[tokio::test]
208    async fn handler_exposes_connection_default_db() {
209        let handler = PostgresHandler::new(&Config {
210            database: DatabaseConfig {
211                name: Some("mydb".into()),
212                ..base_config()
213            },
214            http: None,
215            pii: dbmcp_config::PiiConfig::default(),
216        });
217        assert_eq!(handler.connection.default_database_name(), "mydb");
218    }
219
220    #[tokio::test]
221    async fn router_exposes_all_thirteen_tools_in_read_write_mode() {
222        let router = handler(false).tool_router;
223        for name in [
224            "listDatabases",
225            "listTables",
226            "listViews",
227            "listTriggers",
228            "listFunctions",
229            "listProcedures",
230            "listMaterializedViews",
231            "readQuery",
232            "explainQuery",
233            "createDatabase",
234            "dropDatabase",
235            "dropTable",
236            "writeQuery",
237        ] {
238            assert!(router.has_route(name), "missing tool: {name}");
239        }
240    }
241
242    #[tokio::test]
243    async fn router_hides_write_tools_in_read_only_mode() {
244        let router = handler(true).tool_router;
245        assert!(router.has_route("listDatabases"));
246        assert!(router.has_route("listTables"));
247        assert!(router.has_route("listViews"));
248        assert!(router.has_route("listTriggers"));
249        assert!(router.has_route("listFunctions"));
250        assert!(router.has_route("listProcedures"));
251        assert!(router.has_route("listMaterializedViews"));
252        assert!(router.has_route("readQuery"));
253        assert!(router.has_route("explainQuery"));
254        assert!(!router.has_route("writeQuery"));
255        assert!(!router.has_route("createDatabase"));
256        assert!(!router.has_route("dropDatabase"));
257        assert!(!router.has_route("dropTable"));
258    }
259
260    #[tokio::test]
261    async fn instructions_match_read_only_mode() {
262        let read_write = handler(false).get_info().instructions.expect("instructions present");
263        assert!(
264            read_write.contains("writeQuery"),
265            "read-write instructions mention writeQuery"
266        );
267
268        let read_only = handler(true).get_info().instructions.expect("instructions present");
269        for tool in ["writeQuery", "createDatabase", "dropDatabase", "dropTable"] {
270            assert!(
271                !read_only.contains(tool),
272                "read-only instructions must not mention {tool}"
273            );
274        }
275    }
276
277    #[tokio::test]
278    async fn router_does_not_expose_get_table_schema() {
279        // Postgres drops `getTableSchema` in favour of `listTables({detailed: true})`.
280        // MySQL and SQLite still expose it; their routers are tested separately.
281        let rw = handler(false).tool_router;
282        let ro = handler(true).tool_router;
283        assert!(
284            !rw.has_route("getTableSchema"),
285            "read-write router must not expose getTableSchema"
286        );
287        assert!(
288            !ro.has_route("getTableSchema"),
289            "read-only router must not expose getTableSchema"
290        );
291    }
292
293    #[tokio::test]
294    async fn router_hides_cross_database_tools_when_name_pinned() {
295        let router = pinned_handler(false).tool_router;
296        for present in ["listTables", "readQuery", "explainQuery", "dropTable", "writeQuery"] {
297            assert!(router.has_route(present), "missing tool: {present}");
298        }
299        for absent in ["listDatabases", "createDatabase", "dropDatabase"] {
300            assert!(!router.has_route(absent), "pinned router must not expose {absent}");
301        }
302    }
303
304    #[tokio::test]
305    async fn router_hides_list_databases_when_name_pinned_read_only() {
306        let router = pinned_handler(true).tool_router;
307        assert!(!router.has_route("listDatabases"));
308        assert!(!router.has_route("createDatabase"));
309        assert!(!router.has_route("dropDatabase"));
310        assert!(router.has_route("listTables"));
311        assert!(router.has_route("readQuery"));
312    }
313
314    #[tokio::test]
315    async fn instructions_match_pinned_mode() {
316        for read_only in [false, true] {
317            let instructions = pinned_handler(read_only)
318                .get_info()
319                .instructions
320                .expect("instructions present");
321            for tool in ["listDatabases", "createDatabase", "dropDatabase"] {
322                assert!(
323                    !instructions.contains(tool),
324                    "pinned instructions must not mention {tool} (read_only={read_only})"
325                );
326            }
327        }
328    }
329}
dbmcp_postgres/handler.rs

dbmcp_postgres/
handler.rs
