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

cli/
database.rs
