dbmcp_postgres/tools/
explain_query.rs1use 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
26pub(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
57pub(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 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}