cli/data/import_v2/
executor.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//! DDL execution for import.
16
17use std::collections::HashSet;
18
19use common_catalog::consts::DEFAULT_SCHEMA_NAME;
20use common_telemetry::info;
21use servers::query_handler::sql::{MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS};
22use snafu::ResultExt;
23use sql::dialect::GreptimeDbDialect;
24use sql::parser::{ParseOptions, ParserContext};
25use sql::statements::statement::Statement;
26use store_api::metric_engine_consts::{
27    LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
28};
29
30use crate::data::import_v2::error::{DatabaseSnafu, Result};
31use crate::database::DatabaseClient;
32
33/// A DDL statement with an explicit execution schema context.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct DdlStatement {
36    pub sql: String,
37    pub execution_schema: Option<String>,
38}
39
40impl DdlStatement {
41    pub fn new(sql: String) -> Self {
42        Self {
43            sql,
44            execution_schema: None,
45        }
46    }
47
48    pub fn with_execution_schema(sql: String, schema: String) -> Self {
49        Self {
50            sql,
51            execution_schema: Some(schema),
52        }
53    }
54}
55
56/// Executes DDL statements against the database.
57pub struct DdlExecutor<'a> {
58    client: &'a DatabaseClient,
59}
60
61impl<'a> DdlExecutor<'a> {
62    /// Creates a new DDL executor.
63    pub fn new(client: &'a DatabaseClient) -> Self {
64        Self { client }
65    }
66
67    /// Executes DDL in order, stopping at the first failed request.
68    pub async fn execute_strict(
69        &self,
70        statements: &[DdlStatement],
71        batch_supported: bool,
72    ) -> Result<()> {
73        let mut batch = LogicalBatch::default();
74        for (i, stmt) in statements.iter().enumerate() {
75            info!("Executing DDL ({}/{})", i + 1, statements.len());
76            let schema = stmt
77                .execution_schema
78                .as_deref()
79                .unwrap_or(DEFAULT_SCHEMA_NAME);
80            let candidate = if batch_supported {
81                classify(stmt, self.client.catalog()).context(DatabaseSnafu)?
82            } else {
83                None
84            };
85            if let Some(candidate) = candidate {
86                if !batch.accepts(&candidate, &stmt.sql) {
87                    batch.flush(self.client).await?;
88                }
89                batch.push(candidate, &stmt.sql);
90            } else {
91                batch.flush(self.client).await?;
92                self.client
93                    .sql(&stmt.sql, schema)
94                    .await
95                    .context(DatabaseSnafu)?;
96            }
97        }
98        batch.flush(self.client).await
99    }
100}
101
102struct LogicalTable {
103    catalog: String,
104    execution_schema: String,
105    schema: String,
106    physical: String,
107    name: String,
108}
109
110fn classify(stmt: &DdlStatement, catalog: &str) -> crate::error::Result<Option<LogicalTable>> {
111    if stmt.sql.len() > MAX_LOGICAL_TABLE_DDL_BYTES {
112        return Ok(None);
113    }
114    let parsed = parse_ddl(&stmt.sql)?;
115    let [Statement::CreateTable(create)] = parsed.as_slice() else {
116        return Ok(None);
117    };
118    let Some(physical) = create.options.get(LOGICAL_TABLE_METADATA_KEY) else {
119        return Ok(None);
120    };
121    if create.engine != METRIC_ENGINE_NAME
122        || create.options.value(PHYSICAL_TABLE_METADATA_KEY).is_some()
123    {
124        return Ok(None);
125    }
126    let parts: Vec<_> = create
127        .name
128        .0
129        .iter()
130        .filter_map(|p| p.as_ident().map(|i| i.value.as_str()))
131        .collect();
132    let schema = stmt
133        .execution_schema
134        .as_deref()
135        .unwrap_or(DEFAULT_SCHEMA_NAME);
136    let (catalog, schema, name) = match parts.as_slice() {
137        [name] => (catalog, schema, *name),
138        [schema, name] => (catalog, *schema, *name),
139        [catalog, schema, name] => (*catalog, *schema, *name),
140        _ => return Ok(None),
141    };
142    Ok(Some(LogicalTable {
143        catalog: catalog.into(),
144        execution_schema: stmt
145            .execution_schema
146            .as_deref()
147            .unwrap_or(DEFAULT_SCHEMA_NAME)
148            .into(),
149        schema: schema.into(),
150        physical: physical.into(),
151        name: name.into(),
152    }))
153}
154
155pub(crate) fn parse_ddl(sql: &str) -> crate::error::Result<Vec<Statement>> {
156    ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()).map_err(
157        |_| {
158            crate::error::InvalidArgumentsSnafu {
159                msg: "cannot parse import DDL",
160            }
161            .build()
162        },
163    )
164}
165
166#[derive(Default)]
167struct LogicalBatch {
168    group: Option<LogicalTable>,
169    names: HashSet<String>,
170    sql: String,
171}
172
173impl LogicalBatch {
174    fn accepts(&self, table: &LogicalTable, sql: &str) -> bool {
175        self.group.as_ref().is_none_or(|g| {
176            g.catalog == table.catalog
177                && g.schema == table.schema
178                && g.physical == table.physical
179                && g.execution_schema == table.execution_schema
180        }) && !self.names.contains(&table.name)
181            && self.names.len() < MAX_LOGICAL_TABLE_DDL_STATEMENTS
182            && self.sql.len() + (if self.sql.is_empty() { 0 } else { 2 }) + sql.len()
183                <= MAX_LOGICAL_TABLE_DDL_BYTES
184    }
185
186    fn push(&mut self, table: LogicalTable, sql: &str) {
187        if !self.sql.is_empty() {
188            self.sql.push_str(";\n");
189        }
190        self.sql.push_str(sql);
191        self.names.insert(table.name.clone());
192        self.group = Some(table);
193    }
194
195    async fn flush(&mut self, client: &DatabaseClient) -> Result<()> {
196        if let Some(group) = self.group.as_ref() {
197            client
198                .logical_tables(&self.sql, &group.execution_schema, self.names.len())
199                .await
200                .context(DatabaseSnafu)?;
201            self.sql.clear();
202            self.names.clear();
203            self.group = None;
204        }
205        Ok(())
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    fn logical(name: &str, physical: &str) -> String {
214        format!(
215            "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='{physical}')"
216        )
217    }
218
219    #[test]
220    fn classification_and_batch_boundaries() {
221        let mut batch = LogicalBatch::default();
222        for (name, context, physical, expected) in [
223            ("A", "public", "p.q", true),
224            ("public.b", "other", "p.q", false),
225            ("greptime.public.c", "public", "p.q", true),
226            ("a", "public", "p.q", false),
227            ("\"A\"", "public", "p.q", true),
228            ("d", "public", "p", false),
229            ("d", "other", "p.q", false),
230            ("other.public.d", "public", "p.q", false),
231        ] {
232            let sql = logical(name, physical);
233            let stmt = DdlStatement::with_execution_schema(sql.clone(), context.into());
234            let table = classify(&stmt, "greptime").unwrap().unwrap();
235            assert_eq!(batch.accepts(&table, &sql), expected, "{name}");
236            if expected {
237                batch.push(table, &sql);
238            }
239        }
240        for sql in [
241            "CREATE DATABASE d",
242            "CREATE TABLE t (ts TIMESTAMP TIME INDEX)",
243            "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table='')",
244            "CREATE VIEW v AS SELECT * FROM a",
245        ] {
246            assert!(
247                classify(&DdlStatement::new(sql.into()), "greptime")
248                    .unwrap()
249                    .is_none()
250            );
251        }
252        let mut batch = LogicalBatch::default();
253        for i in 0..128 {
254            let sql = logical(&format!("t{i}"), "p");
255            let table = classify(&DdlStatement::new(sql.clone()), "greptime")
256                .unwrap()
257                .unwrap();
258            assert!(batch.accepts(&table, &sql));
259            batch.push(table, &sql);
260        }
261        let table = classify(&DdlStatement::new(logical("next", "p")), "greptime")
262            .unwrap()
263            .unwrap();
264        assert!(!batch.accepts(&table, &logical("next", "p")));
265        batch.names.clear();
266        batch.sql = "测".repeat(MAX_LOGICAL_TABLE_DDL_BYTES / 3 - 1);
267        assert!(batch.accepts(&table, "ab"));
268        assert!(!batch.accepts(&table, "测"));
269        let oversized =
270            DdlStatement::new(logical("large", &"x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES)));
271        assert!(classify(&oversized, "greptime").unwrap().is_none());
272    }
273
274    #[tokio::test]
275    async fn oversized_ddl_uses_ordinary_sql_in_order() {
276        let large = "x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES);
277        let statements = [
278            DdlStatement::new(logical("before", "p")),
279            DdlStatement::new(format!(
280                "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}')"
281            )),
282            DdlStatement::new(logical("between", "p")),
283            DdlStatement::new(format!("CREATE VIEW v AS SELECT '{large}' AS payload")),
284            DdlStatement::new(logical("oversized", &large)),
285            DdlStatement::new(logical("after", "p")),
286        ];
287        let (client, requests, server) = crate::database::tests::test_server(
288            200,
289            r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#,
290        )
291        .await;
292        DdlExecutor::new(&client)
293            .execute_strict(&statements, true)
294            .await
295            .unwrap();
296        {
297            let requests = requests.lock().unwrap();
298            assert_eq!(requests.len(), statements.len());
299            for (i, (request, statement)) in requests.iter().zip(&statements).enumerate() {
300                let endpoint = if [0, 2, 5].contains(&i) {
301                    "ddl/logical-tables"
302                } else {
303                    "sql"
304                };
305                assert!(request.starts_with(&format!("POST /v1/{endpoint} ")));
306                let body = request.split_once("\r\n\r\n").unwrap().1;
307                let sql = url::form_urlencoded::parse(body.as_bytes())
308                    .find(|(key, _)| key == "sql")
309                    .unwrap()
310                    .1;
311                assert_eq!(sql, statement.sql);
312            }
313        }
314        server.abort();
315        let _ = server.await;
316    }
317
318    #[tokio::test]
319    async fn qualified_targets_preserve_request_context() {
320        let (client, requests, server) = crate::database::tests::test_server(
321            200,
322            r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#,
323        )
324        .await;
325        let statements = [
326            DdlStatement::new(logical("other.metrics.a", "p")),
327            DdlStatement::with_execution_schema(
328                logical("other.metrics.b", "p"),
329                "existing_schema".into(),
330            ),
331        ];
332        DdlExecutor::new(&client)
333            .execute_strict(&statements, true)
334            .await
335            .unwrap();
336        {
337            let requests = requests.lock().unwrap();
338            assert_eq!(requests.len(), 2);
339            for (request, schema) in requests.iter().zip(["public", "existing_schema"]) {
340                assert!(request.starts_with("POST /v1/ddl/logical-tables "));
341                let body = request.split_once("\r\n\r\n").unwrap().1;
342                let db = url::form_urlencoded::parse(body.as_bytes())
343                    .find(|(k, _)| k == "db")
344                    .unwrap()
345                    .1;
346                assert_eq!(db, format!("greptime-{schema}"));
347            }
348        }
349        server.abort();
350        let _ = server.await;
351    }
352
353    #[tokio::test]
354    async fn transport_preserves_order_and_stops_without_fallback() {
355        use crate::database::tests::test_server;
356        let statements = [
357            DdlStatement::new("CREATE DATABASE d".into()),
358            DdlStatement::new(logical("a", "p").replace(
359                "on_physical_table='p'",
360                "on_physical_table='p', access_key_id='original-secret'",
361            )),
362            DdlStatement::new(logical("a", "p")),
363            DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()),
364        ];
365        for supported in [true, false] {
366            let (client, requests, server) = test_server(
367                200,
368                r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#,
369            )
370            .await;
371            DdlExecutor::new(&client)
372                .execute_strict(&statements, supported)
373                .await
374                .unwrap();
375            let requests = requests.lock().unwrap();
376            assert_eq!(requests.len(), 4);
377            for (i, request) in requests.iter().enumerate() {
378                let path = if supported && (i == 1 || i == 2) {
379                    "ddl/logical-tables"
380                } else {
381                    "sql"
382                };
383                assert!(request.starts_with(&format!("POST /v1/{path} ")));
384                assert!(
385                    request
386                        .to_lowercase()
387                        .contains("authorization: basic dxnlcjpwyxnzd29yza==")
388                );
389                let body = request.split_once("\r\n\r\n").unwrap().1;
390                let form: std::collections::HashMap<_, _> =
391                    url::form_urlencoded::parse(body.as_bytes())
392                        .into_owned()
393                        .collect();
394                assert_eq!(form["sql"], statements[i].sql);
395                assert_eq!(form["db"], "greptime-public");
396            }
397            server.abort();
398        }
399        let statements = [
400            DdlStatement::new(logical("a", "p")),
401            DdlStatement::new(logical("b", "p")),
402            DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()),
403        ];
404        for (status, body, success) in [
405            (
406                200,
407                r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":0}]}"#,
408                true,
409            ),
410            (
411                200,
412                r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#,
413                false,
414            ),
415            (
416                200,
417                r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":1}]}"#,
418                false,
419            ),
420            (200, r#"{"execution_time_ms":0,"output":[{},{}]}"#, false),
421            (
422                200,
423                r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"error":"secret"}]}"#,
424                false,
425            ),
426            (500, r#"{"error":"secret"}"#, false),
427            (200, "invalid", false),
428        ] {
429            let (client, requests, server) = test_server(status, body).await;
430            let result = DdlExecutor::new(&client)
431                .execute_strict(&statements, true)
432                .await;
433            assert_eq!(result.is_ok(), success);
434            assert_eq!(requests.lock().unwrap().len(), if success { 2 } else { 1 });
435            if let Err(error) = result {
436                assert!(!format!("{error:?}").contains("secret"));
437            }
438            server.abort();
439            let _ = server.await;
440        }
441    }
442}
cli/data/import_v2/executor.rs

cli/data/import_v2/
executor.rs
