refactor: modified struct

Author: parmesant

src/correlation/correlation_utils.rs

@@ -16,39 +16,24 @@
  *
  */
 
-use datafusion::common::tree_node::TreeNode;
+use itertools::Itertools;
 
-use crate::{
-    query::{TableScanVisitor, QUERY_SESSION},
-    rbac::{
+use crate::rbac::{
         map::SessionKey,
         role::{Action, Permission},
         Users,
-    },
-};
+    };
 
-use super::CorrelationError;
-
-async fn get_tables_from_query(query: &str) -> Result<TableScanVisitor, CorrelationError> {
-    let session_state = QUERY_SESSION.state();
-    let raw_logical_plan = session_state
-        .create_logical_plan(query)
-        .await
-        .map_err(|err| CorrelationError::AnyhowError(err.into()))?;
-
-    let mut visitor = TableScanVisitor::default();
-    let _ = raw_logical_plan.visit(&mut visitor);
-    Ok(visitor)
-}
+use super::{CorrelationError, TableConfig};
 
 pub async fn user_auth_for_query(
     session_key: &SessionKey,
-    query: &str,
+    table_configs: &[TableConfig],
 ) -> Result<(), CorrelationError> {
-    let tables = get_tables_from_query(query).await?;
+    let tables = table_configs.iter().map(|t| &t.table_name).collect_vec();
     let permissions = Users.get_permissions(session_key);
 
-    for table_name in tables.into_inner().iter() {
+    for table_name in tables {
         let mut authorized = false;
 
         // in permission check if user can run query on the stream.

src/correlation/http_handlers.rs

@@ -23,7 +23,7 @@ use relative_path::RelativePathBuf;
 use crate::{
     option::CONFIG,
     storage::{CORRELATION_DIRECTORY, PARSEABLE_ROOT_DIRECTORY},
-    utils::{actix::extract_session_key_from_req, uid::Uid},
+    utils::actix::extract_session_key_from_req,
 };
 
 use super::{
@@ -53,7 +53,7 @@ pub async fn get(req: HttpRequest) -> Result<impl Responder, CorrelationError> {
 
     let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
 
-    if user_auth_for_query(&session_key, &correlation.query)
+    if user_auth_for_query(&session_key, &correlation.table_configs)
         .await
         .is_ok()
     {
@@ -68,10 +68,10 @@ pub async fn post(req: HttpRequest, body: Bytes) -> Result<impl Responder, Corre
         .map_err(|err| CorrelationError::AnyhowError(anyhow::Error::msg(err.to_string())))?;
 
     let correlation_request: CorrelationRequest = serde_json::from_slice(&body)?;
-    let correlation: CorrelationConfig = correlation_request.into();
 
-    // validate user's query auth
-    user_auth_for_query(&session_key, &correlation.query).await?;
+    correlation_request.validate(&session_key).await?;
+
+    let correlation: CorrelationConfig = correlation_request.into();
 
     // Save to disk
     let store = CONFIG.storage().get_object_store();
@@ -82,7 +82,7 @@ pub async fn post(req: HttpRequest, body: Bytes) -> Result<impl Responder, Corre
 
     Ok(format!(
         "Saved correlation with ID- {}",
-        correlation.id.to_string()
+        correlation.id
     ))
 }
 
@@ -95,17 +95,14 @@ pub async fn modify(req: HttpRequest, body: Bytes) -> Result<impl Responder, Cor
         .get("correlation_id")
         .ok_or(CorrelationError::Metadata("No correlation ID Provided"))?;
 
-    let correlation_request: CorrelationRequest = serde_json::from_slice(&body)?;
+    // validate whether user has access to this correlation object or not
+    let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
+    user_auth_for_query(&session_key, &correlation.table_configs).await?;
 
-    // validate user's query auth
-    user_auth_for_query(&session_key, &correlation_request.query).await?;
+    let correlation_request: CorrelationRequest = serde_json::from_slice(&body)?;
+    correlation_request.validate(&session_key).await?;
 
-    let correlation: CorrelationConfig = CorrelationConfig {
-        version: correlation_request.version,
-        id: Uid::from_string(correlation_id)
-            .map_err(|err| CorrelationError::AnyhowError(anyhow::Error::msg(err.to_string())))?,
-        query: correlation_request.query,
-    };
+    let correlation = correlation_request.generate_correlation_config(correlation_id.to_owned());
 
     // Save to disk
     let store = CONFIG.storage().get_object_store();
@@ -116,7 +113,7 @@ pub async fn modify(req: HttpRequest, body: Bytes) -> Result<impl Responder, Cor
 
     Ok(format!(
         "Modified correlation with ID- {}",
-        correlation.id.to_string()
+        correlation.id
     ))
 }
 
@@ -132,14 +129,14 @@ pub async fn delete(req: HttpRequest) -> Result<impl Responder, CorrelationError
     let correlation = CORRELATIONS.get_correlation_by_id(correlation_id).await?;
 
     // validate user's query auth
-    user_auth_for_query(&session_key, &correlation.query).await?;
+    user_auth_for_query(&session_key, &correlation.table_configs).await?;
 
     // Delete from disk
     let store = CONFIG.storage().get_object_store();
     let path = RelativePathBuf::from_iter([
         PARSEABLE_ROOT_DIRECTORY,
         CORRELATION_DIRECTORY,
-        &format!("{}", correlation.id),
+        &correlation.id.to_string(),
     ]);
     store.delete_object(&path).await?;
 

src/correlation/mod.rs

@@ -16,8 +16,12 @@
  *
  */
 
+use std::collections::HashSet;
+
 use actix_web::http::header::ContentType;
+use chrono::Utc;
 use correlation_utils::user_auth_for_query;
+use datafusion::error::DataFusionError;
 use http::StatusCode;
 use itertools::Itertools;
 use once_cell::sync::Lazy;
@@ -27,8 +31,8 @@ use tokio::sync::RwLock;
 use tracing::{trace, warn};
 
 use crate::{
-    handlers::http::rbac::RBACError, option::CONFIG, rbac::map::SessionKey,
-    storage::ObjectStorageError, utils::uid::Uid,
+    handlers::http::rbac::RBACError, option::CONFIG, query::QUERY_SESSION, rbac::map::SessionKey,
+    storage::ObjectStorageError, users::filters::FilterQuery, utils::get_hash,
 };
 
 pub mod correlation_utils;
@@ -69,7 +73,10 @@ impl Correlation {
 
         let mut user_correlations = vec![];
         for c in correlations {
-            if user_auth_for_query(session_key, &c.query).await.is_ok() {
+            if user_auth_for_query(session_key, &c.table_configs)
+                .await
+                .is_ok()
+            {
                 user_correlations.push(c);
             }
         }
@@ -83,7 +90,7 @@ impl Correlation {
         let read = self.0.read().await;
         let correlation = read
             .iter()
-            .find(|c| c.id.to_string() == correlation_id)
+            .find(|c| c.id == correlation_id)
             .cloned();
 
         if let Some(c) = correlation {
@@ -110,7 +117,7 @@ impl Correlation {
         let index = read_access
             .iter()
             .enumerate()
-            .find(|(_, c)| c.id.to_string() == correlation_id)
+            .find(|(_, c)| c.id == correlation_id)
             .to_owned();
 
         if let Some((index, _)) = index {
@@ -126,6 +133,7 @@ impl Correlation {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub enum CorrelationVersion {
     V1,
 }
@@ -134,8 +142,12 @@ pub enum CorrelationVersion {
 #[serde(rename_all = "camelCase")]
 pub struct CorrelationConfig {
     pub version: CorrelationVersion,
-    pub id: Uid,
-    pub query: String,
+    pub id: String,
+    pub table_configs: Vec<TableConfig>,
+    pub join_config: JoinConfig,
+    pub filter: Option<FilterQuery>,
+    pub start_time: Option<String>,
+    pub end_time: Option<String>,
 }
 
 impl CorrelationConfig {}
@@ -144,16 +156,91 @@ impl CorrelationConfig {}
 #[serde(rename_all = "camelCase")]
 pub struct CorrelationRequest {
     pub version: CorrelationVersion,
-    pub query: String,
+    pub table_configs: Vec<TableConfig>,
+    pub join_config: JoinConfig,
+    pub filter: Option<FilterQuery>,
+    pub start_time: Option<String>,
+    pub end_time: Option<String>,
 }
 
 impl From<CorrelationRequest> for CorrelationConfig {
     fn from(val: CorrelationRequest) -> Self {
         Self {
             version: val.version,
-            id: crate::utils::uid::gen(),
-            query: val.query,
+            id: get_hash(Utc::now().timestamp_micros().to_string().as_str()),
+            table_configs: val.table_configs,
+            join_config: val.join_config,
+            filter: val.filter,
+            start_time: val.start_time,
+            end_time: val.end_time,
+        }
+    }
+}
+
+impl CorrelationRequest {
+    pub fn generate_correlation_config(self, id: String) -> CorrelationConfig {
+        CorrelationConfig {
+            version: self.version,
+            id,
+            table_configs: self.table_configs,
+            join_config: self.join_config,
+            filter: self.filter,
+            start_time: self.start_time,
+            end_time: self.end_time,
+        }
+    }
+
+    /// This function will validate the TableConfigs, JoinConfig, and user auth
+    pub async fn validate(&self, session_key: &SessionKey) -> Result<(), CorrelationError> {
+        let ctx = &QUERY_SESSION;
+
+        let mut h1 = HashSet::new();
+
+        h1 = self.table_configs.iter().map(|t| &t.table_name).collect();
+        let h2 = HashSet::from([&self.join_config.table_one, &self.join_config.table_two]);
+
+        // check if table config tables are the same
+        if h1.len() != 2 {
+            return Err(CorrelationError::Metadata(
+                "Must provide config for two unique tables",
+            ));
+        }
+
+        // check that the tables mentioned in join config are
+        // the same as those in table config
+        if h1 != h2 {
+            return Err(CorrelationError::Metadata(
+                "Must provide same tables for join config and table config",
+            ));
+        }
+
+        // check if user has access to table
+        user_auth_for_query(session_key, &self.table_configs).await?;
+
+        // to validate table config, we need to check whether the mentioned fields
+        // are present in the table or not
+        for table_config in self.table_configs.iter() {
+            // table config check
+            let df = ctx.table(&table_config.table_name).await?;
+
+            let mut selected_fields = table_config
+                .selected_fields
+                .iter()
+                .map(|c| c.as_str())
+                .collect_vec();
+            let join_field = if table_config.table_name == self.join_config.table_one {
+                &self.join_config.field_one
+            } else {
+                &self.join_config.field_two
+            };
+
+            selected_fields.push(join_field.as_str());
+
+            // if this errors out then the table config is incorrect or join config is incorrect
+            df.select_columns(selected_fields.as_slice())?;
         }
+
+        Ok(())
     }
 }
 
@@ -171,6 +258,8 @@ pub enum CorrelationError {
     AnyhowError(#[from] anyhow::Error),
     #[error("Unauthorized")]
     Unauthorized,
+    #[error("DataFusion Error: {0}")]
+    DataFusion(#[from] DataFusionError),
 }
 
 impl actix_web::ResponseError for CorrelationError {
@@ -182,6 +271,7 @@ impl actix_web::ResponseError for CorrelationError {
             Self::UserDoesNotExist(_) => StatusCode::NOT_FOUND,
             Self::AnyhowError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             Self::Unauthorized => StatusCode::BAD_REQUEST,
+            Self::DataFusion(_) => StatusCode::INTERNAL_SERVER_ERROR,
         }
     }
 
@@ -191,3 +281,19 @@ impl actix_web::ResponseError for CorrelationError {
             .body(self.to_string())
     }
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct TableConfig {
+    pub selected_fields: Vec<String>,
+    pub table_name: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct JoinConfig {
+    pub table_one: String,
+    pub field_one: String,
+    pub table_two: String,
+    pub field_two: String,
+}

src/handlers/http/modal/server.rs

@@ -194,7 +194,7 @@ impl Server {
                     ),
             )
             .service(
-                web::resource("/correlation/{correlation_id}")
+                web::resource("/{correlation_id}")
                     .route(
                         web::get()
                             .to(correlation::http_handlers::get)

src/storage/object_storage.rs

@@ -640,7 +640,7 @@ pub trait ObjectStorage: Send + Sync + 'static {
         let path = RelativePathBuf::from_iter([
             PARSEABLE_ROOT_DIRECTORY,
             CORRELATION_DIRECTORY,
-            &format!("{}", correlation.id),
+            &format!("{}.json",correlation.id.to_string()),
         ]);
         self.put_object(&path, to_bytes(correlation)).await?;
         Ok(())