Fixed the issue

iPeluwa · iPeluwa · commit af53a195e2d1 · 2025-03-27T15:37:55.000+01:00
diff --git a/src/ingest.rs b/src/ingest.rs
@@ -1,14 +1,11 @@
-// ingest.rs
-
-use std::{
-    collections::HashMap,
-    sync::{Arc, RwLock},
-};
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
 
 use actix_web::{get, post, web, HttpResponse, Responder};
+use datafusion::arrow::record_batch::RecordBatch;
 use futures::future::join_all;
 use serde::Deserialize;
-use serde_json::json;
+use serde_json::{json, Value};
 use tracing::{error, info};
 
 use crate::{
@@ -697,4 +694,123 @@ pub async fn queue_length(queue: web::Data<Arc<PersistentQueue>>) -> impl Respon
             HttpResponse::InternalServerError().body("Error getting queue length")
         }
     }
+}
+
+#[get("/data")]
+pub async fn get_all_data(db: web::Data<Arc<Database>>) -> impl Responder {
+    let query = "SELECT projectId, id, timestamp, traceId, spanId, eventType, durationNs FROM telemetry_events";
+    match db.query(query).await {
+        Ok(df) => {
+            match df.collect().await {
+                Ok(batches) => {
+                    let json_rows = record_batches_to_json_rows(&batches).unwrap_or_default();
+                    HttpResponse::Ok().json(json_rows)
+                }
+                Err(e) => {
+                    error!("Error collecting data: {:?}", e);
+                    HttpResponse::InternalServerError().body("Error collecting data")
+                }
+            }
+        }
+        Err(e) => {
+            error!("Error querying data: {:?}", e);
+            HttpResponse::InternalServerError().body("Error querying data")
+        }
+    }
+}
+
+#[get("/data/{id}")]
+pub async fn get_data_by_id(
+    path: web::Path<String>,
+    db: web::Data<Arc<Database>>,
+) -> impl Responder {
+    let id = path.into_inner();
+    let query = format!(
+        "SELECT projectId, id, timestamp, traceId, spanId, eventType, durationNs FROM telemetry_events WHERE id = '{}'",
+        id
+    );
+    match db.query(&query).await {
+        Ok(df) => {
+            match df.collect().await {
+                Ok(batches) => {
+                    let json_rows = record_batches_to_json_rows(&batches).unwrap_or_default();
+                    if json_rows.is_empty() {
+                        HttpResponse::NotFound().body(format!("No data found for id: {}", id))
+                    } else {
+                        HttpResponse::Ok().json(json_rows)
+                    }
+                }
+                Err(e) => {
+                    error!("Error collecting data for id {}: {:?}", id, e);
+                    HttpResponse::InternalServerError().body("Error collecting data")
+                }
+            }
+        }
+        Err(e) => {
+            error!("Error querying data for id {}: {:?}", id, e);
+            HttpResponse::InternalServerError().body("Error querying data")
+        }
+    }
+}
+
+pub fn record_batches_to_json_rows(batches: &[RecordBatch]) -> Result<Vec<Value>, anyhow::Error> {
+    let mut rows = Vec::new();
+    for batch in batches {
+        let schema = batch.schema();
+        let num_rows = batch.num_rows();
+        for row_idx in 0..num_rows {
+            let mut row = json!({});
+            for (col_idx, field) in schema.fields().iter().enumerate() {
+                let column = batch.column(col_idx);
+                let value = if column.is_null(row_idx) {
+                    Value::Null
+                } else {
+                    match column.data_type() {
+                        datafusion::arrow::datatypes::DataType::Int32 => {
+                            column
+                                .as_any()
+                                .downcast_ref::<datafusion::arrow::array::Int32Array>()
+                                .map_or(Value::Null, |arr| Value::Number(arr.value(row_idx).into()))
+                        }
+                        datafusion::arrow::datatypes::DataType::Int64 => {
+                            column
+                                .as_any()
+                                .downcast_ref::<datafusion::arrow::array::Int64Array>()
+                                .map_or(Value::Null, |arr| Value::Number(arr.value(row_idx).into()))
+                        }
+                        datafusion::arrow::datatypes::DataType::Float64 => {
+                            column
+                                .as_any()
+                                .downcast_ref::<datafusion::arrow::array::Float64Array>()
+                                .map_or(Value::Null, |arr| {
+                                    Value::Number(
+                                        serde_json::Number::from_f64(arr.value(row_idx))
+                                            .unwrap_or_else(|| serde_json::Number::from(0)),
+                                    )
+                                })
+                        }
+                        datafusion::arrow::datatypes::DataType::Utf8 => {
+                            column
+                                .as_any()
+                                .downcast_ref::<datafusion::arrow::array::StringArray>()
+                                .map_or(Value::Null, |arr| Value::String(arr.value(row_idx).to_string()))
+                        }
+                        datafusion::arrow::datatypes::DataType::Timestamp(_, _) => {
+                            column
+                                .as_any()
+                                .downcast_ref::<datafusion::arrow::array::TimestampNanosecondArray>()
+                                .map_or(Value::Null, |arr| Value::String(arr.value(row_idx).to_string()))
+                        }
+                        _ => {
+                            // Fallback for unsupported types
+                            Value::Null
+                        }
+                    }
+                };
+                row[field.name()] = value;
+            }
+            rows.push(row);
+        }
+    }
+    Ok(rows)
 }
diff --git a/src/main.rs b/src/main.rs
@@ -72,8 +72,8 @@ async fn dashboard(
     query: web::Query<HashMap<String, String>>,
 ) -> impl Responder {
     let uptime = Utc::now().signed_duration_since(app_info.start_time).num_seconds() as f64;
-    let http_requests = 0.0;
-    let queue_size = queue.get_ref().len().await.unwrap_or(0);
+    let http_requests = 0.0; // Placeholder; update if metrics_middleware tracks this
+    let queue_size = queue.len().unwrap_or(0);
     let db_status = match db.query("SELECT 1 AS test").await {
         Ok(_) => "success",
         Err(_) => "error",
@@ -238,19 +238,14 @@ async fn main() -> anyhow::Result<()> {
 
     info!("Starting TimeFusion application");
 
-    // Read AWS configuration from environment variables.
-    // AWS_S3_BUCKET should be your bucket name, e.g. "my-aws-bucket"
-    // AWS_S3_ENDPOINT should be your AWS S3 endpoint, e.g. "https://s3.amazonaws.com"
+    // Read AWS configuration from environment variables
     let bucket = env::var("AWS_S3_BUCKET").expect("AWS_S3_BUCKET environment variable not set");
     let aws_endpoint = env::var("AWS_S3_ENDPOINT").unwrap_or_else(|_| "https://s3.amazonaws.com".to_string());
-    // Build the storage URI for AWS S3.
     let storage_uri = format!("s3://{}/?endpoint={}", bucket, aws_endpoint);
     info!("Storage URI configured: {}", storage_uri);
-    // Set AWS_ENDPOINT so that the underlying S3 client uses the specified endpoint.
     unsafe {
         env::set_var("AWS_ENDPOINT", &aws_endpoint);
     }
-    // Parse the AWS endpoint URL and register it.
     let aws_url = Url::parse(&aws_endpoint).expect("AWS endpoint must be a valid URL");
     deltalake::aws::register_handlers(Some(aws_url));
     info!("AWS handlers registered");
@@ -266,7 +261,6 @@ async fn main() -> anyhow::Result<()> {
             return Err(e);
         }
     };
-    // Use the fixed table name "telemetry_events".
     if let Err(e) = db.add_project("telemetry_events", &storage_uri).await {
         error!("Failed to add table 'telemetry_events': {:?}", e);
         return Err(e);
@@ -287,7 +281,7 @@ async fn main() -> anyhow::Result<()> {
     let queue_db_path = env::var("QUEUE_DB_PATH").unwrap_or_else(|_| "/app/queue_db".to_string());
     info!("Using queue DB path: {}", queue_db_path);
 
-    let queue = match PersistentQueue::new(&queue_db_path) {
+    let queue = match PersistentQueue::new(&queue_db_path, db.clone()) {
         Ok(q) => {
             info!("PersistentQueue initialized successfully");
             Arc::new(q)
@@ -387,20 +381,18 @@ async fn main() -> anyhow::Result<()> {
                     }
                     _ = sleep(Duration::from_secs(5)) => {
                         debug!("Checking queue for records to flush");
-                        let records = match queue_clone.dequeue_all().await {
-                            Ok(r) => {
-                                debug!("Dequeued {} records", r.len());
-                                r
-                            },
+                        match queue_clone.dequeue_all().await {
+                            Ok(records) => {
+                                debug!("Dequeued {} records", records.len());
+                                if !records.is_empty() {
+                                    info!("Flushing {} enqueued records", records.len());
+                                    for (key, record) in records {
+                                        process_record(&db_clone, &queue_clone, &status_store_clone, key, record).await;
+                                    }
+                                }
+                            }
                             Err(e) => {
                                 error!("Error during dequeue_all: {:?}", e);
-                                Vec::new()
-                            }
-                        };
-                        if !records.is_empty() {
-                            info!("Flushing {} enqueued records", records.len());
-                            for (key, record) in records {
-                                process_record(&db_clone, &queue_clone, &status_store_clone, key, record).await;
                             }
                         }
                     }
@@ -485,18 +477,34 @@ async fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
-async fn process_record(db: &Arc<Database>, queue: &Arc<PersistentQueue>, status_store: &IngestStatusStore, key: sled::IVec, record: IngestRecord) {
-    use std::str;
-    let id = str::from_utf8(&key).unwrap_or("unknown").to_string();
-    status_store.set_status(id.clone(), "Processing".to_string());
-    if chrono::DateTime::parse_from_rfc3339(&record.timestamp).is_ok() {
+async fn process_record(
+    db: &Arc<Database>,
+    queue: &Arc<PersistentQueue>,
+    status_store: &Arc<IngestStatusStore>,
+    key: String,
+    record: IngestRecord,
+) {
+    status_store.set_status(key.clone(), "Processing".to_string());
+    let timestamp = chrono::DateTime::from_timestamp(
+        record.start_time_unix_nano / 1_000_000_000,
+        (record.start_time_unix_nano % 1_000_000_000) as u32,
+    )
+    .unwrap_or(Utc::now())
+    .to_rfc3339();
+
+    if chrono::DateTime::parse_from_rfc3339(&timestamp).is_ok() {
         match db.write(&record).await {
             Ok(()) => {
                 INGESTION_COUNTER.inc();
-                status_store.set_status(id.clone(), "Ingested".to_string());
+                status_store.set_status(key.clone(), "Ingested".to_string());
                 if let Err(e) = spawn_blocking({
                     let queue = queue.clone();
-                    move || queue.remove_sync(key)
+                    let key = key.clone();
+                    move || {
+                        if let Err(e) = queue.db.remove(key.as_bytes()) {
+                            error!("Failed to remove record from queue: {:?}", e);
+                        }
+                    }
                 })
                 .await
                 {
@@ -506,16 +514,20 @@ async fn process_record(db: &Arc<Database>, queue: &Arc<PersistentQueue>, status
             Err(e) => {
                 ERROR_COUNTER.inc();
                 error!("Error writing record: {:?}", e);
-                status_store.set_status(id, format!("Failed: {:?}", e));
+                status_store.set_status(key, format!("Failed: {:?}", e));
             }
         }
     } else {
         ERROR_COUNTER.inc();
-        error!("Invalid timestamp in record: {}", record.timestamp);
-        status_store.set_status(id, "Invalid timestamp".to_string());
+        error!("Invalid timestamp in record: {}", timestamp);
+        status_store.set_status(key.clone(), "Invalid timestamp".to_string());
         let _ = spawn_blocking({
             let queue = queue.clone();
-            move || queue.remove_sync(key)
+            move || {
+                if let Err(e) = queue.db.remove(key.as_bytes()) {
+                    error!("Failed to remove record from queue: {:?}", e);
+                }
+            }
         })
         .await;
     }
diff --git a/src/persistent_queue.rs b/src/persistent_queue.rs
@@ -1,7 +1,5 @@
-// persistent_queue.rs
-
 use std::path::Path;
-use std::sync::Arc; // Added missing import
+use std::sync::Arc;
 
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
@@ -209,26 +207,29 @@ pub struct IngestRecord {
 }
 
 pub struct PersistentQueue {
-    db: Db,
+    pub db: Db,
     sender: Sender<(String, IngestRecord)>,
-    database: Arc<Database>, // Added to enable writing to the database
+    database: Arc<Database>,
 }
 
 impl PersistentQueue {
-    pub fn new<P: AsRef<Path>>(path: P, database: &Database) -> Result<Self> {
+    pub fn new<P: AsRef<Path>>(path: P, database: Arc<Database>) -> Result<Self> {
         let db = sled::open(path)?;
-        let database = Arc::new(database.clone()); // Clone the database reference
         let (sender, mut receiver) = mpsc::channel::<(String, IngestRecord)>(1000);
 
         let queue = Self {
-            db: db.clone(),
-            sender: sender.clone(),
-            database: database.clone(),
+            db,
+            sender,
+            database,
         };
 
+        // Clone necessary components for the spawned task
+        let db_clone = queue.db.clone();
+        let database_clone = queue.database.clone();
+
         tokio::spawn(async move {
             while let Some((receipt, record)) = receiver.recv().await {
-                if let Err(e) = queue.process_record(receipt.clone(), &record).await {
+                if let Err(e) = Self::process_record_static(&db_clone, &database_clone, receipt.clone(), &record).await {
                     error!("Failed to process record with receipt {}: {:?}", receipt, e);
                 }
             }
@@ -237,6 +238,14 @@ impl PersistentQueue {
         Ok(queue)
     }
 
+    async fn process_record_static(db: &Db, database: &Arc<Database>, receipt: String, record: &IngestRecord) -> Result<()> {
+        tracing::info!("Processing record with receipt: {}", receipt);
+        database.write(record).await?;
+        db.remove(receipt.as_bytes())?;
+        tracing::info!("Record processed and removed from queue: {}", record.trace_id);
+        Ok(())
+    }
+
     pub async fn enqueue(&self, record: &IngestRecord) -> Result<String> {
         let receipt = Uuid::new_v4().to_string();
         let serialized = bincode::serialize(record)?;
@@ -246,13 +255,7 @@ impl PersistentQueue {
     }
 
     pub async fn process_record(&self, receipt: String, record: &IngestRecord) -> Result<()> {
-        tracing::info!("Processing record with receipt: {}", receipt);
-        // Write to the database
-        self.database.write(record).await?;
-        // Remove from the queue after successful write
-        self.db.remove(receipt.as_bytes())?;
-        tracing::info!("Record processed and removed from queue: {}", record.trace_id);
-        Ok(())
+        Self::process_record_static(&self.db, &self.database, receipt, record).await
     }
 
     pub async fn dequeue(&self) -> Result<Option<(String, IngestRecord)>> {
@@ -265,6 +268,16 @@ impl PersistentQueue {
         }
     }
 
+    pub async fn dequeue_all(&self) -> Result<Vec<(String, IngestRecord)>> {
+        let mut records = Vec::new();
+        while let Some((key, value)) = self.db.pop_min()? {
+            let receipt = String::from_utf8(key.to_vec())?;
+            let record: IngestRecord = bincode::deserialize(&value)?;
+            records.push((receipt, record));
+        }
+        Ok(records)
+    }
+
     pub fn len(&self) -> Result<usize> {
         Ok(self.db.len())
     }
