databend_client/
request.rs

1// Copyright 2021 Datafuse Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16
17use crate::session::SessionState;
18use serde::Serialize;
19
20#[derive(Serialize, Debug)]
21pub struct QueryRequest<'a> {
22    #[serde(skip_serializing_if = "Option::is_none")]
23    session: Option<SessionState>,
24    sql: &'a str,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pagination: Option<PaginationConfig>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    stage_attachment: Option<StageAttachmentConfig<'a>>,
29}
30
31#[derive(Serialize, Debug)]
32pub struct PaginationConfig {
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub wait_time_secs: Option<i64>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub max_rows_in_buffer: Option<i64>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub max_rows_per_page: Option<i64>,
39}
40
41#[derive(Serialize, Debug)]
42pub struct StageAttachmentConfig<'a> {
43    pub location: &'a str,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub file_format_options: Option<BTreeMap<&'a str, &'a str>>,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub copy_options: Option<BTreeMap<&'a str, &'a str>>,
48}
49
50impl<'r, 't: 'r> QueryRequest<'r> {
51    pub fn new(sql: &'r str) -> QueryRequest<'r> {
52        QueryRequest {
53            session: None,
54            sql,
55            pagination: None,
56            stage_attachment: None,
57        }
58    }
59
60    pub fn with_session(mut self, session: Option<SessionState>) -> Self {
61        self.session = session;
62        self
63    }
64
65    pub fn with_pagination(mut self, pagination: Option<PaginationConfig>) -> Self {
66        self.pagination = pagination;
67        self
68    }
69
70    pub fn with_stage_attachment(
71        mut self,
72        stage_attachment: Option<StageAttachmentConfig<'t>>,
73    ) -> Self {
74        self.stage_attachment = stage_attachment;
75        self
76    }
77}
78
79#[cfg(test)]
80mod test {
81    use super::*;
82    use crate::error::Result;
83
84    #[test]
85    fn build_request() -> Result<()> {
86        let req = QueryRequest::new("select 1")
87            .with_session(Some(
88                SessionState::default().with_database(Some("default".to_string())),
89            ))
90            .with_pagination(Some(PaginationConfig {
91                wait_time_secs: Some(1),
92                max_rows_in_buffer: Some(1),
93                max_rows_per_page: Some(1),
94            }))
95            .with_stage_attachment(Some(StageAttachmentConfig {
96                location: "@~/my_location",
97                file_format_options: None,
98                copy_options: None,
99            }));
100        assert_eq!(
101            serde_json::to_string(&req)?,
102            r#"{"session":{"database":"default"},"sql":"select 1","pagination":{"wait_time_secs":1,"max_rows_in_buffer":1,"max_rows_per_page":1},"stage_attachment":{"location":"@~/my_location"}}"#
103        );
104        Ok(())
105    }
106}
databend_client/request.rs

databend_client/
request.rs
