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;
@@ -54,3 +57,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/ingest_server.rs

@@ -122,7 +122,9 @@ impl IngestServer {
         config
             .service(
                 // Base path "{url}/api/v1"
-                web::scope(&base_path()).service(Server::get_ingest_factory()),
+                web::scope(&base_path())
+                .service(Server::get_query_factory())
+                .service(Server::get_ingest_factory()),
             )
             .service(Server::get_liveness_factory())
             .service(Server::get_readiness_factory())

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

@@ -23,6 +23,7 @@ use actix_web::web;
 use actix_web::web::ServiceConfig;
 use actix_web::{App, HttpServer};
 use async_trait::async_trait;
+use chrono::Utc;
 use itertools::Itertools;
 use relative_path::RelativePathBuf;
 use std::sync::Arc;
@@ -49,11 +50,12 @@ impl ParseableServer for QueryServer {
         prometheus: actix_web_prometheus::PrometheusMetrics,
         oidc_client: Option<crate::oidc::OpenidConfig>,
     ) -> anyhow::Result<()> {
+
         let data = Self::get_ingestor_info().await?;
 
         // on subsequent runs, the qurier should check if the ingestor is up and running or not
         for ingester in data.iter() {
-            // dbg!(&ingester);
+            dbg!(&ingester);
             // yes the format macro does not need the '/' ingester.origin already
             // has '/' because Url::Parse will add it if it is not present
             // uri should be something like `http://address/api/v1/liveness`
@@ -148,7 +150,7 @@ impl QueryServer {
     }
 
     // update the .query.json file and return the new IngesterMetadataArr
-    async fn get_ingestor_info() -> anyhow::Result<IngesterMetadataArr> {
+    pub async fn get_ingestor_info() -> anyhow::Result<IngesterMetadataArr> {
         let store = CONFIG.storage().get_object_store();
 
         let root_path = RelativePathBuf::from("");
@@ -165,7 +167,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)
     }
 
@@ -182,8 +184,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
@@ -216,20 +219,19 @@ impl QueryServer {
         }
 
         // spawn the sync thread
-        // tokio::spawn(Self::sync_ingestor_metadata());
+        //tokio::spawn(Self::sync_ingestor_metadata());
 
         self.start(prometheus, CONFIG.parseable.openid.clone())
             .await?;
 
         Ok(())
     }
 
-    #[allow(dead_code)]
     async fn sync_ingestor_metadata() {
         let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(60 / 10));
         loop {
             interval.tick().await;
-            // dbg!("Tick");
+            dbg!("Tick");
             Self::get_ingestor_info().await.unwrap();
         }
     }

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

@@ -69,8 +69,7 @@ async fn main() -> anyhow::Result<()> {
 
     // add logic for graceful shutdown if
     // MODE == Query / Ingest and storage = local-store
-    // option.rs ln: 161
-    // CONFIG.run_time_mode_validation()?;
+    // server.validate()?;
 
     server.init().await?;
 

server/src/option.rs

@@ -157,18 +157,6 @@ impl Config {
         }
         "S3 bucket"
     }
-
-    #[allow(dead_code)]
-    pub fn run_time_mode_validation(&self) -> anyhow::Result<()> {
-        let check = (self.parseable.mode == Mode::Ingest || self.parseable.mode == Mode::Query)
-            && self.storage_name == "drive";
-
-        if check {
-            anyhow::bail!(format!("Cannot start the server in {} mode with local storage, please use S3 bucket for storage", self.parseable.mode.to_str()))
-        }
-
-        Ok(())
-    }
 }
 
 fn create_parseable_cli_command() -> Command {

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()