Skip to main content

dbmcp_postgres/tools/
write_query.rs

1//! MCP tool: `writeQuery`.
2
3use std::borrow::Cow;
4
5use dbmcp_server::types::{QueryRequest, QueryResponse};
6use dbmcp_sql::Connection as _;
7use rmcp::handler::server::router::tool::{AsyncTool, ToolBase};
8use rmcp::model::{ErrorData, ToolAnnotations};
9
10use crate::PostgresHandler;
11
12/// Marker type for the `writeQuery` MCP tool.
13pub(crate) struct WriteQueryTool;
14
15impl WriteQueryTool {
16    const NAME: &'static str = "writeQuery";
17    const TITLE: &'static str = "Write Query";
18    const DESCRIPTION: &'static str = r#"Execute a write SQL query (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
19
20<usecase>
21Use when:
22- Inserting, updating, or deleting rows
23- Creating or altering tables, indexes, views, or other schema objects
24- Any data modification operation
25</usecase>
26
27<when_not_to_use>
28- Read-only queries (SELECT, SHOW) → use readQuery
29- Query performance analysis → use explainQuery
30- Creating/dropping entire databases → use createDatabase or dropDatabase
31</when_not_to_use>
32
33<examples>
34✓ "INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com')"
35✓ "UPDATE orders SET status = 'shipped' WHERE id = 42"
36✓ "CREATE TABLE logs (id SERIAL PRIMARY KEY, message TEXT)"
37✗ "SELECT * FROM users" → use readQuery
38</examples>
39
40<what_it_returns>
41A JSON array of affected/returning row objects, each keyed by column name.
42</what_it_returns>"#;
43}
44
45impl ToolBase for WriteQueryTool {
46    type Parameter = QueryRequest;
47    type Output = QueryResponse;
48    type Error = ErrorData;
49
50    fn name() -> Cow<'static, str> {
51        Self::NAME.into()
52    }
53
54    fn title() -> Option<String> {
55        Some(Self::TITLE.into())
56    }
57
58    fn description() -> Option<Cow<'static, str>> {
59        Some(Self::DESCRIPTION.into())
60    }
61
62    fn annotations() -> Option<ToolAnnotations> {
63        Some(
64            ToolAnnotations::new()
65                .read_only(false)
66                .destructive(true)
67                .idempotent(false)
68                .open_world(true),
69        )
70    }
71}
72
73impl AsyncTool<PostgresHandler> for WriteQueryTool {
74    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
75        handler.write_query(params).await
76    }
77}
78
79impl PostgresHandler {
80    /// Executes a write SQL query.
81    ///
82    /// # Errors
83    ///
84    /// Returns [`SqlError`] if the query fails.
85    pub async fn write_query(
86        &self,
87        QueryRequest { query, database }: QueryRequest,
88    ) -> Result<QueryResponse, ErrorData> {
89        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
90        let mut rows = self.connection.fetch_json(query.as_str(), database).await?;
91        if let Some(r) = &self.redactor {
92            r.apply(&mut rows)?;
93        }
94        Ok(QueryResponse { rows })
95    }
96}