Skip to main content

database_mcp_postgres/
handler.rs

1//! MCP handler for the `PostgreSQL` backend.
2//!
3//! [`PostgresHandler`] wraps [`PostgresBackend`] and implements
4//! [`ServerHandler`] using rmcp tool macros.
5
6use database_mcp_backend::types::{CreateDatabaseRequest, GetTableSchemaRequest, ListTablesRequest, QueryRequest};
7use database_mcp_config::DatabaseConfig;
8use database_mcp_server::tools;
9use rmcp::ServerHandler;
10use rmcp::handler::server::router::tool::ToolRouter;
11use rmcp::handler::server::wrapper::Parameters;
12use rmcp::model::ErrorData;
13
14use super::PostgresBackend;
15
16/// MCP handler for `PostgreSQL` databases.
17///
18/// Owns a [`PostgresBackend`] and a pre-filtered [`ToolRouter`].
19/// Write tools are removed when the backend is in read-only mode.
20#[derive(Clone, Debug)]
21pub struct PostgresHandler {
22    backend: PostgresBackend,
23    tool_router: ToolRouter<Self>,
24}
25
26impl PostgresHandler {
27    /// Creates a new `PostgreSQL` handler.
28    ///
29    /// # Errors
30    ///
31    /// Returns an error if the database connection cannot be established.
32    pub async fn new(config: &DatabaseConfig) -> Result<Self, database_mcp_backend::AppError> {
33        let backend = PostgresBackend::new(config).await?;
34        let mut tool_router = Self::tool_router();
35        if backend.read_only {
36            tool_router.remove_route("write_query");
37            tool_router.remove_route("create_database");
38        }
39        Ok(Self { backend, tool_router })
40    }
41}
42
43#[rmcp::tool_router]
44impl PostgresHandler {
45    /// List all accessible databases.
46    #[rmcp::tool(
47        name = "list_databases",
48        description = "List all accessible databases on the connected database server. Call this first to discover available database names.",
49        annotations(
50            read_only_hint = true,
51            destructive_hint = false,
52            idempotent_hint = true,
53            open_world_hint = false
54        )
55    )]
56    async fn list_databases(&self) -> Result<String, ErrorData> {
57        tools::list_databases(self.backend.list_databases()).await
58    }
59
60    /// List all tables in a specific database.
61    #[rmcp::tool(
62        name = "list_tables",
63        description = "List all tables in a specific database. Requires database_name from list_databases.",
64        annotations(
65            read_only_hint = true,
66            destructive_hint = false,
67            idempotent_hint = true,
68            open_world_hint = false
69        )
70    )]
71    async fn list_tables(&self, Parameters(req): Parameters<ListTablesRequest>) -> Result<String, ErrorData> {
72        tools::list_tables(self.backend.list_tables(&req.database_name), &req.database_name).await
73    }
74
75    /// Get column definitions for a table.
76    #[rmcp::tool(
77        name = "get_table_schema",
78        description = "Get column definitions (type, nullable, key, default) and foreign key relationships for a table. Requires database_name and table_name.",
79        annotations(
80            read_only_hint = true,
81            destructive_hint = false,
82            idempotent_hint = true,
83            open_world_hint = false
84        )
85    )]
86    async fn get_table_schema(&self, Parameters(req): Parameters<GetTableSchemaRequest>) -> Result<String, ErrorData> {
87        tools::get_table_schema(
88            self.backend.get_table_schema(&req.database_name, &req.table_name),
89            &req.database_name,
90            &req.table_name,
91        )
92        .await
93    }
94
95    /// Execute a read-only SQL query.
96    #[rmcp::tool(
97        name = "read_query",
98        description = "Execute a read-only SQL query (SELECT, SHOW, DESCRIBE, USE, EXPLAIN).",
99        annotations(
100            read_only_hint = true,
101            destructive_hint = false,
102            idempotent_hint = true,
103            open_world_hint = true
104        )
105    )]
106    async fn read_query(&self, Parameters(req): Parameters<QueryRequest>) -> Result<String, ErrorData> {
107        tools::read_query(
108            self.backend
109                .execute_query(&req.sql_query, tools::resolve_database(&req.database_name)),
110            &req.sql_query,
111            &req.database_name,
112            |sql| {
113                database_mcp_backend::validation::validate_read_only_with_dialect(
114                    sql,
115                    &sqlparser::dialect::PostgreSqlDialect {},
116                )
117            },
118        )
119        .await
120    }
121
122    /// Execute a write SQL query.
123    #[rmcp::tool(
124        name = "write_query",
125        description = "Execute a write SQL query (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).",
126        annotations(
127            read_only_hint = false,
128            destructive_hint = true,
129            idempotent_hint = false,
130            open_world_hint = true
131        )
132    )]
133    async fn write_query(&self, Parameters(req): Parameters<QueryRequest>) -> Result<String, ErrorData> {
134        tools::write_query(
135            self.backend
136                .execute_query(&req.sql_query, tools::resolve_database(&req.database_name)),
137            &req.sql_query,
138            &req.database_name,
139        )
140        .await
141    }
142
143    /// Create a new database.
144    #[rmcp::tool(
145        name = "create_database",
146        description = "Create a new database. Not supported for SQLite.",
147        annotations(
148            read_only_hint = false,
149            destructive_hint = false,
150            idempotent_hint = false,
151            open_world_hint = false
152        )
153    )]
154    async fn create_database(&self, Parameters(req): Parameters<CreateDatabaseRequest>) -> Result<String, ErrorData> {
155        tools::create_database(self.backend.create_database(&req.database_name), &req.database_name).await
156    }
157}
158
159#[rmcp::tool_handler(router = self.tool_router)]
160impl ServerHandler for PostgresHandler {
161    fn get_info(&self) -> rmcp::model::ServerInfo {
162        database_mcp_server::server_info()
163    }
164}