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::{BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result};
30
31#[derive(Debug, Clone)]
32pub struct DatabaseClient {
33    addr: String,
34    catalog: String,
35    auth_header: Option<String>,
36    timeout: Duration,
37    proxy: Option<reqwest::Proxy>,
38    no_proxy: bool,
39    client: Arc<tokio::sync::OnceCell<reqwest::Client>>,
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            client: Arc::new(tokio::sync::OnceCell::new()),
90        }
91    }
92
93    pub fn addr(&self) -> &str {
94        &self.addr
95    }
96
97    pub async fn sql_in_public(&self, sql: &str) -> Result<Option<Vec<Vec<Value>>>> {
98        self.sql(sql, DEFAULT_SCHEMA_NAME).await
99    }
100
101    pub(crate) fn catalog(&self) -> &str {
102        &self.catalog
103    }
104
105    /// Requires packed export support before creating snapshot artifacts.
106    pub async fn require_packed_export(&self) -> Result<()> {
107        self.capabilities().await?.require("metric_packed_export")
108    }
109
110    /// Reads target capabilities once; older targets advertise no extensions.
111    pub(crate) async fn capabilities(&self) -> Result<Capabilities> {
112        let url = format!("http://{}/v1/capabilities", self.addr);
113        let mut request = self.http_client().await?.get(&url).timeout(self.timeout);
114        if let Some(auth) = &self.auth_header {
115            request = request.header("Authorization", auth);
116        }
117        let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
118            reason: "capability request failed",
119        })?;
120        if response.status() == reqwest::StatusCode::NOT_FOUND {
121            return Ok(Capabilities(serde_json::json!({})));
122        }
123        let response = response
124            .error_for_status()
125            .with_context(|_| HttpQuerySqlSnafu {
126                reason: "capability request rejected",
127            })?;
128        let body = response.text().await.with_context(|_| HttpQuerySqlSnafu {
129            reason: "cannot read capability response",
130        })?;
131        let value: Value = serde_json::from_str(&body).map_err(|_| {
132            crate::error::UnexpectedSnafu {
133                msg: "invalid capability response",
134            }
135            .build()
136        })?;
137        if !value.is_object()
138            || value.get("error").is_some()
139            || value.get("code").is_some_and(|v| v.as_u64() != Some(0))
140        {
141            return crate::error::UnexpectedSnafu {
142                msg: "invalid capability response: expected JSON object",
143            }
144            .fail();
145        }
146        for key in [
147            "metric_batch_ddl",
148            "metric_packed_import",
149            "metric_packed_export",
150        ] {
151            if value.get(key).is_some_and(|v| v.as_u64().is_none()) {
152                return crate::error::UnexpectedSnafu {
153                    msg: "invalid capability version",
154                }
155                .fail();
156            }
157        }
158        Ok(Capabilities(value))
159    }
160
161    /// Submits a batch once. An ambiguous failure must not be replayed here.
162    pub(crate) async fn logical_tables(&self, sql: &str, schema: &str, count: usize) -> Result<()> {
163        let db = format!("{}-{}", self.catalog, schema);
164        if 8 + form_encoded_len(&db) + form_encoded_len(sql) > 4 * 1024 * 1024 {
165            return crate::error::InvalidArgumentsSnafu {
166                msg: "logical-table batch exceeds 4 MiB form limit",
167            }
168            .fail();
169        }
170        let response = self.ddl_response(sql, schema, "ddl/logical-tables").await?;
171        if response.output().len() != count
172            || response
173                .output()
174                .iter()
175                .any(|output| !matches!(output, GreptimeQueryOutput::AffectedRows(0)))
176        {
177            return crate::error::UnexpectedSnafu {
178                msg: "invalid logical-table batch output",
179            }
180            .fail();
181        }
182        Ok(())
183    }
184
185    /// Execute sql query.
186    pub async fn sql(&self, sql: &str, schema: &str) -> Result<Option<Vec<Vec<Value>>>> {
187        let body = self.sql_response(sql, schema).await?;
188        Ok(body.output().first().and_then(|output| match output {
189            GreptimeQueryOutput::Records(records) => Some(records.rows().clone()),
190            GreptimeQueryOutput::AffectedRows(_) => None,
191        }))
192    }
193
194    pub(crate) async fn sql_response(
195        &self,
196        sql: &str,
197        schema: &str,
198    ) -> Result<GreptimedbV1Response> {
199        self.ddl_response(sql, schema, "sql").await
200    }
201
202    async fn ddl_response(
203        &self,
204        sql: &str,
205        schema: &str,
206        endpoint: &str,
207    ) -> Result<GreptimedbV1Response> {
208        let url = format!("http://{}/v1/{endpoint}", self.addr);
209        let params = [
210            ("db", format!("{}-{}", self.catalog, schema)),
211            ("sql", sql.to_string()),
212        ];
213        let mut request = self
214            .http_client()
215            .await?
216            .post(&url)
217            .header("Content-Type", "application/x-www-form-urlencoded");
218        if endpoint == "ddl/logical-tables" {
219            let body = url::form_urlencoded::Serializer::new(String::new())
220                .extend_pairs(&params)
221                .finish();
222            // A streaming body cannot be cloned for reqwest's protocol retries.
223            request = request
224                .timeout(self.timeout)
225                .header(reqwest::header::CONTENT_LENGTH, body.len())
226                .body(reqwest::Body::wrap_stream(futures::stream::once(
227                    std::future::ready(Ok::<_, std::io::Error>(body)),
228                )));
229        } else {
230            request = request.form(&params);
231        }
232        if let Some(ref auth) = self.auth_header {
233            request = request.header("Authorization", auth);
234        }
235
236        request = request.header(
237            GREPTIME_DB_HEADER_TIMEOUT,
238            format_duration(self.timeout).to_string(),
239        );
240
241        let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
242            reason: format!("bad url: {}", url),
243        })?;
244        let status = response.status();
245        let text = response.text().await.with_context(|_| HttpQuerySqlSnafu {
246            reason: "cannot get response text".to_string(),
247        })?;
248        let value: Value = serde_json::from_str(&text).map_err(|_| {
249            crate::error::UnexpectedSnafu {
250                msg: format!("invalid SQL response ({status})"),
251            }
252            .build()
253        })?;
254        if !status.is_success()
255            || value.get("error").is_some()
256            || value
257                .get("code")
258                .is_some_and(|code| code.as_u64() != Some(0))
259        {
260            return crate::error::UnexpectedSnafu {
261                msg: format!(
262                    "SQL request failed ({status}, code {:?})",
263                    value.get("code").and_then(Value::as_u64)
264                ),
265            }
266            .fail();
267        }
268        if let Some(outputs) = value.get("output").and_then(Value::as_array) {
269            for (index, output) in outputs.iter().enumerate() {
270                if output.get("error").is_some()
271                    || output
272                        .get("code")
273                        .is_some_and(|code| code.as_u64() != Some(0))
274                {
275                    return crate::error::UnexpectedSnafu {
276                        msg: format!(
277                            "SQL statement {} failed (code {:?})",
278                            index + 1,
279                            output.get("code").and_then(Value::as_u64)
280                        ),
281                    }
282                    .fail();
283                }
284            }
285        }
286        serde_json::from_value(value).map_err(|_| {
287            crate::error::UnexpectedSnafu {
288                msg: format!("invalid SQL response ({status})"),
289            }
290            .build()
291        })
292    }
293
294    async fn http_client(&self) -> Result<&reqwest::Client> {
295        self.client
296            .get_or_try_init(|| async {
297                let mut builder = reqwest::Client::builder().redirect(
298                    reqwest::redirect::Policy::custom(|attempt| {
299                        if attempt
300                            .previous()
301                            .first()
302                            .is_some_and(|url| url.path() == "/v1/ddl/logical-tables")
303                        {
304                            attempt.stop()
305                        } else {
306                            reqwest::redirect::Policy::default().redirect(attempt)
307                        }
308                    }),
309                );
310                if let Some(proxy) = self.proxy.clone() {
311                    builder = builder.proxy(proxy);
312                }
313                if self.no_proxy {
314                    builder = builder.no_proxy();
315                }
316                builder.build().context(BuildClientSnafu)
317            })
318            .await
319    }
320}
321
322fn form_encoded_len(value: &str) -> usize {
323    url::form_urlencoded::byte_serialize(value.as_bytes())
324        .map(str::len)
325        .sum()
326}
327
328pub(crate) struct Capabilities(Value);
329
330impl Capabilities {
331    pub(crate) fn supports(&self, capability: &str) -> bool {
332        self.0.get(capability).and_then(Value::as_u64) == Some(1)
333    }
334
335    pub(crate) fn require(&self, capability: &str) -> Result<()> {
336        if !self.supports(capability) {
337            return crate::error::InvalidArgumentsSnafu {
338                msg: format!("server does not support {capability} version 1"),
339            }
340            .fail();
341        }
342        Ok(())
343    }
344}
345
346/// Split at `-`.
347pub(crate) fn split_database(database: &str) -> Result<(String, Option<String>)> {
348    let (catalog, schema) = match database.split_once('-') {
349        Some((catalog, schema)) => (catalog, schema),
350        None => (DEFAULT_CATALOG_NAME, database),
351    };
352
353    if schema == "*" {
354        Ok((catalog.to_string(), None))
355    } else {
356        Ok((catalog.to_string(), Some(schema.to_string())))
357    }
358}
359
360#[cfg(test)]
361pub(crate) mod tests {
362    #[tokio::test]
363    async fn packed_capability_probe_checks_auth_and_protocol_version() {
364        use tokio::io::{AsyncReadExt, AsyncWriteExt};
365        for export in [false, true] {
366            for (status, body, supported) in [
367                (200, r#"{"metric_packed_import":1,"future":2}"#, true),
368                (200, "{}", false),
369                (200, r#"{"metric_packed_import":2}"#, false),
370                (404, "{}", false),
371                (401, "{}", false),
372                (403, "{}", false),
373                (200, "[]", false),
374                (200, "invalid-json", false),
375            ] {
376                let body = if export {
377                    body.replace("metric_packed_import", "metric_packed_export")
378                } else {
379                    body.to_string()
380                };
381                let response_body = body.clone();
382                let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
383                let address = listener.local_addr().unwrap();
384                let server = tokio::spawn(async move {
385                    let body = response_body;
386                    let (mut socket, _) = listener.accept().await.unwrap();
387                    let mut request = vec![0; 4096];
388                    let n = socket.read(&mut request).await.unwrap();
389                    let request = String::from_utf8_lossy(&request[..n]).to_lowercase();
390                    assert!(request.starts_with("get /v1/capabilities "));
391                    assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza=="));
392                    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();
393                });
394                let client = super::DatabaseClient::new(
395                    address.to_string(),
396                    "greptime".into(),
397                    Some("user:password".into()),
398                    std::time::Duration::from_secs(5),
399                    None,
400                    true,
401                );
402                assert_eq!(
403                    if export {
404                        client.require_packed_export().await.is_ok()
405                    } else {
406                        client
407                            .capabilities()
408                            .await
409                            .and_then(|c| c.require("metric_packed_import"))
410                            .is_ok()
411                    },
412                    supported,
413                    "{status}: {body}"
414                );
415                server.await.unwrap();
416            }
417        }
418    }
419
420    use super::*;
421
422    #[test]
423    fn form_size_counts_utf8_and_escaping() {
424        for (db, sql) in [
425            ("greptime-public", "SELECT 1"),
426            ("测-试", "+&='秘密';\n"),
427            ("a_b.*", "a~b"),
428        ] {
429            let encoded = url::form_urlencoded::Serializer::new(String::new())
430                .append_pair("db", db)
431                .append_pair("sql", sql)
432                .finish();
433            assert_eq!(
434                8 + form_encoded_len(db) + form_encoded_len(sql),
435                encoded.len()
436            );
437        }
438    }
439
440    #[tokio::test]
441    async fn batch_capabilities_distinguish_absence_from_errors() {
442        for (status, body, expected) in [
443            (200, r#"{"metric_batch_ddl":1}"#, Some(true)),
444            (200, r#"{"metric_batch_ddl":2}"#, Some(false)),
445            (200, "{}", Some(false)),
446            (404, "", Some(false)),
447            (401, "{}", None),
448            (403, "{}", None),
449            (500, "{}", None),
450            (200, "[]", None),
451            (200, "bad", None),
452            (200, r#"{"metric_batch_ddl":"1"}"#, None),
453            (200, r#"{"error":"secret"}"#, None),
454        ] {
455            let (client, requests, server) = test_server(status, body).await;
456            let result = client.capabilities().await;
457            assert_eq!(
458                result.as_ref().ok().map(|c| c.supports("metric_batch_ddl")),
459                expected
460            );
461            assert_eq!(requests.lock().unwrap().len(), 1);
462            server.abort();
463            let _ = server.await;
464        }
465    }
466
467    #[tokio::test]
468    async fn batch_redirects_are_not_replayed() {
469        for status in [303, 307, 308] {
470            let (client, requests, server) = test_server(
471                status,
472                r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#,
473            )
474            .await;
475            assert!(
476                client
477                    .logical_tables("CREATE TABLE t", "public", 1)
478                    .await
479                    .is_err()
480            );
481            assert_eq!(requests.lock().unwrap().len(), 1);
482            assert!(requests.lock().unwrap()[0].starts_with("POST /v1/ddl/logical-tables "));
483            // Ordinary SQL retains the shared client's existing redirect behavior.
484            assert!(client.sql("SELECT 1", "public").await.is_err());
485            assert!(requests.lock().unwrap().len() > 2);
486            server.abort();
487            let _ = server.await;
488        }
489    }
490
491    #[tokio::test]
492    async fn batch_disconnect_and_timeout_are_errors() {
493        use tokio::io::AsyncReadExt;
494        for disconnect in [true, false] {
495            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
496            let mut client = DatabaseClient::new(
497                listener.local_addr().unwrap().to_string(),
498                "greptime".into(),
499                None,
500                Duration::from_millis(100),
501                None,
502                true,
503            );
504            let (stop, stopped) = tokio::sync::oneshot::channel::<()>();
505            let server = tokio::spawn(async move {
506                let (mut socket, _) = listener.accept().await.unwrap();
507                let mut request = [0; 4096];
508                assert!(socket.read(&mut request).await.unwrap() > 0);
509                if !disconnect {
510                    let _ = stopped.await;
511                }
512            });
513            assert!(
514                client
515                    .logical_tables("CREATE TABLE t", "public", 1)
516                    .await
517                    .is_err()
518            );
519            let _ = stop.send(());
520            server.await.unwrap();
521            client.catalog = "测".repeat(500_000);
522            let error = client
523                .logical_tables("CREATE TABLE t", "public", 1)
524                .await
525                .unwrap_err();
526            assert!(format!("{error:?}").contains("4 MiB form limit"));
527        }
528    }
529
530    pub(crate) async fn test_server(
531        status: u16,
532        body: &str,
533    ) -> (
534        DatabaseClient,
535        Arc<std::sync::Mutex<Vec<String>>>,
536        tokio::task::JoinHandle<()>,
537    ) {
538        use tokio::io::{AsyncReadExt, AsyncWriteExt};
539        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
540        let client = DatabaseClient::new(
541            listener.local_addr().unwrap().to_string(),
542            "greptime".into(),
543            Some("user:password".into()),
544            Duration::from_secs(5),
545            None,
546            true,
547        );
548        let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
549        let captured = requests.clone();
550        let body = body.to_string();
551        let server = tokio::spawn(async move {
552            loop {
553                let (mut socket, _) = listener.accept().await.unwrap();
554                let mut bytes = Vec::new();
555                loop {
556                    let mut buf = [0; 4096];
557                    let n = socket.read(&mut buf).await.unwrap();
558                    if n == 0 {
559                        break;
560                    }
561                    bytes.extend_from_slice(&buf[..n]);
562                    if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") {
563                        let headers = String::from_utf8_lossy(&bytes[..end]).to_lowercase();
564                        let length = headers
565                            .lines()
566                            .find_map(|l| l.strip_prefix("content-length: "))
567                            .map(|n| n.parse::<usize>().unwrap())
568                            .unwrap_or(0);
569                        if bytes.len() >= end + 4 + length {
570                            break;
571                        }
572                    }
573                }
574                captured
575                    .lock()
576                    .unwrap()
577                    .push(String::from_utf8(bytes).unwrap());
578                let location = if (300..400).contains(&status) {
579                    "Location: /v1/sql\r\n"
580                } else {
581                    ""
582                };
583                socket.write_all(format!("HTTP/1.1 {status} Response\r\n{location}Content-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
584            }
585        });
586        (client, requests, server)
587    }
588
589    #[test]
590    fn test_split_database() {
591        let result = split_database("catalog-schema").unwrap();
592        assert_eq!(result, ("catalog".to_string(), Some("schema".to_string())));
593
594        let result = split_database("schema").unwrap();
595        assert_eq!(result, ("greptime".to_string(), Some("schema".to_string())));
596
597        let result = split_database("catalog-*").unwrap();
598        assert_eq!(result, ("catalog".to_string(), None));
599
600        let result = split_database("*").unwrap();
601        assert_eq!(result, ("greptime".to_string(), None));
602    }
603}
cli/database.rs

cli/
database.rs
