mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-15 16:32:34 +00:00
feat: export Metric snapshots with packed Parquet objects (#9382)
* feat: export Metric snapshots with packed Parquet objects Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: reject empty packed export time ranges Signed-off-by: jeremyhi <fengjiachun@gmail.com> * test: cover cancellation during packed export I/O Signed-off-by: jeremyhi <fengjiachun@gmail.com> --------- Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
@@ -71,3 +71,4 @@ common-test-util.workspace = true
|
||||
object-store = { workspace = true, features = ["testing"] }
|
||||
serde.workspace = true
|
||||
tempfile.workspace = true
|
||||
tokio-util.workspace = true
|
||||
|
||||
@@ -41,11 +41,13 @@
|
||||
//! `--experimental-metric-export` enables shared Metric physical scans for Parquet.
|
||||
//! Enable `experimental_metric_export = true` on the frontend or standalone server
|
||||
//! and use an endpoint whose frontends all support and enable this option.
|
||||
//! Add `--metric-data-layout packed` to write version-2 packed snapshots.
|
||||
//! This requires the server packed-export capability; schema-only exports remain version 1.
|
||||
//! The command checks the server capability before modifying the snapshot.
|
||||
//! Each snapshot path belongs to one export task. Resume requires stable source
|
||||
//! schemas/data and confirmation that the previous export and storage writes ended;
|
||||
//! an HTTP timeout does not establish that. Only unfinished chunks are cleaned and
|
||||
//! rerun. Completed chunks and the V2 snapshot/import format remain unchanged.
|
||||
//! rerun. Resume retains the snapshot version/layout and completed chunks.
|
||||
|
||||
mod chunker;
|
||||
mod command;
|
||||
|
||||
@@ -296,6 +296,10 @@ pub struct ExportCreateCommand {
|
||||
#[clap(long)]
|
||||
experimental_metric_export: bool,
|
||||
|
||||
/// Store independent Metric Parquet streams in shared packed objects.
|
||||
#[clap(long, value_parser = ["packed"], requires = "experimental_metric_export")]
|
||||
metric_data_layout: Option<String>,
|
||||
|
||||
/// Delete existing snapshot and recreate.
|
||||
#[clap(long)]
|
||||
force: bool,
|
||||
@@ -345,6 +349,16 @@ impl ExportCreateCommand {
|
||||
|
||||
let time_range = TimeRange::parse(self.start_time.as_deref(), self.end_time.as_deref())
|
||||
.map_err(BoxedError::new)?;
|
||||
if self.metric_data_layout.is_some()
|
||||
&& time_range.is_bounded()
|
||||
&& time_range.start == time_range.end
|
||||
{
|
||||
return crate::error::InvalidArgumentsSnafu {
|
||||
msg: "Packed export requires --start-time to be earlier than --end-time",
|
||||
}
|
||||
.fail()
|
||||
.map_err(BoxedError::new);
|
||||
}
|
||||
if self.chunk_time_window.is_some() && !time_range.is_bounded() {
|
||||
return ChunkTimeWindowRequiresBoundsSnafu
|
||||
.fail()
|
||||
@@ -425,6 +439,13 @@ impl ExportCreateCommand {
|
||||
.map_err(BoxedError::new);
|
||||
}
|
||||
}
|
||||
if self.metric_data_layout.is_some() && !self.schema_only {
|
||||
database_client
|
||||
.require_packed_export()
|
||||
.await
|
||||
.context(DatabaseSnafu)
|
||||
.map_err(BoxedError::new)?;
|
||||
}
|
||||
let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?;
|
||||
|
||||
Ok(Box::new(ExportCreate {
|
||||
@@ -434,6 +455,7 @@ impl ExportCreateCommand {
|
||||
schema_only: self.schema_only,
|
||||
format: self.format,
|
||||
experimental_metric_export: self.experimental_metric_export,
|
||||
packed: self.metric_data_layout.is_some(),
|
||||
force: self.force,
|
||||
time_range,
|
||||
chunk_time_window: self.chunk_time_window,
|
||||
@@ -462,6 +484,7 @@ struct ExportConfig {
|
||||
schema_only: bool,
|
||||
format: DataFormat,
|
||||
experimental_metric_export: bool,
|
||||
packed: bool,
|
||||
force: bool,
|
||||
time_range: TimeRange,
|
||||
chunk_time_window: Option<Duration>,
|
||||
@@ -504,7 +527,7 @@ impl ExportCreate {
|
||||
let mut manifest = self.storage.read_manifest().await?;
|
||||
|
||||
// Check version compatibility
|
||||
if manifest.version != MANIFEST_VERSION || manifest.data_layout.is_some() {
|
||||
if manifest.validate_layout().is_err() {
|
||||
return ManifestVersionMismatchSnafu {
|
||||
expected: MANIFEST_VERSION,
|
||||
found: manifest.version,
|
||||
@@ -513,6 +536,15 @@ impl ExportCreate {
|
||||
}
|
||||
|
||||
validate_resume_config(&manifest, &self.config)?;
|
||||
if manifest.is_packed() {
|
||||
self.database_client
|
||||
.require_packed_export()
|
||||
.await
|
||||
.context(DatabaseSnafu)?;
|
||||
if !self.config.experimental_metric_export {
|
||||
return crate::data::export_v2::error::MetricExportUnavailableSnafu.fail();
|
||||
}
|
||||
}
|
||||
|
||||
info!(
|
||||
"Resuming existing snapshot: {} (completed: {}/{} chunks)",
|
||||
@@ -571,10 +603,14 @@ impl ExportCreate {
|
||||
self.config.chunk_time_window,
|
||||
)?;
|
||||
|
||||
if self.config.packed && !self.config.schema_only {
|
||||
manifest.version = 2;
|
||||
manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into());
|
||||
}
|
||||
if self.config.experimental_metric_export {
|
||||
for chunk in &manifest.chunks {
|
||||
self.storage
|
||||
.prepare_export_chunk(&schema_names, chunk.id, false)
|
||||
.prepare_export_chunk(&schema_names, chunk.id, false, manifest.is_packed())
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
@@ -796,6 +832,14 @@ fn build_schema_ddl(
|
||||
}
|
||||
|
||||
fn validate_resume_config(manifest: &Manifest, config: &ExportConfig) -> Result<()> {
|
||||
if config.packed && !manifest.schema_only && !manifest.is_packed() {
|
||||
return ResumeConfigMismatchSnafu {
|
||||
field: "metric_data_layout",
|
||||
existing: "standalone".to_string(),
|
||||
requested: "packed".to_string(),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
if manifest.schema_only != config.schema_only {
|
||||
return SchemaOnlyModeMismatchSnafu {
|
||||
existing_schema_only: manifest.schema_only,
|
||||
@@ -1801,6 +1845,7 @@ mod tests {
|
||||
schema_only: false,
|
||||
format: DataFormat::Parquet,
|
||||
experimental_metric_export: false,
|
||||
packed: false,
|
||||
force: false,
|
||||
time_range: TimeRange::unbounded(),
|
||||
chunk_time_window: None,
|
||||
@@ -1820,7 +1865,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_validate_resume_config_accepts_schema_selection_with_different_case_and_order() {
|
||||
let manifest = Manifest::new_for_export(
|
||||
let mut manifest = Manifest::new_for_export(
|
||||
"greptime".to_string(),
|
||||
vec!["public".to_string(), "analytics".to_string()],
|
||||
false,
|
||||
@@ -1829,7 +1874,7 @@ mod tests {
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let config = ExportConfig {
|
||||
let mut config = ExportConfig {
|
||||
catalog: "greptime".to_string(),
|
||||
schemas: Some(vec![
|
||||
"ANALYTICS".to_string(),
|
||||
@@ -1839,6 +1884,7 @@ mod tests {
|
||||
schema_only: false,
|
||||
format: DataFormat::Parquet,
|
||||
experimental_metric_export: false,
|
||||
packed: false,
|
||||
force: false,
|
||||
time_range: TimeRange::unbounded(),
|
||||
chunk_time_window: None,
|
||||
@@ -1850,6 +1896,18 @@ mod tests {
|
||||
};
|
||||
|
||||
assert!(validate_resume_config(&manifest, &config).is_ok());
|
||||
config.packed = true;
|
||||
assert!(
|
||||
validate_resume_config(&manifest, &config)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("metric_data_layout")
|
||||
);
|
||||
manifest.version = 2;
|
||||
manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into());
|
||||
assert!(validate_resume_config(&manifest, &config).is_ok());
|
||||
config.packed = false;
|
||||
assert!(validate_resume_config(&manifest, &config).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1872,6 +1930,7 @@ mod tests {
|
||||
schema_only: false,
|
||||
format: DataFormat::Parquet,
|
||||
experimental_metric_export: false,
|
||||
packed: false,
|
||||
force: false,
|
||||
time_range,
|
||||
chunk_time_window: Some(Duration::from_secs(3600)),
|
||||
@@ -1906,6 +1965,7 @@ mod tests {
|
||||
schema_only: false,
|
||||
format: DataFormat::Csv,
|
||||
experimental_metric_export: false,
|
||||
packed: false,
|
||||
force: false,
|
||||
time_range: TimeRange::unbounded(),
|
||||
chunk_time_window: None,
|
||||
@@ -1942,6 +2002,7 @@ mod tests {
|
||||
schema_only: false,
|
||||
format: DataFormat::Parquet,
|
||||
experimental_metric_export: false,
|
||||
packed: false,
|
||||
force: false,
|
||||
time_range: TimeRange::new(Some(start), Some(start)),
|
||||
chunk_time_window: None,
|
||||
|
||||
@@ -40,6 +40,7 @@ struct ExportContext<'a> {
|
||||
format: DataFormat,
|
||||
parallelism: usize,
|
||||
experimental_metric_export: bool,
|
||||
packed: bool,
|
||||
resume: bool,
|
||||
}
|
||||
|
||||
@@ -71,6 +72,7 @@ pub async fn export_data(
|
||||
catalog: manifest.catalog.clone(),
|
||||
schemas: manifest.schemas.clone(),
|
||||
format: manifest.format,
|
||||
packed: manifest.is_packed(),
|
||||
parallelism: options.parallelism,
|
||||
experimental_metric_export: options.experimental_metric_export,
|
||||
resume: options.resume,
|
||||
@@ -279,7 +281,7 @@ async fn export_chunk(
|
||||
if context.experimental_metric_export {
|
||||
context
|
||||
.storage
|
||||
.prepare_export_chunk(&context.schemas, chunk_id, context.resume)
|
||||
.prepare_export_chunk(&context.schemas, chunk_id, context.resume, context.packed)
|
||||
.await?;
|
||||
}
|
||||
let scheme = StorageScheme::from_uri(context.snapshot_uri)?;
|
||||
@@ -289,6 +291,7 @@ async fn export_chunk(
|
||||
time_range,
|
||||
parallelism: context.parallelism,
|
||||
experimental_metric_export: context.experimental_metric_export,
|
||||
packed: context.packed,
|
||||
};
|
||||
|
||||
for schema in &context.schemas {
|
||||
@@ -314,6 +317,39 @@ async fn export_chunk(
|
||||
}
|
||||
|
||||
let files = list_chunk_files(context.storage, &context.schemas, chunk_id).await?;
|
||||
if context.packed {
|
||||
use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex};
|
||||
let mut expected = Vec::new();
|
||||
for schema in &context.schemas {
|
||||
let prefix = data_dir_for_schema_chunk(schema, chunk_id);
|
||||
let index_path = format!("{prefix}{PACK_INDEX_FILE}");
|
||||
let text = context.storage.read_text(&index_path).await?;
|
||||
let index: PackIndex = serde_json::from_str(&text).map_err(|e| {
|
||||
crate::data::export_v2::error::InvalidUriSnafu {
|
||||
uri: context.snapshot_uri,
|
||||
reason: e.to_string(),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
index.validate().map_err(|e| {
|
||||
crate::data::export_v2::error::InvalidUriSnafu {
|
||||
uri: context.snapshot_uri,
|
||||
reason: e.to_string(),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
expected.push(index_path);
|
||||
expected.extend(index.objects.iter().map(|o| format!("{prefix}{}", o.path)));
|
||||
}
|
||||
expected.sort();
|
||||
if files != expected {
|
||||
return crate::data::export_v2::error::InvalidUriSnafu {
|
||||
uri: context.snapshot_uri,
|
||||
reason: "chunk inventory differs from packed indexes",
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
}
|
||||
info!("Collected {} files for chunk {}", files.len(), chunk_id);
|
||||
Ok(files)
|
||||
}
|
||||
@@ -462,6 +498,158 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct PausedClose {
|
||||
inner: Option<object_store::layers::mock::oio::Writer>,
|
||||
started: std::sync::Arc<tokio::sync::Notify>,
|
||||
release: std::sync::Arc<tokio::sync::Notify>,
|
||||
}
|
||||
|
||||
impl object_store::layers::mock::oio::Write for PausedClose {
|
||||
async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
|
||||
self.inner.as_mut().unwrap().write(bytes).await
|
||||
}
|
||||
|
||||
async fn close(&mut self) -> object_store::Result<object_store::layers::mock::Metadata> {
|
||||
let mut inner = self.inner.take().unwrap();
|
||||
let (started, release) = (self.started.clone(), self.release.clone());
|
||||
// Model blocking storage I/O that survives dropping its caller.
|
||||
let (inner, result) = tokio::spawn(async move {
|
||||
started.notify_one();
|
||||
release.notified().await;
|
||||
let result = inner.close().await;
|
||||
(inner, result)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
self.inner = Some(inner);
|
||||
result
|
||||
}
|
||||
|
||||
async fn abort(&mut self) -> object_store::Result<()> {
|
||||
self.inner.as_mut().unwrap().abort().await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn packed_cancellation_drains_close_before_failing_chunk() {
|
||||
use std::sync::Arc;
|
||||
|
||||
use common_datasource::packed_snapshot::PACK_INDEX_FILE;
|
||||
use common_datasource::packed_writer::{PackedTableWriter, PackedWriter};
|
||||
use common_datasource::parquet_writer::ParquetFileWriter;
|
||||
use datatypes::arrow::datatypes::{DataType, Field, Schema};
|
||||
use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::data::snapshot_storage::OpenDalStorage;
|
||||
|
||||
for paused_path in ["pack-000000.bin", PACK_INDEX_FILE] {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let uri = url::Url::from_directory_path(directory.path()).unwrap();
|
||||
let storage =
|
||||
OpenDalStorage::from_uri(uri.as_str(), &ObjectStoreConfig::default()).unwrap();
|
||||
let started = Arc::new(tokio::sync::Notify::new());
|
||||
let release = Arc::new(tokio::sync::Notify::new());
|
||||
let factory: MockWriterFactory = Arc::new({
|
||||
let (started, release) = (started.clone(), release.clone());
|
||||
move |path, _, inner| {
|
||||
if path.trim_start_matches('/') == paused_path {
|
||||
Box::new(PausedClose {
|
||||
inner: Some(inner),
|
||||
started: started.clone(),
|
||||
release: release.clone(),
|
||||
})
|
||||
} else {
|
||||
inner
|
||||
}
|
||||
}
|
||||
});
|
||||
let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
|
||||
.unwrap()
|
||||
.build_operator()
|
||||
.layer(
|
||||
MockLayerBuilder::default()
|
||||
.writer_factory(factory)
|
||||
.build()
|
||||
.unwrap(),
|
||||
);
|
||||
let packed = PackedWriter::new(store.clone()).unwrap();
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"value",
|
||||
DataType::Int64,
|
||||
true,
|
||||
)]));
|
||||
let mut table = ParquetFileWriter::open_packed(
|
||||
schema,
|
||||
store.clone(),
|
||||
"unused",
|
||||
None,
|
||||
PackedTableWriter::new(packed.clone(), "empty".into(), 1, false),
|
||||
)
|
||||
.unwrap();
|
||||
table.finish(None).await.unwrap();
|
||||
let token = CancellationToken::new();
|
||||
let mut manifest = pending_manifest(1);
|
||||
manifest.version = 2;
|
||||
manifest.data_layout = Some(common_datasource::packed_snapshot::PACKED_LAYOUT.into());
|
||||
let export = export_data_concurrent(
|
||||
&storage,
|
||||
&mut manifest,
|
||||
2,
|
||||
&crate::data::progress::NoopProgress,
|
||||
|_, _| async {
|
||||
let mut packed = packed.lock().await;
|
||||
let result = packed.finish(&token).await;
|
||||
if result.is_err() {
|
||||
packed.abort().await.unwrap();
|
||||
}
|
||||
result.map_err(|error| {
|
||||
crate::data::export_v2::error::IoSnafu {
|
||||
operation: "exporting packed chunk",
|
||||
error: std::io::Error::other(error),
|
||||
}
|
||||
.build()
|
||||
})
|
||||
},
|
||||
);
|
||||
tokio::pin!(export);
|
||||
tokio::select! {
|
||||
_ = started.notified() => {},
|
||||
result = &mut export => panic!("completed while close paused: {result:?}"),
|
||||
}
|
||||
token.cancel();
|
||||
assert!(futures::poll!(&mut export).is_pending());
|
||||
assert_eq!(
|
||||
storage.read_manifest().await.unwrap().chunks[0].status,
|
||||
ChunkStatus::InProgress
|
||||
);
|
||||
release.notify_one();
|
||||
assert!(export.await.unwrap_err().to_string().contains("cancelled"));
|
||||
let persisted = storage.read_manifest().await.unwrap();
|
||||
assert_eq!(persisted.chunks[0].status, ChunkStatus::Failed);
|
||||
assert!(!persisted.is_complete());
|
||||
assert!(persisted.chunks[0].files.is_empty());
|
||||
assert!(
|
||||
store
|
||||
.stat("pack-000000.bin")
|
||||
.await
|
||||
.unwrap()
|
||||
.content_length()
|
||||
> 0
|
||||
);
|
||||
assert_eq!(
|
||||
store.exists(PACK_INDEX_FILE).await.unwrap(),
|
||||
paused_path == PACK_INDEX_FILE
|
||||
);
|
||||
if paused_path == PACK_INDEX_FILE {
|
||||
let index: common_datasource::packed_snapshot::PackIndex =
|
||||
serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes())
|
||||
.unwrap();
|
||||
index.validate_membership(["empty"]).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_next_eligible_chunk_scans_in_order() {
|
||||
let manifest = pending_manifest(3);
|
||||
|
||||
@@ -31,6 +31,7 @@ pub(super) struct CopyOptions {
|
||||
pub(super) time_range: TimeRange,
|
||||
pub(super) parallelism: usize,
|
||||
pub(crate) experimental_metric_export: bool,
|
||||
pub(super) packed: bool,
|
||||
}
|
||||
|
||||
pub(super) struct CopyTarget {
|
||||
@@ -220,6 +221,9 @@ fn build_with_options(options: &CopyOptions) -> String {
|
||||
if options.experimental_metric_export {
|
||||
parts.push("experimental_metric_export='true'".to_string());
|
||||
}
|
||||
if options.packed {
|
||||
parts.push("metric_data_layout='packed'".to_string());
|
||||
}
|
||||
if let Some(start) = options.time_range.start {
|
||||
parts.push(format!(
|
||||
"START_TIME='{}'",
|
||||
|
||||
@@ -28,7 +28,7 @@ use crate::data::export_v2::error::{
|
||||
TimeParseInvalidFormatSnafu,
|
||||
};
|
||||
|
||||
/// Manifest format version produced by the current exporter.
|
||||
/// Default standalone manifest version. Packed data snapshots use version 2.
|
||||
pub const MANIFEST_VERSION: u32 = 1;
|
||||
|
||||
/// Manifest file name within snapshot directory.
|
||||
|
||||
@@ -295,8 +295,9 @@ pub trait SnapshotStorage: Send + Sync {
|
||||
schemas: &[String],
|
||||
chunk_id: u32,
|
||||
resume: bool,
|
||||
packed: bool,
|
||||
) -> Result<()> {
|
||||
let _ = (schemas, chunk_id, resume);
|
||||
let _ = (schemas, chunk_id, resume, packed);
|
||||
InvalidUriSnafu {
|
||||
uri: "snapshot",
|
||||
reason: "storage does not support preparing export chunks",
|
||||
@@ -730,6 +731,7 @@ impl SnapshotStorage for OpenDalStorage {
|
||||
schemas: &[String],
|
||||
chunk_id: u32,
|
||||
resume: bool,
|
||||
packed: bool,
|
||||
) -> Result<()> {
|
||||
let mut files = Vec::new();
|
||||
for schema in schemas {
|
||||
@@ -751,10 +753,17 @@ impl SnapshotStorage for OpenDalStorage {
|
||||
continue;
|
||||
}
|
||||
let name = path.strip_prefix(&prefix).unwrap_or("");
|
||||
if !resume || entry.metadata().is_dir() || !valid_chunk_filename(name) {
|
||||
if !resume
|
||||
|| entry.metadata().is_dir()
|
||||
|| !(if packed {
|
||||
valid_packed_chunk_filename(name)
|
||||
} else {
|
||||
valid_chunk_filename(name)
|
||||
})
|
||||
{
|
||||
return InvalidUriSnafu {
|
||||
uri: path,
|
||||
reason: "expected an empty new chunk or direct Parquet files in an owned unfinished chunk",
|
||||
reason: "expected an empty new chunk or recognized files in an owned unfinished chunk",
|
||||
}.fail();
|
||||
}
|
||||
files.push(path.to_string());
|
||||
@@ -786,6 +795,17 @@ impl SnapshotStorage for OpenDalStorage {
|
||||
}
|
||||
}
|
||||
|
||||
fn valid_packed_chunk_filename(name: &str) -> bool {
|
||||
name == common_datasource::packed_snapshot::PACK_INDEX_FILE
|
||||
|| [("pack-", ".bin"), ("table-", ".parquet")]
|
||||
.iter()
|
||||
.any(|(prefix, suffix)| {
|
||||
name.strip_prefix(prefix)
|
||||
.and_then(|n| n.strip_suffix(suffix))
|
||||
.is_some_and(|id| !id.is_empty() && id.bytes().all(|b| b.is_ascii_digit()))
|
||||
})
|
||||
}
|
||||
|
||||
fn valid_chunk_filename(name: &str) -> bool {
|
||||
name.strip_suffix(".parquet")
|
||||
.is_some_and(|stem| !stem.is_empty())
|
||||
@@ -1074,7 +1094,7 @@ mod tests {
|
||||
let schemas = vec!["public".to_string(), "other".to_string()];
|
||||
assert!(
|
||||
storage
|
||||
.prepare_export_chunk(&schemas, 2, false)
|
||||
.prepare_export_chunk(&schemas, 2, false, false)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
@@ -1085,11 +1105,11 @@ mod tests {
|
||||
.unwrap()
|
||||
);
|
||||
storage
|
||||
.prepare_export_chunk(&schemas, 2, true)
|
||||
.prepare_export_chunk(&schemas, 2, true, false)
|
||||
.await
|
||||
.unwrap();
|
||||
storage
|
||||
.prepare_export_chunk(&schemas, 2, true)
|
||||
.prepare_export_chunk(&schemas, 2, true, false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
@@ -1129,7 +1149,7 @@ mod tests {
|
||||
}
|
||||
assert!(
|
||||
storage
|
||||
.prepare_export_chunk(&["public".into(), "other".into()], 2, true)
|
||||
.prepare_export_chunk(&["public".into(), "other".into()], 2, true, false)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
@@ -1143,6 +1163,70 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn packed_retry_preserves_foreign_files_and_completed_chunks() {
|
||||
let dir = tempdir().unwrap();
|
||||
let storage = make_storage_with_rooted_fs(dir.path());
|
||||
let owned = ["pack-000000.bin", "table-42.parquet", "pack-index.json"];
|
||||
for name in owned {
|
||||
storage
|
||||
.write_text(&format!("data/public/2/{name}"), "partial")
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
for path in [
|
||||
"data/public/1/pack-000000.bin",
|
||||
"data/public/2/foreign.parquet",
|
||||
] {
|
||||
storage.write_text(path, "keep").await.unwrap();
|
||||
}
|
||||
let schemas = ["public".into()];
|
||||
assert!(
|
||||
storage
|
||||
.prepare_export_chunk(&schemas, 2, true, true)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
for name in owned {
|
||||
assert!(
|
||||
storage
|
||||
.file_exists(&format!("data/public/2/{name}"))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
storage
|
||||
.read_text("data/public/2/foreign.parquet")
|
||||
.await
|
||||
.unwrap(),
|
||||
"keep"
|
||||
);
|
||||
storage
|
||||
.object_store
|
||||
.delete("data/public/2/foreign.parquet")
|
||||
.await
|
||||
.unwrap();
|
||||
storage
|
||||
.prepare_export_chunk(&schemas, 2, true, true)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
storage
|
||||
.list_files_recursive("data/public/2/")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
assert_eq!(
|
||||
storage
|
||||
.read_text("data/public/1/pack-000000.bin")
|
||||
.await
|
||||
.unwrap(),
|
||||
"keep"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_prepare_export_chunk_reports_delete_failure() {
|
||||
let dir = tempdir().unwrap();
|
||||
@@ -1159,7 +1243,7 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
let error = storage
|
||||
.prepare_export_chunk(&["public".into()], 2, true)
|
||||
.prepare_export_chunk(&["public".into()], 2, true, false)
|
||||
.await
|
||||
.unwrap_err();
|
||||
let Error::StorageOperation {
|
||||
@@ -1191,7 +1275,7 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
storage
|
||||
.prepare_export_chunk(&[schema], 2, true)
|
||||
.prepare_export_chunk(&[schema], 2, true, false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!storage.file_exists(&owned).await.unwrap());
|
||||
|
||||
+62
-40
@@ -99,6 +99,15 @@ impl DatabaseClient {
|
||||
|
||||
/// 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
|
||||
}
|
||||
|
||||
/// Requires packed export support before creating snapshot artifacts.
|
||||
pub async fn require_packed_export(&self) -> Result<()> {
|
||||
self.require_packed_capability("metric_packed_export").await
|
||||
}
|
||||
|
||||
async fn require_packed_capability(&self, capability: &str) -> Result<()> {
|
||||
let url = format!("http://{}/v1/capabilities", self.addr);
|
||||
let mut builder = reqwest::Client::builder().timeout(self.timeout);
|
||||
if let Some(proxy) = self.proxy.clone() {
|
||||
@@ -113,18 +122,18 @@ impl DatabaseClient {
|
||||
request = request.header("Authorization", auth);
|
||||
}
|
||||
let response = request.send().await.with_context(|_| HttpQuerySqlSnafu {
|
||||
reason: "packed import capability request failed",
|
||||
reason: "packed snapshot capability request failed",
|
||||
})?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
return crate::error::InvalidArgumentsSnafu {
|
||||
msg: "target does not support packed import",
|
||||
msg: format!("server does not support {capability}"),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
let response = response
|
||||
.error_for_status()
|
||||
.with_context(|_| HttpQuerySqlSnafu {
|
||||
reason: "packed import capability request rejected",
|
||||
reason: "packed snapshot capability request rejected",
|
||||
})?;
|
||||
let body = response.text().await.with_context(|_| HttpQuerySqlSnafu {
|
||||
reason: "cannot read capability response",
|
||||
@@ -136,9 +145,9 @@ impl DatabaseClient {
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
if value.get("metric_packed_import").and_then(Value::as_u64) != Some(1) {
|
||||
if value.get(capability).and_then(Value::as_u64) != Some(1) {
|
||||
return crate::error::InvalidArgumentsSnafu {
|
||||
msg: "target does not support packed import version 1",
|
||||
msg: format!("server does not support {capability} version 1"),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
@@ -221,41 +230,54 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn packed_capability_probe_checks_auth_and_protocol_version() {
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
for (status, body, supported) in [
|
||||
(200, r#"{"metric_packed_import":1,"future":2}"#, true),
|
||||
(200, "{}", false),
|
||||
(200, r#"{"metric_packed_import":2}"#, false),
|
||||
(404, "{}", false),
|
||||
(401, "{}", false),
|
||||
(403, "{}", false),
|
||||
(200, "[]", false),
|
||||
(200, "invalid-json", false),
|
||||
] {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0; 4096];
|
||||
let n = socket.read(&mut request).await.unwrap();
|
||||
let request = String::from_utf8_lossy(&request[..n]).to_lowercase();
|
||||
assert!(request.starts_with("get /v1/capabilities "));
|
||||
assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza=="));
|
||||
socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
});
|
||||
let client = super::DatabaseClient::new(
|
||||
address.to_string(),
|
||||
"greptime".into(),
|
||||
Some("user:password".into()),
|
||||
std::time::Duration::from_secs(5),
|
||||
None,
|
||||
true,
|
||||
);
|
||||
assert_eq!(
|
||||
client.require_packed_import().await.is_ok(),
|
||||
supported,
|
||||
"{status}: {body}"
|
||||
);
|
||||
server.await.unwrap();
|
||||
for export in [false, true] {
|
||||
for (status, body, supported) in [
|
||||
(200, r#"{"metric_packed_import":1,"future":2}"#, true),
|
||||
(200, "{}", false),
|
||||
(200, r#"{"metric_packed_import":2}"#, false),
|
||||
(404, "{}", false),
|
||||
(401, "{}", false),
|
||||
(403, "{}", false),
|
||||
(200, "[]", false),
|
||||
(200, "invalid-json", false),
|
||||
] {
|
||||
let body = if export {
|
||||
body.replace("metric_packed_import", "metric_packed_export")
|
||||
} else {
|
||||
body.to_string()
|
||||
};
|
||||
let response_body = body.clone();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let body = response_body;
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0; 4096];
|
||||
let n = socket.read(&mut request).await.unwrap();
|
||||
let request = String::from_utf8_lossy(&request[..n]).to_lowercase();
|
||||
assert!(request.starts_with("get /v1/capabilities "));
|
||||
assert!(request.contains("authorization: basic dxnlcjpwyxnzd29yza=="));
|
||||
socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
});
|
||||
let client = super::DatabaseClient::new(
|
||||
address.to_string(),
|
||||
"greptime".into(),
|
||||
Some("user:password".into()),
|
||||
std::time::Duration::from_secs(5),
|
||||
None,
|
||||
true,
|
||||
);
|
||||
assert_eq!(
|
||||
if export {
|
||||
client.require_packed_export().await.is_ok()
|
||||
} else {
|
||||
client.require_packed_import().await.is_ok()
|
||||
},
|
||||
supported,
|
||||
"{status}: {body}"
|
||||
);
|
||||
server.await.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user