Skip to main content

dbmcp_postgres/tools/
explain_query.rs

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