diff --git a/src/cli/src/data/export_v2/tests.rs b/src/cli/src/data/export_v2/tests.rs index dd972b59b1..91a2c53bbc 100644 --- a/src/cli/src/data/export_v2/tests.rs +++ b/src/cli/src/data/export_v2/tests.rs @@ -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, +} + +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 { + 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 { 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> { + 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",