cli/data/
schema_export.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
15//! Bounded SHOW CREATE requests shared by both export commands.
16
17use common_catalog::consts::DEFAULT_SCHEMA_NAME;
18use servers::http::GreptimeQueryOutput;
19use servers::http::result::greptime_result_v1::GreptimedbV1Response;
20
21use crate::data::sql::escape_sql_identifier;
22use crate::database::DatabaseClient;
23use crate::error::{InvalidArgumentsSnafu, Result, UnexpectedSnafu};
24
25const MAX_STATEMENTS: usize = 128;
26const MAX_SQL_BYTES: usize = 128 * 1024;
27
28/// Appends DDL in dependency order supplied by the caller using bounded requests.
29pub(crate) fn append_schema_ddl<'a>(
30    client: &'a DatabaseClient,
31    catalog: &'a str,
32    schema: &'a str,
33    objects: impl Iterator<Item = (&'static str, Option<&'a str>)> + Send + 'a,
34    ddl: &'a mut String,
35) -> impl std::future::Future<Output = Result<()>> + Send + 'a {
36    let mut statements = objects
37        .map(move |(kind, table)| show_create_sql(kind, catalog, schema, table))
38        .peekable();
39    async move {
40        loop {
41            let batch = next_batch(&mut statements)?;
42            if batch.is_empty() {
43                return Ok(());
44            }
45            let response = client
46                .sql_response(&batch.concat(), DEFAULT_SCHEMA_NAME)
47                .await
48                .map_err(|error| {
49                    UnexpectedSnafu {
50                        msg: format!("SHOW CREATE batch failed for {}: {error}", batch.concat()),
51                    }
52                    .build()
53                })?;
54            append_results(&batch, &response, ddl)?;
55        }
56    }
57}
58
59fn show_create_sql(kind: &str, catalog: &str, schema: &str, table: Option<&str>) -> String {
60    let mut sql = format!(
61        "SHOW CREATE {kind} \"{}\".\"{}\"",
62        escape_sql_identifier(catalog),
63        escape_sql_identifier(schema),
64    );
65    if let Some(table) = table {
66        sql.push_str(".\"");
67        sql.push_str(&escape_sql_identifier(table));
68        sql.push('"');
69    }
70    sql.push_str(";\n");
71    sql
72}
73
74fn next_batch(
75    statements: &mut std::iter::Peekable<impl Iterator<Item = String>>,
76) -> Result<Vec<String>> {
77    let mut batch = Vec::new();
78    let mut bytes = 0;
79    while let Some(sql) = statements.peek() {
80        if sql.len() > MAX_SQL_BYTES {
81            return InvalidArgumentsSnafu {
82                msg: format!(
83                    "SHOW CREATE exceeds {MAX_SQL_BYTES} SQL bytes ({}): {sql}",
84                    sql.len()
85                ),
86            }
87            .fail();
88        }
89        if batch.len() == MAX_STATEMENTS || bytes + sql.len() > MAX_SQL_BYTES {
90            break;
91        }
92        bytes += sql.len();
93        if let Some(sql) = statements.next() {
94            batch.push(sql);
95        }
96    }
97    Ok(batch)
98}
99
100fn append_results(
101    batch: &[String],
102    response: &GreptimedbV1Response,
103    ddl: &mut String,
104) -> Result<()> {
105    let outputs = response.output();
106    if outputs.len() != batch.len() {
107        return UnexpectedSnafu {
108            msg: format!(
109                "SHOW CREATE expected {} results, received {} for {}",
110                batch.len(),
111                outputs.len(),
112                batch.concat()
113            ),
114        }
115        .fail();
116    }
117    for (sql, output) in batch.iter().zip(outputs) {
118        let GreptimeQueryOutput::Records(records) = output else {
119            return UnexpectedSnafu {
120                msg: format!("expected SHOW CREATE records for {sql}"),
121            }
122            .fail();
123        };
124        if records.num_cols() != 2 || records.rows().len() != 1 || records.rows()[0].len() != 2 {
125            return UnexpectedSnafu {
126                msg: format!("expected one two-column SHOW CREATE row for {sql}"),
127            }
128            .fail();
129        }
130        let row = &records.rows()[0];
131        let Some(create) = row[1].as_str().filter(|value| !value.is_empty()) else {
132            return UnexpectedSnafu {
133                msg: format!("expected SHOW CREATE DDL string for {sql}"),
134            }
135            .fail();
136        };
137        if !row[0].is_string() {
138            return UnexpectedSnafu {
139                msg: format!("expected SHOW CREATE object name for {sql}"),
140            }
141            .fail();
142        }
143        ddl.push_str(create);
144        ddl.push_str(";\n");
145    }
146    Ok(())
147}
148
149#[cfg(test)]
150mod tests {
151    use serde_json::{Value, json};
152    use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
153
154    use super::*;
155
156    fn records(name: &str, ddl: &str) -> Value {
157        json!({"records": {"schema": {"column_schemas": [
158            {"name": "Table", "data_type": "String"},
159            {"name": "Create Table", "data_type": "String"}
160        ]}, "rows": [[name, ddl]], "total_rows": 1}})
161    }
162
163    #[test]
164    fn batch_limits_include_quoted_utf8_and_separators() {
165        let special = show_create_sql("TABLE", "cat.name", "s\";库", Some("ta.\";表"));
166        assert_eq!(
167            special,
168            "SHOW CREATE TABLE \"cat.name\".\"s\"\";库\".\"ta.\"\";表\";\n"
169        );
170        let mut statements = std::iter::repeat_n(special.clone(), 129).peekable();
171        assert_eq!(
172            next_batch(&mut statements).unwrap(),
173            vec![special.clone(); 128]
174        );
175        assert_eq!(next_batch(&mut statements).unwrap(), vec![special]);
176        assert!(next_batch(&mut statements).unwrap().is_empty());
177
178        let prefix = show_create_sql("TABLE", "c", "s", Some(""));
179        // Every quote expands to two bytes; this also exercises a multibyte name.
180        let name = format!("表{}", "\"".repeat((MAX_SQL_BYTES - prefix.len() - 3) / 2));
181        let mut exact = show_create_sql("TABLE", "c", "s", Some(&name));
182        let padding = MAX_SQL_BYTES - exact.len();
183        exact = show_create_sql(
184            "TABLE",
185            "c",
186            "s",
187            Some(&format!("{name}{}", "a".repeat(padding))),
188        );
189        assert_eq!(exact.len(), MAX_SQL_BYTES);
190        let small = show_create_sql("TABLE", "c", "s", Some("next"));
191        let mut statements = [exact.clone(), small.clone()].into_iter().peekable();
192        assert_eq!(next_batch(&mut statements).unwrap(), vec![exact.clone()]);
193        assert_eq!(next_batch(&mut statements).unwrap(), vec![small]);
194        let mut oversized = std::iter::once(format!("{exact} ")).peekable();
195        assert!(
196            next_batch(&mut oversized)
197                .unwrap_err()
198                .to_string()
199                .contains("exceeds 131072")
200        );
201
202        let half = show_create_sql(
203            "TABLE",
204            "c",
205            "s",
206            Some(&"a".repeat(MAX_SQL_BYTES / 2 - prefix.len())),
207        );
208        let mut statements = [half.clone(), half.clone(), prefix.clone()]
209            .into_iter()
210            .peekable();
211        assert_eq!(
212            next_batch(&mut statements).unwrap(),
213            vec![half.clone(), half]
214        );
215        assert_eq!(next_batch(&mut statements).unwrap(), vec![prefix]);
216    }
217
218    #[test]
219    fn show_create_requires_exact_record_shape_and_count() {
220        let batch = vec![
221            "SHOW CREATE TABLE first".into(),
222            "SHOW CREATE TABLE middle".into(),
223        ];
224        let valid = records("first", "CREATE TABLE first");
225        for output in [
226            vec![],
227            vec![valid.clone()],
228            vec![valid.clone(); 3],
229            vec![valid.clone(), json!({"affectedrows": 1})],
230            vec![valid.clone(), records("middle", "")],
231            vec![valid.clone(), {
232                let mut r = valid.clone();
233                r["records"]["rows"] = json!([]);
234                r
235            }],
236            vec![valid.clone(), {
237                let mut r = valid.clone();
238                r["records"]["rows"] = json!([["middle", 4]]);
239                r
240            }],
241            vec![valid.clone(), {
242                let mut r = valid.clone();
243                r["records"]["rows"] = json!([["middle", "ddl", "extra"]]);
244                r
245            }],
246        ] {
247            let response =
248                serde_json::from_value(json!({"output": output, "execution_time_ms": 0})).unwrap();
249            let error = append_results(&batch, &response, &mut String::new()).unwrap_err();
250            assert!(error.to_string().contains("middle"), "{error}");
251        }
252    }
253
254    async fn start_mock_sql_server(
255        bodies: Vec<(u16, String)>,
256    ) -> (DatabaseClient, tokio::task::JoinHandle<Vec<String>>) {
257        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
258        let address = listener.local_addr().unwrap();
259        let server = tokio::spawn(async move {
260            tokio::time::timeout(std::time::Duration::from_secs(5), async move {
261                let (socket, _) = listener.accept().await.unwrap();
262                let mut socket = BufReader::new(socket);
263                let mut queries = Vec::new();
264                for (status, body) in bodies {
265                    let mut headers = String::new();
266                    loop {
267                        let mut line = String::new();
268                        if socket.read_line(&mut line).await.unwrap() == 0 {
269                            assert!(headers.is_empty());
270                            return queries;
271                        }
272                        headers.push_str(&line);
273                        if line == "\r\n" {
274                            break;
275                        }
276                    }
277                    let headers = headers.to_lowercase();
278                    assert!(headers.starts_with("post /v1/sql "));
279                    assert!(headers.contains("authorization: basic dxnlcjpwyxnzd29yza=="));
280                    assert!(headers.contains("x-greptime-timeout: 7s"));
281                    let length: usize = headers
282                        .lines()
283                        .find_map(|line| line.strip_prefix("content-length: "))
284                        .unwrap()
285                        .parse()
286                        .unwrap();
287                    let mut form = vec![0; length];
288                    socket.read_exact(&mut form).await.unwrap();
289                    let params: std::collections::HashMap<_, _> =
290                        url::form_urlencoded::parse(&form).into_owned().collect();
291                    assert_eq!(params["db"], "catalog-public");
292                    queries.push(params["sql"].clone());
293                    socket
294                        .get_mut()
295                        .write_all(
296                            format!(
297                                "HTTP/1.1 {status} Response\r\nContent-Length: {}\r\n\r\n{body}",
298                                body.len()
299                            )
300                            .as_bytes(),
301                        )
302                        .await
303                        .unwrap();
304                }
305                queries
306            })
307            .await
308            .expect("mock SQL exchange timed out")
309        });
310        let client = DatabaseClient::new(
311            address.to_string(),
312            "catalog".into(),
313            Some("user:password".into()),
314            std::time::Duration::from_secs(7),
315            None,
316            true,
317        );
318        (client, server)
319    }
320
321    #[tokio::test]
322    async fn batches_preserve_order_context_and_reuse_connection() {
323        let names: Vec<_> = (0..129).map(|i| format!("表.{i}\";x")).collect();
324        let outputs: Vec<_> = names
325            .iter()
326            .map(|name| records(name, &format!("DDL {name}")))
327            .collect();
328        let bodies = outputs
329            .chunks(128)
330            .map(|chunk| {
331                (
332                    200,
333                    json!({"output": chunk, "execution_time_ms": 0}).to_string(),
334                )
335            })
336            .collect();
337        let (client, server) = start_mock_sql_server(bodies).await;
338        let mut ddl = String::new();
339        append_schema_ddl(
340            &client.clone(),
341            "catalog",
342            "schema",
343            names.iter().map(|name| ("TABLE", Some(name.as_str()))),
344            &mut ddl,
345        )
346        .await
347        .unwrap();
348        assert_eq!(
349            ddl,
350            names
351                .iter()
352                .map(|name| format!("DDL {name};\n"))
353                .collect::<String>()
354        );
355        let queries = server.await.unwrap();
356        let expected: Vec<_> = (0..129)
357            .map(|i| format!("SHOW CREATE TABLE \"catalog\".\"schema\".\"表.{i}\"\";x\";\n"))
358            .collect();
359        assert_eq!(
360            queries,
361            vec![expected[..128].concat(), expected[128..].concat()]
362        );
363    }
364
365    #[tokio::test]
366    async fn legacy_command_stops_before_data_after_schema_failure() {
367        use clap::Parser;
368
369        use crate::data::export::ExportCommand;
370
371        let output = |rows: Value, columns: usize| {
372            json!({"records": {"schema": {"column_schemas": (0..columns)
373                .map(|i| json!({"name": i.to_string(), "data_type": "String"}))
374                .collect::<Vec<_>>()}, "rows": rows, "total_rows": rows.as_array().unwrap().len()}})
375        };
376        let databases = output(json!([["public"]]), 1);
377        let mut malformed = records("middle", "CREATE TABLE middle");
378        malformed["records"]["rows"][0][1] = json!(42);
379        let bodies = vec![
380            vec![databases.clone()],
381            vec![records("public", "CREATE DATABASE public")],
382            vec![databases.clone()],
383            vec![output(json!([]), 3)],
384            vec![output(
385                json!([
386                    ["catalog", "public", "first", "BASE TABLE"],
387                    ["catalog", "public", "middle", "BASE TABLE"],
388                    ["catalog", "public", "last", "BASE TABLE"]
389                ]),
390                4,
391            )],
392            vec![
393                records("first", "CREATE TABLE first"),
394                malformed,
395                records("last", "CREATE TABLE last"),
396            ],
397            vec![databases],
398            vec![json!({"affectedrows": 0})],
399        ]
400        .into_iter()
401        .map(|output| {
402            (
403                200,
404                json!({"output": output, "execution_time_ms": 0}).to_string(),
405            )
406        })
407        .collect::<Vec<_>>();
408        for target in ["schema", "all"] {
409            let (client, server) = start_mock_sql_server(bodies.clone()).await;
410            let dir = tempfile::tempdir().unwrap();
411            let command = ExportCommand::parse_from([
412                "export",
413                "--addr",
414                client.addr(),
415                "--database",
416                "catalog-public",
417                "--output-dir",
418                dir.path().to_str().unwrap(),
419                "--target",
420                target,
421                "--auth-basic",
422                "user:password",
423                "--timeout",
424                "7s",
425                "--no-proxy",
426            ]);
427            let tool = command.build().await.unwrap();
428            let result = tool.do_work().await;
429            drop(tool);
430            let queries = server.await.unwrap();
431            assert!(result.unwrap_err().to_string().contains("middle"));
432            assert!(
433                dir.path()
434                    .join("catalog/public/create_database.sql")
435                    .is_file()
436            );
437            assert!(!dir.path().join("catalog/public/create_tables.sql").exists());
438            assert_eq!(queries.len(), 6);
439            assert!(queries[5].contains("SHOW CREATE TABLE \"catalog\".\"public\".\"middle\""));
440            assert!(!queries.iter().any(|sql| sql.starts_with("COPY")));
441        }
442    }
443
444    #[tokio::test]
445    async fn sql_response_errors_do_not_expose_copy_credentials() {
446        let secret = "sentinel-connection-secret";
447        let echoed = "sentinel-response-secret";
448        let sql = format!(
449            "COPY DATABASE public TO 's3://bucket/' CONNECTION (SECRET_ACCESS_KEY='{secret}')"
450        );
451        let bodies = vec![
452            (200, format!("not json {echoed}")),
453            (200, json!({"output": [{"records": echoed}], "execution_time_ms": 0}).to_string()),
454            (400, json!({"code": 3001, "error": echoed}).to_string()),
455            (200, json!({"error": echoed}).to_string()),
456            (200, json!({"code": 3001, "output": [records("first", echoed)], "execution_time_ms": 0}).to_string()),
457            (200, json!({"output": [{"error": echoed}]}).to_string()),
458            (200, json!({"output": [{"code": 3001, "records": records("first", echoed)["records"]}], "execution_time_ms": 0}).to_string()),
459        ];
460        let count = bodies.len();
461        let (client, server) = start_mock_sql_server(bodies).await;
462        for _ in 0..count {
463            let error = client
464                .sql_response(&sql, "public")
465                .await
466                .unwrap_err()
467                .to_string();
468            assert!(error.contains("SQL"), "{error}");
469            assert!(!error.contains(secret), "{error}");
470            assert!(!error.contains(echoed), "{error}");
471            assert!(!error.contains("COPY DATABASE"), "{error}");
472        }
473        assert_eq!(server.await.unwrap().len(), count);
474    }
475
476    #[tokio::test]
477    async fn rejects_http_and_statement_errors_without_retry() {
478        let valid = records("first", "CREATE TABLE first");
479        let bodies = vec![
480            (400, json!({"code": 3001, "error": "middle not found", "execution_time_ms": 0}).to_string()),
481            (200, json!({"code": 3001, "output": [valid.clone(), valid.clone(), valid.clone()], "execution_time_ms": 0}).to_string()),
482            (200, json!({"output": [valid.clone(), {"error": "middle denied"}, valid.clone()], "execution_time_ms": 0}).to_string()),
483            (200, json!({"output": [valid.clone(), {"error": "middle denied", "records": valid["records"]}, valid], "execution_time_ms": 0}).to_string()),
484            (200, json!({"output": [records("first", "DDL"), {"code": 3001, "records": records("middle", "DDL")["records"]}, records("last", "DDL")], "execution_time_ms": 0}).to_string()),
485            (200, r#"{"execution_time_ms":0}"#.into()),
486            (200, "not json".into()),
487        ];
488        let count = bodies.len();
489        let (client, server) = start_mock_sql_server(bodies).await;
490        for _ in 0..count {
491            let mut ddl = String::new();
492            let error = append_schema_ddl(
493                &client,
494                "catalog",
495                "schema",
496                ["first", "middle", "last"]
497                    .into_iter()
498                    .map(|name| ("TABLE", Some(name))),
499                &mut ddl,
500            )
501            .await
502            .unwrap_err();
503            assert!(error.to_string().contains("middle"), "{error}");
504            assert!(ddl.is_empty());
505        }
506        assert_eq!(server.await.unwrap().len(), count);
507    }
508}
cli/data/schema_export.rs

cli/data/
schema_export.rs
