Skip to main content

dbmcp_postgres/tools/
explain_query.rs

1//! MCP tool: `explainQuery`.
2
3use dbmcp_pii::MaybeRedact as _;
4use dbmcp_server::types::{ExplainQueryRequest, QueryResponse};
5use dbmcp_sql::validation::{validate_read_only, validate_single_statement};
6
7use super::prelude::*;
8
9const NAME: &str = "explainQuery";
10const TITLE: &str = "Explain Query";
11const DESCRIPTION_PINNED: &str = include_str!("../../assets/tools/explain_query/pinned.md");
12const DESCRIPTION_UNPINNED: &str = include_str!("../../assets/tools/explain_query/unpinned.md");
13
14fn annotations() -> ToolAnnotations {
15    ToolAnnotations::new()
16        .read_only(true)
17        .destructive(false)
18        .idempotent(true)
19        .open_world(true)
20}
21
22/// Marker type for the `explainQuery` MCP tool (pinned variant — no `database` field).
23pub(crate) struct PinnedExplainQueryTool;
24
25impl ToolBase for PinnedExplainQueryTool {
26    type Parameter = ExplainQueryRequest;
27    type Output = QueryResponse;
28    type Error = ErrorData;
29
30    fn name() -> Cow<'static, str> {
31        NAME.into()
32    }
33
34    fn title() -> Option<String> {
35        Some(TITLE.into())
36    }
37
38    fn description() -> Option<Cow<'static, str>> {
39        Some(DESCRIPTION_PINNED.into())
40    }
41
42    fn annotations() -> Option<ToolAnnotations> {
43        Some(annotations())
44    }
45
46    fn input_schema() -> Option<Arc<JsonObject>> {
47        Some(input_schema::<Self::Parameter>(true))
48    }
49
50    fn output_schema() -> Option<Arc<JsonObject>> {
51        Some(output_schema::<Self::Output>())
52    }
53}
54
55impl AsyncTool<PostgresHandler> for PinnedExplainQueryTool {
56    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
57        handler.explain_query(None, params.query, params.analyze).await
58    }
59}
60
61/// Marker type for the `explainQuery` MCP tool (unpinned variant — carries `database`).
62pub(crate) struct UnpinnedExplainQueryTool;
63
64impl ToolBase for UnpinnedExplainQueryTool {
65    type Parameter = ExplainQueryRequest;
66    type Output = QueryResponse;
67    type Error = ErrorData;
68
69    fn name() -> Cow<'static, str> {
70        NAME.into()
71    }
72
73    fn title() -> Option<String> {
74        Some(TITLE.into())
75    }
76
77    fn description() -> Option<Cow<'static, str>> {
78        Some(DESCRIPTION_UNPINNED.into())
79    }
80
81    fn annotations() -> Option<ToolAnnotations> {
82        Some(annotations())
83    }
84
85    fn input_schema() -> Option<Arc<JsonObject>> {
86        Some(input_schema::<Self::Parameter>(false))
87    }
88
89    fn output_schema() -> Option<Arc<JsonObject>> {
90        Some(output_schema::<Self::Output>())
91    }
92}
93
94impl AsyncTool<PostgresHandler> for UnpinnedExplainQueryTool {
95    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
96        handler
97            .explain_query(params.database, params.query, params.analyze)
98            .await
99    }
100}
101
102impl PostgresHandler {
103    /// Returns the execution plan for a query.
104    ///
105    /// The inner query must be a single statement. When `analyze` is true
106    /// and read-only mode is enabled, it must also be read-only.
107    ///
108    /// # Errors
109    ///
110    /// Returns [`SqlError::MultiStatement`] if the query holds more than
111    /// one statement.
112    /// Returns [`SqlError::ReadOnlyViolation`] if `analyze` is true,
113    /// read-only mode is enabled, and the query is a write statement.
114    /// Returns [`SqlError::Query`] if the backend reports an error.
115    pub async fn explain_query(
116        &self,
117        database: Option<String>,
118        query: String,
119        analyze: bool,
120    ) -> Result<QueryResponse, ErrorData> {
121        if analyze && self.config.read_only {
122            let _ = validate_read_only(&query, &sqlparser::dialect::PostgreSqlDialect {})?;
123        } else {
124            validate_single_statement(&query, &sqlparser::dialect::PostgreSqlDialect {})?;
125        }
126
127        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
128
129        let explain_sql = if analyze {
130            format!("EXPLAIN (ANALYZE, FORMAT JSON) {query}")
131        } else {
132            format!("EXPLAIN (FORMAT JSON) {query}")
133        };
134
135        let rows = self.connection.fetch_json(explain_sql.as_str(), database).await?;
136        let rows = self.redactor.redact_rows(rows).await?;
137
138        Ok(QueryResponse { rows })
139    }
140}