dbmcp_postgres/tools/
write_query.rs

1//! MCP tool: `writeQuery`.
2
3use dbmcp_pii::MaybeRedact as _;
4use dbmcp_server::types::{QueryRequest, QueryResponse};
5
6use super::prelude::*;
7
8const NAME: &str = "writeQuery";
9const TITLE: &str = "Write Query";
10const DESCRIPTION_PINNED: &str = include_str!("../../assets/tools/write_query/pinned.md");
11const DESCRIPTION_UNPINNED: &str = include_str!("../../assets/tools/write_query/unpinned.md");
12
13fn annotations() -> ToolAnnotations {
14    ToolAnnotations::new()
15        .read_only(false)
16        .destructive(true)
17        .idempotent(false)
18        .open_world(true)
19}
20
21/// Marker type for the `writeQuery` MCP tool (pinned variant — no `database` field).
22pub(crate) struct PinnedWriteQueryTool;
23
24impl ToolBase for PinnedWriteQueryTool {
25    type Parameter = QueryRequest;
26    type Output = QueryResponse;
27    type Error = ErrorData;
28
29    fn name() -> Cow<'static, str> {
30        NAME.into()
31    }
32
33    fn title() -> Option<String> {
34        Some(TITLE.into())
35    }
36
37    fn description() -> Option<Cow<'static, str>> {
38        Some(DESCRIPTION_PINNED.into())
39    }
40
41    fn annotations() -> Option<ToolAnnotations> {
42        Some(annotations())
43    }
44
45    fn input_schema() -> Option<Arc<JsonObject>> {
46        Some(input_schema::<Self::Parameter>(true))
47    }
48
49    fn output_schema() -> Option<Arc<JsonObject>> {
50        Some(output_schema::<Self::Output>())
51    }
52}
53
54impl AsyncTool<PostgresHandler> for PinnedWriteQueryTool {
55    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
56        handler.write_query(params.query, None).await
57    }
58}
59
60/// Marker type for the `writeQuery` MCP tool (unpinned variant — carries `database`).
61pub(crate) struct UnpinnedWriteQueryTool;
62
63impl ToolBase for UnpinnedWriteQueryTool {
64    type Parameter = QueryRequest;
65    type Output = QueryResponse;
66    type Error = ErrorData;
67
68    fn name() -> Cow<'static, str> {
69        NAME.into()
70    }
71
72    fn title() -> Option<String> {
73        Some(TITLE.into())
74    }
75
76    fn description() -> Option<Cow<'static, str>> {
77        Some(DESCRIPTION_UNPINNED.into())
78    }
79
80    fn annotations() -> Option<ToolAnnotations> {
81        Some(annotations())
82    }
83
84    fn input_schema() -> Option<Arc<JsonObject>> {
85        Some(input_schema::<Self::Parameter>(false))
86    }
87
88    fn output_schema() -> Option<Arc<JsonObject>> {
89        Some(output_schema::<Self::Output>())
90    }
91}
92
93impl AsyncTool<PostgresHandler> for UnpinnedWriteQueryTool {
94    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
95        handler.write_query(params.query, params.database).await
96    }
97}
98
99impl PostgresHandler {
100    /// Executes a write SQL query.
101    ///
102    /// # Errors
103    ///
104    /// Returns [`SqlError`] if the query fails.
105    pub async fn write_query(&self, query: String, database: Option<String>) -> Result<QueryResponse, ErrorData> {
106        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
107        let rows = self.connection.fetch_json(query.as_str(), database).await?;
108        let rows = self.redactor.redact_rows(rows).await?;
109        Ok(QueryResponse { rows })
110    }
111}
dbmcp_postgres/tools/write_query.rs

dbmcp_postgres/tools/
write_query.rs
