dbmcp_sqlite/tools/
read_query.rs1use dbmcp_pii::MaybeRedact as _;
4use dbmcp_server::pagination::Pager;
5use dbmcp_server::types::ReadQueryResponse;
6
7use dbmcp_sql::StatementKind;
8use dbmcp_sql::pagination::with_limit_offset;
9use dbmcp_sql::validation::validate_read_only;
10
11use super::prelude::*;
12use crate::types::ReadQueryRequest;
13
14const NAME: &str = "readQuery";
15const TITLE: &str = "Read Query";
16const DESCRIPTION: &str = include_str!("../../assets/tools/read_query.md");
17
18pub(crate) struct ReadQueryTool;
20
21impl ToolBase for ReadQueryTool {
22 type Parameter = ReadQueryRequest;
23 type Output = ReadQueryResponse;
24 type Error = ErrorData;
25
26 fn name() -> Cow<'static, str> {
27 NAME.into()
28 }
29
30 fn title() -> Option<String> {
31 Some(TITLE.into())
32 }
33
34 fn description() -> Option<Cow<'static, str>> {
35 Some(DESCRIPTION.into())
36 }
37
38 fn annotations() -> Option<ToolAnnotations> {
39 Some(
40 ToolAnnotations::new()
41 .read_only(true)
42 .destructive(false)
43 .idempotent(true)
44 .open_world(true),
45 )
46 }
47
48 fn input_schema() -> Option<Arc<JsonObject>> {
49 Some(input_schema::<Self::Parameter>(true))
50 }
51
52 fn output_schema() -> Option<Arc<JsonObject>> {
53 Some(output_schema::<Self::Output>())
54 }
55}
56
57impl AsyncTool<SqliteHandler> for ReadQueryTool {
58 async fn invoke(handler: &SqliteHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
59 handler.read_query(params).await
60 }
61}
62
63impl SqliteHandler {
64 pub async fn read_query(
78 &self,
79 ReadQueryRequest { query, cursor }: ReadQueryRequest,
80 ) -> Result<ReadQueryResponse, ErrorData> {
81 let kind = validate_read_only(&query, &sqlparser::dialect::SQLiteDialect {})?;
82
83 let (rows, next_cursor) = match kind {
84 StatementKind::Select => {
85 let pager = Pager::new(cursor, self.config.page_size);
86 let wrapped = with_limit_offset(&query, pager.limit(), pager.offset());
87 let rows = self.connection.fetch_json(wrapped.as_str(), None).await?;
88 pager.paginate(rows)
89 }
90 StatementKind::NonSelect => {
91 let rows = self.connection.fetch_json(query.as_str(), None).await?;
92 (rows, None)
93 }
94 };
95 let rows = self.redactor.redact_rows(rows).await?;
96 Ok(ReadQueryResponse { rows, next_cursor })
97 }
98}