Query Server Impl WIP

Author: Eshanatnight

server/src/handlers/http.rs

@@ -17,6 +17,9 @@
  */
 
 use actix_cors::Cors;
+use serde_json::Value;
+
+use self::{modal::query_server::QueryServer, query::Query};
 
 pub(crate) mod about;
 pub(crate) mod health_check;
@@ -51,3 +54,26 @@ pub(crate) fn cross_origin_config() -> Cors {
         Cors::default().block_on_origin_mismatch(false)
     }
 }
+
+pub async fn send_query_request_to_ingestor(query: &Query) -> anyhow::Result<Vec<Value>> {
+    // send the query request to the ingestor
+    let mut res = vec![];
+    let ima = QueryServer::get_ingestor_info().await.unwrap();
+
+    for im in ima {
+        let uri = format!("{}{}/{}",im.domain_name, base_path(), "query");
+        let reqw = reqwest::Client::new()
+            .post(uri)
+            .json(query)
+            .basic_auth("admin", Some("admin"))
+            .send()
+            .await?;
+
+        if reqw.status().is_success() {
+            let v: Value = serde_json::from_slice(&reqw.bytes().await?)?;
+            res.push(v);
+        }
+    }
+
+    Ok(res)
+}

server/src/handlers/http/modal/query_server.rs

@@ -158,7 +158,7 @@ impl QueryServer {
 
         let mut f = Self::get_meta_file().await;
         // writer the arr in f
-        let _ = f.write(serde_json::to_string(&arr)?.as_bytes()).await?;
+        let write_size = f.write(serde_json::to_string(&arr)?.as_bytes()).await?;
         Ok(arr)
     }
 
@@ -177,8 +177,9 @@ impl QueryServer {
     /// initialize the server, run migrations as needed and start the server
     async fn initialize(&self) -> anyhow::Result<()> {
         migration::run_metadata_migration(&CONFIG).await?;
-        let metadata = storage::resolve_parseable_metadata().await?;
         tokio::fs::File::create(CONFIG.staging_dir().join(".query.json")).await?;
+
+        let metadata = storage::resolve_parseable_metadata().await?;
         banner::print(&CONFIG, &metadata).await;
 
         // initialize the rbac map

server/src/handlers/http/query.rs

@@ -19,7 +19,7 @@
 use actix_web::http::header::ContentType;
 use actix_web::web::{self, Json};
 use actix_web::{FromRequest, HttpRequest, Responder};
-use chrono::{DateTime, Utc};
+use chrono::{DateTime, Timelike, Utc};
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::SessionState;
 use futures_util::Future;
@@ -36,8 +36,10 @@ use crate::rbac::Users;
 use crate::response::QueryResponse;
 use crate::utils::actix::extract_session_key_from_req;
 
+use super::send_query_request_to_ingestor;
+
 /// Query Request through http endpoint.
-#[derive(Debug, serde::Deserialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct Query {
     query: String,
@@ -52,6 +54,15 @@ pub struct Query {
 }
 
 pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Responder, QueryError> {
+    // create a new query to send to the ingestors
+    let mut mmem= vec![];
+    if let Some(query) = transform_query_for_ingestor(&query_request) {
+        mmem = send_query_request_to_ingestor(&query)
+            .await
+            .map_err(|err| QueryError::Custom(err.to_string()))?;
+    }
+
+    // let rbj = arrow_json::ReaderBuilder::new();
     let creds = extract_session_key_from_req(&req).expect("expects basic auth");
     let permissions = Users.get_permissions(&creds);
     let session_state = QUERY_SESSION.state();
@@ -94,7 +105,7 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
 
     let time = Instant::now();
 
-    let (records, fields) = query.execute().await?;
+    let (records, fields) = query.execute(Some(mmem)).await?;
     let response = QueryResponse {
         records,
         fields,
@@ -183,6 +194,24 @@ async fn into_query(
     })
 }
 
+fn transform_query_for_ingestor(query: &Query) -> Option<Query> {
+    let end_time = DateTime::parse_from_rfc3339(&query.end_time).ok()?;
+    let start_time = end_time - chrono::Duration::minutes(1);
+
+    dbg!(start_time.minute());
+
+    let q = Query {
+        query: query.query.clone(),
+        fields: query.fields,
+        filter_tags: query.filter_tags.clone(),
+        send_null: query.send_null,
+        start_time: start_time.to_rfc3339(),
+        end_time: end_time.to_rfc3339(),
+    };
+
+    Some(q)
+}
+
 #[derive(Debug, thiserror::Error)]
 pub enum QueryError {
     #[error("Query cannot be empty")]
@@ -207,6 +236,8 @@ pub enum QueryError {
     Datafusion(#[from] DataFusionError),
     #[error("Execution Error: {0}")]
     Execute(#[from] ExecuteError),
+    #[error("Query Error: {0}")]
+    Custom(String),
 }
 
 impl actix_web::ResponseError for QueryError {

server/src/main.rs

@@ -68,7 +68,8 @@ async fn main() -> anyhow::Result<()> {
     };
 
     // MODE == Query / Ingest and storage = local-store
-    server.validate()?;
+    // server.validate()?;
+
     server.init().await?;
 
     Ok(())

server/src/query.rs

@@ -33,6 +33,7 @@ use datafusion::logical_expr::{Explain, Filter, LogicalPlan, PlanType, ToStringi
 use datafusion::prelude::*;
 use itertools::Itertools;
 use once_cell::sync::Lazy;
+use serde_json::Value;
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -102,11 +103,18 @@ impl Query {
         SessionContext::new_with_state(state)
     }
 
-    pub async fn execute(&self) -> Result<(Vec<RecordBatch>, Vec<String>), ExecuteError> {
+    pub async fn execute(&self, mem: Option<Vec<Value>>) -> Result<(Vec<RecordBatch>, Vec<String>), ExecuteError> {
         let df = QUERY_SESSION
             .execute_logical_plan(self.final_logical_plan())
             .await?;
 
+        let schema = df.schema();
+        if let Some(mem) = mem {
+            let mem = arrow_json::ReaderBuilder::new(schema.clone())
+            .build(mem.iter().map(|x| serde_json::to_string(x).unwrap()).collect());
+
+        }
+
         let fields = df
             .schema()
             .fields()