rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 06:52:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: batch logical table DDL during import-v2 (#9467)

Browse Source 
* feat: batch logical table DDL during import-v2

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: preserve one-shot batch DDL transport and context

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test: check schema-only import state at its own path

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: preserve oversized import DDL compatibility

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* refactor: reuse import DDL encoding and metric constants

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-10-08 13:14:06 +00:00
committed by GitHub
parent 5126623e63
commit 6265c31e4a
5 changed files with 904 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cli/src/data/import_v2/command.rs
+92 -74
View File
@@ -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<String> {
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<Item = String> + '_ {
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 {
src/cli/src/data/import_v2/executor.rs
+363 -43
View File
@@ -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<Option<LogicalTable>> {
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<Vec<Statement>> {
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<LogicalTable>,
names: HashSet<String>,
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;
}
}
}
src/cli/src/data/import_v2/packed.rs
+41 -44
View File
@@ -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()]);
src/cli/src/database.rs
+285 -23
View File
@@ -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<Capabilities> {
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<GreptimedbV1Response> {
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<GreptimedbV1Response> {
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(&params)
.header("Content-Type", "application/x-www-form-urlencoded");
if endpoint == "ddl/logical-tables" {
let body = url::form_urlencoded::Serializer::new(String::new())
.extend_pairs(&params)
.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(&params);
}
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<String>)> {
let (catalog, schema) = match database.split_once('-') {
@@ -267,7 +358,7 @@ pub(crate) fn split_database(database: &str) -> Result<(String, Option<String>)>
}
#[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<std::sync::Mutex<Vec<String>>>,
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::<usize>().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();
tests-integration/tests/export_logical_tables.rs
+123 -6
View File
@@ -1416,6 +1416,59 @@ async fn export_http(instance: Arc<Instance>) -> (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<Instance>,
older: bool,
) -> (String, tokio::task::JoinHandle<()>, Arc<ImportRequests>) {
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();