From 9d073c05ef1bd752d6c3baae1eacc91e8514f1b2 Mon Sep 17 00:00:00 2001 From: ByteBaker <42913098+ByteBaker@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:03:54 +0530 Subject: [PATCH 1/2] refactor: remove repeated storage and metadata code - Share cloud and local storage code so fixes stay in one place. - Reuse metadata helpers and migration steps without changing paths or errors. - Use existing derives and reduce repeated telemetry setup and assertions. - Remove 1,508 net lines while keeping existing behavior and input cases. --- src/alerts/alert_enums.rs | 115 +-- src/alerts/alert_structs.rs | 15 +- src/handlers/http/oidc.rs | 24 +- src/handlers/http/role.rs | 26 +- src/livetail.rs | 9 +- .../metastores/object_store_metastore.rs | 148 ++-- src/migration/stream_metadata_migration.rs | 101 +-- src/otel/logs.rs | 71 +- src/otel/metrics.rs | 201 +++-- src/otel/otel_utils.rs | 10 +- src/otel/traces.rs | 389 +++------- src/storage/azure_blob.rs | 628 ++------------- src/storage/gcs.rs | 651 ++-------------- src/storage/localfs.rs | 251 +++--- src/storage/mod.rs | 728 ++++++++++++++++- src/storage/s3.rs | 733 ++---------------- src/utils/json/flatten.rs | 348 +++------ 17 files changed, 1470 insertions(+), 2978 deletions(-) diff --git a/src/alerts/alert_enums.rs b/src/alerts/alert_enums.rs index 5824d8b84..03daea1a6 100644 --- a/src/alerts/alert_enums.rs +++ b/src/alerts/alert_enums.rs @@ -51,6 +51,7 @@ impl From<&str> for AlertVersion { } #[derive( + derive_more::Display, Debug, serde::Serialize, serde::Deserialize, @@ -71,101 +72,72 @@ pub enum Severity { Low, } -impl Display for Severity { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Severity::Critical => write!(f, "Critical"), - Severity::High => write!(f, "High"), - Severity::Medium => write!(f, "Medium"), - Severity::Low => write!(f, "Low"), - } - } -} - -#[derive(Debug, serde::Serialize, serde::Deserialize, Clone)] +#[derive(derive_more::Display, Debug, serde::Serialize, serde::Deserialize, Clone)] #[serde(rename_all = "camelCase")] pub enum LogicalOperator { + #[display("AND")] And, + #[display("OR")] Or, } -impl Display for LogicalOperator { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - LogicalOperator::And => write!(f, "AND"), - LogicalOperator::Or => write!(f, "OR"), - } - } -} - -#[derive(Debug, serde::Serialize, serde::Deserialize, Clone, PartialEq)] +#[derive(derive_more::Display, Debug, serde::Serialize, serde::Deserialize, Clone, PartialEq)] #[serde(rename_all = "camelCase")] pub enum AlertType { + #[display("threshold")] Threshold, + #[display("anomaly")] Anomaly(AnomalyConfig), + #[display("forecast")] Forecast(ForecastConfig), } -impl Display for AlertType { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AlertType::Threshold => write!(f, "threshold"), - AlertType::Anomaly(_) => write!(f, "anomaly"), - AlertType::Forecast(_) => write!(f, "forecast"), - } - } -} - -#[derive(Debug, serde::Serialize, serde::Deserialize, Clone, Copy, PartialEq, Eq, Default)] +#[derive( + derive_more::Display, + Debug, + serde::Serialize, + serde::Deserialize, + Clone, + Copy, + PartialEq, + Eq, + Default, +)] #[serde(rename_all = "camelCase")] pub enum AlertQueryType { #[default] + #[display("builder")] Builder, #[serde(alias = "sql")] + #[display("code")] Code, + #[display("promql")] Promql, } -impl Display for AlertQueryType { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AlertQueryType::Builder => write!(f, "builder"), - AlertQueryType::Code => write!(f, "code"), - AlertQueryType::Promql => write!(f, "promql"), - } - } -} - -#[derive(Debug, serde::Serialize, serde::Deserialize, Clone)] +#[derive(derive_more::Display, Debug, serde::Serialize, serde::Deserialize, Clone)] #[serde(rename_all = "camelCase")] pub enum AlertOperator { #[serde(rename = ">")] + #[display(">")] GreaterThan, #[serde(rename = "<")] + #[display("<")] LessThan, #[serde(rename = "=")] + #[display("=")] Equal, #[serde(rename = "!=")] + #[display("!=")] NotEqual, #[serde(rename = ">=")] + #[display(">=")] GreaterThanOrEqual, #[serde(rename = "<=")] + #[display("<=")] LessThanOrEqual, } -impl Display for AlertOperator { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AlertOperator::GreaterThan => write!(f, ">"), - AlertOperator::LessThan => write!(f, "<"), - AlertOperator::Equal => write!(f, "="), - AlertOperator::NotEqual => write!(f, "!="), - AlertOperator::GreaterThanOrEqual => write!(f, ">="), - AlertOperator::LessThanOrEqual => write!(f, "<="), - } - } -} - #[derive(Debug, serde::Serialize, serde::Deserialize, Clone, FromStr, PartialEq, Eq)] #[serde(rename_all = "camelCase")] pub enum WhereConfigOperator { @@ -292,7 +264,7 @@ mod tests { } } -#[derive(Debug, serde::Serialize, serde::Deserialize, Clone)] +#[derive(derive_more::Display, Debug, serde::Serialize, serde::Deserialize, Clone)] #[serde(rename_all = "camelCase")] pub enum AggregateFunction { Avg, @@ -303,19 +275,6 @@ pub enum AggregateFunction { Sum, } -impl Display for AggregateFunction { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AggregateFunction::Avg => write!(f, "Avg"), - AggregateFunction::Count => write!(f, "Count"), - AggregateFunction::CountDistinct => write!(f, "CountDistinct"), - AggregateFunction::Min => write!(f, "Min"), - AggregateFunction::Max => write!(f, "Max"), - AggregateFunction::Sum => write!(f, "Sum"), - } - } -} - #[derive(Debug, serde::Serialize, serde::Deserialize, Clone)] #[serde(rename_all = "camelCase")] pub enum EvalConfig { @@ -323,6 +282,7 @@ pub enum EvalConfig { } #[derive( + derive_more::Display, Debug, serde::Serialize, serde::Deserialize, @@ -337,23 +297,16 @@ pub enum EvalConfig { )] #[serde(rename_all = "camelCase")] pub enum AlertState { + #[display("triggered")] Triggered, #[default] #[serde(rename = "not-triggered")] + #[display("not-triggered")] NotTriggered, + #[display("disabled")] Disabled, } -impl Display for AlertState { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - AlertState::Triggered => write!(f, "triggered"), - AlertState::Disabled => write!(f, "disabled"), - AlertState::NotTriggered => write!(f, "not-triggered"), - } - } -} - #[derive(Debug, serde::Serialize, serde::Deserialize, Clone, PartialEq, Default)] #[serde(rename_all = "camelCase")] pub enum NotificationState { diff --git a/src/alerts/alert_structs.rs b/src/alerts/alert_structs.rs index 52407d1fd..0647dbcc9 100644 --- a/src/alerts/alert_structs.rs +++ b/src/alerts/alert_structs.rs @@ -647,7 +647,7 @@ pub struct NotificationStateRequest { } /// MTTR (Mean Time To Recovery) statistics -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize, Default)] #[serde(rename_all = "camelCase")] pub struct MTTRStats { /// Total number of incidents (triggered -> not-triggered cycles) @@ -709,19 +709,6 @@ impl MTTRStats { } } -impl Default for MTTRStats { - fn default() -> Self { - Self { - total_incidents: 0, - mean_seconds: 0.0, - median_seconds: 0.0, - min_seconds: 0.0, - max_seconds: 0.0, - recovery_times_seconds: Vec::new(), - } - } -} - /// Aggregated MTTR statistics across multiple alerts #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/src/handlers/http/oidc.rs b/src/handlers/http/oidc.rs index 1ce09c69e..32b8e0aeb 100644 --- a/src/handlers/http/oidc.rs +++ b/src/handlers/http/oidc.rs @@ -40,6 +40,7 @@ use serde::Deserialize; use ulid::Ulid; use url::Url; +use crate::handlers::http::modal::utils::rbac_utils::{get_metadata, put_metadata}; use crate::{ handlers::{ COOKIE_AGE_DAYS, SESSION_COOKIE_NAME, USER_COOKIE_NAME, USER_ID_COOKIE_NAME, @@ -52,7 +53,7 @@ use crate::{ map::{DEFAULT_ROLE, SessionKey}, user::{self, GroupUser, User, UserType}, }, - storage::{self, ObjectStorageError, StorageMetadata}, + storage::{ObjectStorageError, StorageMetadata}, utils::{ actix::extract_session_key_from_req, get_tenant_id_from_key, get_tenant_id_from_request, }, @@ -615,27 +616,6 @@ pub async fn update_user_if_changed( Ok(user) } -async fn get_metadata( - tenant_id: &Option, -) -> Result { - let metadata = PARSEABLE - .metastore - .get_parseable_metadata(tenant_id) - .await - .map_err(|e| ObjectStorageError::MetastoreError(Box::new(e.to_detail())))? - .ok_or_else(|| ObjectStorageError::Custom("parseable metadata not initialized".into()))?; - Ok(serde_json::from_slice::(&metadata)?) -} - -async fn put_metadata( - metadata: &StorageMetadata, - tenant_id: &Option, -) -> Result<(), ObjectStorageError> { - storage::put_remote_metadata(metadata, tenant_id).await?; - storage::put_staging_metadata(metadata, tenant_id)?; - Ok(()) -} - #[derive(Debug, thiserror::Error)] pub enum OIDCError { #[error("Failed to connect to storage: {0}")] diff --git a/src/handlers/http/role.rs b/src/handlers/http/role.rs index 969020c24..cb9caad73 100644 --- a/src/handlers/http/role.rs +++ b/src/handlers/http/role.rs @@ -25,12 +25,13 @@ use actix_web::{ web::{self, Json}, }; +use crate::handlers::http::modal::utils::rbac_utils::{get_metadata, put_metadata}; use crate::rbac::map::roles; use crate::rbac::role::model::{Role, RoleType, RoleUI}; use crate::{ - parseable::{DEFAULT_TENANT, PARSEABLE}, + parseable::DEFAULT_TENANT, rbac::map::{DEFAULT_ROLE, mut_roles, mut_sessions, read_user_groups, users}, - storage::{self, ObjectStorageError, StorageMetadata}, + storage::ObjectStorageError, utils::get_tenant_id_from_request, validator::{self, error::UsernameValidationError}, }; @@ -228,27 +229,6 @@ pub fn get_default_internal(tenant_id: &Option) -> serde_json::Value { res } -async fn get_metadata( - tenant_id: &Option, -) -> Result { - let metadata = PARSEABLE - .metastore - .get_parseable_metadata(tenant_id) - .await - .map_err(|e| ObjectStorageError::MetastoreError(Box::new(e.to_detail())))? - .ok_or_else(|| ObjectStorageError::Custom("parseable metadata not initialized".into()))?; - Ok(serde_json::from_slice::(&metadata)?) -} - -async fn put_metadata( - metadata: &StorageMetadata, - tenant_id: &Option, -) -> Result<(), ObjectStorageError> { - storage::put_remote_metadata(metadata, tenant_id).await?; - storage::put_staging_metadata(metadata, tenant_id)?; - Ok(()) -} - #[derive(Debug, thiserror::Error)] pub enum RoleError { #[error("Failed to connect to storage: {0}")] diff --git a/src/livetail.rs b/src/livetail.rs index 6ce7445e1..50416b802 100644 --- a/src/livetail.rs +++ b/src/livetail.rs @@ -34,6 +34,7 @@ pub static LIVETAIL: Lazy = Lazy::new(LiveTail::default); pub type LiveTailRegistry = RwLock>>; +#[derive(Default)] pub struct LiveTail { pipes: Arc, } @@ -61,14 +62,6 @@ impl LiveTail { } } -impl Default for LiveTail { - fn default() -> Self { - Self { - pipes: Arc::new(RwLock::default()), - } - } -} - #[derive(Debug, PartialEq)] pub enum Message { Record(RecordBatch), diff --git a/src/metastore/metastores/object_store_metastore.rs b/src/metastore/metastores/object_store_metastore.rs index 867544601..2a9c7bda5 100644 --- a/src/metastore/metastores/object_store_metastore.rs +++ b/src/metastore/metastores/object_store_metastore.rs @@ -71,6 +71,33 @@ pub struct ObjectStoreMetastore { pub storage: Arc, } +impl ObjectStoreMetastore { + async fn put_object( + &self, + obj: &dyn MetastoreObject, + tenant_id: &Option, + ) -> Result<(), MetastoreError> { + let path = obj.get_object_path(); + + Ok(self + .storage + .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) + .await?) + } + + async fn delete_object( + &self, + obj: &dyn MetastoreObject, + tenant_id: &Option, + ) -> Result<(), MetastoreError> { + let path = obj.get_object_path(); + Ok(self + .storage + .delete_object(&RelativePathBuf::from(path), tenant_id) + .await?) + } +} + pub fn is_missing_optional_dir(err: &ObjectStorageError) -> bool { match err { ObjectStorageError::NoSuchKey(_) => true, @@ -433,11 +460,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = RelativePathBuf::from(obj.get_object_path()); - Ok(self - .storage - .put_object(&path, to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } /// This function fetches all the llmconfigs from the underlying object store @@ -469,12 +492,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - - Ok(self - .storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } /// Delete an llmconfig @@ -483,11 +501,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - Ok(self - .storage - .delete_object(&RelativePathBuf::from(path), tenant_id) - .await?) + self.delete_object(obj, tenant_id).await } /// Fetch all dashboards @@ -533,13 +547,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - // we need the path to store in obj store - let path = obj.get_object_path(); - - Ok(self - .storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } /// Delete a dashboard @@ -548,11 +556,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - Ok(self - .storage - .delete_object(&RelativePathBuf::from(path), tenant_id) - .await?) + self.delete_object(obj, tenant_id).await } /// Fetch all chats @@ -607,13 +611,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - // we need the path to store in obj store - let path = obj.get_object_path(); - - Ok(self - .storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } /// Delete a chat @@ -622,11 +620,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - Ok(self - .storage - .delete_object(&RelativePathBuf::from(path), tenant_id) - .await?) + self.delete_object(obj, tenant_id).await } // for get filters, take care of migration and removal of incorrect/old filters @@ -744,13 +738,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - // we need the path to store in obj store - let path = obj.get_object_path(); - - Ok(self - .storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } /// Delete a filter @@ -759,12 +747,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - - Ok(self - .storage - .delete_object(&RelativePathBuf::from(path), tenant_id) - .await?) + self.delete_object(obj, tenant_id).await } /// Fetch an `ObjectStoreFormat` file @@ -1075,13 +1058,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - // we need the path to store in obj store - let path = obj.get_object_path(); - - Ok(self - .storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), tenant_id) - .await?) + self.put_object(obj, tenant_id).await } async fn delete_target( @@ -1089,13 +1066,7 @@ impl Metastore for ObjectStoreMetastore { obj: &dyn MetastoreObject, tenant_id: &Option, ) -> Result<(), MetastoreError> { - // we need the path to store in obj store - let path = obj.get_object_path(); - - Ok(self - .storage - .delete_object(&RelativePathBuf::from(path), tenant_id) - .await?) + self.delete_object(obj, tenant_id).await } async fn get_outbound_policies( @@ -1281,11 +1252,7 @@ impl Metastore for ObjectStoreMetastore { } async fn put_node_metadata(&self, obj: &dyn MetastoreObject) -> Result<(), MetastoreError> { - let path = obj.get_object_path(); - self.storage - .put_object(&RelativePathBuf::from(path), to_bytes(obj), &None) - .await?; - Ok(()) + self.put_object(obj, &None).await } async fn delete_node_metadata( @@ -1408,6 +1375,43 @@ mod tests { } } + #[tokio::test] + async fn object_path_persistence_keeps_serialization_tenants_and_errors() { + let dir = tempfile::tempdir().unwrap(); + let storage = Arc::new(LocalFS::new(dir.path().to_path_buf())); + let metastore = ObjectStoreMetastore { + storage: storage.clone(), + }; + for tenant in [None, Some("tenant".to_owned())] { + let mut object = TestMetastoreObject { + id: Ulid::new(), + path: format!("{}/object.json", tenant.as_deref().unwrap_or("default")), + }; + let path = RelativePathBuf::from(&object.path); + metastore.put_object(&object, &tenant).await.unwrap(); + let bytes = storage.get_object(&path, &tenant).await.unwrap(); + assert_eq!( + serde_json::from_slice::(&bytes).unwrap(), + serde_json::json!({"id": object.id.to_string(), "path": object.path}) + ); + metastore.delete_object(&object, &tenant).await.unwrap(); + assert!(matches!( + storage.get_object(&path, &tenant).await, + Err(ObjectStorageError::NoSuchKey(_)) + )); + assert!(matches!( + metastore.delete_object(&object, &tenant).await, + Err(MetastoreError::ObjectStorageError(_)) + )); + std::fs::write(dir.path().join("blocked"), "data").unwrap(); + object.path = "blocked/object.json".into(); + assert!(matches!( + metastore.put_object(&object, &tenant).await, + Err(MetastoreError::ObjectStorageError(_)) + )); + } + } + #[tokio::test] async fn delete_alert_uses_request_tenant_path_and_is_idempotent() { let temp_dir = tempfile::tempdir().unwrap(); diff --git a/src/migration/stream_metadata_migration.rs b/src/migration/stream_metadata_migration.rs index 13be5bec8..9c0cea0d9 100644 --- a/src/migration/stream_metadata_migration.rs +++ b/src/migration/stream_metadata_migration.rs @@ -24,8 +24,7 @@ use crate::{ }; use serde_json::{Value, json}; -pub fn v1_v4(mut stream_metadata: Value) -> Value { - let stream_metadata_map = stream_metadata.as_object_mut().unwrap(); +fn migrate_stats_and_version(stream_metadata_map: &mut serde_json::Map) { let stats = stream_metadata_map.get("stats").unwrap().clone(); let default_stats = json!({ "lifetime_stats": { @@ -53,46 +52,11 @@ pub fn v1_v4(mut stream_metadata: Value) -> Value { "objectstore-format".to_owned(), Value::String(storage::CURRENT_OBJECT_STORE_VERSION.into()), ); - stream_metadata_map.insert( - "snapshot".to_owned(), - json!({ - "version": CURRENT_SNAPSHOT_VERSION, - "manifest_list": [] - }), - ); - stream_metadata } -pub fn v2_v4(mut stream_metadata: Value) -> Value { +pub fn v1_v4(mut stream_metadata: Value) -> Value { let stream_metadata_map = stream_metadata.as_object_mut().unwrap(); - let stats = stream_metadata_map.get("stats").unwrap().clone(); - let default_stats = json!({ - "lifetime_stats": { - "events": stats.get("events").unwrap(), - "ingestion": stats.get("ingestion").unwrap(), - "storage": stats.get("storage").unwrap() - }, - "current_stats": { - "events": stats.get("events").unwrap(), - "ingestion": stats.get("ingestion").unwrap(), - "storage": stats.get("storage").unwrap() - }, - "deleted_stats": { - "events": 0, - "ingestion": 0, - "storage": 0 - } - }); - stream_metadata_map.insert("stats".to_owned(), default_stats); - stream_metadata_map.insert( - "version".to_owned(), - Value::String(storage::CURRENT_SCHEMA_VERSION.into()), - ); - stream_metadata_map.insert( - "objectstore-format".to_owned(), - Value::String(storage::CURRENT_OBJECT_STORE_VERSION.into()), - ); - + migrate_stats_and_version(stream_metadata_map); stream_metadata_map.insert( "snapshot".to_owned(), json!({ @@ -103,37 +67,14 @@ pub fn v2_v4(mut stream_metadata: Value) -> Value { stream_metadata } +pub fn v2_v4(stream_metadata: Value) -> Value { + v1_v4(stream_metadata) +} + pub fn v3_v4(mut stream_metadata: Value) -> Value { let stream_metadata_map: &mut serde_json::Map = stream_metadata.as_object_mut().unwrap(); - let stats = stream_metadata_map.get("stats").unwrap().clone(); - let default_stats = json!({ - "lifetime_stats": { - "events": stats.get("events").unwrap(), - "ingestion": stats.get("ingestion").unwrap(), - "storage": stats.get("storage").unwrap() - }, - "current_stats": { - "events": stats.get("events").unwrap(), - "ingestion": stats.get("ingestion").unwrap(), - "storage": stats.get("storage").unwrap() - }, - "deleted_stats": { - "events": 0, - "ingestion": 0, - "storage": 0 - } - }); - stream_metadata_map.insert("stats".to_owned(), default_stats); - stream_metadata_map.insert( - "version".to_owned(), - Value::String(storage::CURRENT_SCHEMA_VERSION.into()), - ); - stream_metadata_map.insert( - "objectstore-format".to_owned(), - Value::String(storage::CURRENT_OBJECT_STORE_VERSION.into()), - ); - + migrate_stats_and_version(stream_metadata_map); let snapshot = stream_metadata_map.get("snapshot").unwrap().clone(); let version = snapshot .as_object() @@ -297,6 +238,32 @@ fn v1_v2_snapshot_migration(mut snapshot: Value) -> Value { #[cfg(test)] mod tests { + use super::{v1_v4, v2_v4, v3_v4}; + + #[rstest::rstest] + #[case::v1(v1_v4)] + #[case::v2(v2_v4)] + #[case::v3(v3_v4)] + fn legacy_stats_and_unrelated_metadata_are_preserved( + #[case] migrate: fn(serde_json::Value) -> serde_json::Value, + ) { + let stats = serde_json::json!({"events": 7, "ingestion": 11, "storage": 13}); + let snapshot = serde_json::json!({"version": crate::catalog::snapshot::CURRENT_SNAPSHOT_VERSION, "manifest_list": []}); + let result = migrate( + serde_json::json!({"stats": stats, "snapshot": snapshot, "time_partition": "legacy_timestamp"}), + ); + assert_eq!( + result, + serde_json::json!({ + "version": crate::storage::CURRENT_SCHEMA_VERSION, + "objectstore-format": crate::storage::CURRENT_OBJECT_STORE_VERSION, + "stats": {"lifetime_stats": stats, "current_stats": stats, "deleted_stats": {"events": 0, "ingestion": 0, "storage": 0}}, + "snapshot": snapshot, + "time_partition": "legacy_timestamp" + }) + ); + } + #[test] fn test_v5_v6_with_log_source() { let stream_metadata = serde_json::json!({"version":"v5","schema_version":"v0","objectstore-format":"v5","created-at":"2025-03-10T14:38:29.355131524-04:00","first-event-at":"2025-03-10T14:38:29.356-04:00","owner":{"id":"admin","group":"admin"},"permissions":[{"id":"admin","group":"admin","access":["all"]}],"stats":{"lifetime_stats":{"events":3,"ingestion":70,"storage":1969},"current_stats":{"events":3,"ingestion":70,"storage":1969},"deleted_stats":{"events":0,"ingestion":0,"storage":0}},"snapshot":{"version":"v2","manifest_list":[{"manifest_path":"home/nikhilsinha/Parseable/parseable/data/test10/date=2025-03-10/manifest.json","time_lower_bound":"2025-03-10T00:00:00Z","time_upper_bound":"2025-03-10T23:59:59.999999999Z","events_ingested":3,"ingestion_size":70,"storage_size":1969}]},"hot_tier_enabled":false,"stream_type":"UserDefined","log_source":"OtelLogs"}); diff --git a/src/otel/logs.rs b/src/otel/logs.rs index a7cbc0b42..bb88ef757 100644 --- a/src/otel/logs.rs +++ b/src/otel/logs.rs @@ -113,16 +113,10 @@ fn severity_number_from_text(severity_text: &str) -> Option { /// Add the normalized severity fields to the flattened JSON. fn flatten_severity(severity_number: i32, severity_text: &str) -> Map { - let mut severity_json: Map = Map::new(); - severity_json.insert( - "severity_number".to_string(), - Value::Number(severity_number.into()), - ); - severity_json.insert( - "severity_text".to_string(), - Value::String(severity_text.to_string()), - ); - severity_json + Map::from_iter([ + ("severity_number".into(), severity_number.into()), + ("severity_text".into(), severity_text.to_string().into()), + ]) } /// Maps OTel severity_number (0–24) to a log category. @@ -171,19 +165,16 @@ fn category_from_body(body_str: &str) -> &'static str { /// and returns a `Map` of the flattened json /// this function is called recursively for each log record object in the otel logs pub fn flatten_log_record(log_record: &LogRecord) -> Map { - let mut log_record_json: Map = Map::new(); - log_record_json.insert( - "time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - log_record.time_unix_nano as i64, - )), - ); - log_record_json.insert( - "observed_time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - log_record.observed_time_unix_nano as i64, - )), - ); + let mut log_record_json = Map::from_iter([ + ( + "time_unix_nano".into(), + convert_epoch_nano_to_timestamp(log_record.time_unix_nano as i64).into(), + ), + ( + "observed_time_unix_nano".into(), + convert_epoch_nano_to_timestamp(log_record.observed_time_unix_nano as i64).into(), + ), + ]); let (severity_number, severity_text) = normalize_severity(log_record.severity_number, &log_record.severity_text); @@ -233,27 +224,19 @@ pub fn flatten_log_record(log_record: &LogRecord) -> Map { insert_attributes(&mut log_record_json, &log_record.attributes); // Insert after attributes so a client-sent "p_log_category" cannot override - log_record_json.insert( - "p_log_category".to_string(), - Value::String(log_category.unwrap_or("UNSPECIFIED").to_string()), - ); - log_record_json.insert( - "log_record_dropped_attributes_count".to_string(), - Value::Number(log_record.dropped_attributes_count.into()), - ); - - log_record_json.insert( - "flags".to_string(), - Value::Number((log_record.flags).into()), - ); - log_record_json.insert( - "span_id".to_string(), - Value::String(hex::encode(&log_record.span_id)), - ); - log_record_json.insert( - "trace_id".to_string(), - Value::String(hex::encode(&log_record.trace_id)), - ); + log_record_json.extend([ + ( + "p_log_category".into(), + log_category.unwrap_or("UNSPECIFIED").to_string().into(), + ), + ( + "log_record_dropped_attributes_count".into(), + log_record.dropped_attributes_count.into(), + ), + ("flags".into(), (log_record.flags).into()), + ("span_id".into(), hex::encode(&log_record.span_id).into()), + ("trace_id".into(), hex::encode(&log_record.trace_id).into()), + ]); log_record_json } diff --git a/src/otel/metrics.rs b/src/otel/metrics.rs index 6684dd1bf..b96841217 100644 --- a/src/otel/metrics.rs +++ b/src/otel/metrics.rs @@ -137,21 +137,17 @@ fn target_info_record( timestamp: String, ) -> Map { let mut target = resource_fields.clone(); - target.insert( - "metric_name".to_string(), - Value::String("target_info".to_string()), - ); - target.insert( - "metric_description".to_string(), - Value::String("Target metadata".to_string()), - ); - target.insert("metric_unit".to_string(), Value::String(String::default())); - target.insert( - "metric_type".to_string(), - Value::String("gauge".to_string()), - ); - target.insert("time_unix_nano".to_string(), Value::String(timestamp)); - target.insert("data_point_value".to_string(), Value::Number(1.into())); + target.extend([ + ("metric_name".into(), "target_info".to_string().into()), + ( + "metric_description".into(), + "Target metadata".to_string().into(), + ), + ("metric_unit".into(), String::default().into()), + ("metric_type".into(), "gauge".to_string().into()), + ("time_unix_nano".into(), timestamp.into()), + ("data_point_value".into(), 1.into()), + ]); let series_hash = compute_series_hash(&target); target.insert( SERIES_HASH_COLUMN.to_string(), @@ -306,18 +302,16 @@ fn flatten_number_data_points( .map(|data_point| { let mut data_point_json = Map::with_capacity(data_point.attributes.len() + 8); insert_attributes(&mut data_point_json, &data_point.attributes); - data_point_json.insert( - "start_time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - data_point.start_time_unix_nano as i64, - )), - ); - data_point_json.insert( - "time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - data_point.time_unix_nano as i64, - )), - ); + data_point_json.extend([ + ( + "start_time_unix_nano".into(), + convert_epoch_nano_to_timestamp(data_point.start_time_unix_nano as i64).into(), + ), + ( + "time_unix_nano".into(), + convert_epoch_nano_to_timestamp(data_point.time_unix_nano as i64).into(), + ), + ]); insert_exemplars( &mut data_point_json, @@ -379,22 +373,17 @@ fn flatten_histogram(histogram: &Histogram, flatten_exemplars: bool) -> Vec Vec> { for data_point in &summary.data_points { let mut data_point_json = Map::with_capacity(data_point.attributes.len() + 6); insert_attributes(&mut data_point_json, &data_point.attributes); - data_point_json.insert( - "start_time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - data_point.start_time_unix_nano as i64, - )), - ); - data_point_json.insert( - "time_unix_nano".to_string(), - Value::String(convert_epoch_nano_to_timestamp( - data_point.time_unix_nano as i64, - )), - ); - data_point_json.insert( - "data_point_count".to_string(), - Value::Number(data_point.count.into()), - ); - data_point_json.insert( - "data_point_sum".to_string(), - serde_json::Number::from_f64(data_point.sum) - .map(Value::Number) - .unwrap_or(Value::Null), - ); - data_point_json.insert( - "data_point_quantile_values".to_string(), - Value::Array( - data_point - .quantile_values - .iter() - .map(|quantile_value| { - let mut quantile_map = Map::with_capacity(2); - quantile_map.insert( - "quantile".to_string(), - serde_json::Number::from_f64(quantile_value.quantile) - .map(Value::Number) - .unwrap_or(Value::Null), - ); - quantile_map.insert( - "value".to_string(), - serde_json::Number::from_f64(quantile_value.value) - .map(Value::Number) - .unwrap_or(Value::Null), - ); - Value::Object(quantile_map) - }) - .collect(), + data_point_json.extend([ + ( + "start_time_unix_nano".into(), + convert_epoch_nano_to_timestamp(data_point.start_time_unix_nano as i64).into(), ), - ); + ( + "time_unix_nano".into(), + convert_epoch_nano_to_timestamp(data_point.time_unix_nano as i64).into(), + ), + ("data_point_count".into(), data_point.count.into()), + ( + "data_point_sum".into(), + serde_json::Number::from_f64(data_point.sum) + .map(Value::Number) + .unwrap_or(Value::Null), + ), + ( + "data_point_quantile_values".into(), + Value::Array( + data_point + .quantile_values + .iter() + .map(|quantile_value| { + let mut quantile_map = Map::with_capacity(2); + quantile_map.insert( + "quantile".to_string(), + serde_json::Number::from_f64(quantile_value.quantile) + .map(Value::Number) + .unwrap_or(Value::Null), + ); + quantile_map.insert( + "value".to_string(), + serde_json::Number::from_f64(quantile_value.value) + .map(Value::Number) + .unwrap_or(Value::Null), + ); + Value::Object(quantile_map) + }) + .collect(), + ), + ), + ]); data_points_json.push(data_point_json); } diff --git a/src/otel/otel_utils.rs b/src/otel/otel_utils.rs index c267235a1..1a022da8c 100644 --- a/src/otel/otel_utils.rs +++ b/src/otel/otel_utils.rs @@ -219,14 +219,8 @@ pub fn value_to_string(value: serde_json::Value) -> String { } pub fn flatten_attributes(attributes: &[KeyValue]) -> Map { - let mut attributes_json: Map = Map::with_capacity(attributes.len()); - for attribute in attributes { - if let Some(value) = &attribute.value - && let Some(value) = value.value.clone() - { - insert_json_from_value(&mut attributes_json, &attribute.key, value); - } - } + let mut attributes_json = Map::with_capacity(attributes.len()); + insert_attributes(&mut attributes_json, attributes); attributes_json } diff --git a/src/otel/traces.rs b/src/otel/traces.rs index 5aba0b248..6ca6fc8c4 100644 --- a/src/otel/traces.rs +++ b/src/otel/traces.rs @@ -228,15 +228,10 @@ fn flatten_links(links: &[Link]) -> Vec> { links .iter() .map(|link| { - let mut link_json = Map::new(); - link_json.insert( - "link_span_id".to_string(), - Value::String(hex::encode(&link.span_id)), - ); - link_json.insert( - "link_trace_id".to_string(), - Value::String(hex::encode(&link.trace_id)), - ); + let mut link_json = Map::from_iter([ + ("link_span_id".into(), hex::encode(&link.span_id).into()), + ("link_trace_id".into(), hex::encode(&link.trace_id).into()), + ]); insert_links_attributes(&mut link_json, &link.attributes); link_json.insert( @@ -253,27 +248,17 @@ fn flatten_links(links: &[Link]) -> Vec> { /// this function fetches the status description from the status code /// and adds it to the flattened json fn flatten_status(status: &Status) -> Map { - let mut status_json = Map::new(); - status_json.insert( - "span_status_message".to_string(), - Value::String(status.message.clone()), - ); - status_json.insert( - "span_status_code".to_string(), - Value::Number(status.code.into()), - ); let description = match status.code { 0 => "UNSET", 1 => "OK", 2 => "ERROR", _ => "", }; - status_json.insert( - "span_status_description".to_string(), - Value::String(description.to_string()), - ); - - status_json + Map::from_iter([ + ("span_status_message".into(), status.message.clone().into()), + ("span_status_code".into(), status.code.into()), + ("span_status_description".into(), description.into()), + ]) } /// otel log event has flags @@ -281,8 +266,6 @@ fn flatten_status(status: &Status) -> Map { /// this function fetches the flags description from the flags /// and adds it to the flattened json fn flatten_flags(flags: u32) -> Map { - let mut flags_json = Map::new(); - flags_json.insert("span_flags".to_string(), Value::Number(flags.into())); let description = match flags { 0 => "DO_NOT_USE", 255 => "TRACE_FLAGS_MASK", @@ -290,12 +273,10 @@ fn flatten_flags(flags: u32) -> Map { 512 => "CONTEXT_IS_REMOTE_MASK", _ => "", }; - flags_json.insert( - "span_flags_description".to_string(), - Value::String(description.to_string()), - ); - - flags_json + Map::from_iter([ + ("span_flags".into(), flags.into()), + ("span_flags_description".into(), description.into()), + ]) } /// otel span event has kind @@ -303,8 +284,6 @@ fn flatten_flags(flags: u32) -> Map { /// this function fetches the kind description from the kind /// and adds it to the flattened json fn flatten_kind(kind: i32) -> Map { - let mut kind_json = Map::new(); - kind_json.insert("span_kind".to_string(), Value::Number(kind.into())); let description = match kind { 0 => "UNSPECIFIED", 1 => "INTERNAL", @@ -314,12 +293,10 @@ fn flatten_kind(kind: i32) -> Map { 5 => "CONSUMER", _ => "", }; - kind_json.insert( - "span_kind_description".to_string(), - Value::String(description.to_string()), - ); - - kind_json + Map::from_iter([ + ("span_kind".into(), kind.into()), + ("span_kind_description".into(), description.into()), + ]) } /// this function flattens the `Span` object @@ -497,21 +474,14 @@ mod tests { let result = flatten_status(&status); + assert_eq!(result["span_status_code"], serde_json::json!(code)); assert_eq!( - result.get("span_status_code").unwrap(), - &Value::Number(code.into()), - "Status code should be preserved" - ); - assert_eq!( - result.get("span_status_description").unwrap(), - &Value::String(expected_description.to_string()), - "Status description should match expected value for code {}", - code + result["span_status_description"], + serde_json::json!(expected_description) ); assert_eq!( - result.get("span_status_message").unwrap(), - &Value::String("Test message".to_string()), - "Status message should be preserved" + result["span_status_message"], + serde_json::json!("Test message") ); } } @@ -532,16 +502,10 @@ mod tests { for (kind, expected_description) in test_cases { let result = flatten_kind(kind); + assert_eq!(result["span_kind"], serde_json::json!(kind)); assert_eq!( - result.get("span_kind").unwrap(), - &Value::Number(kind.into()), - "Span kind should be preserved" - ); - assert_eq!( - result.get("span_kind_description").unwrap(), - &Value::String(expected_description.to_string()), - "Span kind description should match expected value for kind {}", - kind + result["span_kind_description"], + serde_json::json!(expected_description) ); } } @@ -560,16 +524,10 @@ mod tests { for (flags, expected_description) in test_cases { let result = flatten_flags(flags); + assert_eq!(result["span_flags"], serde_json::json!(flags)); assert_eq!( - result.get("span_flags").unwrap(), - &Value::Number(flags.into()), - "Span flags should be preserved" - ); - assert_eq!( - result.get("span_flags_description").unwrap(), - &Value::String(expected_description.to_string()), - "Span flags description should match expected value for flags {}", - flags + result["span_flags_description"], + serde_json::json!(expected_description) ); } } @@ -587,8 +545,7 @@ mod tests { Event { time_unix_nano: 1640995201000000000, name: "request.end".to_string(), - attributes: vec![], - dropped_attributes_count: 0, + ..Default::default() }, ]; @@ -598,36 +555,23 @@ mod tests { // Check first event let first_event = &result[0]; - assert!( - first_event.contains_key("event_time_unix_nano"), - "Should contain timestamp" - ); + assert!(first_event.contains_key("event_time_unix_nano")); assert_eq!( - first_event.get("event_name").unwrap(), - &Value::String("request.start".to_string()), - "Event name should be preserved" + first_event["event_name"], + serde_json::json!("request.start") ); assert_eq!( - first_event.get("event_dropped_attributes_count").unwrap(), - &Value::Number(2.into()), - "Dropped attributes count should be preserved" - ); - assert!( - first_event.contains_key("event_service.name"), - "Should contain flattened attributes" + first_event["event_dropped_attributes_count"], + serde_json::json!(2) ); + assert!(first_event.contains_key("event_service.name")); // Check second event let second_event = &result[1]; + assert_eq!(second_event["event_name"], serde_json::json!("request.end")); assert_eq!( - second_event.get("event_name").unwrap(), - &Value::String("request.end".to_string()), - "Second event name should be preserved" - ); - assert_eq!( - second_event.get("event_dropped_attributes_count").unwrap(), - &Value::Number(0.into()), - "Second event dropped count should be zero" + second_event["event_dropped_attributes_count"], + serde_json::json!(0) ); } @@ -649,24 +593,12 @@ mod tests { let link = &result[0]; assert_eq!( - link.get("link_trace_id").unwrap(), - &Value::String("0102030405060708090a0b0c0d0e0f10".to_string()), - "Trace ID should be hex encoded" - ); - assert_eq!( - link.get("link_span_id").unwrap(), - &Value::String("0102030405060708".to_string()), - "Span ID should be hex encoded" - ); - assert_eq!( - link.get("link_dropped_attributes_count").unwrap(), - &Value::Number(1.into()), - "Dropped attributes count should be preserved" - ); - assert!( - link.contains_key("link_service.name"), - "Should contain flattened attributes" + link["link_trace_id"], + serde_json::json!("0102030405060708090a0b0c0d0e0f10") ); + assert_eq!(link["link_span_id"], serde_json::json!("0102030405060708")); + assert_eq!(link["link_dropped_attributes_count"], serde_json::json!(1)); + assert!(link.contains_key("link_service.name")); } #[test] @@ -676,34 +608,26 @@ mod tests { trace_id: sample_trace_id(), span_id: sample_span_id(), trace_state: "test-state".to_string(), - parent_span_id: vec![], - flags: 0, name: "test-span".to_string(), kind: 2, // SERVER start_time_unix_nano: 1640995200000000000, end_time_unix_nano: 1640995201000000000, attributes: sample_attributes(), - dropped_attributes_count: 0, events: vec![Event { time_unix_nano: 1640995200500000000, name: "middleware".to_string(), - attributes: vec![], - dropped_attributes_count: 0, + ..Default::default() }], - dropped_events_count: 0, links: vec![Link { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - attributes: vec![], - dropped_attributes_count: 0, - flags: 0, + ..Default::default() }], - dropped_links_count: 0, status: Some(Status { message: "OK".to_string(), code: 1, }), + ..Default::default() }; let result = flatten_span_record(&span); @@ -713,45 +637,15 @@ mod tests { // Both records should contain span-level information for record in &result { - assert_eq!( - record.get("span_name").unwrap(), - &Value::String("test-span".to_string()), - "All records should contain span name" - ); - assert_eq!( - record.get("span_kind").unwrap(), - &Value::Number(2.into()), - "All records should contain span kind" - ); - assert_eq!( - record.get("span_kind_description").unwrap(), - &Value::String("SERVER".to_string()), - "All records should contain span kind description" - ); - assert!( - record.contains_key("span_trace_id"), - "Should contain trace ID" - ); - assert!( - record.contains_key("span_span_id"), - "Should contain span ID" - ); - assert!( - record.contains_key("span_start_time_unix_nano"), - "Should contain start time" - ); - assert!( - record.contains_key("span_end_time_unix_nano"), - "Should contain end time" - ); - assert!( - record.contains_key("service.name"), - "Should contain span attributes" - ); - assert!( - record.contains_key("span_status_code"), - "Should contain status" - ); + assert_eq!(record["span_name"], serde_json::json!("test-span")); + assert_eq!(record["span_kind"], serde_json::json!(2)); + assert_eq!(record["span_kind_description"], serde_json::json!("SERVER")); + assert!(record.contains_key("span_trace_id")); + assert!(record.contains_key("span_span_id")); + assert!(record.contains_key("span_start_time_unix_nano")); + assert!(record.contains_key("span_end_time_unix_nano")); + assert!(record.contains_key("service.name")); + assert!(record.contains_key("span_status_code")); } // One record should be an event, one should be a link @@ -767,20 +661,11 @@ mod tests { let span = Span { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - parent_span_id: vec![], - flags: 0, name: "simple-span".to_string(), kind: 1, // INTERNAL start_time_unix_nano: 1640995200000000000, end_time_unix_nano: 1640995201000000000, - attributes: vec![], - dropped_attributes_count: 0, - events: vec![], - dropped_events_count: 0, - links: vec![], - dropped_links_count: 0, - status: None, + ..Default::default() }; let result = flatten_span_record(&span); @@ -792,23 +677,10 @@ mod tests { ); let record = &result[0]; - assert_eq!( - record.get("span_name").unwrap(), - &Value::String("simple-span".to_string()), - "Should contain span name" - ); - assert!( - !record.contains_key("event_name"), - "Should not contain event fields" - ); - assert!( - !record.contains_key("link_trace_id"), - "Should not contain link fields" - ); - assert!( - !record.contains_key("span_status_code"), - "Should not contain status when none provided" - ); + assert_eq!(record["span_name"], serde_json::json!("simple-span")); + assert!(!record.contains_key("event_name")); + assert!(!record.contains_key("link_trace_id")); + assert!(!record.contains_key("span_status_code")); } #[test] @@ -820,27 +692,14 @@ mod tests { let span = Span { trace_id: trace_id.clone(), span_id: span_id.clone(), - trace_state: "".to_string(), parent_span_id: vec![0x87, 0x65, 0x43, 0x21], - flags: 0, name: "hex-test".to_string(), - kind: 0, - start_time_unix_nano: 0, - end_time_unix_nano: 0, - attributes: vec![], - dropped_attributes_count: 0, - events: vec![], - dropped_events_count: 0, links: vec![Link { trace_id: trace_id.clone(), span_id: span_id.clone(), - trace_state: "".to_string(), - attributes: vec![], - dropped_attributes_count: 0, - flags: 0, + ..Default::default() }], - dropped_links_count: 0, - status: None, + ..Default::default() }; let result = flatten_span_record(&span); @@ -894,30 +753,22 @@ mod tests { scope: Some(opentelemetry_proto::tonic::common::v1::InstrumentationScope { name: "test-tracer".to_string(), version: "1.0.0".to_string(), - attributes: vec![], - dropped_attributes_count: 0, - }), +..Default::default() +}), spans: vec![Span { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - parent_span_id: vec![], - flags: 0, name: "integration-test-span".to_string(), kind: 3, // CLIENT start_time_unix_nano: 1640995200000000000, end_time_unix_nano: 1640995201000000000, attributes: sample_attributes(), - dropped_attributes_count: 0, - events: vec![], - dropped_events_count: 0, - links: vec![], - dropped_links_count: 0, status: Some(Status { message: "Success".to_string(), code: 1, }), - }], +..Default::default() +}], schema_url: "https://opentelemetry.io/schemas/1.21.0".to_string(), }], schema_url: "https://opentelemetry.io/schemas/1.21.0".to_string(), @@ -932,54 +783,30 @@ mod tests { // Verify resource-level fields assert_eq!( - record.get("deployment.environment").unwrap(), - &Value::String("production".to_string()), - "Should contain resource attributes" + record["deployment.environment"], + serde_json::json!("production") ); assert_eq!( - record.get("resource_schema_url").unwrap(), - &Value::String("https://opentelemetry.io/schemas/1.21.0".to_string()), - "Should contain resource schema URL" + record["resource_schema_url"], + serde_json::json!("https://opentelemetry.io/schemas/1.21.0") ); // Verify scope-level fields + assert_eq!(record["scope_name"], serde_json::json!("test-tracer")); + assert_eq!(record["scope_version"], serde_json::json!("1.0.0")); assert_eq!( - record.get("scope_name").unwrap(), - &Value::String("test-tracer".to_string()), - "Should contain scope name" - ); - assert_eq!( - record.get("scope_version").unwrap(), - &Value::String("1.0.0".to_string()), - "Should contain scope version" - ); - assert_eq!( - record.get("scope_schema_url").unwrap(), - &Value::String("https://opentelemetry.io/schemas/1.21.0".to_string()), - "Should contain scope schema URL" + record["scope_schema_url"], + serde_json::json!("https://opentelemetry.io/schemas/1.21.0") ); // Verify span-level fields assert_eq!( - record.get("span_name").unwrap(), - &Value::String("integration-test-span".to_string()), - "Should contain span name" - ); - assert_eq!( - record.get("span_kind_description").unwrap(), - &Value::String("CLIENT".to_string()), - "Should contain span kind description" - ); - assert_eq!( - record.get("service.name").unwrap(), - &Value::String("test-service".to_string()), - "Should contain span attributes" - ); - assert_eq!( - record.get("span_status_description").unwrap(), - &Value::String("OK".to_string()), - "Should contain status description" + record["span_name"], + serde_json::json!("integration-test-span") ); + assert_eq!(record["span_kind_description"], serde_json::json!("CLIENT")); + assert_eq!(record["service.name"], serde_json::json!("test-service")); + assert_eq!(record["span_status_description"], serde_json::json!("OK")); } #[test] @@ -1066,53 +893,33 @@ mod tests { let span = Span { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - parent_span_id: vec![], - flags: 0, name: "events-only-span".to_string(), kind: 1, start_time_unix_nano: 1640995200000000000, end_time_unix_nano: 1640995201000000000, - attributes: vec![], - dropped_attributes_count: 0, events: vec![ Event { time_unix_nano: 1640995200100000000, name: "evt-a".to_string(), - attributes: vec![], - dropped_attributes_count: 0, + ..Default::default() }, Event { time_unix_nano: 1640995200200000000, name: "evt-b".to_string(), - attributes: vec![], - dropped_attributes_count: 0, + ..Default::default() }, ], - dropped_events_count: 0, - links: vec![], - dropped_links_count: 0, - status: None, + ..Default::default() }; let result = flatten_span_record(&span); assert_eq!(result.len(), 2, "One record per event when links are empty"); - assert_eq!( - result[0].get("event_name").unwrap(), - &Value::String("evt-a".to_string()), - ); - assert_eq!( - result[1].get("event_name").unwrap(), - &Value::String("evt-b".to_string()), - ); + assert_eq!(result[0]["event_name"], serde_json::json!("evt-a")); + assert_eq!(result[1]["event_name"], serde_json::json!("evt-b")); // Each event record must also carry the span-level fields. for record in &result { - assert_eq!( - record.get("span_name").unwrap(), - &Value::String("events-only-span".to_string()), - "Event records must be enriched with span fields" - ); + assert_eq!(record["span_name"], serde_json::json!("events-only-span")); } } @@ -1123,40 +930,22 @@ mod tests { let span = Span { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - parent_span_id: vec![], - flags: 0, name: "links-only-span".to_string(), kind: 3, start_time_unix_nano: 1640995200000000000, end_time_unix_nano: 1640995201000000000, - attributes: vec![], - dropped_attributes_count: 0, - events: vec![], - dropped_events_count: 0, links: vec![Link { trace_id: sample_trace_id(), span_id: sample_span_id(), - trace_state: "".to_string(), - attributes: vec![], - dropped_attributes_count: 0, - flags: 0, + ..Default::default() }], - dropped_links_count: 0, - status: None, + ..Default::default() }; let result = flatten_span_record(&span); assert_eq!(result.len(), 1, "One record per link when events are empty"); - assert!( - result[0].contains_key("link_trace_id"), - "Link record must contain link fields" - ); - assert_eq!( - result[0].get("span_name").unwrap(), - &Value::String("links-only-span".to_string()), - "Link records must be enriched with span fields" - ); + assert!(result[0].contains_key("link_trace_id")); + assert_eq!(result[0]["span_name"], serde_json::json!("links-only-span")); } } diff --git a/src/storage/azure_blob.rs b/src/storage/azure_blob.rs index 6c0bc776b..09ea4734d 100644 --- a/src/storage/azure_blob.rs +++ b/src/storage/azure_blob.rs @@ -19,7 +19,6 @@ use std::{ collections::HashSet, num::NonZeroUsize, - ops::Range, path::{Path, PathBuf}, sync::Arc, time::Duration, @@ -35,7 +34,7 @@ use datafusion::{ runtime_env::RuntimeEnvBuilder, }, }; -use futures::{StreamExt, TryStreamExt, stream::FuturesUnordered}; + use object_store::{ BackoffConfig, ClientOptions, ListResult, ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload, RetryConfig, @@ -50,23 +49,16 @@ use tracing::error; use url::Url; use crate::{ - metrics::{ - increment_bytes_scanned_in_object_store_calls_by_date, - increment_files_scanned_in_object_store_calls_by_date, - increment_object_store_calls_by_date, - increment_partial_file_scans_in_object_store_calls_by_date, - }, + metrics::increment_object_store_calls_by_date, parseable::{DEFAULT_TENANT, LogStream, PARSEABLE}, storage::RETRY_TIMEOUT_SECS, }; use super::{ CONNECT_TIMEOUT_SECS, DEFAULT_MAX_OBJECT_STORE_REQUESTS, ObjectStorage, ObjectStorageError, - ObjectStorageProvider, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, - STREAM_METADATA_FILE_NAME, metrics_layer::MetricLayer, object_storage::parseable_json_path, - partial_path, to_object_store_path, + ObjectStorageProvider, REQUEST_TIMEOUT_SECS, metrics_layer::MetricLayer, + object_storage::parseable_json_path, partial_path, to_object_store_path, }; -use crate::storage::GET_OBJECTS_CONCURRENCY; #[derive(Debug, Clone, clap::Args)] #[command( @@ -262,66 +254,15 @@ impl BlobStore { tenant_id: &Option, partial_path: PathBuf, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let src = to_object_store_path(path); - - let meta = self.client.head(&src).await?; - increment_object_store_calls_by_date("HEAD", &date, tenant_str); - let total = meta.size; - - if let Some(parent) = partial_path.parent() { - tokio::fs::create_dir_all(parent).await?; - } - let file = tokio::fs::File::create(&partial_path).await?; - file.set_len(total).await?; - let std_file = Arc::new(file.into_std().await); - - let chunk = PARSEABLE.options.hot_tier_download_chunk_size; - let concurrency = PARSEABLE.options.hot_tier_download_concurrency; - let ranges: Vec> = (0..total) - .step_by(chunk as usize) - .map(|s| s..(s + chunk).min(total)) - .collect(); - let chunk_count = ranges.len() as u64; - tracing::Span::current() - .record("total_bytes", total) - .record("chunks", chunk_count); - let semaphore = Arc::new(tokio::sync::Semaphore::new(concurrency as usize)); - - futures::stream::iter(ranges) - .map(|r| { - let src = src.clone(); - let std_file = std_file.clone(); - let semaphore = semaphore.clone(); - async move { - let _permit = semaphore.acquire_owned().await.map_err(|e| { - ObjectStorageError::Custom(format!("semaphore closed: {e}")) - })?; - let bytes = self.client.get_range(&src, r.clone()).await?; - let offset = r.start; - tokio::task::spawn_blocking(move || -> std::io::Result<()> { - crate::storage::write_all_at(&std_file, &bytes, offset) - }) - .await - .map_err(|e| ObjectStorageError::Custom(format!("join: {e}")))??; - Ok::<_, ObjectStorageError>(()) - } - }) - .buffer_unordered(concurrency as usize) - .try_collect::>() - .await?; - - increment_object_store_calls_by_date("GET", &date, tenant_str); - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant_str); - increment_bytes_scanned_in_object_store_calls_by_date("GET", total, &date, tenant_str); - increment_partial_file_scans_in_object_store_calls_by_date( - "GET", - chunk_count, - &date, - tenant_str, - ); - Ok(()) + super::parallel_download( + self.client.as_ref(), + path, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + partial_path, + PARSEABLE.options.hot_tier_download_chunk_size, + PARSEABLE.options.hot_tier_download_concurrency, + ) + .await } #[tracing::instrument( @@ -334,30 +275,7 @@ impl BlobStore { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.get(&to_object_store_path(path)).await; - increment_object_store_calls_by_date("GET", &Utc::now().date_naive().to_string(), tenant); - - match resp { - Ok(resp) => { - let body: Bytes = resp.bytes().await?; - tracing::Span::current().record("bytes", body.len()); - increment_files_scanned_in_object_store_calls_by_date( - "GET", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(body) - } - Err(err) => Err(err.into()), - } + super::get_object(self.client.as_ref(), path, tenant_id).await } #[tracing::instrument( @@ -371,21 +289,8 @@ impl BlobStore { resource: PutPayload, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.put(&to_object_store_path(path), resource).await; - increment_object_store_calls_by_date("PUT", &Utc::now().date_naive().to_string(), tenant); - match resp { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + let result = self.client.put(&to_object_store_path(path), resource).await; + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -399,54 +304,12 @@ impl BlobStore { key: &str, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - - // List all objects under the prefix - let object_stream = self.client.list(Some(&(key.into()))); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - // Map listed objects to their locations for delete_stream - let locations_stream = object_stream.map(|result| result.map(|obj| obj.location)); - - // Use delete_stream which batches deletes via the bulk delete API - let mut delete_results = self.client.delete_stream(Box::pin(locations_stream)); - - let mut files_deleted: u64 = 0; - let mut failed_deletes: u64 = 0; - while let Some(result) = delete_results.next().await { - match result { - Ok(_) => files_deleted += 1, - Err(err) => { - failed_deletes += 1; - error!("Failed to delete object during prefix deletion: {:?}", err); - } - } - } - - let total_files = files_deleted + failed_deletes; - tracing::Span::current() - .record("deleted", files_deleted) - .record("failed", failed_deletes); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - total_files, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - files_deleted, - &Utc::now().date_naive().to_string(), - tenant, - ); - - if failed_deletes > 0 { - return Err(ObjectStorageError::UnhandledError( - anyhow::anyhow!("Failed to delete {failed_deletes} out of {total_files} objects under prefix '{key}'").into(), - )); - } - - Ok(()) + super::delete_prefix( + self.client.as_ref(), + key, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + ) + .await } #[tracing::instrument( @@ -459,38 +322,13 @@ impl BlobStore { stream: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let resp: Result = self - .client - .list_with_delimiter(Some(&(stream.into()))) - .await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - let resp = match resp { - Ok(resp) => resp, - Err(err) => { - return Err(err.into()); - } - }; - - let common_prefixes = resp.common_prefixes; - - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - // return prefixes at the root level - let dates: Vec<_> = common_prefixes - .iter() - .filter_map(|path| path.as_ref().strip_prefix(&format!("{stream}/"))) - .map(String::from) - .collect(); - tracing::Span::current().record("dates", dates.len()); - - Ok(dates) + super::list_dates( + self.client.as_ref(), + &StorePath::from(stream), + &format!("{stream}/"), + tenant_id, + ) + .await } #[tracing::instrument( @@ -505,21 +343,8 @@ impl BlobStore { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let bytes = tokio::fs::read(path).await?; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.put(&key.into(), bytes.into()).await; - increment_object_store_calls_by_date("PUT", &Utc::now().date_naive().to_string(), tenant); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -545,63 +370,17 @@ impl BlobStore { let mut data = Vec::new(); file.read_to_end(&mut data).await?; let result = self.client.put(location, data.into()).await; - increment_object_store_calls_by_date( - "PUT", - &Utc::now().date_naive().to_string(), - tenant, - ); - - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - Err(err) => { - return Err(err.into()); - } - } + super::record_request(result, "PUT", tenant_id)?; return Ok(()); } else { - let async_writer = self.client.put_multipart(location).await; - let mut async_writer = match async_writer { - Ok(writer) => writer, - Err(err) => { - return Err(err.into()); - } - }; + let mut async_writer = self.client.put_multipart(location).await?; let mut data = Vec::new(); file.read_to_end(&mut data).await?; - // let mut upload_parts = Vec::new(); - - let has_final_partial_part = !total_size.is_multiple_of(min_multipart_size); - let num_full_parts = total_size / min_multipart_size; - let total_parts = num_full_parts + if has_final_partial_part { 1 } else { 0 }; - - // Upload each part with metrics - for part_number in 0..(total_parts) { - let start_pos = part_number * min_multipart_size; - let end_pos = if part_number == num_full_parts && has_final_partial_part { - // Last part might be smaller than 5MB (which is allowed) - total_size - } else { - // All other parts must be at least 5MB - start_pos + min_multipart_size - }; - - // Extract this part's data - let part_data = data[start_pos..end_pos].to_vec(); - - let result = async_writer.put_part(part_data.into()).await; - if result.is_err() { - return Err(result.err().unwrap().into()); - } + for part_data in data[..total_size].chunks(min_multipart_size) { + async_writer.put_part(part_data.to_vec().into()).await?; increment_object_store_calls_by_date( "PUT_MULTIPART", &Utc::now().date_naive().to_string(), @@ -663,19 +442,8 @@ impl ObjectStorage for BlobStore { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.head(&to_object_store_path(path)).await; - increment_object_store_calls_by_date("HEAD", &Utc::now().date_naive().to_string(), tenant); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result?) + super::record_request(result, "HEAD", tenant_id) } async fn get_object( @@ -683,7 +451,7 @@ impl ObjectStorage for BlobStore { path: &RelativePath, tenant_id: &Option, ) -> Result { - Ok(self._get_object(path, tenant_id).await?) + self._get_object(path, tenant_id).await } async fn get_object_ranged( @@ -691,23 +459,7 @@ impl ObjectStorage for BlobStore { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let body = crate::storage::get_object_ranged( - &self.client, - &to_object_store_path(path), - path, - || increment_object_store_calls_by_date("GET", &date, tenant), - ) - .await?; - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &date, - tenant, - ); - Ok(body) + super::get_object_ranged(self.client.as_ref(), path, tenant_id).await } async fn get_objects( @@ -721,89 +473,14 @@ impl ObjectStorage for BlobStore { } else { self.root.clone() }; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let mut list_stream = self.client.list(Some(&prefix)); - - let mut paths = vec![]; - let mut files_scanned = 0; - - // Drain the listing first, then fetch. Awaiting each GET inside the - // listing loop serialises every download against the one before it. - // Note: We track each streaming list item retrieval - while let Some(meta_result) = list_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let ingestor_file = filter_func(meta.location.filename().unwrap().to_string()); - - if !ingestor_file { - continue; - } - - paths.push( - RelativePath::from_path(meta.location.as_ref()) - .map_err(ObjectStorageError::PathError)? - .to_relative_path_buf(), - ); - } - - // `buffered` rather than `buffer_unordered`: callers of this generic - // helper may rely on results following listing order. - let res: Vec = futures::stream::iter(paths) - .map(|path| async move { self.get_object(&path, tenant_id).await }) - .buffered(GET_OBJECTS_CONCURRENCY) - .try_collect() - .await?; - - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - Ok(res) + super::get_objects(self.client.as_ref(), self, &prefix, filter_func, tenant_id).await } async fn get_ingestor_meta_file_paths( &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let mut path_arr = vec![]; - let mut files_scanned = 0; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let mut object_stream = self.client.list(Some(&self.root)); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - while let Some(meta_result) = object_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let flag = meta.location.filename().unwrap().starts_with("ingestor"); - - if flag { - path_arr.push(RelativePathBuf::from(meta.location.as_ref())); - } - } - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(path_arr) + super::get_ingestor_meta_file_paths(self.client.as_ref(), &self.root, tenant_id).await } async fn put_object( @@ -824,9 +501,7 @@ impl ObjectStorage for BlobStore { path: &RelativePath, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - self._delete_prefix(path.as_ref(), tenant_id).await?; - - Ok(()) + self._delete_prefix(path.as_ref(), tenant_id).await } #[tracing::instrument( @@ -841,22 +516,7 @@ impl ObjectStorage for BlobStore { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let result = self.client.delete(&to_object_store_path(path)).await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant, - ); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result?) + super::record_request(result, "DELETE", tenant_id) } #[tracing::instrument( @@ -870,18 +530,7 @@ impl ObjectStorage for BlobStore { .head(&to_object_store_path(&parseable_json_path())) .await; tracing::Span::current().record("ok", result.is_ok()); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("HEAD", &Utc::now().date_naive().to_string(), tenant); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result.map(|_| ())?) + super::record_request(result.map(|_| ()), "HEAD", tenant_id) } async fn delete_stream( @@ -894,9 +543,7 @@ impl ObjectStorage for BlobStore { } else { stream_name }; - self._delete_prefix(prefix, tenant_id).await?; - - Ok(()) + self._delete_prefix(prefix, tenant_id).await } async fn try_delete_node_meta( @@ -905,25 +552,8 @@ impl ObjectStorage for BlobStore { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let file = RelativePathBuf::from(&node_filename); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.delete(&to_object_store_path(&file)).await; - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "DELETE", tenant_id) } async fn list_streams(&self) -> Result, ObjectStorageError> { @@ -934,52 +564,7 @@ impl ObjectStorage for BlobStore { } async fn list_old_streams(&self) -> Result, ObjectStorageError> { - let resp = self.client.list_with_delimiter(None).await?; - - let common_prefixes = resp.common_prefixes; // get all dirs - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - // return prefixes at the root level - let dirs: HashSet<_> = common_prefixes - .iter() - .filter_map(|path| path.parts().next()) - .map(|name| name.as_ref().to_string()) - .filter(|x| x != PARSEABLE_ROOT_DIRECTORY) - .collect(); - - let stream_json_check = FuturesUnordered::new(); - - for dir in &dirs { - let key = format!("{dir}/{STREAM_METADATA_FILE_NAME}"); - let task = async move { - let result = self.client.head(&StorePath::from(key)).await; - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - result.map(|_| ()) - }; - stream_json_check.push(task); - } - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - dirs.len() as u64, - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - stream_json_check.try_collect::<()>().await?; - - Ok(dirs) + super::list_old_streams(self.client.as_ref(), DEFAULT_TENANT).await } async fn list_dates( @@ -987,9 +572,7 @@ impl ObjectStorage for BlobStore { stream_name: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let streams = self._list_dates(stream_name, tenant_id).await?; - - Ok(streams) + self._list_dates(stream_name, tenant_id).await } async fn list_hours( @@ -998,35 +581,8 @@ impl ObjectStorage for BlobStore { date: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from(format!("{}/{}/", stream_name, date)); - let resp = self.client.list_with_delimiter(Some(&pre)).await?; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - let hours: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = path_str.strip_prefix(&format!("{}/{}/", stream_name, date)) - { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("hour=")) - .collect(); - - Ok(hours) + let prefix = format!("{stream_name}/{date}/"); + super::list_partitions(self.client.as_ref(), &prefix, "hour=", tenant_id).await } async fn list_minutes( @@ -1036,35 +592,8 @@ impl ObjectStorage for BlobStore { hour: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from(format!("{}/{}/{}/", stream_name, date, hour)); - let resp = self.client.list_with_delimiter(Some(&pre)).await?; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let minutes: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = - path_str.strip_prefix(&format!("{}/{}/{}/", stream_name, date, hour)) - { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("minute=")) - .collect(); - - Ok(minutes) + let prefix = format!("{stream_name}/{date}/{hour}/"); + super::list_partitions(self.client.as_ref(), &prefix, "minute=", tenant_id).await } async fn upload_file( @@ -1073,7 +602,7 @@ impl ObjectStorage for BlobStore { path: &Path, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - Ok(self._upload_file(key, path, tenant_id).await?) + self._upload_file(key, path, tenant_id).await } fn absolute_url(&self, prefix: &RelativePath) -> object_store::path::Path { @@ -1102,32 +631,7 @@ impl ObjectStorage for BlobStore { &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from("/"); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.list_with_delimiter(Some(&pre)).await; - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + super::list_dirs(self.client.as_ref(), &StorePath::from("/"), tenant_id).await } async fn list_dirs_relative( @@ -1135,32 +639,8 @@ impl ObjectStorage for BlobStore { relative_path: &RelativePath, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let prefix = object_store::path::Path::from(relative_path.as_str()); - let resp = self.client.list_with_delimiter(Some(&prefix)).await; - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + let prefix = StorePath::from(relative_path.as_str()); + super::list_dirs(self.client.as_ref(), &prefix, tenant_id).await } async fn list_with_delimiter( diff --git a/src/storage/gcs.rs b/src/storage/gcs.rs index 8ac87f763..3cba2a50c 100644 --- a/src/storage/gcs.rs +++ b/src/storage/gcs.rs @@ -19,19 +19,13 @@ use std::{ collections::HashSet, num::NonZeroUsize, - ops::Range, path::{Path, PathBuf}, sync::Arc, time::Duration, }; use crate::{ - metrics::{ - increment_bytes_scanned_in_object_store_calls_by_date, - increment_files_scanned_in_object_store_calls_by_date, - increment_object_store_calls_by_date, - increment_partial_file_scans_in_object_store_calls_by_date, - }, + metrics::increment_object_store_calls_by_date, parseable::{DEFAULT_TENANT, LogStream, PARSEABLE}, storage::RETRY_TIMEOUT_SECS, }; @@ -45,7 +39,7 @@ use datafusion::{ runtime_env::RuntimeEnvBuilder, }, }; -use futures::{StreamExt, TryStreamExt, stream::FuturesUnordered}; + use object_store::{ BackoffConfig, ClientOptions, ListResult, ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload, RetryConfig, @@ -61,11 +55,9 @@ use tracing::error; use super::{ CONNECT_TIMEOUT_SECS, DEFAULT_MAX_OBJECT_STORE_REQUESTS, ObjectStorage, ObjectStorageError, - ObjectStorageProvider, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, - STREAM_METADATA_FILE_NAME, metrics_layer::MetricLayer, object_storage::parseable_json_path, - partial_path, to_object_store_path, + ObjectStorageProvider, REQUEST_TIMEOUT_SECS, metrics_layer::MetricLayer, + object_storage::parseable_json_path, partial_path, to_object_store_path, }; -use crate::storage::GET_OBJECTS_CONCURRENCY; #[derive(Debug, Clone, clap::Args)] #[command( @@ -229,66 +221,15 @@ impl Gcs { tenant_id: &Option, partial_path: PathBuf, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let src = to_object_store_path(path); - - let meta = self.client.head(&src).await?; - increment_object_store_calls_by_date("HEAD", &date, tenant_str); - let total = meta.size; - - if let Some(parent) = partial_path.parent() { - tokio::fs::create_dir_all(parent).await?; - } - let file = tokio::fs::File::create(&partial_path).await?; - file.set_len(total).await?; - let std_file = Arc::new(file.into_std().await); - - let chunk = PARSEABLE.options.hot_tier_download_chunk_size; - let concurrency = PARSEABLE.options.hot_tier_download_concurrency; - let ranges: Vec> = (0..total) - .step_by(chunk as usize) - .map(|s| s..(s + chunk).min(total)) - .collect(); - let chunk_count = ranges.len() as u64; - tracing::Span::current() - .record("total_bytes", total) - .record("chunks", chunk_count); - let semaphore = Arc::new(tokio::sync::Semaphore::new(concurrency as usize)); - - futures::stream::iter(ranges) - .map(|r| { - let src = src.clone(); - let std_file = std_file.clone(); - let semaphore = semaphore.clone(); - async move { - let _permit = semaphore.acquire_owned().await.map_err(|e| { - ObjectStorageError::Custom(format!("semaphore closed: {e}")) - })?; - let bytes = self.client.get_range(&src, r.clone()).await?; - let offset = r.start; - tokio::task::spawn_blocking(move || -> std::io::Result<()> { - crate::storage::write_all_at(&std_file, &bytes, offset) - }) - .await - .map_err(|e| ObjectStorageError::Custom(format!("join: {e}")))??; - Ok::<_, ObjectStorageError>(()) - } - }) - .buffer_unordered(concurrency as usize) - .try_collect::>() - .await?; - - increment_object_store_calls_by_date("GET", &date, tenant_str); - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant_str); - increment_bytes_scanned_in_object_store_calls_by_date("GET", total, &date, tenant_str); - increment_partial_file_scans_in_object_store_calls_by_date( - "GET", - chunk_count, - &date, - tenant_str, - ); - Ok(()) + super::parallel_download( + self.client.as_ref(), + path, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + partial_path, + PARSEABLE.options.hot_tier_download_chunk_size, + PARSEABLE.options.hot_tier_download_concurrency, + ) + .await } #[tracing::instrument( @@ -301,29 +242,7 @@ impl Gcs { path: &RelativePath, tenant_id: &Option, ) -> Result { - let resp = self.client.get(&to_object_store_path(path)).await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("GET", &Utc::now().date_naive().to_string(), tenant); - match resp { - Ok(resp) => { - let body: Bytes = resp.bytes().await?; - tracing::Span::current().record("bytes", body.len()); - increment_files_scanned_in_object_store_calls_by_date( - "GET", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(body) - } - Err(err) => Err(err.into()), - } + super::get_object(self.client.as_ref(), path, tenant_id).await } #[tracing::instrument( @@ -337,21 +256,8 @@ impl Gcs { resource: PutPayload, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.put(&to_object_store_path(path), resource).await; - increment_object_store_calls_by_date("PUT", &Utc::now().date_naive().to_string(), tenant); - match resp { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + let result = self.client.put(&to_object_store_path(path), resource).await; + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -365,54 +271,12 @@ impl Gcs { key: &str, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - - // List all objects under the prefix - let object_stream = self.client.list(Some(&(key.into()))); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - // Map listed objects to their locations for delete_stream - let locations_stream = object_stream.map(|result| result.map(|obj| obj.location)); - - // Use delete_stream which batches deletes via the bulk delete API - let mut delete_results = self.client.delete_stream(Box::pin(locations_stream)); - - let mut files_deleted: u64 = 0; - let mut failed_deletes: u64 = 0; - while let Some(result) = delete_results.next().await { - match result { - Ok(_) => files_deleted += 1, - Err(err) => { - failed_deletes += 1; - error!("Failed to delete object during prefix deletion: {:?}", err); - } - } - } - - let total_files = files_deleted + failed_deletes; - tracing::Span::current() - .record("deleted", files_deleted) - .record("failed", failed_deletes); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - total_files, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - files_deleted, - &Utc::now().date_naive().to_string(), - tenant, - ); - - if failed_deletes > 0 { - return Err(ObjectStorageError::UnhandledError( - anyhow::anyhow!("Failed to delete {failed_deletes} out of {total_files} objects under prefix '{key}'").into(), - )); - } - - Ok(()) + super::delete_prefix( + self.client.as_ref(), + key, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + ) + .await } #[tracing::instrument( @@ -425,38 +289,13 @@ impl Gcs { stream: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let resp: Result = self - .client - .list_with_delimiter(Some(&(stream.into()))) - .await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - let resp = match resp { - Ok(resp) => resp, - Err(err) => { - return Err(err.into()); - } - }; - - let common_prefixes = resp.common_prefixes; - - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - // return prefixes at the root level - let dates: Vec<_> = common_prefixes - .iter() - .filter_map(|path| path.as_ref().strip_prefix(&format!("{stream}/"))) - .map(String::from) - .collect(); - tracing::Span::current().record("dates", dates.len()); - - Ok(dates) + super::list_dates( + self.client.as_ref(), + &StorePath::from(stream), + &format!("{stream}/"), + tenant_id, + ) + .await } #[tracing::instrument( @@ -471,21 +310,8 @@ impl Gcs { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let bytes = tokio::fs::read(path).await?; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.put(&key.into(), bytes.into()).await; - increment_object_store_calls_by_date("PUT", &Utc::now().date_naive().to_string(), tenant); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -512,59 +338,16 @@ impl Gcs { // Track single PUT operation for small files let result = self.client.put(location, data.into()).await; - increment_object_store_calls_by_date( - "PUT", - &Utc::now().date_naive().to_string(), - tenant, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - Err(err) => { - return Err(err.into()); - } - } + super::record_request(result, "PUT", tenant_id)?; + return Ok(()); } else { - let async_writer = self.client.put_multipart(location).await; - let mut async_writer = match async_writer { - Ok(writer) => writer, - Err(err) => { - return Err(err.into()); - } - }; + let mut async_writer = self.client.put_multipart(location).await?; let mut data = Vec::new(); file.read_to_end(&mut data).await?; - let has_final_partial_part = !total_size.is_multiple_of(min_multipart_size); - let num_full_parts = total_size / min_multipart_size; - let total_parts = num_full_parts + if has_final_partial_part { 1 } else { 0 }; - - // Upload each part with metrics - for part_number in 0..(total_parts) { - let start_pos = part_number * min_multipart_size; - let end_pos = if part_number == num_full_parts && has_final_partial_part { - // Last part might be smaller than 5MB (which is allowed) - total_size - } else { - // All other parts must be at least 5MB - start_pos + min_multipart_size - }; - - // Extract this part's data - let part_data = data[start_pos..end_pos].to_vec(); - - // Track individual part upload - let result = async_writer.put_part(part_data.into()).await; - if result.is_err() { - return Err(result.err().unwrap().into()); - } + for part_data in data[..total_size].chunks(min_multipart_size) { + async_writer.put_part(part_data.to_vec().into()).await?; increment_object_store_calls_by_date( "PUT_MULTIPART", &Utc::now().date_naive().to_string(), @@ -603,27 +386,9 @@ impl ObjectStorage for Gcs { path: &RelativePath, tenant_id: &Option, ) -> Result { - let path = &to_object_store_path(path); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let meta = self.client.head(path).await; - increment_object_store_calls_by_date("HEAD", &Utc::now().date_naive().to_string(), tenant); - let meta = match meta { - Ok(meta) => { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - meta - } - Err(err) => { - return Err(err.into()); - } - }; - - let buf = object_store::buffered::BufReader::new(self.client.clone(), &meta); - Ok(buf) + let result = self.client.head(&to_object_store_path(path)).await; + let meta = super::record_request(result, "HEAD", tenant_id)?; + Ok(BufReader::new(self.client.clone(), &meta)) } async fn upload_multipart( @@ -646,18 +411,7 @@ impl ObjectStorage for Gcs { tenant_id: &Option, ) -> Result { let result = self.client.head(&to_object_store_path(path)).await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("HEAD", &Utc::now().date_naive().to_string(), tenant); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result?) + super::record_request(result, "HEAD", tenant_id) } async fn get_object( @@ -665,7 +419,7 @@ impl ObjectStorage for Gcs { path: &RelativePath, tenant_id: &Option, ) -> Result { - Ok(self._get_object(path, tenant_id).await?) + self._get_object(path, tenant_id).await } async fn get_object_ranged( @@ -673,23 +427,7 @@ impl ObjectStorage for Gcs { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let body = crate::storage::get_object_ranged( - &self.client, - &to_object_store_path(path), - path, - || increment_object_store_calls_by_date("GET", &date, tenant), - ) - .await?; - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &date, - tenant, - ); - Ok(body) + super::get_object_ranged(self.client.as_ref(), path, tenant_id).await } async fn get_objects( @@ -703,89 +441,14 @@ impl ObjectStorage for Gcs { } else { self.root.clone() }; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let mut list_stream = self.client.list(Some(&prefix)); - - let mut paths = vec![]; - let mut files_scanned = 0; - - // Drain the listing first, then fetch. Awaiting each GET inside the - // listing loop serialises every download against the one before it. - // Note: We track each streaming list item retrieval - while let Some(meta_result) = list_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let ingestor_file = filter_func(meta.location.filename().unwrap().to_string()); - - if !ingestor_file { - continue; - } - - paths.push( - RelativePath::from_path(meta.location.as_ref()) - .map_err(ObjectStorageError::PathError)? - .to_relative_path_buf(), - ); - } - - // `buffered` rather than `buffer_unordered`: callers of this generic - // helper may rely on results following listing order. - let res: Vec = futures::stream::iter(paths) - .map(|path| async move { self.get_object(&path, tenant_id).await }) - .buffered(GET_OBJECTS_CONCURRENCY) - .try_collect() - .await?; - - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - Ok(res) + super::get_objects(self.client.as_ref(), self, &prefix, filter_func, tenant_id).await } async fn get_ingestor_meta_file_paths( &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let mut path_arr = vec![]; - let mut files_scanned = 0; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let mut object_stream = self.client.list(Some(&self.root)); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - while let Some(meta_result) = object_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let flag = meta.location.filename().unwrap().starts_with("ingestor"); - - if flag { - path_arr.push(RelativePathBuf::from(meta.location.as_ref())); - } - } - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(path_arr) + super::get_ingestor_meta_file_paths(self.client.as_ref(), &self.root, tenant_id).await } async fn put_object( @@ -806,9 +469,7 @@ impl ObjectStorage for Gcs { path: &RelativePath, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - self._delete_prefix(path.as_ref(), tenant_id).await?; - - Ok(()) + self._delete_prefix(path.as_ref(), tenant_id).await } #[tracing::instrument( @@ -823,22 +484,7 @@ impl ObjectStorage for Gcs { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let result = self.client.delete(&to_object_store_path(path)).await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant, - ); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result?) + super::record_request(result, "DELETE", tenant_id) } #[tracing::instrument( @@ -852,19 +498,7 @@ impl ObjectStorage for Gcs { .head(&to_object_store_path(&parseable_json_path())) .await; tracing::Span::current().record("ok", result.is_ok()); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("HEAD", &Utc::now().date_naive().to_string(), tenant); - - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - } - - Ok(result.map(|_| ())?) + super::record_request(result.map(|_| ()), "HEAD", tenant_id) } async fn delete_stream( @@ -877,9 +511,7 @@ impl ObjectStorage for Gcs { } else { stream_name }; - self._delete_prefix(prefix, tenant_id).await?; - - Ok(()) + self._delete_prefix(prefix, tenant_id).await } async fn try_delete_node_meta( @@ -888,25 +520,8 @@ impl ObjectStorage for Gcs { tenant_id: &Option, ) -> Result<(), ObjectStorageError> { let file = RelativePathBuf::from(&node_filename); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.delete(&to_object_store_path(&file)).await; - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "DELETE", tenant_id) } async fn list_streams(&self) -> Result, ObjectStorageError> { @@ -917,51 +532,7 @@ impl ObjectStorage for Gcs { } async fn list_old_streams(&self) -> Result, ObjectStorageError> { - let resp = self.client.list_with_delimiter(None).await?; - let common_prefixes = resp.common_prefixes; // get all dirs - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - // return prefixes at the root level - let dirs: HashSet<_> = common_prefixes - .iter() - .filter_map(|path| path.parts().next()) - .map(|name| name.as_ref().to_string()) - .filter(|x| x != PARSEABLE_ROOT_DIRECTORY) - .collect(); - - let stream_json_check = FuturesUnordered::new(); - - for dir in &dirs { - let key = format!("{dir}/{STREAM_METADATA_FILE_NAME}"); - let task = async move { - let result = self.client.head(&StorePath::from(key)).await; - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - result.map(|_| ()) - }; - stream_json_check.push(task); - } - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - dirs.len() as u64, - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - stream_json_check.try_collect::<()>().await?; - - Ok(dirs) + super::list_old_streams(self.client.as_ref(), DEFAULT_TENANT).await } async fn list_dates( @@ -969,9 +540,7 @@ impl ObjectStorage for Gcs { stream_name: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let streams = self._list_dates(stream_name, tenant_id).await?; - - Ok(streams) + self._list_dates(stream_name, tenant_id).await } async fn list_hours( @@ -980,35 +549,8 @@ impl ObjectStorage for Gcs { date: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from(format!("{}/{}/", stream_name, date)); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.list_with_delimiter(Some(&pre)).await?; - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - - let hours: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = path_str.strip_prefix(&format!("{}/{}/", stream_name, date)) - { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("hour=")) - .collect(); - - Ok(hours) + let prefix = format!("{stream_name}/{date}/"); + super::list_partitions(self.client.as_ref(), &prefix, "hour=", tenant_id).await } async fn list_minutes( @@ -1018,35 +560,8 @@ impl ObjectStorage for Gcs { hour: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from(format!("{}/{}/{}/", stream_name, date, hour)); - let resp = self.client.list_with_delimiter(Some(&pre)).await?; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let minutes: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = - path_str.strip_prefix(&format!("{}/{}/{}/", stream_name, date, hour)) - { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("minute=")) - .collect(); - - Ok(minutes) + let prefix = format!("{stream_name}/{date}/{hour}/"); + super::list_partitions(self.client.as_ref(), &prefix, "minute=", tenant_id).await } async fn upload_file( @@ -1055,7 +570,7 @@ impl ObjectStorage for Gcs { path: &Path, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - Ok(self._upload_file(key, path, tenant_id).await?) + self._upload_file(key, path, tenant_id).await } fn absolute_url(&self, prefix: &RelativePath) -> object_store::path::Path { @@ -1080,33 +595,7 @@ impl ObjectStorage for Gcs { &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let pre = object_store::path::Path::from("/"); - - let resp = self.client.list_with_delimiter(Some(&pre)).await; - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + super::list_dirs(self.client.as_ref(), &StorePath::from("/"), tenant_id).await } async fn list_dirs_relative( @@ -1114,32 +603,8 @@ impl ObjectStorage for Gcs { relative_path: &RelativePath, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let prefix = object_store::path::Path::from(relative_path.as_str()); - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.list_with_delimiter(Some(&prefix)).await; - increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + let prefix = StorePath::from(relative_path.as_str()); + super::list_dirs(self.client.as_ref(), &prefix, tenant_id).await } async fn list_with_delimiter( diff --git a/src/storage/localfs.rs b/src/storage/localfs.rs index d66714023..b77c831a6 100644 --- a/src/storage/localfs.rs +++ b/src/storage/localfs.rs @@ -262,13 +262,7 @@ impl ObjectStorage for LocalFS { let mut files_scanned = 0u64; let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let entries_result = fs::read_dir(&self.root).await; - let mut entries = match entries_result { - Ok(entries) => entries, - Err(err) => { - return Err(err.into()); - } - }; + let mut entries = fs::read_dir(&self.root).await?; while let Some(entry) = entries.next_entry().await? { files_scanned += 1; @@ -317,13 +311,7 @@ impl ObjectStorage for LocalFS { }; let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let entries_result = fs::read_dir(&prefix).await; - let mut entries = match entries_result { - Ok(entries) => entries, - Err(err) => { - return Err(err.into()); - } - }; + let mut entries = fs::read_dir(&prefix).await?; let mut res = Vec::new(); let mut files_scanned = 0; @@ -537,32 +525,8 @@ impl ObjectStorage for LocalFS { TOMBSTONE_ROOT_DIRECTORY, ]; - let result = fs::read_dir(&self.root).await; - let directories = match result { - Ok(read_dir) => { - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - ReadDirStream::new(read_dir) - } - Err(err) => { - return Err(err.into()); - } - }; - - let entries: Vec = directories.try_collect().await?; - let entries = entries - .into_iter() - .map(|entry| dir_with_stream(entry, ignore_dir)); - - let logstream_dirs: Vec> = - FuturesUnordered::from_iter(entries).try_collect().await?; - - let logstreams = logstream_dirs.into_iter().flatten().collect(); - - Ok(logstreams) + let metadata_path = PathBuf::from(STREAM_ROOT_DIRECTORY).join(STREAM_METADATA_FILE_NAME); + stream_directories(&self.root, ignore_dir, &metadata_path).await } async fn list_old_streams(&self) -> Result, ObjectStorageError> { @@ -575,32 +539,8 @@ impl ObjectStorage for LocalFS { TOMBSTONE_ROOT_DIRECTORY, ]; - let result = fs::read_dir(&self.root).await; - let directories = match result { - Ok(read_dir) => { - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - DEFAULT_TENANT, - ); - ReadDirStream::new(read_dir) - } - Err(err) => { - return Err(err.into()); - } - }; - - let entries: Vec = directories.try_collect().await?; - let entries = entries - .into_iter() - .map(|entry| dir_with_old_stream(entry, ignore_dir)); - - let logstream_dirs: Vec> = - FuturesUnordered::from_iter(entries).try_collect().await?; - - let logstreams = logstream_dirs.into_iter().flatten().collect(); - - Ok(logstreams) + let metadata_path = PathBuf::from(STREAM_METADATA_FILE_NAME); + stream_directories(&self.root, ignore_dir, &metadata_path).await } async fn list_dirs( @@ -627,20 +567,7 @@ impl ObjectStorage for LocalFS { } }; - let dirs = ReadDirStream::new(read_dir) - .try_collect::>() - .await? - .into_iter() - .map(dir_name); - - let dirs = FuturesUnordered::from_iter(dirs) - .try_collect::>() - .await? - .into_iter() - .flatten() - .collect::>(); - - Ok(dirs) + directory_names(read_dir).await } async fn list_dirs_relative( @@ -662,20 +589,7 @@ impl ObjectStorage for LocalFS { } }; - let dirs = ReadDirStream::new(read_dir) - .try_collect::>() - .await? - .into_iter() - .map(dir_name); - - let dirs = FuturesUnordered::from_iter(dirs) - .try_collect::>() - .await? - .into_iter() - .flatten() - .collect::>(); - - Ok(dirs) + directory_names(read_dir).await } async fn list_dates( @@ -705,12 +619,7 @@ impl ObjectStorage for LocalFS { } }; - let directories = ReadDirStream::new(read_dir); - let entries: Vec = directories.try_collect().await?; - let entries = entries.into_iter().map(dir_name); - let dates: Vec<_> = FuturesUnordered::from_iter(entries).try_collect().await?; - - Ok(dates.into_iter().flatten().collect()) + directory_names(read_dir).await } async fn list_hours( @@ -720,13 +629,9 @@ impl ObjectStorage for LocalFS { _tenant_id: &Option, ) -> Result, ObjectStorageError> { let path = self.root.join(stream_name).join(date); - let directories = ReadDirStream::new(fs::read_dir(&path).await?); - let entries: Vec = directories.try_collect().await?; - let entries = entries.into_iter().map(dir_name); - let hours: Vec<_> = FuturesUnordered::from_iter(entries).try_collect().await?; - Ok(hours + Ok(directory_names(fs::read_dir(&path).await?) + .await? .into_iter() - .flatten() .filter(|dir| dir.starts_with("hour=")) .collect()) } @@ -740,14 +645,9 @@ impl ObjectStorage for LocalFS { ) -> Result, ObjectStorageError> { let path = self.root.join(stream_name).join(date).join(hour); // Propagate any read_dir errors instead of swallowing them - let directories = ReadDirStream::new(fs::read_dir(&path).await?); - let entries: Vec = directories.try_collect().await?; - let entries = entries.into_iter().map(dir_name); - let minutes: Vec<_> = FuturesUnordered::from_iter(entries).try_collect().await?; - // Filter down to only the "minute=" prefixed directories - Ok(minutes + Ok(directory_names(fs::read_dir(&path).await?) + .await? .into_iter() - .flatten() .filter(|dir| dir.starts_with("minute=")) .collect()) } @@ -826,43 +726,50 @@ impl ObjectStorage for LocalFS { } } -async fn dir_with_old_stream( - entry: DirEntry, - ignore_dirs: &[&str], -) -> Result, ObjectStorageError> { - let dir_name = entry - .path() - .file_name() - .expect("valid path") - .to_str() - .expect("valid unicode") - .to_owned(); - - if ignore_dirs.contains(&dir_name.as_str()) { - return Ok(None); - } - - if entry.file_type().await?.is_dir() { - let path = entry.path(); - - // even in ingest mode, we should only look for the global stream metadata file - let stream_json_path = path.join(STREAM_METADATA_FILE_NAME); +async fn directory_names(read_dir: fs::ReadDir) -> Result, ObjectStorageError> { + let entries = ReadDirStream::new(read_dir) + .try_collect::>() + .await?; + Ok( + FuturesUnordered::from_iter(entries.into_iter().map(dir_name)) + .try_collect::>() + .await? + .into_iter() + .flatten() + .collect(), + ) +} - if stream_json_path.exists() { - Ok(Some(dir_name)) - } else { - let err: Box = - format!("found {}", entry.path().display()).into(); - Err(ObjectStorageError::UnhandledError(err)) - } - } else { - Ok(None) - } +async fn stream_directories( + root: &Path, + ignore_dirs: &[&str], + metadata_path: &Path, +) -> Result, ObjectStorageError> { + let read_dir = fs::read_dir(root).await?; + increment_object_store_calls_by_date( + "LIST", + &Utc::now().date_naive().to_string(), + DEFAULT_TENANT, + ); + let entries = ReadDirStream::new(read_dir) + .try_collect::>() + .await?; + Ok(FuturesUnordered::from_iter( + entries + .into_iter() + .map(|entry| dir_with_stream(entry, ignore_dirs, metadata_path)), + ) + .try_collect::>() + .await? + .into_iter() + .flatten() + .collect()) } async fn dir_with_stream( entry: DirEntry, ignore_dirs: &[&str], + metadata_path: &Path, ) -> Result, ObjectStorageError> { let dir_name = entry .path() @@ -880,9 +787,7 @@ async fn dir_with_stream( let path = entry.path(); // even in ingest mode, we should only look for the global stream metadata file - let stream_json_path = path - .join(STREAM_ROOT_DIRECTORY) - .join(STREAM_METADATA_FILE_NAME); + let stream_json_path = path.join(metadata_path); if stream_json_path.exists() { Ok(Some(dir_name)) @@ -916,3 +821,55 @@ impl From for ObjectStorageError { ObjectStorageError::UnhandledError(Box::new(e)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn directory_listing_keeps_metadata_layouts_and_missing_path_rules() { + let root = tempfile::tempdir().unwrap(); + let store = LocalFS::new(root.path().to_path_buf()); + let new_metadata = root + .path() + .join("new") + .join(STREAM_ROOT_DIRECTORY) + .join(STREAM_METADATA_FILE_NAME); + let old_metadata = root.path().join("old").join(STREAM_METADATA_FILE_NAME); + std::fs::create_dir_all(new_metadata.parent().unwrap()).unwrap(); + std::fs::create_dir_all(old_metadata.parent().unwrap()).unwrap(); + std::fs::write(new_metadata, "{}").unwrap(); + std::fs::write(old_metadata, "{}").unwrap(); + std::fs::write(root.path().join("file"), "data").unwrap(); + let new = stream_directories( + root.path(), + &["old"], + &PathBuf::from(STREAM_ROOT_DIRECTORY).join(STREAM_METADATA_FILE_NAME), + ) + .await + .unwrap(); + let old = stream_directories(root.path(), &["new"], Path::new(STREAM_METADATA_FILE_NAME)) + .await + .unwrap(); + assert_eq!(new, ["new".to_owned()].into()); + assert_eq!(old, ["old".to_owned()].into()); + let mut names = store.list_dirs(&None).await.unwrap(); + names.sort(); + assert_eq!(names, ["new", "old"]); + assert!(store.list_dates("missing", &None).await.unwrap().is_empty()); + assert!( + store + .list_dirs_relative(RelativePath::new("missing"), &None) + .await + .unwrap() + .is_empty() + ); + assert!(store.list_hours("missing", "date", &None).await.is_err()); + assert!( + store + .list_minutes("missing", "date", "hour", &None) + .await + .is_err() + ); + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index c376a487f..6b71c5c3a 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -16,7 +16,8 @@ * */ -use object_store::path::Path; +use futures::{StreamExt, TryStreamExt}; +use object_store::{ObjectStoreExt, path::Path}; use relative_path::RelativePath; use serde::{Deserialize, Serialize}; use tokio::task::JoinError; @@ -28,16 +29,24 @@ use crate::{ hottier::StreamHotTier, metadata::SchemaVersion, metastore::{MetastoreErrorDetail, metastore_traits::MetastoreObject}, + metrics::{ + increment_bytes_scanned_in_object_store_calls_by_date, + increment_files_scanned_in_object_store_calls_by_date, + increment_object_store_calls_by_date, + increment_partial_file_scans_in_object_store_calls_by_date, + }, option::StandaloneWithDistributed, - parseable::StreamNotFound, + parseable::{DEFAULT_TENANT, LogStream, StreamNotFound}, stats::FullStats, utils::json::{deserialize_string_as_true, serialize_bool_as_true}, }; use chrono::Utc; +use std::collections::HashSet; use std::fmt::Debug; use std::num::NonZeroUsize; +use std::sync::Arc; mod azure_blob; pub mod field_stats; @@ -85,17 +94,28 @@ pub(crate) const GET_OBJECTS_CONCURRENCY: usize = 32; /// full length, so an object smaller than [`RANGED_GET_PROBE_BYTES`] costs /// exactly one request and needs no HEAD. /// -/// `on_request` is invoked once per issued GET so callers can keep their own -/// request accounting accurate. +/// Counts each issued GET and records file/byte metrics only after a successful read. pub(crate) async fn get_object_ranged( client: &S, - src: &Path, log_path: &RelativePath, - on_request: impl Fn() + Send + Sync, + tenant_id: &Option, ) -> Result { + let src = to_object_store_path(log_path); + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let on_request = || increment_object_store_calls_by_date("GET", &date, tenant); for attempt in 1..=RANGED_GET_MAX_ATTEMPTS { - match get_object_ranged_once(client, src, log_path, &on_request).await { - Ok(body) => return Ok(body), + match get_object_ranged_once(client, &src, log_path, &on_request).await { + Ok(body) => { + increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); + increment_bytes_scanned_in_object_store_calls_by_date( + "GET", + body.len() as u64, + &date, + tenant, + ); + return Ok(body); + } // The object was rewritten while its ranges were being read. Every // range fetched so far may belong to a different generation, so the // assembled buffer is discarded and the read restarts. @@ -611,6 +631,394 @@ pub fn to_object_store_path(path: &RelativePath) -> Path { Path::from(path.as_str()) } +async fn list_dates( + client: &S, + prefix: &Path, + strip_prefix: &str, + tenant_id: &Option, +) -> Result, ObjectStorageError> { + let resp = client.list_with_delimiter(Some(prefix)).await; + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + + let resp = resp?; + + let common_prefixes = resp.common_prefixes; + + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + common_prefixes.len() as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + + // return prefixes at the root level + let dates: Vec<_> = common_prefixes + .iter() + .filter_map(|path| path.as_ref().strip_prefix(strip_prefix)) + .map(String::from) + .collect(); + tracing::Span::current().record("dates", dates.len()); + + Ok(dates) +} + +async fn list_old_streams( + client: &S, + head_tenant: &str, +) -> Result, ObjectStorageError> { + let resp = client.list_with_delimiter(None).await?; + let common_prefixes = resp.common_prefixes; // get all dirs + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + common_prefixes.len() as u64, + &Utc::now().date_naive().to_string(), + DEFAULT_TENANT, + ); + increment_object_store_calls_by_date( + "LIST", + &Utc::now().date_naive().to_string(), + DEFAULT_TENANT, + ); + // return prefixes at the root level + let dirs: HashSet<_> = common_prefixes + .iter() + .filter_map(|path| path.parts().next()) + .map(|name| name.as_ref().to_string()) + .filter(|x| x != PARSEABLE_ROOT_DIRECTORY) + .collect(); + + let stream_json_check = futures::stream::FuturesUnordered::new(); + + for dir in &dirs { + let key = format!("{dir}/{STREAM_METADATA_FILE_NAME}"); + let task = async move { + let result = client.head(&Path::from(key)).await; + increment_object_store_calls_by_date( + "HEAD", + &Utc::now().date_naive().to_string(), + head_tenant, + ); + result.map(|_| ()) + }; + stream_json_check.push(task); + } + increment_files_scanned_in_object_store_calls_by_date( + "HEAD", + dirs.len() as u64, + &Utc::now().date_naive().to_string(), + DEFAULT_TENANT, + ); + stream_json_check.try_collect::<()>().await?; + + Ok(dirs) +} + +async fn list_partitions( + client: &S, + prefix: &str, + partition: &str, + tenant_id: &Option, +) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let resp = client + .list_with_delimiter(Some(&Path::from(prefix))) + .await?; + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + resp.common_prefixes.len() as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + + let partitions: Vec = resp + .common_prefixes + .iter() + .filter_map(|path| { + let path_str = path.as_ref(); + if let Some(stripped) = path_str.strip_prefix(prefix) { + // Remove trailing slash if present, otherwise use as is + let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); + Some(clean_path.to_string()) + } else { + None + } + }) + .filter(|dir| dir.starts_with(partition)) + .collect(); + + Ok(partitions) +} + +async fn get_objects( + client: &S, + provider: &impl ObjectStorage, + prefix: &Path, + filter_func: Box bool + Send>, + tenant_id: &Option, +) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let mut list_stream = client.list(Some(prefix)); + + let mut paths = vec![]; + let mut files_scanned = 0; + + // Drain the listing first, then fetch. Awaiting each GET inside the + // listing loop serialises every download against the one before it. + // Note: We track each streaming list item retrieval + while let Some(meta_result) = list_stream.next().await { + let meta = meta_result?; + + files_scanned += 1; + let ingestor_file = filter_func(meta.location.filename().unwrap().to_string()); + + if !ingestor_file { + continue; + } + + paths.push( + RelativePath::from_path(meta.location.as_ref()) + .map_err(ObjectStorageError::PathError)? + .to_relative_path_buf(), + ); + } + + // `buffered` rather than `buffer_unordered`: callers of this generic + // helper may rely on results following listing order. + let res: Vec = futures::stream::iter(paths) + .map(|path| async move { provider.get_object(&path, tenant_id).await }) + .buffered(GET_OBJECTS_CONCURRENCY) + .try_collect() + .await?; + + // Record total files scanned + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + files_scanned as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + Ok(res) +} + +async fn get_object( + client: &S, + path: &RelativePath, + tenant_id: &Option, +) -> Result { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let resp = client.get(&to_object_store_path(path)).await; + increment_object_store_calls_by_date("GET", &Utc::now().date_naive().to_string(), tenant); + + let resp = resp?; + let body = resp.bytes().await?; + tracing::Span::current().record("bytes", body.len()); + increment_files_scanned_in_object_store_calls_by_date( + "GET", + 1, + &Utc::now().date_naive().to_string(), + tenant, + ); + increment_bytes_scanned_in_object_store_calls_by_date( + "GET", + body.len() as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + Ok(body) +} + +fn record_request( + result: Result, + method: &str, + tenant_id: &Option, +) -> Result { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + increment_object_store_calls_by_date(method, &Utc::now().date_naive().to_string(), tenant); + let value = result?; + increment_files_scanned_in_object_store_calls_by_date( + method, + 1, + &Utc::now().date_naive().to_string(), + tenant, + ); + Ok(value) +} + +async fn list_dirs( + client: &S, + prefix: &Path, + tenant_id: &Option, +) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + + let resp = client.list_with_delimiter(Some(prefix)).await; + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + let resp = resp?; + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + resp.common_prefixes.len() as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + + Ok(resp + .common_prefixes + .iter() + .flat_map(|path| path.parts()) + .map(|name| name.as_ref().to_string()) + .collect::>()) +} + +async fn get_ingestor_meta_file_paths( + client: &S, + prefix: &Path, + tenant_id: &Option, +) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let mut path_arr = vec![]; + let mut files_scanned = 0; + + let mut object_stream = client.list(Some(prefix)); + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + + while let Some(meta_result) = object_stream.next().await { + let meta = meta_result?; + + files_scanned += 1; + let flag = meta.location.filename().unwrap().starts_with("ingestor"); + + if flag { + path_arr.push(relative_path::RelativePathBuf::from(meta.location.as_ref())); + } + } + // Record total files scanned + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + files_scanned as u64, + &Utc::now().date_naive().to_string(), + tenant, + ); + Ok(path_arr) +} + +async fn parallel_download( + client: &S, + path: &RelativePath, + tenant: &str, + partial_path: std::path::PathBuf, + chunk: u64, + concurrency: u64, +) -> Result<(), ObjectStorageError> { + let date = Utc::now().date_naive().to_string(); + let src = to_object_store_path(path); + + let meta = client.head(&src).await?; + increment_object_store_calls_by_date("HEAD", &date, tenant); + let total = meta.size; + + if let Some(parent) = partial_path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + let file = tokio::fs::File::create(&partial_path).await?; + file.set_len(total).await?; + let std_file = Arc::new(file.into_std().await); + + let ranges: Vec> = (0..total) + .step_by(chunk as usize) + .map(|s| s..(s + chunk).min(total)) + .collect(); + let chunk_count = ranges.len() as u64; + tracing::Span::current() + .record("total_bytes", total) + .record("chunks", chunk_count); + + let semaphore = Arc::new(tokio::sync::Semaphore::new(concurrency as usize)); + + futures::stream::iter(ranges) + .map(|r| { + let src = src.clone(); + let std_file = std_file.clone(); + let semaphore = semaphore.clone(); + async move { + let _permit = semaphore + .acquire_owned() + .await + .map_err(|e| ObjectStorageError::Custom(format!("semaphore closed: {e}")))?; + let bytes = client.get_range(&src, r.clone()).await?; + let offset = r.start; + tokio::task::spawn_blocking(move || -> std::io::Result<()> { + write_all_at(&std_file, &bytes, offset) + }) + .await + .map_err(|e| ObjectStorageError::Custom(format!("join: {e}")))??; + Ok::<_, ObjectStorageError>(()) + } + }) + .buffer_unordered(concurrency as usize) + .try_collect::>() + .await?; + + increment_object_store_calls_by_date("GET", &date, tenant); + increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); + increment_bytes_scanned_in_object_store_calls_by_date("GET", total, &date, tenant); + increment_partial_file_scans_in_object_store_calls_by_date("GET", chunk_count, &date, tenant); + Ok(()) +} + +async fn delete_prefix( + client: &S, + key: &str, + tenant: &str, +) -> Result<(), ObjectStorageError> { + let object_stream = client.list(Some(&(key.into()))); + increment_object_store_calls_by_date("LIST", &Utc::now().date_naive().to_string(), tenant); + + let locations_stream = object_stream.map(|result| result.map(|obj| obj.location)); + + // Use delete_stream which batches deletes via the bulk delete API + let mut delete_results = client.delete_stream(Box::pin(locations_stream)); + + let mut files_deleted: u64 = 0; + let mut failed_deletes: u64 = 0; + while let Some(result) = delete_results.next().await { + match result { + Ok(_) => files_deleted += 1, + Err(err) => { + failed_deletes += 1; + tracing::error!("Failed to delete object during prefix deletion: {:?}", err); + } + } + } + + let total_files = files_deleted + failed_deletes; + tracing::Span::current() + .record("deleted", files_deleted) + .record("failed", failed_deletes); + increment_files_scanned_in_object_store_calls_by_date( + "LIST", + total_files, + &Utc::now().date_naive().to_string(), + tenant, + ); + increment_files_scanned_in_object_store_calls_by_date( + "DELETE", + files_deleted, + &Utc::now().date_naive().to_string(), + tenant, + ); + + if failed_deletes > 0 { + return Err(ObjectStorageError::UnhandledError( + anyhow::anyhow!("Failed to delete {failed_deletes} out of {total_files} objects under prefix '{key}'").into(), + )); + } + + Ok(()) +} + /// Append `.partial` to the file name of a local path. Used by hot-tier /// downloaders to write to a sibling path and atomically rename on success. pub fn partial_path( @@ -629,8 +1037,310 @@ pub fn partial_path( #[cfg(test)] mod tests { use clap::Parser; + use object_store::{ + ObjectStore, ObjectStoreExt, local::LocalFileSystem, memory::InMemory, path::Path, + }; + use relative_path::RelativePath; + + use super::{ + DEFAULT_MAX_OBJECT_STORE_REQUESTS, GcsConfig, LocalFS, ObjectStorageError, + PARSEABLE_ROOT_DIRECTORY, STREAM_METADATA_FILE_NAME, delete_prefix, + get_ingestor_meta_file_paths, get_object, get_object_ranged, get_objects, list_dates, + list_dirs, list_old_streams, list_partitions, parallel_download, record_request, + }; + + #[tokio::test] + async fn partition_listings_keep_matching_children_with_and_without_tenant_prefixes() { + let store = InMemory::new(); + for prefix in ["stream/date=2026-11-01/", "tenant/stream/date=2026-11-01/"] { + for suffix in [ + "hour=00/minute=01/data", + "hour=00/other/data", + "hour=invalid/data", + "other/data", + "file", + ] { + store + .put(&Path::from(format!("{prefix}{suffix}")), "data".into()) + .await + .unwrap(); + } + let hours = list_partitions(&store, prefix, "hour=", &None) + .await + .unwrap(); + assert_eq!(hours, ["hour=00", "hour=invalid"]); + let prefix = format!("{prefix}hour=00/"); + let minutes = list_partitions(&store, &prefix, "minute=", &None) + .await + .unwrap(); + assert_eq!(minutes, ["minute=01"]); + } + assert!( + list_partitions(&store, "missing/", "hour=", &None) + .await + .unwrap() + .is_empty() + ); + } + + #[tokio::test] + async fn old_stream_discovery_excludes_metadata_and_propagates_missing_stream_files() { + let store = InMemory::new(); + for key in [ + format!("{PARSEABLE_ROOT_DIRECTORY}/ignored"), + format!("one/{STREAM_METADATA_FILE_NAME}"), + format!("two/{STREAM_METADATA_FILE_NAME}"), + ] { + store.put(&Path::from(key), "data".into()).await.unwrap(); + } + let streams = list_old_streams(&store, "old-stream-test").await.unwrap(); + assert_eq!(streams, ["one".to_owned(), "two".to_owned()].into()); + store + .put(&Path::from("missing/other"), "data".into()) + .await + .unwrap(); + assert!(matches!( + list_old_streams(&store, "old-stream-test").await, + Err(ObjectStorageError::NoSuchKey(_)) + )); + } + + #[tokio::test] + async fn get_objects_filters_filenames_and_preserves_listing_order() { + let dir = tempfile::tempdir().unwrap(); + let store = LocalFileSystem::new_with_prefix(dir.path()).unwrap(); + let provider = LocalFS::new(dir.path().to_path_buf()); + for key in [ + "nested/a.json", + "nested/b.skip", + "nested/c.json", + "outside.json", + ] { + store + .put(&Path::from(key), key.as_bytes().to_vec().into()) + .await + .unwrap(); + } + let seen = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let accepted = seen.clone(); + let filter = Box::new(move |name: String| { + let keep = name.ends_with(".json"); + if keep { + accepted.lock().unwrap().push(format!("nested/{name}")); + } + keep + }); + let results = get_objects(&store, &provider, &Path::from("nested"), filter, &None) + .await + .unwrap(); + let results: Vec<_> = results + .into_iter() + .map(|bytes| String::from_utf8(bytes.to_vec()).unwrap()) + .collect(); + assert_eq!(results, *seen.lock().unwrap()); + assert_eq!(results.len(), 2); + + let target = dir.path().join("nested/a.json"); + let filter = Box::new(move |name: String| { + if name != "a.json" { + return false; + } + std::fs::remove_file(&target).unwrap(); + true + }); + let result = get_objects(&store, &provider, &Path::from("nested"), filter, &None).await; + assert!(matches!(result, Err(ObjectStorageError::NoSuchKey(_)))); + } + + #[test] + fn record_request_counts_attempts_and_successful_files() { + use crate::metrics::{ + TOTAL_FILES_SCANNED_IN_OBJECT_STORE_CALLS_BY_DATE as FILES, + TOTAL_OBJECT_STORE_CALLS_BY_DATE as CALLS, + }; + let tenant = Some("storage-request-test".to_owned()); + let date = chrono::Utc::now().date_naive().to_string(); + for method in ["PUT", "HEAD", "DELETE"] { + let labels = [method, &date, tenant.as_deref().unwrap()]; + let calls = CALLS.with_label_values(&labels); + let files = FILES.with_label_values(&labels); + let before = (calls.get(), files.get()); + assert_eq!(record_request(Ok(7), method, &tenant).unwrap(), 7); + let err = object_store::Error::NotFound { + path: "missing".into(), + source: std::io::Error::from(std::io::ErrorKind::NotFound).into(), + }; + assert!(matches!( + record_request::<()>(Err(err), method, &tenant), + Err(ObjectStorageError::NoSuchKey(path)) if path == "missing" + )); + assert_eq!((calls.get() - before.0, files.get() - before.1), (2, 1)); + } + } + + #[tokio::test] + async fn cloud_listings_preserve_prefixes_and_match_ingestor_filenames() { + let store = InMemory::new(); + let tenant = Some("storage-listing-test".to_owned()); + for key in [ + "tenant/stream/file", + "tenant/stream/date=2026-11-01/file", + "tenant-other/file", + "ingestor-one", + "nested/ingestor-two", + "not-ingestor", + ] { + store.put(&Path::from(key), "data".into()).await.unwrap(); + } + let prefix = "tenant/stream/"; + let dates = list_dates(&store, &Path::from(prefix), prefix, &tenant) + .await + .unwrap(); + assert_eq!(dates, ["date=2026-11-01"]); + let root = Path::from("/"); + let dirs = list_dirs(&store, &root, &tenant).await.unwrap(); + assert_eq!(dirs, ["nested", "tenant", "tenant-other"]); + let nested = list_dirs(&store, &Path::from("tenant"), &tenant) + .await + .unwrap(); + assert_eq!(nested, ["tenant", "stream"]); + let paths = get_ingestor_meta_file_paths(&store, &root, &tenant) + .await + .unwrap(); + let paths: Vec<_> = paths.iter().map(|path| path.as_str()).collect(); + assert_eq!(paths, ["ingestor-one", "nested/ingestor-two"]); + let path = RelativePath::new("ingestor-one"); + assert_eq!(get_object(&store, path, &tenant).await.unwrap(), "data"); + assert_eq!( + get_object_ranged(&store, path, &tenant).await.unwrap(), + "data" + ); + let path = RelativePath::new("missing"); + assert!(matches!( + get_object(&store, path, &tenant).await, + Err(ObjectStorageError::NoSuchKey(_)) + )); + assert!(matches!( + get_object_ranged(&store, path, &tenant).await, + Err(ObjectStorageError::NoSuchKey(_)) + )); + } + + #[tokio::test] + async fn parallel_download_writes_all_chunks_and_truncates_existing_file() { + let store = InMemory::new(); + let dir = tempfile::tempdir().unwrap(); + let partial = dir.path().join("nested/file.partial"); + std::fs::create_dir_all(partial.parent().unwrap()).unwrap(); + for size in [0, 1, 4, 5, 12, 13] { + let data: Vec = (0..size).collect(); + store + .put(&Path::from("file"), data.clone().into()) + .await + .unwrap(); + for concurrency in [1, 3] { + std::fs::write(&partial, [255; 32]).unwrap(); + parallel_download( + &store, + RelativePath::new("file"), + "download-test", + partial.clone(), + 4, + concurrency, + ) + .await + .unwrap(); + assert_eq!(std::fs::read(&partial).unwrap(), data); + } + } + } + + #[tokio::test] + async fn parallel_download_propagates_store_and_file_errors() { + let store = InMemory::new(); + let dir = tempfile::tempdir().unwrap(); + let partial = dir.path().join("file.partial"); + let result = parallel_download( + &store, + RelativePath::new("missing"), + "download-error-test", + partial.clone(), + 4, + 2, + ) + .await; + assert!(matches!(result, Err(ObjectStorageError::NoSuchKey(_)))); + assert!(!partial.exists()); + + store.put(&Path::from("file"), "data".into()).await.unwrap(); + std::fs::write(&partial, "existing").unwrap(); + let result = parallel_download( + &store, + RelativePath::new("file"), + "download-error-test", + partial.join("child"), + 4, + 2, + ) + .await; + assert!(matches!(result, Err(ObjectStorageError::IoError(_)))); + assert_eq!(std::fs::read(&partial).unwrap(), b"existing"); + } + + #[tokio::test] + async fn delete_prefix_preserves_other_prefixes_and_accepts_empty_results() { + let store = InMemory::new(); + for key in ["stream/a", "stream/nested/b", "stream-other/a"] { + store.put(&Path::from(key), "data".into()).await.unwrap(); + } + delete_prefix(&store, "stream/", "delete-test") + .await + .unwrap(); + let remaining = store.list_with_delimiter(None).await.unwrap(); + assert_eq!(remaining.common_prefixes, [Path::from("stream-other")]); + assert!(remaining.objects.is_empty()); + delete_prefix(&store, "stream/", "delete-test") + .await + .unwrap(); + } - use super::{DEFAULT_MAX_OBJECT_STORE_REQUESTS, GcsConfig}; + #[tokio::test] + async fn delete_prefix_reports_listing_errors() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("blocked"), "data").unwrap(); + let store = LocalFileSystem::new_with_prefix(dir.path()).unwrap(); + let prefix = Path::from("blocked/child"); + assert!(list_dirs(&store, &prefix, &None).await.is_err()); + assert!( + list_dates(&store, &prefix, "blocked/child/", &None) + .await + .is_err() + ); + assert!( + list_partitions(&store, "blocked/child/", "hour=", &None) + .await + .is_err() + ); + let provider = LocalFS::new(dir.path().to_path_buf()); + assert!( + get_objects(&store, &provider, &prefix, Box::new(|_| true), &None) + .await + .is_err() + ); + assert!( + get_ingestor_meta_file_paths(&store, &prefix, &None) + .await + .is_err() + ); + let err = delete_prefix(&store, "blocked/child", "delete-error-test") + .await + .unwrap_err(); + assert!(matches!(err, ObjectStorageError::UnhandledError(_))); + assert_eq!( + err.to_string(), + "Unhandled Error: Failed to delete 1 out of 1 objects under prefix 'blocked/child'" + ); + } #[derive(Debug, Parser)] struct TestArgs { diff --git a/src/storage/s3.rs b/src/storage/s3.rs index a92923172..61033903f 100644 --- a/src/storage/s3.rs +++ b/src/storage/s3.rs @@ -20,7 +20,6 @@ use std::{ collections::HashSet, fmt::Display, num::NonZeroUsize, - ops::Range, path::{Path, PathBuf}, str::FromStr, sync::Arc, @@ -37,7 +36,7 @@ use datafusion::{ runtime_env::RuntimeEnvBuilder, }, }; -use futures::{StreamExt, TryStreamExt, stream::FuturesUnordered}; + use object_store::{ BackoffConfig, ClientOptions, ListResult, ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload, RetryConfig, @@ -51,23 +50,16 @@ use tokio::{fs::OpenOptions, io::AsyncReadExt}; use tracing::error; use crate::{ - metrics::{ - increment_bytes_scanned_in_object_store_calls_by_date, - increment_files_scanned_in_object_store_calls_by_date, - increment_object_store_calls_by_date, - increment_partial_file_scans_in_object_store_calls_by_date, - }, + metrics::increment_object_store_calls_by_date, parseable::{DEFAULT_TENANT, LogStream, PARSEABLE}, storage::RETRY_TIMEOUT_SECS, }; use super::{ CONNECT_TIMEOUT_SECS, DEFAULT_MAX_OBJECT_STORE_REQUESTS, ObjectStorage, ObjectStorageError, - ObjectStorageProvider, PARSEABLE_ROOT_DIRECTORY, REQUEST_TIMEOUT_SECS, - STREAM_METADATA_FILE_NAME, metrics_layer::MetricLayer, object_storage::parseable_json_path, - partial_path, to_object_store_path, + ObjectStorageProvider, REQUEST_TIMEOUT_SECS, metrics_layer::MetricLayer, + object_storage::parseable_json_path, partial_path, to_object_store_path, }; -use crate::storage::GET_OBJECTS_CONCURRENCY; // in bytes // const MULTIPART_UPLOAD_SIZE: usize = 1024 * 1024 * 100; @@ -429,67 +421,15 @@ impl S3 { tenant_id: &Option, partial_path: PathBuf, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let src = to_object_store_path(path); - - let meta = self.client.head(&src).await?; - increment_object_store_calls_by_date("HEAD", &date, tenant_str); - let total = meta.size; - - if let Some(parent) = partial_path.parent() { - tokio::fs::create_dir_all(parent).await?; - } - let file = tokio::fs::File::create(&partial_path).await?; - file.set_len(total).await?; - let std_file = Arc::new(file.into_std().await); - - let chunk = PARSEABLE.options.hot_tier_download_chunk_size; - let concurrency = PARSEABLE.options.hot_tier_download_concurrency; - let ranges: Vec> = (0..total) - .step_by(chunk as usize) - .map(|s| s..(s + chunk).min(total)) - .collect(); - let chunk_count = ranges.len() as u64; - tracing::Span::current() - .record("total_bytes", total) - .record("chunks", chunk_count); - - let semaphore = Arc::new(tokio::sync::Semaphore::new(concurrency as usize)); - - futures::stream::iter(ranges) - .map(|r| { - let src = src.clone(); - let std_file = std_file.clone(); - let semaphore = semaphore.clone(); - async move { - let _permit = semaphore.acquire_owned().await.map_err(|e| { - ObjectStorageError::Custom(format!("semaphore closed: {e}")) - })?; - let bytes = self.client.get_range(&src, r.clone()).await?; - let offset = r.start; - tokio::task::spawn_blocking(move || -> std::io::Result<()> { - crate::storage::write_all_at(&std_file, &bytes, offset) - }) - .await - .map_err(|e| ObjectStorageError::Custom(format!("join: {e}")))??; - Ok::<_, ObjectStorageError>(()) - } - }) - .buffer_unordered(concurrency as usize) - .try_collect::>() - .await?; - - increment_object_store_calls_by_date("GET", &date, tenant_str); - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant_str); - increment_bytes_scanned_in_object_store_calls_by_date("GET", total, &date, tenant_str); - increment_partial_file_scans_in_object_store_calls_by_date( - "GET", - chunk_count, - &date, - tenant_str, - ); - Ok(()) + super::parallel_download( + self.client.as_ref(), + path, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + partial_path, + PARSEABLE.options.hot_tier_download_chunk_size, + PARSEABLE.options.hot_tier_download_concurrency, + ) + .await } #[tracing::instrument( @@ -502,34 +442,7 @@ impl S3 { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.get(&to_object_store_path(path)).await; - increment_object_store_calls_by_date( - "GET", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - match resp { - Ok(resp) => { - let body = resp.bytes().await?; - tracing::Span::current().record("bytes", body.len()); - increment_files_scanned_in_object_store_calls_by_date( - "GET", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(body) - } - Err(err) => Err(err.into()), - } + super::get_object(self.client.as_ref(), path, tenant_id).await } #[tracing::instrument( @@ -543,25 +456,8 @@ impl S3 { resource: PutPayload, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let resp = self.client.put(&to_object_store_path(path), resource).await; - increment_object_store_calls_by_date( - "PUT", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - match resp { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + let result = self.client.put(&to_object_store_path(path), resource).await; + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -575,58 +471,12 @@ impl S3 { key: &str, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - - // List all objects under the prefix - let object_stream = self.client.list(Some(&(key.into()))); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - // Map listed objects to their locations for delete_stream - let locations_stream = object_stream.map(|result| result.map(|obj| obj.location)); - - // Use delete_stream which batches deletes via the S3 DeleteObjects API - let mut delete_results = self.client.delete_stream(Box::pin(locations_stream)); - - let mut files_deleted: u64 = 0; - let mut failed_deletes: u64 = 0; - while let Some(result) = delete_results.next().await { - match result { - Ok(_) => files_deleted += 1, - Err(err) => { - failed_deletes += 1; - error!("Failed to delete object during prefix deletion: {:?}", err); - } - } - } - - let total_files = files_deleted + failed_deletes; - tracing::Span::current() - .record("deleted", files_deleted) - .record("failed", failed_deletes); - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - total_files, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - files_deleted, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - if failed_deletes > 0 { - return Err(ObjectStorageError::UnhandledError( - anyhow::anyhow!("Failed to delete {failed_deletes} out of {total_files} objects under prefix '{key}'").into(), - )); - } - - Ok(()) + super::delete_prefix( + self.client.as_ref(), + key, + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + ) + .await } #[tracing::instrument( @@ -639,46 +489,18 @@ impl S3 { stream: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let (prefix, tenant_str) = if let Some(tenant) = tenant_id { - (format!("{}/{}/", tenant, stream), tenant.as_str()) + let prefix = if let Some(tenant) = tenant_id { + format!("{tenant}/{stream}/") } else { - (format!("{}/", stream), DEFAULT_TENANT) - }; - let resp: Result = self - .client - .list_with_delimiter(Some(&(prefix.as_str().into()))) - .await; - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - let resp = match resp { - Ok(resp) => resp, - Err(err) => { - return Err(err.into()); - } + format!("{stream}/") }; - - let common_prefixes = resp.common_prefixes; - - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - // return prefixes at the root level - let dates: Vec<_> = common_prefixes - .iter() - .filter_map(|path| path.as_ref().strip_prefix(&prefix)) - .map(String::from) - .collect(); - tracing::Span::current().record("dates", dates.len()); - - Ok(dates) + super::list_dates( + self.client.as_ref(), + &StorePath::from(prefix.as_str()), + &prefix, + tenant_id, + ) + .await } #[tracing::instrument( @@ -692,27 +514,9 @@ impl S3 { path: &Path, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let bytes = tokio::fs::read(path).await?; - let result = self.client.put(&key.into(), bytes.into()).await; - increment_object_store_calls_by_date( - "PUT", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "PUT", tenant_id).map(|_| ()) } #[tracing::instrument( @@ -739,66 +543,17 @@ impl S3 { // Track single PUT operation for small files let result = self.client.put(location, data.into()).await; - increment_object_store_calls_by_date( - "PUT", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "PUT", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - } - Err(err) => { - return Err(err.into()); - } - } + super::record_request(result, "PUT", tenant_id)?; return Ok(()); } else { - let async_writer = self.client.put_multipart(location).await; - let mut async_writer = match async_writer { - Ok(writer) => writer, - Err(err) => { - return Err(err.into()); - } - }; + let mut async_writer = self.client.put_multipart(location).await?; let mut data = Vec::new(); file.read_to_end(&mut data).await?; - // let mut upload_parts = Vec::new(); - - let has_final_partial_part = !total_size.is_multiple_of(min_multipart_size); - let num_full_parts = total_size / min_multipart_size; - let total_parts = num_full_parts + if has_final_partial_part { 1 } else { 0 }; - tracing::Span::current() - .record("total_bytes", total_size) - .record("parts", total_parts); - - // Upload each part with metrics - for part_number in 0..(total_parts) { - let start_pos = part_number * min_multipart_size; - let end_pos = if part_number == num_full_parts && has_final_partial_part { - // Last part might be smaller than 5MB (which is allowed) - total_size - } else { - // All other parts must be at least 5MB - start_pos + min_multipart_size - }; - - // Extract this part's data - let part_data = data[start_pos..end_pos].to_vec(); - - // Track individual part upload - let result = async_writer.put_part(part_data.into()).await; - if result.is_err() { - return Err(result.err().unwrap().into()); - } + for part_data in data[..total_size].chunks(min_multipart_size) { + async_writer.put_part(part_data.to_vec().into()).await?; increment_object_store_calls_by_date( "PUT_MULTIPART", &Utc::now().date_naive().to_string(), @@ -825,31 +580,9 @@ impl ObjectStorage for S3 { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let path = &to_object_store_path(path); - let meta = self.client.head(path).await; - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - let meta = match meta { - Ok(meta) => { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - meta - } - Err(err) => { - return Err(err.into()); - } - }; - - let buf = object_store::buffered::BufReader::new(self.client.clone(), &meta); - Ok(buf) + let result = self.client.head(&to_object_store_path(path)).await; + let meta = super::record_request(result, "HEAD", tenant_id)?; + Ok(BufReader::new(self.client.clone(), &meta)) } async fn upload_multipart( @@ -871,23 +604,8 @@ impl ObjectStorage for S3 { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.head(&to_object_store_path(path)).await; - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - } - - Ok(result?) + super::record_request(result, "HEAD", tenant_id) } async fn parallel_chunked_download( @@ -912,23 +630,7 @@ impl ObjectStorage for S3 { path: &RelativePath, tenant_id: &Option, ) -> Result { - let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let date = Utc::now().date_naive().to_string(); - let body = crate::storage::get_object_ranged( - &self.client, - &to_object_store_path(path), - path, - || increment_object_store_calls_by_date("GET", &date, tenant), - ) - .await?; - increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); - increment_bytes_scanned_in_object_store_calls_by_date( - "GET", - body.len() as u64, - &date, - tenant, - ); - Ok(body) + super::get_object_ranged(self.client.as_ref(), path, tenant_id).await } async fn get_objects( @@ -937,104 +639,19 @@ impl ObjectStorage for S3 { filter_func: Box bool + Send>, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let prefix = if let Some(base_path) = base_path { to_object_store_path(base_path) } else { self.root.clone() }; - - let mut list_stream = self.client.list(Some(&prefix)); - - let mut paths = vec![]; - let mut files_scanned = 0; - - // Drain the listing first, then fetch. Awaiting each GET inside the - // listing loop serialises every download against the one before it. - // Note: We track each streaming list item retrieval - while let Some(meta_result) = list_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let ingestor_file = filter_func(meta.location.filename().unwrap().to_string()); - - if !ingestor_file { - continue; - } - - paths.push( - RelativePath::from_path(meta.location.as_ref()) - .map_err(ObjectStorageError::PathError)? - .to_relative_path_buf(), - ); - } - - // `buffered` rather than `buffer_unordered`: callers of this generic - // helper may rely on results following listing order. - let res: Vec = futures::stream::iter(paths) - .map(|path| async move { self.get_object(&path, tenant_id).await }) - .buffered(GET_OBJECTS_CONCURRENCY) - .try_collect() - .await?; - - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(res) + super::get_objects(self.client.as_ref(), self, &prefix, filter_func, tenant_id).await } async fn get_ingestor_meta_file_paths( &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let mut path_arr = vec![]; - let mut files_scanned = 0; - - let mut object_stream = self.client.list(Some(&self.root)); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - while let Some(meta_result) = object_stream.next().await { - let meta = match meta_result { - Ok(meta) => meta, - Err(err) => { - return Err(err.into()); - } - }; - - files_scanned += 1; - let flag = meta.location.filename().unwrap().starts_with("ingestor"); - - if flag { - path_arr.push(RelativePathBuf::from(meta.location.as_ref())); - } - } - // Record total files scanned - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - files_scanned as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(path_arr) + super::get_ingestor_meta_file_paths(self.client.as_ref(), &self.root, tenant_id).await } async fn put_object( @@ -1055,9 +672,7 @@ impl ObjectStorage for S3 { path: &RelativePath, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - self._delete_prefix(path.as_ref(), tenant_id).await?; - - Ok(()) + self._delete_prefix(path.as_ref(), tenant_id).await } #[tracing::instrument( @@ -1071,23 +686,8 @@ impl ObjectStorage for S3 { path: &RelativePath, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self.client.delete(&to_object_store_path(path)).await; - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - } - - Ok(result?) + super::record_request(result, "DELETE", tenant_id) } #[tracing::instrument( @@ -1096,28 +696,12 @@ impl ObjectStorage for S3 { fields(tenant = ?tenant_id, ok = tracing::field::Empty) )] async fn check(&self, tenant_id: &Option) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let result = self .client .head(&to_object_store_path(&parseable_json_path())) .await; tracing::Span::current().record("ok", result.is_ok()); - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - if result.is_ok() { - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - } - - Ok(result.map(|_| ())?) + super::record_request(result.map(|_| ()), "HEAD", tenant_id) } async fn delete_stream( @@ -1130,9 +714,7 @@ impl ObjectStorage for S3 { } else { stream_name }; - self._delete_prefix(prefix, tenant_id).await?; - - Ok(()) + self._delete_prefix(prefix, tenant_id).await } async fn try_delete_node_meta( @@ -1140,27 +722,9 @@ impl ObjectStorage for S3 { node_filename: String, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); let file = RelativePathBuf::from(&node_filename); - let result = self.client.delete(&to_object_store_path(&file)).await; - increment_object_store_calls_by_date( - "DELETE", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - match result { - Ok(_) => { - increment_files_scanned_in_object_store_calls_by_date( - "DELETE", - 1, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - Ok(()) - } - Err(err) => Err(err.into()), - } + super::record_request(result, "DELETE", tenant_id) } async fn list_streams(&self) -> Result, ObjectStorageError> { @@ -1171,52 +735,7 @@ impl ObjectStorage for S3 { } async fn list_old_streams(&self) -> Result, ObjectStorageError> { - let tenant_str = DEFAULT_TENANT; - let resp = self.client.list_with_delimiter(None).await?; - let common_prefixes = resp.common_prefixes; // get all dirs - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - // return prefixes at the root level - let dirs: HashSet<_> = common_prefixes - .iter() - .filter_map(|path| path.parts().next()) - .map(|name| name.as_ref().to_string()) - .filter(|x| x != PARSEABLE_ROOT_DIRECTORY) - .collect(); - - let stream_json_check = FuturesUnordered::new(); - - for dir in &dirs { - let key = format!("{dir}/{STREAM_METADATA_FILE_NAME}"); - let task = async move { - let result = self.client.head(&StorePath::from(key)).await; - increment_object_store_calls_by_date( - "HEAD", - &Utc::now().date_naive().to_string(), - "", - ); - result.map(|_| ()) - }; - stream_json_check.push(task); - } - increment_files_scanned_in_object_store_calls_by_date( - "HEAD", - dirs.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - stream_json_check.try_collect::<()>().await?; - - Ok(dirs) + super::list_old_streams(self.client.as_ref(), "").await } async fn list_dates( @@ -1224,9 +743,7 @@ impl ObjectStorage for S3 { stream_name: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let streams = self._list_dates(stream_name, tenant_id).await?; - - Ok(streams) + self._list_dates(stream_name, tenant_id).await } async fn list_hours( @@ -1235,44 +752,12 @@ impl ObjectStorage for S3 { date: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let (pre, tenant_str) = if let Some(tenant) = tenant_id { - (format!("{tenant}/{stream_name}/{date}/"), tenant.as_str()) + let prefix = if let Some(tenant) = tenant_id { + format!("{tenant}/{stream_name}/{date}/") } else { - (format!("{stream_name}/{date}/"), DEFAULT_TENANT) + format!("{stream_name}/{date}/") }; - let resp = self - .client - .list_with_delimiter(Some(&object_store::path::Path::from(pre.as_str()))) - .await?; - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - let hours: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = path_str.strip_prefix(&pre) { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("hour=")) - .collect(); - - Ok(hours) + super::list_partitions(self.client.as_ref(), &prefix, "hour=", tenant_id).await } async fn list_minutes( @@ -1282,46 +767,12 @@ impl ObjectStorage for S3 { hour: &str, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let (pre, tenant_str) = if let Some(tenant) = tenant_id { - ( - format!("{tenant}/{stream_name}/{date}/{hour}/"), - tenant.as_str(), - ) + let prefix = if let Some(tenant) = tenant_id { + format!("{tenant}/{stream_name}/{date}/{hour}/") } else { - (format!("{stream_name}/{date}/{hour}/"), DEFAULT_TENANT) + format!("{stream_name}/{date}/{hour}/") }; - let resp = self - .client - .list_with_delimiter(Some(&object_store::path::Path::from(pre.as_str()))) - .await?; - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - let minutes: Vec = resp - .common_prefixes - .iter() - .filter_map(|path| { - let path_str = path.as_ref(); - if let Some(stripped) = path_str.strip_prefix(&pre) { - // Remove trailing slash if present, otherwise use as is - let clean_path = stripped.strip_suffix('/').unwrap_or(stripped); - Some(clean_path.to_string()) - } else { - None - } - }) - .filter(|dir| dir.starts_with("minute=")) - .collect(); - - Ok(minutes) + super::list_partitions(self.client.as_ref(), &prefix, "minute=", tenant_id).await } async fn upload_file( @@ -1330,7 +781,7 @@ impl ObjectStorage for S3 { path: &Path, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { - Ok(self._upload_file(key, path, tenant_id).await?) + self._upload_file(key, path, tenant_id).await } fn absolute_url(&self, prefix: &RelativePath) -> object_store::path::Path { @@ -1355,36 +806,7 @@ impl ObjectStorage for S3 { &self, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let pre = object_store::path::Path::from("/"); - let resp = self.client.list_with_delimiter(Some(&pre)).await; - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + super::list_dirs(self.client.as_ref(), &StorePath::from("/"), tenant_id).await } async fn list_dirs_relative( @@ -1392,37 +814,8 @@ impl ObjectStorage for S3 { relative_path: &RelativePath, tenant_id: &Option, ) -> Result, ObjectStorageError> { - let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); - let prefix = object_store::path::Path::from(relative_path.as_str()); - - let resp = self.client.list_with_delimiter(Some(&prefix)).await; - increment_object_store_calls_by_date( - "LIST", - &Utc::now().date_naive().to_string(), - tenant_str, - ); - let resp = match resp { - Ok(resp) => { - increment_files_scanned_in_object_store_calls_by_date( - "LIST", - resp.common_prefixes.len() as u64, - &Utc::now().date_naive().to_string(), - tenant_str, - ); - - resp - } - Err(err) => { - return Err(err.into()); - } - }; - - Ok(resp - .common_prefixes - .iter() - .flat_map(|path| path.parts()) - .map(|name| name.as_ref().to_string()) - .collect::>()) + let prefix = StorePath::from(relative_path.as_str()); + super::list_dirs(self.client.as_ref(), &prefix, tenant_id).await } fn get_bucket_name(&self) -> String { diff --git a/src/utils/json/flatten.rs b/src/utils/json/flatten.rs index 2052ccdd4..16e84d772 100644 --- a/src/utils/json/flatten.rs +++ b/src/utils/json/flatten.rs @@ -417,92 +417,68 @@ mod tests { use super::{JsonFlattenError, flatten}; use serde_json::{Map, Value, json}; - #[test] - fn flatten_single_key_string() { - let mut obj = json!({"key": "value"}); - let expected = obj.clone(); - flatten(&mut obj, "_", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn flatten_single_key_int() { - let mut obj = json!({"key": 1}); - let expected = obj.clone(); - flatten(&mut obj, "_", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn flatten_multiple_key_value() { - let mut obj = json!({"key1": 1, "key2": "value2"}); - let expected = obj.clone(); - flatten(&mut obj, "_", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn flatten_nested_single_key_value() { - let mut obj = json!({"key": "value", "nested_key": {"key":"value"}}); - let expected = json!({"key": "value", "nested_key.key": "value"}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_multiple_key_value() { - let mut obj = json!({"key": "value", "nested_key": {"key1":"value1", "key2": "value2"}}); - let expected = - json!({"key": "value", "nested_key.key1": "value1", "nested_key.key2": "value2"}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_key_value_with_array() { - let mut obj = json!({"key": "value", "nested_key": {"key1":[1,2,3]}}); - let expected = json!({"key": "value", "nested_key.key1": [1,2,3]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_obj_array() { - let mut obj = json!({"key": [{"a": "value0"}, {"a": "value1"}]}); - let expected = json!({"key.a": ["value0", "value1"]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_obj_array_nulls() { - let mut obj = json!({"key": [{"a": "value0"}, {"a": "value1", "b": "value1"}]}); - let expected = json!({"key.a": ["value0", "value1"], "key.b": [null, "value1"]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_obj_array_nulls_reversed() { - let mut obj = json!({"key": [{"a": "value0", "b": "value0"}, {"a": "value1"}]}); - let expected = json!({"key.a": ["value0", "value1"], "key.b": ["value0", null]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_obj_array_nested_obj() { - let mut obj = json!({"key": [{"a": {"p": 0}, "b": "value0"}, {"b": "value1"}]}); - let expected = json!({"key.a.p": [0, null], "key.b": ["value0", "value1"]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); - assert_eq!(obj, expected); - } - - #[test] - fn nested_obj_array_nested_obj_array() { - let mut obj = json!({"key": [{"a": [{"p": "value0", "q": "value0"}, {"p": "value1", "q": null}], "b": "value0"}, {"b": "value1"}]}); - let expected = json!({"key.a.p": [["value0", "value1"], null], "key.a.q": [["value0", null], null], "key.b": ["value0", "value1"]}); - flatten(&mut obj, ".", None, None, None, false).unwrap(); + #[rstest::rstest] + #[case::flatten_single_key_string( + json!({"key": "value"}), + "_", + json!({"key": "value"}) + )] + #[case::flatten_single_key_int( + json!({"key": 1}), + "_", + json!({"key": 1}) + )] + #[case::flatten_multiple_key_value( + json!({"key1": 1, "key2": "value2"}), + "_", + json!({"key1": 1, "key2": "value2"}) + )] + #[case::flatten_nested_single_key_value( + json!({"key": "value", "nested_key": {"key":"value"}}), + ".", + json!({"key": "value", "nested_key.key": "value"}) + )] + #[case::nested_multiple_key_value( + json!({"key": "value", "nested_key": {"key1":"value1", "key2": "value2"}}), + ".", + json!({"key": "value", "nested_key.key1": "value1", "nested_key.key2": "value2"}) + )] + #[case::nested_key_value_with_array( + json!({"key": "value", "nested_key": {"key1":[1,2,3]}}), + ".", + json!({"key": "value", "nested_key.key1": [1,2,3]}) + )] + #[case::nested_obj_array( + json!({"key": [{"a": "value0"}, {"a": "value1"}]}), + ".", + json!({"key.a": ["value0", "value1"]}) + )] + #[case::nested_obj_array_nulls( + json!({"key": [{"a": "value0"}, {"a": "value1", "b": "value1"}]}), + ".", + json!({"key.a": ["value0", "value1"], "key.b": [null, "value1"]}) + )] + #[case::nested_obj_array_nulls_reversed( + json!({"key": [{"a": "value0", "b": "value0"}, {"a": "value1"}]}), + ".", + json!({"key.a": ["value0", "value1"], "key.b": ["value0", null]}) + )] + #[case::nested_obj_array_nested_obj( + json!({"key": [{"a": {"p": 0}, "b": "value0"}, {"b": "value1"}]}), + ".", + json!({"key.a.p": [0, null], "key.b": ["value0", "value1"]}) + )] + #[case::nested_obj_array_nested_obj_array( + json!({"key": [{"a": [{"p": "value0", "q": "value0"}, {"p": "value1", "q": null}], "b": "value0"}, {"b": "value1"}]}), + ".", + json!({"key.a.p": [["value0", "value1"], null], "key.a.q": [["value0", null], null], "key.b": ["value0", "value1"]}) + )] + fn flatten_object_cases( + #[case] mut obj: Value, + #[case] separator: &str, + #[case] expected: Value, + ) { + flatten(&mut obj, separator, None, None, None, false).unwrap(); assert_eq!(obj, expected); } @@ -512,163 +488,69 @@ mod tests { assert!(flatten(&mut obj, ".", None, None, None, false).is_err()); } - #[test] - fn flatten_array_nulls_at_start() { - let Value::Array(mut arr) = json!([ + #[rstest::rstest] + #[case::flatten_array_nulls_at_start( + json!([ null, {"p": 2, "q": 2}, {"q": 3}, - ]) else { - unreachable!() - }; - - let mut map = Map::new(); - flatten_array_objects(&mut map, "key", &mut arr, ".").unwrap(); - - assert_eq!(map.len(), 2); - assert_eq!(map.get("key.p").unwrap(), &json!([null, 2, null])); - assert_eq!(map.get("key.q").unwrap(), &json!([null, 2, 3])); - } - - #[test] - fn flatten_array_objects_nulls_at_end() { - let Value::Array(mut arr) = json!([{"a": 1, "b": 1}, {"a": 2}, null]) else { - unreachable!() - }; - - let mut map = Map::new(); - flatten_array_objects(&mut map, "key", &mut arr, ".").unwrap(); - - assert_eq!(map.len(), 2); - assert_eq!(map.get("key.a").unwrap(), &json!([1, 2, null])); - assert_eq!(map.get("key.b").unwrap(), &json!([1, null, null])); - } - - #[test] - fn flatten_array_objects_nulls_in_middle() { - let Value::Array(mut arr) = json!([{"a": 1, "b": 1}, null, {"a": 3, "c": 3}]) else { - unreachable!() - }; - - let mut map = Map::new(); - flatten_array_objects(&mut map, "key", &mut arr, ".").unwrap(); - - assert_eq!(map.len(), 3); - assert_eq!(map.get("key.a").unwrap(), &json!([1, null, 3])); - assert_eq!(map.get("key.b").unwrap(), &json!([1, null, null])); - assert_eq!(map.get("key.c").unwrap(), &json!([null, null, 3])); - } - - #[test] - fn flatten_array_test() { - let Value::Array(mut arr) = json!([ + ]), + json!({"key.p": [null, 2, null], "key.q": [null, 2, 3]}) + )] + #[case::flatten_array_objects_nulls_at_end( + json!([{"a": 1, "b": 1}, {"a": 2}, null]), + json!({"key.a": [1, 2, null], "key.b": [1, null, null]}) + )] + #[case::flatten_array_objects_nulls_in_middle( + json!([{"a": 1, "b": 1}, null, {"a": 3, "c": 3}]), + json!({"key.a": [1, null, 3], "key.b": [1, null, null], "key.c": [null, null, 3]}) + )] + #[case::flatten_array_test( + json!([ {"p": 1, "q": 1}, {"r": 2, "q": 2}, {"p": 3, "r": 3} - ]) else { - unreachable!() - }; - - let mut map = Map::new(); - flatten_array_objects(&mut map, "key", &mut arr, ".").unwrap(); - - assert_eq!(map.len(), 3); - assert_eq!(map.get("key.p").unwrap(), &json!([1, null, 3])); - assert_eq!(map.get("key.q").unwrap(), &json!([1, 2, null])); - assert_eq!(map.get("key.r").unwrap(), &json!([null, 2, 3])); - } - - #[test] - fn flatten_array_nested_test() { - let Value::Array(mut arr) = json!([ + ]), + json!({"key.p": [1, null, 3], "key.q": [1, 2, null], "key.r": [null, 2, 3]}) + )] + #[case::flatten_array_nested_test( + json!([ {"p": 1, "q": [{"x": 1}, {"x": 2}]}, {"r": 2, "q": [{"x": 1}]}, {"p": 3, "r": 3} - ]) else { - unreachable!() - }; - + ]), + json!({"key.p": [1, null, 3], "key.q.x": [[1, 2], [1], null], "key.r": [null, 2, 3]}) + )] + fn flatten_array_objects_cases(#[case] mut value: Value, #[case] expected: Value) { let mut map = Map::new(); - flatten_array_objects(&mut map, "key", &mut arr, ".").unwrap(); - - assert_eq!(map.len(), 3); - assert_eq!(map.get("key.p").unwrap(), &json!([1, null, 3])); - assert_eq!(map.get("key.q.x").unwrap(), &json!([[1, 2], [1], null])); - assert_eq!(map.get("key.r").unwrap(), &json!([null, 2, 3])); - } - - #[test] - fn acceptable_value_custom_parition_test() { - let mut value = json!({ - "a": 1, - }); - assert!(flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).is_ok()); - - let mut value = json!({ - "a": true, - }); - assert!(flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).is_ok()); - - let mut value = json!({ - "a": "yes", - }); - assert!(flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).is_ok()); - - let mut value = json!({ - "a": -1, - }); - assert!(flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).is_ok()); - } - - #[test] - fn unacceptable_value_custom_partition_test() { - let mut value = json!({ - "a": null, - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldEmptyOrNull(_) - ); - - let mut value = json!({ - "a": "", - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldEmptyOrNull(_) - ); - - let mut value = json!({ - "a": {"b": 1}, - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldIsObject(_) - ); - - let mut value = json!({ - "a": ["b", "c"], - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldIsArray(_) - ); - - let mut value = json!({ - "a": "b.c", - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldContainsPeriod(_) - ); - - let mut value = json!({ - "a": 1.0, - }); - matches!( - flatten(&mut value, "_", None, None, Some(&"a".to_string()), true).unwrap_err(), - JsonFlattenError::FieldContainsPeriod(_) - ); + flatten_array_objects(&mut map, "key", value.as_array_mut().unwrap(), ".").unwrap(); + assert_eq!(Value::Object(map), expected); + } + + #[rstest::rstest] + #[case::integer(json!(1))] + #[case::boolean(json!(true))] + #[case::string(json!("yes"))] + #[case::negative_integer(json!(-1))] + fn acceptable_value_custom_parition_test(#[case] value: Value) { + let mut value = json!({"a": value}); + assert!(flatten(&mut value, "_", None, None, Some(&"a".to_owned()), true).is_ok()); + } + + #[rstest::rstest] + #[case::null(json!(null), JsonFlattenError::FieldEmptyOrNull("a".into()))] + #[case::empty_string(json!(""), JsonFlattenError::FieldEmptyOrNull("a".into()))] + #[case::object(json!({"b": 1}), JsonFlattenError::FieldIsObject("a".into()))] + #[case::array(json!(["b", "c"]), JsonFlattenError::FieldIsArray("a".into()))] + #[case::period(json!("b.c"), JsonFlattenError::FieldContainsPeriod("a".into()))] + #[case::float(json!(1.0), JsonFlattenError::FieldContainsPeriod("a".into()))] + fn unacceptable_value_custom_partition_test( + #[case] value: Value, + #[case] expected: JsonFlattenError, + ) { + let mut value = json!({"a": value}); + let error = flatten(&mut value, "_", None, None, Some(&"a".to_owned()), true).unwrap_err(); + assert_eq!(error.to_string(), expected.to_string()); } #[test] From ec2f02660d9e3a69b8806130eb20d9ef37af4724 Mon Sep 17 00:00:00 2001 From: ByteBaker <42913098+ByteBaker@users.noreply.github.com> Date: Sun, 1 Nov 2026 11:11:07 +0530 Subject: [PATCH 2/2] test: fix metastore persistence checks on Windows - Check files under the temporary root to avoid platform-specific read paths. - Document shared storage and migration contracts without changing runtime behavior. --- src/metastore/metastores/object_store_metastore.rs | 11 +++++------ src/migration/stream_metadata_migration.rs | 4 ++++ src/storage/localfs.rs | 3 +++ src/storage/mod.rs | 10 ++++++++++ 4 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/metastore/metastores/object_store_metastore.rs b/src/metastore/metastores/object_store_metastore.rs index 2a9c7bda5..5b933b495 100644 --- a/src/metastore/metastores/object_store_metastore.rs +++ b/src/metastore/metastores/object_store_metastore.rs @@ -72,6 +72,7 @@ pub struct ObjectStoreMetastore { } impl ObjectStoreMetastore { + /// Serialize at the object's own path; the tenant is forwarded without changing the path. async fn put_object( &self, obj: &dyn MetastoreObject, @@ -85,6 +86,7 @@ impl ObjectStoreMetastore { .await?) } + /// Delete at the object's own path, preserving storage errors for missing objects. async fn delete_object( &self, obj: &dyn MetastoreObject, @@ -1387,18 +1389,15 @@ mod tests { id: Ulid::new(), path: format!("{}/object.json", tenant.as_deref().unwrap_or("default")), }; - let path = RelativePathBuf::from(&object.path); + let path = dir.path().join(&object.path); metastore.put_object(&object, &tenant).await.unwrap(); - let bytes = storage.get_object(&path, &tenant).await.unwrap(); + let bytes = tokio::fs::read(&path).await.unwrap(); assert_eq!( serde_json::from_slice::(&bytes).unwrap(), serde_json::json!({"id": object.id.to_string(), "path": object.path}) ); metastore.delete_object(&object, &tenant).await.unwrap(); - assert!(matches!( - storage.get_object(&path, &tenant).await, - Err(ObjectStorageError::NoSuchKey(_)) - )); + assert!(!path.exists()); assert!(matches!( metastore.delete_object(&object, &tenant).await, Err(MetastoreError::ObjectStorageError(_)) diff --git a/src/migration/stream_metadata_migration.rs b/src/migration/stream_metadata_migration.rs index 9c0cea0d9..05017e8f9 100644 --- a/src/migration/stream_metadata_migration.rs +++ b/src/migration/stream_metadata_migration.rs @@ -24,6 +24,7 @@ use crate::{ }; use serde_json::{Value, json}; +/// Expand required legacy stats and update versions without replacing unrelated metadata. fn migrate_stats_and_version(stream_metadata_map: &mut serde_json::Map) { let stats = stream_metadata_map.get("stats").unwrap().clone(); let default_stats = json!({ @@ -54,6 +55,7 @@ fn migrate_stats_and_version(stream_metadata_map: &mut serde_json::Map Value { let stream_metadata_map = stream_metadata.as_object_mut().unwrap(); migrate_stats_and_version(stream_metadata_map); @@ -67,10 +69,12 @@ pub fn v1_v4(mut stream_metadata: Value) -> Value { stream_metadata } +/// Apply the same stats and snapshot upgrade as v1; v2 needs no separate conversion. pub fn v2_v4(stream_metadata: Value) -> Value { v1_v4(stream_metadata) } +/// Upgrade legacy stats and v1 snapshots, leaving newer snapshots and other fields intact. pub fn v3_v4(mut stream_metadata: Value) -> Value { let stream_metadata_map: &mut serde_json::Map = stream_metadata.as_object_mut().unwrap(); diff --git a/src/storage/localfs.rs b/src/storage/localfs.rs index b77c831a6..b4a09d42e 100644 --- a/src/storage/localfs.rs +++ b/src/storage/localfs.rs @@ -726,6 +726,7 @@ impl ObjectStorage for LocalFS { } } +/// Collect directory names in completion order, propagating entry and file-type errors. async fn directory_names(read_dir: fs::ReadDir) -> Result, ObjectStorageError> { let entries = ReadDirStream::new(read_dir) .try_collect::>() @@ -740,6 +741,7 @@ async fn directory_names(read_dir: fs::ReadDir) -> Result, ObjectSto ) } +/// Discover streams using the caller's ignore list and metadata layout; missing roots are errors. async fn stream_directories( root: &Path, ignore_dirs: &[&str], @@ -766,6 +768,7 @@ async fn stream_directories( .collect()) } +/// Skip ignored entries and files; a stream directory without its metadata is an error. async fn dir_with_stream( entry: DirEntry, ignore_dirs: &[&str], diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 6b71c5c3a..3b68c413a 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -631,6 +631,7 @@ pub fn to_object_store_path(path: &RelativePath) -> Path { Path::from(path.as_str()) } +/// Strip the caller's prefix from date directories, counting LIST attempts even on failure. async fn list_dates( client: &S, prefix: &Path, @@ -663,6 +664,7 @@ async fn list_dates( Ok(dates) } +/// Require metadata in every legacy stream directory; retain the provider's HEAD metric tenant. async fn list_old_streams( client: &S, head_tenant: &str, @@ -714,6 +716,7 @@ async fn list_old_streams( Ok(dirs) } +/// Return matching partition names without the parent prefix or trailing slash. async fn list_partitions( client: &S, prefix: &str, @@ -751,6 +754,7 @@ async fn list_partitions( Ok(partitions) } +/// Fetch filtered filenames concurrently in listing order through the provider's traced GET. async fn get_objects( client: &S, provider: &impl ObjectStorage, @@ -803,6 +807,7 @@ async fn get_objects( Ok(res) } +/// Count each GET attempt, recording file and byte metrics only after the full body is read. async fn get_object( client: &S, path: &RelativePath, @@ -830,6 +835,7 @@ async fn get_object( Ok(body) } +/// Count every request attempt and count a scanned file only when the request succeeds. fn record_request( result: Result, method: &str, @@ -847,6 +853,7 @@ fn record_request( Ok(value) } +/// Return every component of each child prefix, retaining the providers' directory listing format. async fn list_dirs( client: &S, prefix: &Path, @@ -872,6 +879,7 @@ async fn list_dirs( .collect::>()) } +/// Select filenames starting with `ingestor`, counting all listed files in scan metrics. async fn get_ingestor_meta_file_paths( client: &S, prefix: &Path, @@ -904,6 +912,7 @@ async fn get_ingestor_meta_file_paths( Ok(path_arr) } +/// Truncate the destination and fill it with bounded concurrent ranges; count successful downloads. async fn parallel_download( client: &S, path: &RelativePath, @@ -968,6 +977,7 @@ async fn parallel_download( Ok(()) } +/// Drain bulk deletion results before reporting failures, counting only successful deletes. async fn delete_prefix( client: &S, key: &str,