Merge remote-tracking branch 'origin/main' into query-param

Author: de-sh

Diff for: Cargo.lock

@@ -106,7 +106,6 @@ bytes = "1.4"
 clokwerk = "0.4"
 derive_more = { version = "1", features = ["full"] }
 itertools = "0.14"
-lazy_static = "1.4"
 once_cell = "1.20"
 rand = "0.8.5"
 regex = "1.7.3"

Diff for: Cargo.toml

@@ -24,6 +24,7 @@ use datafusion::common::tree_node::TreeNode;
 use http::StatusCode;
 use itertools::Itertools;
 use once_cell::sync::Lazy;
+use serde::Serialize;
 use serde_json::Error as SerdeError;
 use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Display};
@@ -873,3 +874,52 @@ impl Alerts {
         Ok(())
     }
 }
+
+#[derive(Debug, Serialize)]
+pub struct AlertsInfo {
+    total: u64,
+    silenced: u64,
+    resolved: u64,
+    triggered: u64,
+    low: u64,
+    medium: u64,
+    high: u64,
+}
+
+// TODO: add RBAC
+pub async fn get_alerts_info() -> Result<AlertsInfo, AlertError> {
+    let alerts = ALERTS.alerts.read().await;
+    let mut total = 0;
+    let mut silenced = 0;
+    let mut resolved = 0;
+    let mut triggered = 0;
+    let mut low = 0;
+    let mut medium = 0;
+    let mut high = 0;
+
+    for (_, alert) in alerts.iter() {
+        total += 1;
+        match alert.state {
+            AlertState::Silenced => silenced += 1,
+            AlertState::Resolved => resolved += 1,
+            AlertState::Triggered => triggered += 1,
+        }
+
+        match alert.severity {
+            Severity::Low => low += 1,
+            Severity::Medium => medium += 1,
+            Severity::High => high += 1,
+            _ => {}
+        }
+    }
+
+    Ok(AlertsInfo {
+        total,
+        silenced,
+        resolved,
+        triggered,
+        low,
+        medium,
+        high,
+    })
+}

Diff for: src/alerts/mod.rs

@@ -340,6 +340,12 @@ pub async fn remove_manifest_from_snapshot(
             Ok(get_first_event(storage.clone(), stream_name, Vec::new()).await?)
         }
         Mode::Query => Ok(get_first_event(storage, stream_name, dates).await?),
+        Mode::Index => Err(ObjectStorageError::UnhandledError(Box::new(
+            std::io::Error::new(
+                std::io::ErrorKind::Unsupported,
+                "Can't remove manifest from within Index server",
+            ),
+        ))),
     }
 }
 
@@ -350,6 +356,7 @@ pub async fn get_first_event(
 ) -> Result<Option<String>, ObjectStorageError> {
     let mut first_event_at: String = String::default();
     match PARSEABLE.options.mode {
+        Mode::Index => unimplemented!(),
         Mode::All | Mode::Ingest => {
             // get current snapshot
             let stream_first_event = PARSEABLE.get_stream(stream_name)?.get_first_event();

Diff for: src/catalog/mod.rs

@@ -16,10 +16,9 @@
  *
  */
 
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use async_trait::async_trait;
-use chrono::Utc;
 use futures_util::StreamExt;
 use rdkafka::consumer::{CommitMode, Consumer};
 use serde_json::Value;
@@ -58,37 +57,10 @@ impl ParseableSinkProcessor {
         let stream = PARSEABLE.get_stream(stream_name)?;
         let schema = stream.get_schema_raw();
         let time_partition = stream.get_time_partition();
+        let custom_partition = stream.get_custom_partition();
         let static_schema_flag = stream.get_static_schema_flag();
         let schema_version = stream.get_schema_version();
 
-        let (json_vec, total_payload_size) = Self::json_vec(records);
-        let batch_json_event = json::Event {
-            data: Value::Array(json_vec),
-        };
-
-        let (rb, is_first) = batch_json_event.into_recordbatch(
-            &schema,
-            static_schema_flag,
-            time_partition.as_ref(),
-            schema_version,
-        )?;
-
-        let p_event = ParseableEvent {
-            rb,
-            stream_name: stream_name.to_string(),
-            origin_format: "json",
-            origin_size: total_payload_size,
-            is_first_event: is_first,
-            parsed_timestamp: Utc::now().naive_utc(),
-            time_partition: None,
-            custom_partition_values: HashMap::new(),
-            stream_type: StreamType::UserDefined,
-        };
-
-        Ok(p_event)
-    }
-
-    fn json_vec(records: &[ConsumerRecord]) -> (Vec<Value>, u64) {
         let mut json_vec = Vec::with_capacity(records.len());
         let mut total_payload_size = 0u64;
 
@@ -99,22 +71,30 @@ impl ParseableSinkProcessor {
             }
         }
 
-        (json_vec, total_payload_size)
+        let p_event = json::Event::new(Value::Array(json_vec)).into_event(
+            stream_name.to_string(),
+            total_payload_size,
+            &schema,
+            static_schema_flag,
+            custom_partition.as_ref(),
+            time_partition.as_ref(),
+            schema_version,
+            StreamType::UserDefined,
+        )?;
+
+        Ok(p_event)
     }
 }
 
 #[async_trait]
 impl Processor<Vec<ConsumerRecord>, ()> for ParseableSinkProcessor {
     async fn process(&self, records: Vec<ConsumerRecord>) -> anyhow::Result<()> {
         let len = records.len();
-        debug!("Processing {} records", len);
+        debug!("Processing {len} records");
 
-        self.build_event_from_chunk(&records)
-            .await?
-            .process()
-            .await?;
+        self.build_event_from_chunk(&records).await?.process()?;
 
-        debug!("Processed {} records", len);
+        debug!("Processed {len} records");
         Ok(())
     }
 }

Diff for: src/connectors/kafka/processor.rs

@@ -0,0 +1 @@
+pub mod utils;

Diff for: src/enterprise/mod.rs

@@ -0,0 +1,159 @@
+use std::{collections::HashMap, path::PathBuf, sync::Arc};
+
+use datafusion::{common::Column, prelude::Expr};
+use itertools::Itertools;
+use relative_path::RelativePathBuf;
+
+use crate::query::stream_schema_provider::extract_primary_filter;
+use crate::{
+    catalog::{
+        manifest::{File, Manifest},
+        snapshot, Snapshot,
+    },
+    event,
+    parseable::PARSEABLE,
+    query::{stream_schema_provider::ManifestExt, PartialTimeFilter},
+    storage::{ObjectStorage, ObjectStorageError, ObjectStoreFormat, STREAM_ROOT_DIRECTORY},
+    utils::time::TimeRange,
+};
+
+pub fn create_time_filter(
+    time_range: &TimeRange,
+    time_partition: Option<String>,
+    table_name: &str,
+) -> Vec<Expr> {
+    let mut new_filters = vec![];
+    let start_time = time_range.start.naive_utc();
+    let end_time = time_range.end.naive_utc();
+    let mut _start_time_filter: Expr;
+    let mut _end_time_filter: Expr;
+
+    match time_partition {
+        Some(time_partition) => {
+            _start_time_filter = PartialTimeFilter::Low(std::ops::Bound::Included(start_time))
+                .binary_expr(Expr::Column(Column::new(
+                    Some(table_name.to_owned()),
+                    time_partition.clone(),
+                )));
+            _end_time_filter =
+                PartialTimeFilter::High(std::ops::Bound::Excluded(end_time)).binary_expr(
+                    Expr::Column(Column::new(Some(table_name.to_owned()), time_partition)),
+                );
+        }
+        None => {
+            _start_time_filter = PartialTimeFilter::Low(std::ops::Bound::Included(start_time))
+                .binary_expr(Expr::Column(Column::new(
+                    Some(table_name.to_owned()),
+                    event::DEFAULT_TIMESTAMP_KEY,
+                )));
+            _end_time_filter = PartialTimeFilter::High(std::ops::Bound::Excluded(end_time))
+                .binary_expr(Expr::Column(Column::new(
+                    Some(table_name.to_owned()),
+                    event::DEFAULT_TIMESTAMP_KEY,
+                )));
+        }
+    }
+
+    new_filters.push(_start_time_filter);
+    new_filters.push(_end_time_filter);
+
+    new_filters
+}
+
+pub async fn fetch_parquet_file_paths(
+    stream: &str,
+    time_range: &TimeRange,
+) -> Result<HashMap<RelativePathBuf, Vec<File>>, ObjectStorageError> {
+    let glob_storage = PARSEABLE.storage.get_object_store();
+
+    let object_store_format = glob_storage.get_object_store_format(stream).await?;
+
+    let time_partition = object_store_format.time_partition;
+
+    let time_filter_expr = create_time_filter(time_range, time_partition.clone(), stream);
+
+    let time_filters = extract_primary_filter(&time_filter_expr, &time_partition);
+
+    let mut merged_snapshot: snapshot::Snapshot = snapshot::Snapshot::default();
+
+    let path = RelativePathBuf::from_iter([stream, STREAM_ROOT_DIRECTORY]);
+    let obs = glob_storage
+        .get_objects(
+            Some(&path),
+            Box::new(|file_name| file_name.ends_with("stream.json")),
+        )
+        .await;
+    if let Ok(obs) = obs {
+        for ob in obs {
+            if let Ok(object_store_format) = serde_json::from_slice::<ObjectStoreFormat>(&ob) {
+                let snapshot = object_store_format.snapshot;
+                for manifest in snapshot.manifest_list {
+                    merged_snapshot.manifest_list.push(manifest);
+                }
+            }
+        }
+    }
+
+    let manifest_files = collect_manifest_files(
+        glob_storage,
+        merged_snapshot
+            .manifests(&time_filters)
+            .into_iter()
+            .sorted_by_key(|file| file.time_lower_bound)
+            .map(|item| item.manifest_path)
+            .collect(),
+    )
+    .await?;
+
+    let mut parquet_files: HashMap<RelativePathBuf, Vec<File>> = HashMap::new();
+
+    let mut selected_files = manifest_files
+        .into_iter()
+        .flat_map(|file| file.files)
+        .rev()
+        .collect_vec();
+
+    for filter in time_filter_expr {
+        selected_files.retain(|file| !file.can_be_pruned(&filter))
+    }
+
+    selected_files
+        .into_iter()
+        .map(|file| {
+            let date = file.file_path.split("/").collect_vec();
+
+            let date = date.as_slice()[1..4].iter().map(|s| s.to_string());
+
+            let date = RelativePathBuf::from_iter(date);
+
+            parquet_files.entry(date).or_default().push(file);
+        })
+        .for_each(|_| {});
+
+    Ok(parquet_files)
+}
+
+async fn collect_manifest_files(
+    storage: Arc<dyn ObjectStorage>,
+    manifest_urls: Vec<String>,
+) -> Result<Vec<Manifest>, ObjectStorageError> {
+    let mut tasks = Vec::new();
+    manifest_urls.into_iter().for_each(|path| {
+        let path = RelativePathBuf::from_path(PathBuf::from(path)).expect("Invalid path");
+        let storage = Arc::clone(&storage);
+        tasks.push(tokio::task::spawn(async move {
+            storage.get_object(&path).await
+        }));
+    });
+
+    let mut op = Vec::new();
+    for task in tasks {
+        let file = task.await??;
+        op.push(file);
+    }
+
+    Ok(op
+        .into_iter()
+        .map(|res| serde_json::from_slice(&res).expect("Data is invalid for Manifest"))
+        .collect())
+}