Skip to main content

databend_client/
settings.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 crate::error::Result;
16use crate::Error;
17use jiff::tz::TimeZone;
18use serde::Deserialize;
19use std::str::FromStr;
20
21#[derive(Debug, Clone)]
22pub struct ResultFormatSettings {
23    pub geometry_output_format: GeometryDataType,
24    pub timezone: TimeZone,
25    pub arrow_result_version: Option<i64>,
26    pub binary_output_format: BinaryFormat,
27}
28
29impl Default for ResultFormatSettings {
30    fn default() -> Self {
31        Self {
32            geometry_output_format: GeometryDataType::default(),
33            binary_output_format: BinaryFormat::default(),
34            timezone: TimeZone::UTC,
35            arrow_result_version: None,
36        }
37    }
38}
39
40impl TryFrom<&Option<QueryResultFormatSettings>> for ResultFormatSettings {
41    type Error = Error;
42
43    fn try_from(settings: &Option<QueryResultFormatSettings>) -> Result<Self> {
44        let settings = settings.clone().unwrap_or_default();
45        let timezone = match settings.timezone {
46            None => TimeZone::UTC,
47            Some(t) => TimeZone::get(&t).map_err(|e| Error::Decode(e.to_string()))?,
48        };
49
50        let geometry_output_format = match settings.geometry_output_format {
51            None => GeometryDataType::default(),
52            Some(t) => GeometryDataType::from_str(&t).map_err(|e| Error::Decode(e.to_string()))?,
53        };
54
55        let binary_output_format = match settings.binary_output_format {
56            None => BinaryFormat::default(),
57            Some(t) => BinaryFormat::from_str(&t).map_err(|e| Error::Decode(e.to_string()))?,
58        };
59
60        Ok(Self {
61            geometry_output_format,
62            timezone,
63            arrow_result_version: settings.arrow_result_version,
64            binary_output_format,
65        })
66    }
67}
68
69#[derive(Debug, Clone, Default, Deserialize)]
70pub struct QueryResultFormatSettings {
71    pub timezone: Option<String>,
72    pub geometry_output_format: Option<String>,
73    pub arrow_result_version: Option<i64>,
74    pub binary_output_format: Option<String>,
75}
76
77#[derive(Debug, Clone, Copy, Default, Deserialize, PartialEq, Eq)]
78pub enum GeometryDataType {
79    WKB,
80    WKT,
81    EWKB,
82    EWKT,
83    #[default]
84    GEOJSON,
85}
86
87impl FromStr for GeometryDataType {
88    type Err = Error;
89
90    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
91        match s.to_uppercase().as_str() {
92            "WKB" => Ok(GeometryDataType::WKB),
93            "WKT" => Ok(GeometryDataType::WKT),
94            "EWKB" => Ok(GeometryDataType::EWKB),
95            "EWKT" => Ok(GeometryDataType::EWKT),
96            "GEOJSON" => Ok(GeometryDataType::GEOJSON),
97            _ => Err(Error::Decode("Invalid geometry type format".to_string())),
98        }
99    }
100}
101
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
103pub enum BinaryFormat {
104    #[default]
105    Hex,
106    Base64,
107    Utf8,
108    Utf8Lossy,
109}
110
111impl FromStr for BinaryFormat {
112    type Err = Error;
113
114    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
115        match s.to_ascii_lowercase().as_str() {
116            "hex" => Ok(BinaryFormat::Hex),
117            "base64" => Ok(BinaryFormat::Base64),
118            "utf-8" | "utf8" => Ok(BinaryFormat::Utf8),
119            "utf-8-lossy" | "utf8-lossy" => Ok(BinaryFormat::Utf8Lossy),
120            other => Err(Error::Decode(format!(
121                "Invalid binary format '{other}', valid values: HEX | BASE64 | UTF-8 | UTF-8-LOSSY"
122            ))),
123        }
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn deserialize_query_result_format_settings_from_strings() {
133        let settings: QueryResultFormatSettings = serde_json::from_str(
134            r#"{
135                "timezone": "Asia/Shanghai",
136                "geometry_output_format": "wkt",
137                "arrow_result_version": 2,
138                "binary_output_format": "utf-8"
139            }"#,
140        )
141        .unwrap();
142
143        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
144        assert_eq!(settings.geometry_output_format, GeometryDataType::WKT);
145        assert_eq!(settings.arrow_result_version, Some(2));
146        assert_eq!(settings.binary_output_format, BinaryFormat::Utf8);
147        assert_eq!(settings.timezone.iana_name(), Some("Asia/Shanghai"));
148    }
149
150    #[test]
151    fn deserialize_query_result_format_settings_with_defaults() {
152        let settings: QueryResultFormatSettings = serde_json::from_str(r#"{}"#).unwrap();
153
154        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
155        assert_eq!(settings.geometry_output_format, GeometryDataType::default());
156        assert_eq!(settings.arrow_result_version, None);
157        assert_eq!(settings.binary_output_format, BinaryFormat::default());
158        assert_eq!(settings.timezone.iana_name(), Some("UTC"));
159    }
160
161    #[test]
162    fn deserialize_query_result_format_settings_accepts_numeric_arrow_version() {
163        let settings: QueryResultFormatSettings =
164            serde_json::from_str(r#"{"arrow_result_version": 2}"#).unwrap();
165
166        let settings = ResultFormatSettings::try_from(&Some(settings)).unwrap();
167        assert_eq!(settings.arrow_result_version, Some(2));
168    }
169}