rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-06 05:50:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(cli): add import-v2 resume e2e coverage (#8311)

Browse Source 
* test(cli): harden export import v2 e2e roundtrip

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

* test(cli): add import-v2 resume e2e coverage

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

* test(cli): make import-v2 resume replay check non-idempotent

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

* test(cli): share export import e2e setup helpers

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

* test(cli): assert import-v2 resume row identities

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

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-06-18 02:02:24 -07:00
committed by GitHub
parent fbdc35a109
commit a3e01e0018
1 changed files with 264 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

312
src/cli/src/data/export_v2/tests.rs
View File

@@ -26,11 +26,51 @@ use super::command::{ExportCreateCommand, ExportVerifyCommand};
use crate::common::ObjectStoreConfig;
use crate::data::export_v2::manifest::ChunkStatus;
use crate::data::import_v2::ImportV2Command;
use crate::data::import_v2::coordinator::build_import_tasks;
use crate::data::import_v2::state::{ImportState, ImportTaskStatus, save_import_state};
use crate::data::snapshot_storage::{OpenDalStorage, SnapshotStorage};
use crate::data::sql::escape_sql_identifier;
use crate::database::DatabaseClient;
use crate::error::{FileIoSnafu, InvalidArgumentsSnafu, OtherSnafu, Result};
struct TestConnection {
addr: String,
catalog: String,
auth_basic: Option<String>,
}
impl TestConnection {
fn from_env() -> Self {
Self {
addr: env::var("GREPTIME_ADDR").unwrap_or_else(|_| "127.0.0.1:4000".to_string()),
catalog: env::var("GREPTIME_CATALOG").unwrap_or_else(|_| "greptime".to_string()),
auth_basic: env::var("GREPTIME_AUTH_BASIC").ok(),
}
}
fn client(&self) -> DatabaseClient {
DatabaseClient::new(
self.addr.clone(),
self.catalog.clone(),
self.auth_basic.clone(),
Duration::from_secs(60),
None,
false,
)
}
}
fn path_to_uri(path: &std::path::Path) -> Result<String> {
Url::from_directory_path(path)
.map(|url| url.to_string())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})
}
async fn query_count(database_client: &DatabaseClient, schema: &str, table: &str) -> Result<u64> {
let sql = format!("SELECT COUNT(*) FROM {}", escape_sql_identifier(table));
let rows = database_client.sql(&sql, schema).await?;
@@ -60,6 +100,22 @@ async fn query_count(database_client: &DatabaseClient, schema: &str, table: &str
}
}
async fn query_hosts(database_client: &DatabaseClient, schema: &str) -> Result<Vec<String>> {
let rows = database_client
.sql("SELECT host FROM metrics ORDER BY host", schema)
.await?
.unwrap_or_default();
rows.into_iter()
.map(|row| match row.first() {
Some(Value::String(value)) => Ok(value.clone()),
_ => InvalidArgumentsSnafu {
msg: "unexpected host value".to_string(),
}
.fail(),
})
.collect()
}
#[tokio::test]
#[ignore]
async fn export_import_v2_schema_parity_e2e() -> Result<()> {
@@ -136,14 +192,7 @@ async fn export_import_v2_schema_parity_e2e() -> Result<()> {
.await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
@@ -306,14 +355,7 @@ async fn import_v2_ddl_dry_run_e2e() -> Result<()> {
.await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
@@ -417,14 +459,7 @@ async fn export_import_v2_data_roundtrip_e2e() -> Result<()> {
let expected_rows = query_count(&database_client, schema, "metrics").await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
@@ -516,6 +551,208 @@ async fn export_import_v2_data_roundtrip_e2e() -> Result<()> {
Ok(())
}
#[tokio::test]
#[ignore]
async fn import_v2_resume_from_completed_chunk_e2e() -> Result<()> {
let conn = TestConnection::from_env();
let schema = "test_db_import_resume_completed_chunk";
let database_client = conn.client();
database_client
.sql_in_public(&format!("DROP DATABASE IF EXISTS {schema}"))
.await?;
database_client
.sql_in_public(&format!("CREATE DATABASE {schema}"))
.await?;
database_client
.sql(
"CREATE TABLE metrics (\
ts TIMESTAMP TIME INDEX, \
host STRING, \
cpu DOUBLE \
) ENGINE=mito",
schema,
)
.await?;
database_client
.sql(
"INSERT INTO metrics (ts, host, cpu) VALUES \
('2025-01-01T00:00:00Z', 'h1', 1.0), \
('2025-01-01T01:00:00Z', 'h2', 2.0)",
schema,
)
.await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
"--addr",
&conn.addr,
"--to",
&src_uri,
"--catalog",
&conn.catalog,
"--schemas",
schema,
"--start-time",
"2025-01-01T00:00:00Z",
"--end-time",
"2025-01-01T02:00:00Z",
"--chunk-time-window",
"1h",
"--chunk-parallelism",
"1",
"--progress",
"never",
];
if let Some(auth) = &conn.auth_basic {
export_args.push("--auth-basic");
export_args.push(auth);
}
let export_cmd = ExportCreateCommand::parse_from(export_args);
export_cmd
.build()
.await
.context(OtherSnafu)?
.do_work()
.await
.context(OtherSnafu)?;
let storage_config = ObjectStoreConfig::default();
let storage = OpenDalStorage::from_uri(&src_uri, &storage_config)
.map_err(BoxedError::new)
.context(OtherSnafu)?;
let full_manifest = storage
.read_manifest()
.await
.map_err(BoxedError::new)
.context(OtherSnafu)?;
assert_eq!(full_manifest.chunks.len(), 2);
database_client
.sql_in_public(&format!("DROP DATABASE IF EXISTS {schema}"))
.await?;
let mut partial_manifest = full_manifest.clone();
partial_manifest.chunks.truncate(1);
storage
.write_manifest(&partial_manifest)
.await
.map_err(BoxedError::new)
.context(OtherSnafu)?;
let mut partial_import_args = vec![
"import-v2",
"--addr",
&conn.addr,
"--from",
&src_uri,
"--catalog",
&conn.catalog,
"--schemas",
schema,
"--task-parallelism",
"1",
"--progress",
"never",
];
if let Some(auth) = &conn.auth_basic {
partial_import_args.push("--auth-basic");
partial_import_args.push(auth);
}
let partial_import_cmd = ImportV2Command::parse_from(partial_import_args);
partial_import_cmd
.build()
.await
.context(OtherSnafu)?
.do_work()
.await
.context(OtherSnafu)?;
let rows_after_partial_import = query_count(&database_client, schema, "metrics").await?;
assert_eq!(rows_after_partial_import, 1);
storage
.write_manifest(&full_manifest)
.await
.map_err(BoxedError::new)
.context(OtherSnafu)?;
let import_state_dir = tempdir().context(FileIoSnafu)?;
let import_state_path = import_state_dir.path().join("import-state.json");
let schemas = vec![schema.to_string()];
let tasks = build_import_tasks(&full_manifest.chunks, &schemas);
assert_eq!(tasks.len(), 2);
let mut import_state = ImportState::new(
full_manifest.snapshot_id.to_string(),
conn.addr.clone(),
conn.catalog.clone(),
&schemas,
tasks,
);
import_state.mark_ddl_completed();
import_state
.set_task_status(
full_manifest.chunks[0].id,
schema,
ImportTaskStatus::Completed,
None,
)
.map_err(BoxedError::new)
.context(OtherSnafu)?;
save_import_state(&import_state_path, &import_state)
.await
.map_err(BoxedError::new)
.context(OtherSnafu)?;
let import_state_path_arg = import_state_path.to_string_lossy().into_owned();
let mut import_args = vec![
"import-v2",
"--addr",
&conn.addr,
"--from",
&src_uri,
"--catalog",
&conn.catalog,
"--schemas",
schema,
"--state-path",
&import_state_path_arg,
"--task-parallelism",
"1",
"--progress",
"never",
];
if let Some(auth) = &conn.auth_basic {
import_args.push("--auth-basic");
import_args.push(auth);
}
let import_cmd = ImportV2Command::parse_from(import_args);
import_cmd
.build()
.await
.context(OtherSnafu)?
.do_work()
.await
.context(OtherSnafu)?;
let actual_rows = query_count(&database_client, schema, "metrics").await?;
assert_eq!(actual_rows, 2);
let hosts = query_hosts(&database_client, schema).await?;
assert_eq!(hosts, vec!["h1".to_string(), "h2".to_string()]);
assert!(!import_state_path.exists());
database_client
.sql_in_public(&format!("DROP DATABASE IF EXISTS {schema}"))
.await?;
Ok(())
}
#[tokio::test]
#[ignore]
async fn import_v2_fails_on_incomplete_snapshot_e2e() -> Result<()> {
@@ -557,14 +794,7 @@ async fn import_v2_fails_on_incomplete_snapshot_e2e() -> Result<()> {
.await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
@@ -692,14 +922,7 @@ async fn import_v2_schema_filter_data_e2e() -> Result<()> {
let expected_rows_a = query_count(&database_client, schema_a, "metrics").await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",
@@ -817,14 +1040,7 @@ async fn export_import_v2_skipped_chunk_e2e() -> Result<()> {
.await?;
let src_dir = tempdir().context(FileIoSnafu)?;
let src_uri = Url::from_directory_path(src_dir.path())
.map_err(|_| {
InvalidArgumentsSnafu {
msg: "invalid temp dir path".to_string(),
}
.build()
})?
.to_string();
let src_uri = path_to_uri(src_dir.path())?;
let mut export_args = vec![
"export-v2-create",