Skip to main content

dbmcp_postgres/tools/
explain_query.rs

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