dbmcp_postgres/tools/
read_query.rs

1//! MCP tool: `readQuery`.
2
3use dbmcp_pii::MaybeRedact as _;
4use dbmcp_server::pagination::{Cursor, Pager};
5use dbmcp_server::types::{ReadQueryRequest, ReadQueryResponse};
6use dbmcp_sql::StatementKind;
7use dbmcp_sql::pagination::with_limit_offset;
8use dbmcp_sql::validation::validate_read_only;
9
10use super::prelude::*;
11
12const NAME: &str = "readQuery";
13const TITLE: &str = "Read Query";
14const DESCRIPTION_PINNED: &str = include_str!("../../assets/tools/read_query/pinned.md");
15const DESCRIPTION_UNPINNED: &str = include_str!("../../assets/tools/read_query/unpinned.md");
16
17fn annotations() -> ToolAnnotations {
18    ToolAnnotations::new()
19        .read_only(true)
20        .destructive(false)
21        .idempotent(true)
22        .open_world(true)
23}
24
25/// Marker type for the `readQuery` MCP tool (pinned variant — no `database` field).
26pub(crate) struct PinnedReadQueryTool;
27
28impl ToolBase for PinnedReadQueryTool {
29    type Parameter = ReadQueryRequest;
30    type Output = ReadQueryResponse;
31    type Error = ErrorData;
32
33    fn name() -> Cow<'static, str> {
34        NAME.into()
35    }
36
37    fn title() -> Option<String> {
38        Some(TITLE.into())
39    }
40
41    fn description() -> Option<Cow<'static, str>> {
42        Some(DESCRIPTION_PINNED.into())
43    }
44
45    fn annotations() -> Option<ToolAnnotations> {
46        Some(annotations())
47    }
48
49    fn input_schema() -> Option<Arc<JsonObject>> {
50        Some(input_schema::<Self::Parameter>(true))
51    }
52
53    fn output_schema() -> Option<Arc<JsonObject>> {
54        Some(output_schema::<Self::Output>())
55    }
56}
57
58impl AsyncTool<PostgresHandler> for PinnedReadQueryTool {
59    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
60        handler.read_query(params.query, None, params.cursor).await
61    }
62}
63
64/// Marker type for the `readQuery` MCP tool (unpinned variant — carries `database`).
65pub(crate) struct UnpinnedReadQueryTool;
66
67impl ToolBase for UnpinnedReadQueryTool {
68    type Parameter = ReadQueryRequest;
69    type Output = ReadQueryResponse;
70    type Error = ErrorData;
71
72    fn name() -> Cow<'static, str> {
73        NAME.into()
74    }
75
76    fn title() -> Option<String> {
77        Some(TITLE.into())
78    }
79
80    fn description() -> Option<Cow<'static, str>> {
81        Some(DESCRIPTION_UNPINNED.into())
82    }
83
84    fn annotations() -> Option<ToolAnnotations> {
85        Some(annotations())
86    }
87
88    fn input_schema() -> Option<Arc<JsonObject>> {
89        Some(input_schema::<Self::Parameter>(false))
90    }
91
92    fn output_schema() -> Option<Arc<JsonObject>> {
93        Some(output_schema::<Self::Output>())
94    }
95}
96
97impl AsyncTool<PostgresHandler> for UnpinnedReadQueryTool {
98    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
99        handler.read_query(params.query, params.database, params.cursor).await
100    }
101}
102
103impl PostgresHandler {
104    /// Executes a read-only SQL query, paginating `SELECT` result rows.
105    ///
106    /// Validates that the query is read-only, then dispatches on the
107    /// classified [`StatementKind`]: `Select` is wrapped in a subquery with
108    /// a server-controlled `LIMIT`/`OFFSET`; `NonSelect` (`SHOW`, `EXPLAIN`
109    /// under the PG dialect) is executed as-is and returned in a single
110    /// page. A malformed `cursor` is rejected by the serde deserializer
111    /// before this method is called, producing JSON-RPC `-32602`.
112    ///
113    /// # Errors
114    ///
115    /// Returns [`SqlError::ReadOnlyViolation`] if the query is not
116    /// read-only, or [`SqlError::Query`] if the backend reports an error.
117    pub async fn read_query(
118        &self,
119        query: String,
120        database: Option<String>,
121        cursor: Option<Cursor>,
122    ) -> Result<ReadQueryResponse, ErrorData> {
123        let kind = validate_read_only(&query, &sqlparser::dialect::PostgreSqlDialect {})?;
124        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
125
126        let (rows, next_cursor) = match kind {
127            StatementKind::Select => {
128                let pager = Pager::new(cursor, self.config.page_size);
129                let wrapped = with_limit_offset(&query, pager.limit(), pager.offset());
130                let rows = self.connection.fetch_json(wrapped.as_str(), database).await?;
131                pager.paginate(rows)
132            }
133            StatementKind::NonSelect => {
134                let rows = self.connection.fetch_json(query.as_str(), database).await?;
135                (rows, None)
136            }
137        };
138        let rows = self.redactor.redact_rows(rows).await?;
139        Ok(ReadQueryResponse { rows, next_cursor })
140    }
141}
dbmcp_postgres/tools/read_query.rs

dbmcp_postgres/tools/
read_query.rs
