cli/
database.rs

1// Copyright 2023 Greptime Team
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::time::Duration;
16
17use base64::Engine;
18use base64::engine::general_purpose;
19use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
20use common_error::ext::BoxedError;
21use humantime::format_duration;
22use serde_json::Value;
23use servers::http::GreptimeQueryOutput;
24use servers::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT;
25use servers::http::result::greptime_result_v1::GreptimedbV1Response;
26use snafu::ResultExt;
27
28use crate::error::{
29    BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result, SerdeJsonSnafu,
30};
31
32#[derive(Debug, Clone)]
33pub struct DatabaseClient {
34    addr: String,
35    catalog: String,
36    auth_header: Option<String>,
37    timeout: Duration,
38    proxy: Option<reqwest::Proxy>,
39    no_proxy: bool,
40}
41
42pub fn parse_proxy_opts(
43    proxy: Option<String>,
44    no_proxy: bool,
45) -> std::result::Result<Option<reqwest::Proxy>, BoxedError> {
46    if no_proxy {
47        return Ok(None);
48    }
49    proxy
50        .map(|proxy| {
51            reqwest::Proxy::all(proxy)
52                .context(ParseProxyOptsSnafu)
53                .map_err(BoxedError::new)
54        })
55        .transpose()
56}
57
58impl DatabaseClient {
59    pub fn new(
60        addr: String,
61        catalog: String,
62        auth_basic: Option<String>,
63        timeout: Duration,
64        proxy: Option<reqwest::Proxy>,
65        no_proxy: bool,
66    ) -> Self {
67        let auth_header = if let Some(basic) = auth_basic {
68            let encoded = general_purpose::STANDARD.encode(basic);
69            Some(format!("basic {}", encoded))
70        } else {
71            None
72        };
73
74        if no_proxy {
75            common_telemetry::info!("Proxy disabled");
76        } else if let Some(ref proxy) = proxy {
77            common_telemetry::info!("Using proxy: {:?}", proxy);
78        } else {
79            common_telemetry::info!("Using system proxy(if any)");
80        }
81
82        Self {
83            addr,
84            catalog,
85            auth_header,
86            timeout,
87            proxy,
88            no_proxy,
89        }
90    }
91
92    pub fn addr(&self) -> &str {
93        &self.addr
94    }
95
96    pub async fn sql_in_public(&self, sql: &str) -> Result<Option<Vec<Vec<Value>>>> {
97        self.sql(sql, DEFAULT_SCHEMA_NAME).await
98    }
99
100    /// Requires the explicit packed-import protocol before any restore mutation.
101    pub async fn require_packed_import(&self) -> Result<()> {
102        self.require_packed_capability("metric_packed_import").await
103    }
104
105    /// Requires packed export support before creating snapshot artifacts.
106    pub async fn require_packed_export(&self) -> Result<()> {
107        self.require_packed_capability("metric_packed_export").await
108    }
109
110    async fn require_packed_capability(&self, capability: &str) -> Result<()> {
111        let url = format!("http://{}/v1/capabilities", self.addr);
112        let mut builder = reqwest::Client::builder().timeout(self.timeout);
113        if let Some(proxy) = self.proxy.clone() {
114            builder = builder.proxy(proxy);
115        }
116        if self.no_proxy {
117            builder = builder.no_proxy();
118        }
119        let client = builder.build().context(BuildClientSnafu)?;
120        let mut request = client.get(&url);
121        if let Some(auth) = &self.auth_header {
122            request = request.header("Authorization", auth);
123        }
124        let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
125            reason: "packed snapshot capability request failed",
126        })?;
127        if response.status() == reqwest::StatusCode::NOT_FOUND {
128            return crate::error::InvalidArgumentsSnafu {
129                msg: format!("server does not support {capability}"),
130            }
131            .fail();
132        }
133        let response = response
134            .error_for_status()
135            .with_context(|_| HttpQuerySqlSnafu {
136                reason: "packed snapshot capability request rejected",
137            })?;
138        let body = response.text().await.with_context(|_| HttpQuerySqlSnafu {
139            reason: "cannot read capability response",
140        })?;
141        let value: Value = serde_json::from_str(&body).context(SerdeJsonSnafu)?;
142        if !value.is_object() {
143            return crate::error::UnexpectedSnafu {
144                msg: "invalid capability response: expected JSON object",
145            }
146            .fail();
147        }
148        if value.get(capability).and_then(Value::as_u64) != Some(1) {
149            return crate::error::InvalidArgumentsSnafu {
150                msg: format!("server does not support {capability} version 1"),
151            }
152            .fail();
153        }
154        Ok(())
155    }
156
157    /// Execute sql query.
158    pub async fn sql(&self, sql: &str, schema: &str) -> Result<Option<Vec<Vec<Value>>>> {
159        let body = self.sql_response(sql, schema).await?;
160        Ok(body.output().first().and_then(|output| match output {
161            GreptimeQueryOutput::Records(records) => Some(records.rows().clone()),
162            GreptimeQueryOutput::AffectedRows(_) => None,
163        }))
164    }
165
166    pub(crate) async fn sql_response(
167        &self,
168        sql: &str,
169        schema: &str,
170    ) -> Result<GreptimedbV1Response> {
171        let url = format!("http://{}/v1/sql", self.addr);
172        let params = [
173            ("db", format!("{}-{}", self.catalog, schema)),
174            ("sql", sql.to_string()),
175        ];
176        let mut builder = reqwest::Client::builder();
177        if let Some(proxy) = self.proxy.clone() {
178            builder = builder.proxy(proxy);
179        }
180        if self.no_proxy {
181            builder = builder.no_proxy();
182        }
183        let client = builder.build().context(BuildClientSnafu)?;
184        let mut request = client
185            .post(&url)
186            .form(&params)
187            .header("Content-Type", "application/x-www-form-urlencoded");
188        if let Some(ref auth) = self.auth_header {
189            request = request.header("Authorization", auth);
190        }
191
192        request = request.header(
193            GREPTIME_DB_HEADER_TIMEOUT,
194            format_duration(self.timeout).to_string(),
195        );
196
197        let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
198            reason: format!("bad url: {}", url),
199        })?;
200        let response = response
201            .error_for_status()
202            .with_context(|_| HttpQuerySqlSnafu {
203                reason: format!("query failed: {}", sql),
204            })?;
205
206        let text = response.text().await.with_context(|_| HttpQuerySqlSnafu {
207            reason: "cannot get response text".to_string(),
208        })?;
209
210        serde_json::from_str::<GreptimedbV1Response>(&text).context(SerdeJsonSnafu)
211    }
212}
213
214/// Split at `-`.
215pub(crate) fn split_database(database: &str) -> Result<(String, Option<String>)> {
216    let (catalog, schema) = match database.split_once('-') {
217        Some((catalog, schema)) => (catalog, schema),
218        None => (DEFAULT_CATALOG_NAME, database),
219    };
220
221    if schema == "*" {
222        Ok((catalog.to_string(), None))
223    } else {
224        Ok((catalog.to_string(), Some(schema.to_string())))
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    #[tokio::test]
231    async fn packed_capability_probe_checks_auth_and_protocol_version() {
232        use tokio::io::{AsyncReadExt, AsyncWriteExt};
233        for export in [false, true] {
234            for (status, body, supported) in [
235                (200, r#"{"metric_packed_import":1,"future":2}"#, true),
236                (200, "{}", false),
237                (200, r#"{"metric_packed_import":2}"#, false),
238                (404, "{}", false),
239                (401, "{}", false),
240                (403, "{}", false),
241                (200, "[]", false),
242                (200, "invalid-json", false),
243            ] {
244                let body = if export {
245                    body.replace("metric_packed_import", "metric_packed_export")
246                } else {
247                    body.to_string()
248                };
249                let response_body = body.clone();
250                let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
251                let address = listener.local_addr().unwrap();
252                let server = tokio::spawn(async move {
253                    let body = response_body;
254                    let (mut socket, _) = listener.accept().await.unwrap();
255                    let mut request = vec![0; 4096];
256                    let n = socket.read(&mut request).await.unwrap();
257                    let request = String::from_utf8_lossy(&request[..n]).to_lowercase();
258                    assert!(request.starts_with("get /v1/capabilities "));
259                    assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza=="));
260                    socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
261                });
262                let client = super::DatabaseClient::new(
263                    address.to_string(),
264                    "greptime".into(),
265                    Some("user:password".into()),
266                    std::time::Duration::from_secs(5),
267                    None,
268                    true,
269                );
270                assert_eq!(
271                    if export {
272                        client.require_packed_export().await.is_ok()
273                    } else {
274                        client.require_packed_import().await.is_ok()
275                    },
276                    supported,
277                    "{status}: {body}"
278                );
279                server.await.unwrap();
280            }
281        }
282    }
283
284    use super::*;
285
286    #[test]
287    fn test_split_database() {
288        let result = split_database("catalog-schema").unwrap();
289        assert_eq!(result, ("catalog".to_string(), Some("schema".to_string())));
290
291        let result = split_database("schema").unwrap();
292        assert_eq!(result, ("greptime".to_string(), Some("schema".to_string())));
293
294        let result = split_database("catalog-*").unwrap();
295        assert_eq!(result, ("catalog".to_string(), None));
296
297        let result = split_database("*").unwrap();
298        assert_eq!(result, ("greptime".to_string(), None));
299    }
300}
cli/database.rs

cli/
database.rs
