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, 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` in read-write mode.
30const INSTRUCTIONS: &str = include_str!("../assets/instructions.md");
31
32/// Backend-specific instructions for `PostgreSQL` in read-only mode.
33const INSTRUCTIONS_READ_ONLY: &str = include_str!("../assets/instructions.readonly.md");
34
35/// `PostgreSQL` database handler.
36///
37/// Composes one [`PostgresConnection`] (which owns every pool
38/// and the pool initialization logic) with the per-backend MCP tool
39/// router.
40#[derive(Clone)]
41pub struct PostgresHandler {
42    pub(crate) config: DatabaseConfig,
43    pub(crate) connection: PostgresConnection,
44    pub(crate) redactor: Option<Redactor>,
45    tool_router: ToolRouter<Self>,
46}
47
48impl std::fmt::Debug for PostgresHandler {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("PostgresHandler")
51            .field("read_only", &self.config.read_only)
52            .field("redact_pii", &self.redactor.is_some())
53            .field("connection", &self.connection)
54            .finish_non_exhaustive()
55    }
56}
57
58impl PostgresHandler {
59    /// Creates a new `PostgreSQL` handler.
60    ///
61    /// Constructs the [`PostgresConnection`] (which builds the
62    /// lazy default pool and the per-database pool cache) and the MCP
63    /// tool router. No network I/O happens here.
64    #[must_use]
65    pub fn new(config: &Config) -> Self {
66        Self {
67            config: config.database.clone(),
68            connection: PostgresConnection::new(&config.database),
69            redactor: Redactor::from_config(&config.pii),
70            tool_router: build_tool_router(config.database.read_only),
71        }
72    }
73}
74
75impl From<PostgresHandler> for Server {
76    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
77    fn from(handler: PostgresHandler) -> Self {
78        Self::new(handler)
79    }
80}
81
82/// Builds the tool router, including write tools only when not in read-only mode.
83fn build_tool_router(read_only: bool) -> ToolRouter<PostgresHandler> {
84    let mut router = ToolRouter::new()
85        .with_async_tool::<ListDatabasesTool>()
86        .with_async_tool::<ListTablesTool>()
87        .with_async_tool::<ListViewsTool>()
88        .with_async_tool::<ListTriggersTool>()
89        .with_async_tool::<ListFunctionsTool>()
90        .with_async_tool::<ListProceduresTool>()
91        .with_async_tool::<ListMaterializedViewsTool>()
92        .with_async_tool::<ReadQueryTool>()
93        .with_async_tool::<ExplainQueryTool>();
94
95    if !read_only {
96        router = router
97            .with_async_tool::<CreateDatabaseTool>()
98            .with_async_tool::<DropDatabaseTool>()
99            .with_async_tool::<DropTableTool>()
100            .with_async_tool::<WriteQueryTool>();
101    }
102    router
103}
104
105impl ServerHandler for PostgresHandler {
106    fn get_info(&self) -> ServerInfo {
107        let mut info = server_info();
108        info.server_info.description = Some(DESCRIPTION.into());
109        info.instructions = Some(if self.config.read_only {
110            INSTRUCTIONS_READ_ONLY.into()
111        } else {
112            INSTRUCTIONS.into()
113        });
114        info
115    }
116
117    async fn call_tool(
118        &self,
119        request: CallToolRequestParams,
120        context: RequestContext<RoleServer>,
121    ) -> Result<CallToolResult, ErrorData> {
122        let tcc = ToolCallContext::new(self, request, context);
123        self.tool_router.call(tcc).await
124    }
125
126    async fn list_tools(
127        &self,
128        _request: Option<PaginatedRequestParams>,
129        _context: RequestContext<RoleServer>,
130    ) -> Result<ListToolsResult, ErrorData> {
131        Ok(ListToolsResult {
132            tools: self.tool_router.list_all(),
133            next_cursor: None,
134            meta: None,
135        })
136    }
137
138    fn get_tool(&self, name: &str) -> Option<Tool> {
139        self.tool_router.get(name).cloned()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use dbmcp_config::DatabaseBackend;
147
148    fn base_config() -> DatabaseConfig {
149        DatabaseConfig {
150            backend: DatabaseBackend::Postgres,
151            host: "pg.example.com".into(),
152            port: 5433,
153            user: "pgadmin".into(),
154            password: Some("pgpass".into()),
155            name: Some("mydb".into()),
156            ..DatabaseConfig::default()
157        }
158    }
159
160    fn handler(read_only: bool) -> PostgresHandler {
161        PostgresHandler::new(&Config {
162            database: DatabaseConfig {
163                read_only,
164                ..base_config()
165            },
166            http: None,
167            pii: dbmcp_config::PiiConfig::default(),
168        })
169    }
170
171    #[tokio::test]
172    async fn handler_exposes_connection_default_db() {
173        let handler = PostgresHandler::new(&Config {
174            database: base_config(),
175            http: None,
176            pii: dbmcp_config::PiiConfig::default(),
177        });
178        assert_eq!(handler.connection.default_database_name(), "mydb");
179    }
180
181    #[tokio::test]
182    async fn router_exposes_all_thirteen_tools_in_read_write_mode() {
183        let router = handler(false).tool_router;
184        for name in [
185            "listDatabases",
186            "listTables",
187            "listViews",
188            "listTriggers",
189            "listFunctions",
190            "listProcedures",
191            "listMaterializedViews",
192            "readQuery",
193            "explainQuery",
194            "createDatabase",
195            "dropDatabase",
196            "dropTable",
197            "writeQuery",
198        ] {
199            assert!(router.has_route(name), "missing tool: {name}");
200        }
201    }
202
203    #[tokio::test]
204    async fn router_hides_write_tools_in_read_only_mode() {
205        let router = handler(true).tool_router;
206        assert!(router.has_route("listDatabases"));
207        assert!(router.has_route("listTables"));
208        assert!(router.has_route("listViews"));
209        assert!(router.has_route("listTriggers"));
210        assert!(router.has_route("listFunctions"));
211        assert!(router.has_route("listProcedures"));
212        assert!(router.has_route("listMaterializedViews"));
213        assert!(router.has_route("readQuery"));
214        assert!(router.has_route("explainQuery"));
215        assert!(!router.has_route("writeQuery"));
216        assert!(!router.has_route("createDatabase"));
217        assert!(!router.has_route("dropDatabase"));
218        assert!(!router.has_route("dropTable"));
219    }
220
221    #[tokio::test]
222    async fn instructions_match_read_only_mode() {
223        let read_write = handler(false).get_info().instructions.expect("instructions present");
224        assert!(
225            read_write.contains("writeQuery"),
226            "read-write instructions mention writeQuery"
227        );
228
229        let read_only = handler(true).get_info().instructions.expect("instructions present");
230        for tool in ["writeQuery", "createDatabase", "dropDatabase", "dropTable"] {
231            assert!(
232                !read_only.contains(tool),
233                "read-only instructions must not mention {tool}"
234            );
235        }
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
