Skip to main content

dbmcp_postgres/tools/
explain_query.rs

1//! MCP tool: `explainQuery`.
2
3use std::borrow::Cow;
4
5use dbmcp_server::types::{ExplainQueryRequest, QueryResponse};
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
13/// Marker type for the `explainQuery` MCP tool.
14pub(crate) struct ExplainQueryTool;
15
16impl ExplainQueryTool {
17    const NAME: &'static str = "explainQuery";
18    const TITLE: &'static str = "Explain Query";
19    const DESCRIPTION: &'static str = include_str!("../../assets/tools/explain_query.md");
20}
21
22impl ToolBase for ExplainQueryTool {
23    type Parameter = ExplainQueryRequest;
24    type Output = QueryResponse;
25    type Error = ErrorData;
26
27    fn name() -> Cow<'static, str> {
28        Self::NAME.into()
29    }
30
31    fn title() -> Option<String> {
32        Some(Self::TITLE.into())
33    }
34
35    fn description() -> Option<Cow<'static, str>> {
36        Some(Self::DESCRIPTION.into())
37    }
38
39    fn annotations() -> Option<ToolAnnotations> {
40        Some(
41            ToolAnnotations::new()
42                .read_only(true)
43                .destructive(false)
44                .idempotent(true)
45                .open_world(true),
46        )
47    }
48}
49
50impl AsyncTool<PostgresHandler> for ExplainQueryTool {
51    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
52        handler.explain_query(params).await
53    }
54}
55
56impl PostgresHandler {
57    /// Returns the execution plan for a query.
58    ///
59    /// When `analyze` is true and read-only mode is enabled, the inner
60    /// query is validated to be read-only before executing.
61    ///
62    /// # Errors
63    ///
64    /// Returns [`SqlError::ReadOnlyViolation`] if `analyze` is true,
65    /// read-only mode is enabled, and the query is a write statement.
66    /// Returns [`SqlError::Query`] if the backend reports an error.
67    pub async fn explain_query(
68        &self,
69        ExplainQueryRequest {
70            database,
71            query,
72            analyze,
73        }: ExplainQueryRequest,
74    ) -> Result<QueryResponse, ErrorData> {
75        if analyze && self.config.read_only {
76            let _ = validate_read_only(&query, &sqlparser::dialect::PostgreSqlDialect {})?;
77        }
78
79        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
80
81        let explain_sql = if analyze {
82            format!("EXPLAIN (ANALYZE, FORMAT JSON) {query}")
83        } else {
84            format!("EXPLAIN (FORMAT JSON) {query}")
85        };
86
87        let mut rows = self.connection.fetch_json(explain_sql.as_str(), database).await?;
88        if let Some(r) = &self.redactor {
89            r.apply(&mut rows)?;
90        }
91
92        Ok(QueryResponse { rows })
93    }
94}