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

dbmcp_postgres/
handler.rs
