From 6265c31e4afa4e064fd0f266f6756cae4853a1db Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 8 Oct 2026 13:14:06 +0000 Subject: [PATCH] feat: batch logical table DDL during import-v2 (#9467) * feat: batch logical table DDL during import-v2 Signed-off-by: jeremyhi * fix: preserve one-shot batch DDL transport and context Signed-off-by: jeremyhi * test: check schema-only import state at its own path Signed-off-by: jeremyhi * fix: preserve oversized import DDL compatibility Signed-off-by: jeremyhi * refactor: reuse import DDL encoding and metric constants Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- src/cli/src/data/import_v2/command.rs | 166 +++---- src/cli/src/data/import_v2/executor.rs | 406 ++++++++++++++++-- src/cli/src/data/import_v2/packed.rs | 85 ++-- src/cli/src/database.rs | 308 ++++++++++++- .../tests/export_logical_tables.rs | 129 +++++- 5 files changed, 904 insertions(+), 190 deletions(-) diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 33cf9cb1ffb..62b9de08816 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -222,12 +222,19 @@ impl Import { manifest.validate_layout().map_err(|reason| { crate::data::import_v2::error::InvalidPackedSnapshotSnafu { reason }.build() })?; - if manifest.is_packed() { - self.database_client - .require_packed_import() + let capabilities = if manifest.is_packed() { + let capabilities = self + .database_client + .capabilities() .await .context(crate::data::import_v2::error::DatabaseSnafu)?; - } + capabilities + .require("metric_packed_import") + .context(crate::data::import_v2::error::DatabaseSnafu)?; + Some(capabilities) + } else { + None + }; info!("Snapshot contains {} schema(s)", manifest.schemas.len()); @@ -279,6 +286,15 @@ impl Import { return Ok(()); } + let capabilities = match capabilities { + Some(capabilities) => capabilities, + None => self + .database_client + .capabilities() + .await + .context(crate::data::import_v2::error::DatabaseSnafu)?, + }; + let mut resume_session = if !data_tasks.is_empty() { let state_path = resolve_state_path( self.state_path.as_deref(), @@ -316,7 +332,9 @@ impl Import { false } else { let executor = DdlExecutor::new(&self.database_client); - executor.execute_strict(&ddl_statements).await?; + executor + .execute_strict(&ddl_statements, capabilities.supports("metric_batch_ddl")) + .await?; if let Some(session) = resume_session.as_mut() { session.mark_ddl_completed().await?; } @@ -356,9 +374,7 @@ impl Import { .await .context(SnapshotStorageSnafu)?; statements.extend( - parse_ddl_statements(&content) - .into_iter() - .map(|sql| ddl_statement_for_schema(schema, sql)), + iter_ddl_statements(&content).map(|sql| ddl_statement_for_schema(schema, sql)), ); } @@ -402,90 +418,92 @@ impl ImportTaskExecutor for CopyDatabaseImportTaskExecutor<'_> { } } +#[cfg(test)] fn parse_ddl_statements(content: &str) -> Vec { - let mut statements = Vec::new(); - let mut current = String::new(); + iter_ddl_statements(content).collect() +} + +pub(crate) fn iter_ddl_statements(content: &str) -> impl Iterator + '_ { let mut chars = content.chars().peekable(); let mut in_single_quote = false; let mut in_double_quote = false; let mut in_line_comment = false; let mut in_block_comment = false; - while let Some(ch) = chars.next() { - if in_line_comment { - if ch == '\n' { - in_line_comment = false; - current.push('\n'); - } - continue; - } - - if in_block_comment { - if ch == '*' && chars.peek() == Some(&'/') { - chars.next(); - in_block_comment = false; - } - continue; - } - - if in_single_quote { - current.push(ch); - if ch == '\'' { - if chars.peek() == Some(&'\'') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_single_quote = false; + std::iter::from_fn(move || { + let mut current = String::new(); + while let Some(ch) = chars.next() { + if in_line_comment { + if ch == '\n' { + in_line_comment = false; + current.push('\n'); } + continue; } - continue; - } - if in_double_quote { - current.push(ch); - if ch == '"' { - if chars.peek() == Some(&'"') { - current.push(chars.next().expect("peeked quote must exist")); - } else { - in_double_quote = false; + if in_block_comment { + if ch == '*' && chars.peek() == Some(&'/') { + chars.next(); + in_block_comment = false; } + continue; } - continue; - } - match ch { - '-' if chars.peek() == Some(&'-') => { - chars.next(); - in_line_comment = true; - } - '/' if chars.peek() == Some(&'*') => { - chars.next(); - in_block_comment = true; - } - '\'' => { - in_single_quote = true; + if in_single_quote { current.push(ch); - } - '"' => { - in_double_quote = true; - current.push(ch); - } - ';' => { - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); + if ch == '\'' { + if chars.peek() == Some(&'\'') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_single_quote = false; + } } - current.clear(); + continue; + } + + if in_double_quote { + current.push(ch); + if ch == '"' { + if chars.peek() == Some(&'"') { + current.push(chars.next().expect("peeked quote must exist")); + } else { + in_double_quote = false; + } + } + continue; + } + + match ch { + '-' if chars.peek() == Some(&'-') => { + chars.next(); + in_line_comment = true; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + in_block_comment = true; + } + '\'' => { + in_single_quote = true; + current.push(ch); + } + '"' => { + in_double_quote = true; + current.push(ch); + } + ';' => { + let statement = current.trim(); + if !statement.is_empty() { + return Some(statement.to_string()); + } + current.clear(); + } + _ => current.push(ch), } - _ => current.push(ch), } - } - let statement = current.trim(); - if !statement.is_empty() { - statements.push(statement.to_string()); - } - - statements + let statement = current.trim(); + (!statement.is_empty()).then(|| statement.to_string()) + }) } fn ddl_statement_for_schema(schema: &str, sql: String) -> DdlStatement { diff --git a/src/cli/src/data/import_v2/executor.rs b/src/cli/src/data/import_v2/executor.rs index 3f2bf66ae68..6f152bc5c6d 100644 --- a/src/cli/src/data/import_v2/executor.rs +++ b/src/cli/src/data/import_v2/executor.rs @@ -14,8 +14,18 @@ //! DDL execution for import. +use std::collections::HashSet; + +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_telemetry::info; +use servers::query_handler::sql::{MAX_LOGICAL_TABLE_DDL_BYTES, MAX_LOGICAL_TABLE_DDL_STATEMENTS}; use snafu::ResultExt; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::statement::Statement; +use store_api::metric_engine_consts::{ + LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, +}; use crate::data::import_v2::error::{DatabaseSnafu, Result}; use crate::database::DatabaseClient; @@ -54,69 +64,379 @@ impl<'a> DdlExecutor<'a> { Self { client } } - /// Executes a list of DDL statements, stopping on first error. - pub async fn execute_strict(&self, statements: &[DdlStatement]) -> Result<()> { - let total = statements.len(); - + /// Executes DDL in order, stopping at the first failed request. + pub async fn execute_strict( + &self, + statements: &[DdlStatement], + batch_supported: bool, + ) -> Result<()> { + let mut batch = LogicalBatch::default(); for (i, stmt) in statements.iter().enumerate() { - let preview = preview_sql(&stmt.sql); - - info!("Executing DDL ({}/{}): {}", i + 1, total, preview); - - if let Some(schema) = stmt.execution_schema.as_deref() { + info!("Executing DDL ({}/{})", i + 1, statements.len()); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let candidate = if batch_supported { + classify(stmt, self.client.catalog()).context(DatabaseSnafu)? + } else { + None + }; + if let Some(candidate) = candidate { + if !batch.accepts(&candidate, &stmt.sql) { + batch.flush(self.client).await?; + } + batch.push(candidate, &stmt.sql); + } else { + batch.flush(self.client).await?; self.client .sql(&stmt.sql, schema) .await .context(DatabaseSnafu)?; - } else { - self.client - .sql_in_public(&stmt.sql) - .await - .context(DatabaseSnafu)?; } } + batch.flush(self.client).await + } +} +struct LogicalTable { + catalog: String, + execution_schema: String, + schema: String, + physical: String, + name: String, +} + +fn classify(stmt: &DdlStatement, catalog: &str) -> crate::error::Result> { + if stmt.sql.len() > MAX_LOGICAL_TABLE_DDL_BYTES { + return Ok(None); + } + let parsed = parse_ddl(&stmt.sql)?; + let [Statement::CreateTable(create)] = parsed.as_slice() else { + return Ok(None); + }; + let Some(physical) = create.options.get(LOGICAL_TABLE_METADATA_KEY) else { + return Ok(None); + }; + if create.engine != METRIC_ENGINE_NAME + || create.options.value(PHYSICAL_TABLE_METADATA_KEY).is_some() + { + return Ok(None); + } + let parts: Vec<_> = create + .name + .0 + .iter() + .filter_map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let schema = stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME); + let (catalog, schema, name) = match parts.as_slice() { + [name] => (catalog, schema, *name), + [schema, name] => (catalog, *schema, *name), + [catalog, schema, name] => (*catalog, *schema, *name), + _ => return Ok(None), + }; + Ok(Some(LogicalTable { + catalog: catalog.into(), + execution_schema: stmt + .execution_schema + .as_deref() + .unwrap_or(DEFAULT_SCHEMA_NAME) + .into(), + schema: schema.into(), + physical: physical.into(), + name: name.into(), + })) +} + +pub(crate) fn parse_ddl(sql: &str) -> crate::error::Result> { + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()).map_err( + |_| { + crate::error::InvalidArgumentsSnafu { + msg: "cannot parse import DDL", + } + .build() + }, + ) +} + +#[derive(Default)] +struct LogicalBatch { + group: Option, + names: HashSet, + sql: String, +} + +impl LogicalBatch { + fn accepts(&self, table: &LogicalTable, sql: &str) -> bool { + self.group.as_ref().is_none_or(|g| { + g.catalog == table.catalog + && g.schema == table.schema + && g.physical == table.physical + && g.execution_schema == table.execution_schema + }) && !self.names.contains(&table.name) + && self.names.len() < MAX_LOGICAL_TABLE_DDL_STATEMENTS + && self.sql.len() + (if self.sql.is_empty() { 0 } else { 2 }) + sql.len() + <= MAX_LOGICAL_TABLE_DDL_BYTES + } + + fn push(&mut self, table: LogicalTable, sql: &str) { + if !self.sql.is_empty() { + self.sql.push_str(";\n"); + } + self.sql.push_str(sql); + self.names.insert(table.name.clone()); + self.group = Some(table); + } + + async fn flush(&mut self, client: &DatabaseClient) -> Result<()> { + if let Some(group) = self.group.as_ref() { + client + .logical_tables(&self.sql, &group.execution_schema, self.names.len()) + .await + .context(DatabaseSnafu)?; + self.sql.clear(); + self.names.clear(); + self.group = None; + } Ok(()) } } -fn preview_sql(sql: &str) -> String { - let mut chars = sql.chars(); - let preview: String = chars.by_ref().take(80).collect(); - if chars.next().is_some() { - format!("{preview}...") - } else { - preview - } -} #[cfg(test)] mod tests { use super::*; - #[test] - fn test_statement_without_execution_schema_uses_public() { - let stmt = DdlStatement::new("CREATE DATABASE IF NOT EXISTS test_db".to_string()); - assert_eq!(stmt.execution_schema, None); + fn logical(name: &str, physical: &str) -> String { + format!( + "CREATE TABLE IF NOT EXISTS {name} (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='{physical}')" + ) } #[test] - fn test_statement_with_execution_schema_preserves_context() { - let stmt = DdlStatement::with_execution_schema( - r#"CREATE TABLE IF NOT EXISTS "my""schema"."metrics" (ts TIMESTAMP TIME INDEX)"# - .to_string(), - r#"my"schema"#.to_string(), - ); - assert_eq!(stmt.execution_schema.as_deref(), Some(r#"my"schema"#)); + fn classification_and_batch_boundaries() { + let mut batch = LogicalBatch::default(); + for (name, context, physical, expected) in [ + ("A", "public", "p.q", true), + ("public.b", "other", "p.q", false), + ("greptime.public.c", "public", "p.q", true), + ("a", "public", "p.q", false), + ("\"A\"", "public", "p.q", true), + ("d", "public", "p", false), + ("d", "other", "p.q", false), + ("other.public.d", "public", "p.q", false), + ] { + let sql = logical(name, physical); + let stmt = DdlStatement::with_execution_schema(sql.clone(), context.into()); + let table = classify(&stmt, "greptime").unwrap().unwrap(); + assert_eq!(batch.accepts(&table, &sql), expected, "{name}"); + if expected { + batch.push(table, &sql); + } + } + for sql in [ + "CREATE DATABASE d", + "CREATE TABLE t (ts TIMESTAMP TIME INDEX)", + "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table='')", + "CREATE VIEW v AS SELECT * FROM a", + ] { + assert!( + classify(&DdlStatement::new(sql.into()), "greptime") + .unwrap() + .is_none() + ); + } + let mut batch = LogicalBatch::default(); + for i in 0..128 { + let sql = logical(&format!("t{i}"), "p"); + let table = classify(&DdlStatement::new(sql.clone()), "greptime") + .unwrap() + .unwrap(); + assert!(batch.accepts(&table, &sql)); + batch.push(table, &sql); + } + let table = classify(&DdlStatement::new(logical("next", "p")), "greptime") + .unwrap() + .unwrap(); + assert!(!batch.accepts(&table, &logical("next", "p"))); + batch.names.clear(); + batch.sql = "测".repeat(MAX_LOGICAL_TABLE_DDL_BYTES / 3 - 1); + assert!(batch.accepts(&table, "ab")); + assert!(!batch.accepts(&table, "测")); + let oversized = + DdlStatement::new(logical("large", &"x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES))); + assert!(classify(&oversized, "greptime").unwrap().is_none()); } - #[test] - fn test_preview_sql_truncates_at_char_boundary() { - let sql = format!( - "CREATE TABLE {} (ts TIMESTAMP TIME INDEX)", - "测".repeat(100) - ); - let preview = preview_sql(&sql); - assert!(preview.ends_with("...")); - assert!(preview.is_char_boundary(preview.len())); + #[tokio::test] + async fn oversized_ddl_uses_ordinary_sql_in_order() { + let large = "x".repeat(MAX_LOGICAL_TABLE_DDL_BYTES); + let statements = [ + DdlStatement::new(logical("before", "p")), + DdlStatement::new(format!( + "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}')" + )), + DdlStatement::new(logical("between", "p")), + DdlStatement::new(format!("CREATE VIEW v AS SELECT '{large}' AS payload")), + DdlStatement::new(logical("oversized", &large)), + DdlStatement::new(logical("after", "p")), + ]; + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), statements.len()); + for (i, (request, statement)) in requests.iter().zip(&statements).enumerate() { + let endpoint = if [0, 2, 5].contains(&i) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{endpoint} "))); + let body = request.split_once("\r\n\r\n").unwrap().1; + let sql = url::form_urlencoded::parse(body.as_bytes()) + .find(|(key, _)| key == "sql") + .unwrap() + .1; + assert_eq!(sql, statement.sql); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn qualified_targets_preserve_request_context() { + let (client, requests, server) = crate::database::tests::test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + let statements = [ + DdlStatement::new(logical("other.metrics.a", "p")), + DdlStatement::with_execution_schema( + logical("other.metrics.b", "p"), + "existing_schema".into(), + ), + ]; + DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await + .unwrap(); + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + for (request, schema) in requests.iter().zip(["public", "existing_schema"]) { + assert!(request.starts_with("POST /v1/ddl/logical-tables ")); + let body = request.split_once("\r\n\r\n").unwrap().1; + let db = url::form_urlencoded::parse(body.as_bytes()) + .find(|(k, _)| k == "db") + .unwrap() + .1; + assert_eq!(db, format!("greptime-{schema}")); + } + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn transport_preserves_order_and_stops_without_fallback() { + use crate::database::tests::test_server; + let statements = [ + DdlStatement::new("CREATE DATABASE d".into()), + DdlStatement::new(logical("a", "p").replace( + "on_physical_table='p'", + "on_physical_table='p', access_key_id='original-secret'", + )), + DdlStatement::new(logical("a", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for supported in [true, false] { + let (client, requests, server) = test_server( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + DdlExecutor::new(&client) + .execute_strict(&statements, supported) + .await + .unwrap(); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 4); + for (i, request) in requests.iter().enumerate() { + let path = if supported && (i == 1 || i == 2) { + "ddl/logical-tables" + } else { + "sql" + }; + assert!(request.starts_with(&format!("POST /v1/{path} "))); + assert!( + request + .to_lowercase() + .contains("authorization: basic dxnlcjpwyxnzd29yza==") + ); + let body = request.split_once("\r\n\r\n").unwrap().1; + let form: std::collections::HashMap<_, _> = + url::form_urlencoded::parse(body.as_bytes()) + .into_owned() + .collect(); + assert_eq!(form["sql"], statements[i].sql); + assert_eq!(form["db"], "greptime-public"); + } + server.abort(); + } + let statements = [ + DdlStatement::new(logical("a", "p")), + DdlStatement::new(logical("b", "p")), + DdlStatement::new("CREATE VIEW v AS SELECT * FROM a".into()), + ]; + for (status, body, success) in [ + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":0}]}"#, + true, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + false, + ), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"affectedrows":1}]}"#, + false, + ), + (200, r#"{"execution_time_ms":0,"output":[{},{}]}"#, false), + ( + 200, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0},{"error":"secret"}]}"#, + false, + ), + (500, r#"{"error":"secret"}"#, false), + (200, "invalid", false), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = DdlExecutor::new(&client) + .execute_strict(&statements, true) + .await; + assert_eq!(result.is_ok(), success); + assert_eq!(requests.lock().unwrap().len(), if success { 2 } else { 1 }); + if let Err(error) = result { + assert!(!format!("{error:?}").contains("secret")); + } + server.abort(); + let _ = server.await; + } } } diff --git a/src/cli/src/data/import_v2/packed.rs b/src/cli/src/data/import_v2/packed.rs index 303c866ba69..c2f1ec35889 100644 --- a/src/cli/src/data/import_v2/packed.rs +++ b/src/cli/src/data/import_v2/packed.rs @@ -16,12 +16,11 @@ use std::collections::HashSet; use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; use snafu::ResultExt; -use sql::dialect::GreptimeDbDialect; -use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use crate::data::export_v2::manifest::{ChunkStatus, Manifest}; use crate::data::import_v2::error::{InvalidPackedSnapshotSnafu, Result, SnapshotStorageSnafu}; +use crate::data::import_v2::{command, executor}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; use crate::data::snapshot_storage::SnapshotStorage; @@ -51,49 +50,46 @@ pub(crate) async fn validate_snapshot( .read_text(&ddl_path_for_schema(schema)) .await .context(SnapshotStorageSnafu)?; - let statements = ParserContext::create_with_dialect( - &ddl, - &GreptimeDbDialect {}, - ParseOptions::default(), - ) - .map_err(|e| { - InvalidPackedSnapshotSnafu { - reason: e.to_string(), - } - .build() - })?; let mut names = HashSet::new(); - for statement in statements { - if let Statement::CreateTable(create) = statement { - if create.engine.eq_ignore_ascii_case("metric") - && create.options.get("on_physical_table").is_none() - { - continue; + for sql in command::iter_ddl_statements(&ddl) { + let statements = executor::parse_ddl(&sql).map_err(|_| { + InvalidPackedSnapshotSnafu { + reason: "invalid snapshot DDL", } - let parts: Vec<_> = create - .name - .0 - .iter() - .map(|p| p.as_ident().map(|i| i.value.as_str())) - .collect(); - let allowed = match parts.as_slice() { - [Some(_)] => true, - [Some(s), Some(_)] => *s == schema, - [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, - _ => false, - }; - if !allowed { - return InvalidPackedSnapshotSnafu { - reason: "snapshot table DDL escapes its catalog/schema", + .build() + })?; + for statement in statements { + if let Statement::CreateTable(create) = statement { + if create.engine.eq_ignore_ascii_case("metric") + && create.options.get("on_physical_table").is_none() + { + continue; } - .fail(); - } - let name = parts.last().and_then(|p| *p).unwrap_or_default(); - if !names.insert(name.to_string()) { - return InvalidPackedSnapshotSnafu { - reason: "duplicate snapshot table DDL", + let parts: Vec<_> = create + .name + .0 + .iter() + .map(|p| p.as_ident().map(|i| i.value.as_str())) + .collect(); + let allowed = match parts.as_slice() { + [Some(_)] => true, + [Some(s), Some(_)] => *s == schema, + [Some(c), Some(s), Some(_)] => *c == manifest.catalog && *s == schema, + _ => false, + }; + if !allowed { + return InvalidPackedSnapshotSnafu { + reason: "snapshot table DDL escapes its catalog/schema", + } + .fail(); + } + let name = parts.last().and_then(|p| *p).unwrap_or_default(); + if !names.insert(name.to_string()) { + return InvalidPackedSnapshotSnafu { + reason: "duplicate snapshot table DDL", + } + .fail(); } - .fail(); } } } @@ -180,13 +176,14 @@ mod tests { .unwrap(); let uri = url::Url::from_file_path(dir.path()).unwrap(); let storage = OpenDalStorage::from_uri(uri.as_str(), &Default::default()).unwrap(); - storage.write_text("schema/ddl/public.sql", "CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\";").await.unwrap(); + let large = "x".repeat(servers::query_handler::sql::MAX_LOGICAL_TABLE_DDL_BYTES); + storage.write_text("schema/ddl/public.sql", &format!("CREATE TABLE p (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(physical_metric_table=''); CREATE TABLE \"logical.name\" (ts TIMESTAMP TIME INDEX) ENGINE=metric WITH(on_physical_table='p'); CREATE VIEW v AS SELECT * FROM \"logical.name\"; CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, payload STRING DEFAULT '{large}'); CREATE VIEW large_view AS SELECT '{large}' AS payload;")).await.unwrap(); let path = "data/public/1/pack-index.json"; - let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":12}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0}]}); + let index = serde_json::json!({"version":1,"objects":[{"path":"pack-0.bin","kind":"pack","length":24}],"tables":[{"table_name":"logical.name","object":"pack-0.bin","offset":0,"length":12,"row_count":0},{"table_name":"ordinary","object":"pack-0.bin","offset":12,"length":12,"row_count":0}]}); storage.write_text(path, &index.to_string()).await.unwrap(); assert!(dir.path().join(path).is_file()); storage - .write_text("data/public/1/pack-0.bin", "abcdefghijkl") + .write_text("data/public/1/pack-0.bin", "abcdefghijklmnopqrstuvwx") .await .unwrap(); let mut manifest = Manifest::new_schema_only("greptime".into(), vec!["public".into()]); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index da277e166be..1efe6d7d0a6 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -26,9 +26,7 @@ use servers::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use snafu::ResultExt; -use crate::error::{ - BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result, SerdeJsonSnafu, -}; +use crate::error::{BuildClientSnafu, HttpQuerySqlSnafu, ParseProxyOptsSnafu, Result}; #[derive(Debug, Clone)] pub struct DatabaseClient { @@ -100,49 +98,84 @@ impl DatabaseClient { self.sql(sql, DEFAULT_SCHEMA_NAME).await } - /// Requires the explicit packed-import protocol before any restore mutation. - pub async fn require_packed_import(&self) -> Result<()> { - self.require_packed_capability("metric_packed_import").await + pub(crate) fn catalog(&self) -> &str { + &self.catalog } /// Requires packed export support before creating snapshot artifacts. pub async fn require_packed_export(&self) -> Result<()> { - self.require_packed_capability("metric_packed_export").await + self.capabilities().await?.require("metric_packed_export") } - async fn require_packed_capability(&self, capability: &str) -> Result<()> { + /// Reads target capabilities once; older targets advertise no extensions. + pub(crate) async fn capabilities(&self) -> Result { let url = format!("http://{}/v1/capabilities", self.addr); let mut request = self.http_client().await?.get(&url).timeout(self.timeout); if let Some(auth) = &self.auth_header { request = request.header("Authorization", auth); } let response = request.send().await.with_context(|_| HttpQuerySqlSnafu { - reason: "packed snapshot capability request failed", + reason: "capability request failed", })?; if response.status() == reqwest::StatusCode::NOT_FOUND { - return crate::error::InvalidArgumentsSnafu { - msg: format!("server does not support {capability}"), - } - .fail(); + return Ok(Capabilities(serde_json::json!({}))); } let response = response .error_for_status() .with_context(|_| HttpQuerySqlSnafu { - reason: "packed snapshot capability request rejected", + reason: "capability request rejected", })?; let body = response.text().await.with_context(|_| HttpQuerySqlSnafu { reason: "cannot read capability response", })?; - let value: Value = serde_json::from_str(&body).context(SerdeJsonSnafu)?; - if !value.is_object() { + let value: Value = serde_json::from_str(&body).map_err(|_| { + crate::error::UnexpectedSnafu { + msg: "invalid capability response", + } + .build() + })?; + if !value.is_object() + || value.get("error").is_some() + || value.get("code").is_some_and(|v| v.as_u64() != Some(0)) + { return crate::error::UnexpectedSnafu { msg: "invalid capability response: expected JSON object", } .fail(); } - if value.get(capability).and_then(Value::as_u64) != Some(1) { + for key in [ + "metric_batch_ddl", + "metric_packed_import", + "metric_packed_export", + ] { + if value.get(key).is_some_and(|v| v.as_u64().is_none()) { + return crate::error::UnexpectedSnafu { + msg: "invalid capability version", + } + .fail(); + } + } + Ok(Capabilities(value)) + } + + /// Submits a batch once. An ambiguous failure must not be replayed here. + pub(crate) async fn logical_tables(&self, sql: &str, schema: &str, count: usize) -> Result<()> { + let db = format!("{}-{}", self.catalog, schema); + if 8 + form_encoded_len(&db) + form_encoded_len(sql) > 4 * 1024 * 1024 { return crate::error::InvalidArgumentsSnafu { - msg: format!("server does not support {capability} version 1"), + msg: "logical-table batch exceeds 4 MiB form limit", + } + .fail(); + } + let response = self.ddl_response(sql, schema, "ddl/logical-tables").await?; + if response.output().len() != count + || response + .output() + .iter() + .any(|output| !matches!(output, GreptimeQueryOutput::AffectedRows(0))) + { + return crate::error::UnexpectedSnafu { + msg: "invalid logical-table batch output", } .fail(); } @@ -163,7 +196,16 @@ impl DatabaseClient { sql: &str, schema: &str, ) -> Result { - let url = format!("http://{}/v1/sql", self.addr); + self.ddl_response(sql, schema, "sql").await + } + + async fn ddl_response( + &self, + sql: &str, + schema: &str, + endpoint: &str, + ) -> Result { + let url = format!("http://{}/v1/{endpoint}", self.addr); let params = [ ("db", format!("{}-{}", self.catalog, schema)), ("sql", sql.to_string()), @@ -172,8 +214,21 @@ impl DatabaseClient { .http_client() .await? .post(&url) - .form(¶ms) .header("Content-Type", "application/x-www-form-urlencoded"); + if endpoint == "ddl/logical-tables" { + let body = url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(¶ms) + .finish(); + // A streaming body cannot be cloned for reqwest's protocol retries. + request = request + .timeout(self.timeout) + .header(reqwest::header::CONTENT_LENGTH, body.len()) + .body(reqwest::Body::wrap_stream(futures::stream::once( + std::future::ready(Ok::<_, std::io::Error>(body)), + ))); + } else { + request = request.form(¶ms); + } if let Some(ref auth) = self.auth_header { request = request.header("Authorization", auth); } @@ -239,7 +294,19 @@ impl DatabaseClient { async fn http_client(&self) -> Result<&reqwest::Client> { self.client .get_or_try_init(|| async { - let mut builder = reqwest::Client::builder(); + let mut builder = reqwest::Client::builder().redirect( + reqwest::redirect::Policy::custom(|attempt| { + if attempt + .previous() + .first() + .is_some_and(|url| url.path() == "/v1/ddl/logical-tables") + { + attempt.stop() + } else { + reqwest::redirect::Policy::default().redirect(attempt) + } + }), + ); if let Some(proxy) = self.proxy.clone() { builder = builder.proxy(proxy); } @@ -252,6 +319,30 @@ impl DatabaseClient { } } +fn form_encoded_len(value: &str) -> usize { + url::form_urlencoded::byte_serialize(value.as_bytes()) + .map(str::len) + .sum() +} + +pub(crate) struct Capabilities(Value); + +impl Capabilities { + pub(crate) fn supports(&self, capability: &str) -> bool { + self.0.get(capability).and_then(Value::as_u64) == Some(1) + } + + pub(crate) fn require(&self, capability: &str) -> Result<()> { + if !self.supports(capability) { + return crate::error::InvalidArgumentsSnafu { + msg: format!("server does not support {capability} version 1"), + } + .fail(); + } + Ok(()) + } +} + /// Split at `-`. pub(crate) fn split_database(database: &str) -> Result<(String, Option)> { let (catalog, schema) = match database.split_once('-') { @@ -267,7 +358,7 @@ pub(crate) fn split_database(database: &str) -> Result<(String, Option)> } #[cfg(test)] -mod tests { +pub(crate) mod tests { #[tokio::test] async fn packed_capability_probe_checks_auth_and_protocol_version() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -312,7 +403,11 @@ mod tests { if export { client.require_packed_export().await.is_ok() } else { - client.require_packed_import().await.is_ok() + client + .capabilities() + .await + .and_then(|c| c.require("metric_packed_import")) + .is_ok() }, supported, "{status}: {body}" @@ -324,6 +419,173 @@ mod tests { use super::*; + #[test] + fn form_size_counts_utf8_and_escaping() { + for (db, sql) in [ + ("greptime-public", "SELECT 1"), + ("测-试", "+&='秘密';\n"), + ("a_b.*", "a~b"), + ] { + let encoded = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("db", db) + .append_pair("sql", sql) + .finish(); + assert_eq!( + 8 + form_encoded_len(db) + form_encoded_len(sql), + encoded.len() + ); + } + } + + #[tokio::test] + async fn batch_capabilities_distinguish_absence_from_errors() { + for (status, body, expected) in [ + (200, r#"{"metric_batch_ddl":1}"#, Some(true)), + (200, r#"{"metric_batch_ddl":2}"#, Some(false)), + (200, "{}", Some(false)), + (404, "", Some(false)), + (401, "{}", None), + (403, "{}", None), + (500, "{}", None), + (200, "[]", None), + (200, "bad", None), + (200, r#"{"metric_batch_ddl":"1"}"#, None), + (200, r#"{"error":"secret"}"#, None), + ] { + let (client, requests, server) = test_server(status, body).await; + let result = client.capabilities().await; + assert_eq!( + result.as_ref().ok().map(|c| c.supports("metric_batch_ddl")), + expected + ); + assert_eq!(requests.lock().unwrap().len(), 1); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_redirects_are_not_replayed() { + for status in [303, 307, 308] { + let (client, requests, server) = test_server( + status, + r#"{"execution_time_ms":0,"output":[{"affectedrows":0}]}"#, + ) + .await; + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + assert_eq!(requests.lock().unwrap().len(), 1); + assert!(requests.lock().unwrap()[0].starts_with("POST /v1/ddl/logical-tables ")); + // Ordinary SQL retains the shared client's existing redirect behavior. + assert!(client.sql("SELECT 1", "public").await.is_err()); + assert!(requests.lock().unwrap().len() > 2); + server.abort(); + let _ = server.await; + } + } + + #[tokio::test] + async fn batch_disconnect_and_timeout_are_errors() { + use tokio::io::AsyncReadExt; + for disconnect in [true, false] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + None, + Duration::from_millis(100), + None, + true, + ); + let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + if !disconnect { + let _ = stopped.await; + } + }); + assert!( + client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .is_err() + ); + let _ = stop.send(()); + server.await.unwrap(); + client.catalog = "测".repeat(500_000); + let error = client + .logical_tables("CREATE TABLE t", "public", 1) + .await + .unwrap_err(); + assert!(format!("{error:?}").contains("4 MiB form limit")); + } + } + + pub(crate) async fn test_server( + status: u16, + body: &str, + ) -> ( + DatabaseClient, + Arc>>, + tokio::task::JoinHandle<()>, + ) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let client = DatabaseClient::new( + listener.local_addr().unwrap().to_string(), + "greptime".into(), + Some("user:password".into()), + Duration::from_secs(5), + None, + true, + ); + let requests = Arc::new(std::sync::Mutex::new(Vec::new())); + let captured = requests.clone(); + let body = body.to_string(); + let server = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + loop { + let mut buf = [0; 4096]; + let n = socket.read(&mut buf).await.unwrap(); + if n == 0 { + break; + } + bytes.extend_from_slice(&buf[..n]); + if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&bytes[..end]).to_lowercase(); + let length = headers + .lines() + .find_map(|l| l.strip_prefix("content-length: ")) + .map(|n| n.parse::().unwrap()) + .unwrap_or(0); + if bytes.len() >= end + 4 + length { + break; + } + } + } + captured + .lock() + .unwrap() + .push(String::from_utf8(bytes).unwrap()); + let location = if (300..400).contains(&status) { + "Location: /v1/sql\r\n" + } else { + "" + }; + 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(); + } + }); + (client, requests, server) + } + #[test] fn test_split_database() { let result = split_database("catalog-schema").unwrap(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index da1232a243b..8e3057bb7aa 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1416,6 +1416,59 @@ async fn export_http(instance: Arc) -> (String, tokio::task::JoinHandl (addr, task) } +#[derive(Default)] +struct ImportRequests { + batches: std::sync::atomic::AtomicUsize, + fail_batch: std::sync::atomic::AtomicBool, + capabilities: std::sync::atomic::AtomicUsize, +} + +async fn import_http( + instance: Arc, + older: bool, +) -> (String, tokio::task::JoinHandle<()>, Arc) { + use axum::response::IntoResponse; + use std::sync::atomic::Ordering; + let requests = Arc::new(ImportRequests::default()); + requests.fail_batch.store(!older, Ordering::SeqCst); + let observed = requests.clone(); + let server = servers::http::HttpServerBuilder::new(Default::default()) + .with_sql_handler(instance) + .build(); + let app = server + .build(server.make_app()) + .unwrap() + .layer(axum::middleware::from_fn( + move |request: axum::extract::Request, next: axum::middleware::Next| { + let observed = observed.clone(); + async move { + let batch = request.uri().path() == "/v1/ddl/logical-tables"; + if request.uri().path() == "/v1/capabilities" { + observed.capabilities.fetch_add(1, Ordering::SeqCst); + if older { + return axum::http::StatusCode::NOT_FOUND.into_response(); + } + } + if batch { + observed.batches.fetch_add(1, Ordering::SeqCst); + } + let response = next.run(request).await; + if batch && observed.fail_batch.swap(false, Ordering::SeqCst) { + assert!(response.status().is_success()); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + response + } + }, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (addr, task, requests) +} + struct FailSecondChunk { fail: std::sync::atomic::AtomicBool, copies: std::sync::atomic::AtomicUsize, @@ -1451,7 +1504,12 @@ impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { #[tokio::test] async fn metric_export_v2_cli_resume_roundtrip() { - metric_export_v2_cli_roundtrip(false, false).await; + metric_export_v2_cli_roundtrip(false, false, false).await; +} + +#[tokio::test] +async fn metric_import_v2_older_target_roundtrip() { + metric_export_v2_cli_roundtrip(false, false, true).await; } #[tokio::test] @@ -1461,13 +1519,13 @@ async fn metric_export_v2_cli_s3_resume_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) { - metric_export_v2_cli_roundtrip(true, false).await; + metric_export_v2_cli_roundtrip(true, false, false).await; } } #[tokio::test(flavor = "multi_thread")] async fn packed_export_v2_cli_local_roundtrip() { - metric_export_v2_cli_roundtrip(false, true).await; + metric_export_v2_cli_roundtrip(false, true, false).await; } #[tokio::test(flavor = "multi_thread")] @@ -1475,7 +1533,7 @@ async fn packed_export_v2_cli_s3_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) && std::env::var("GT_S3_BUCKET").is_ok_and(|b| !b.is_empty()) { - metric_export_v2_cli_roundtrip(true, true).await; + metric_export_v2_cli_roundtrip(true, true, false).await; } } @@ -1508,7 +1566,7 @@ async fn exported_table_path( ) } -async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { +async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { use servers::interceptor::SqlQueryInterceptorRef; let plugins = common_base::Plugins::new(); let faults = Arc::new(FailSecondChunk { @@ -1864,7 +1922,8 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { ) .await; } - let (target_addr, target_server) = export_http(target.fe_instance().clone()).await; + let (target_addr, target_server, requests) = + import_http(target.fe_instance().clone(), older).await; let state = destination.path().join("restore-state.json"); let mut import = vec![ "import-v2", @@ -1879,7 +1938,29 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { "never", ]; import.extend(storage_args.iter().map(String::as_str)); + use std::sync::atomic::Ordering; + if !older { + assert!(run_data_cli(&import).await.is_err()); + assert_eq!(requests.batches.load(Ordering::SeqCst), 1); + assert_eq!(requests.capabilities.load(Ordering::SeqCst), 1); + let failed: serde_json::Value = + serde_json::from_slice(&std::fs::read(&state).unwrap()).unwrap(); + assert_eq!(failed["ddl_completed"], false); + assert!( + failed["tasks"] + .as_array() + .unwrap() + .iter() + .all(|t| t["status"] == "pending") + ); + } run_data_cli(&import).await.unwrap(); + assert_eq!( + requests.capabilities.load(Ordering::SeqCst), + if older { 1 } else { 2 } + ); + assert_eq!(requests.batches.load(Ordering::SeqCst) > 1, !older); + assert!(!state.exists()); for name in names { let source_table = table(instance, &name).await; let target_table = table(target.fe_instance(), &name).await; @@ -1954,6 +2035,42 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool) { assert!(schema_manifest.schema_only && schema_manifest.chunks.is_empty()); assert_eq!(schema_manifest.version, 1); assert!(schema_manifest.data_layout.is_none()); + let schema_target = GreptimeDbStandaloneBuilder::new("schema_only_target") + .build() + .await; + let (schema_addr, schema_server, schema_requests) = + import_http(schema_target.fe_instance().clone(), false).await; + schema_requests.fail_batch.store(false, Ordering::SeqCst); + let schema_state = destination.path().join("schema-only-state.json"); + let mut schema_import = vec![ + "import-v2", + "--addr", + &schema_addr, + "--from", + &schema_uri, + "--state-path", + schema_state.to_str().unwrap(), + "--no-proxy", + "--progress", + "never", + ]; + schema_import.extend(storage_args.iter().map(String::as_str)); + let probes = schema_requests.capabilities.load(Ordering::SeqCst); + schema_import.push("--dry-run"); + run_data_cli(&schema_import).await.unwrap(); + assert_eq!(schema_requests.capabilities.load(Ordering::SeqCst), probes); + schema_import.pop(); + let batches = schema_requests.batches.load(Ordering::SeqCst); + run_data_cli(&schema_import).await.unwrap(); + assert!(schema_requests.batches.load(Ordering::SeqCst) > batches); + assert!(!schema_state.exists()); + assert!( + values(schema_target.fe_instance(), "SELECT * FROM dashboard") + .await + .is_empty() + ); + schema_server.abort(); + let _ = schema_server.await; } server.abort(); target_server.abort();