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