merge schema for tables in query

Author: nikhilsinhaparseable

src/alerts/alerts_utils.rs

@@ -34,6 +34,7 @@ use tracing::{trace, warn};
 
 use crate::{
     alerts::LogicalOperator,
+    handlers::http::query::update_schema_when_distributed,
     parseable::PARSEABLE,
     query::{resolve_stream_names, QUERY_SESSION},
     utils::time::TimeRange,
@@ -117,12 +118,18 @@ async fn execute_base_query(
     query: &crate::query::Query,
     original_query: &str,
 ) -> Result<DataFrame, AlertError> {
-    let streams = resolve_stream_names(&original_query)?;
+    let streams = resolve_stream_names(original_query)?;
     let stream_name = streams.first().ok_or_else(|| {
         AlertError::CustomError(format!("Table name not found in query- {original_query}"))
     })?;
-
-    let time_partition = PARSEABLE.get_stream(&stream_name)?.get_time_partition();
+    update_schema_when_distributed(&streams)
+        .await
+        .map_err(|err| {
+            AlertError::CustomError(format!(
+                "Failed to update schema for distributed streams: {err}"
+            ))
+        })?;
+    let time_partition = PARSEABLE.get_stream(stream_name)?.get_time_partition();
     query
         .get_dataframe(time_partition.as_ref())
         .await

src/handlers/airplane.rs

@@ -143,8 +143,7 @@ impl FlightService for AirServiceImpl {
         // create a visitor to extract the table name
 
         let stream_name = streams
-            .iter()
-            .next()
+            .first()
             .ok_or_else(|| Status::aborted("Malformed SQL Provided, Table Name Not Found"))?
             .to_owned();
 
@@ -199,9 +198,11 @@ impl FlightService for AirServiceImpl {
 
         let permissions = Users.get_permissions(&key);
 
-        user_auth_for_datasets(&permissions, &streams).await.map_err(|_| {
-            Status::permission_denied("User Does not have permission to access this")
-        })?;
+        user_auth_for_datasets(&permissions, &streams)
+            .await
+            .map_err(|_| {
+                Status::permission_denied("User Does not have permission to access this")
+            })?;
         let time = Instant::now();
 
         let (records, _) = execute(query, &stream_name, false)

src/handlers/http/query.rs

@@ -46,8 +46,8 @@ use crate::event::commit_schema;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::parseable::{StreamNotFound, PARSEABLE};
 use crate::query::error::ExecuteError;
-use crate::query::{resolve_stream_names, QUERY_SESSION};
 use crate::query::{execute, CountsRequest, Query as LogicalQuery};
+use crate::query::{resolve_stream_names, QUERY_SESSION};
 use crate::rbac::Users;
 use crate::response::QueryResponse;
 use crate::storage::ObjectStorageError;
@@ -95,8 +95,8 @@ pub async fn get_records_and_fields(
         .first()
         .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
     user_auth_for_datasets(&permissions, &tables).await?;
-
-    let (records, fields) = execute(query, &table_name, false).await?;
+    update_schema_when_distributed(&tables).await?;
+    let (records, fields) = execute(query, table_name, false).await?;
 
     let records = match records {
         Either::Left(vec_rb) => vec_rb,
@@ -122,25 +122,25 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
         .first()
         .ok_or_else(|| QueryError::MalformedQuery("No table name found in query"))?;
     user_auth_for_datasets(&permissions, &tables).await?;
-
+    update_schema_when_distributed(&tables).await?;
     let time = Instant::now();
 
     // if the query is `select count(*) from <dataset>`
     // we use the `get_bin_density` method to get the count of records in the dataset
     // instead of executing the query using datafusion
     if let Some(column_name) = query.is_logical_plan_count_without_filters() {
-        return handle_count_query(&query_request, &table_name, column_name, time).await;
+        return handle_count_query(&query_request, table_name, column_name, time).await;
     }
 
     // if the query request has streaming = false (default)
     // we use datafusion's `execute` method to get the records
     if !query_request.streaming {
-        return handle_non_streaming_query(query, &table_name, &query_request, time).await;
+        return handle_non_streaming_query(query, table_name, &query_request, time).await;
     }
 
     // if the query request has streaming = true
     // we use datafusion's `execute_stream` method to get the records
-    handle_streaming_query(query, &table_name, &query_request, time).await
+    handle_streaming_query(query, table_name, &query_request, time).await
 }
 
 /// Handles count queries (e.g., `SELECT COUNT(*) FROM <dataset-name>`)

src/users/filters.rs

@@ -193,7 +193,10 @@ impl Filters {
                 }
             } else if *filter_type == FilterType::Search || *filter_type == FilterType::Filter {
                 let dataset_name = &f.stream_name;
-                if user_auth_for_datasets(&permissions, &[dataset_name.to_string()]).await.is_ok() {
+                if user_auth_for_datasets(&permissions, &[dataset_name.to_string()])
+                    .await
+                    .is_ok()
+                {
                     filters.push(f.clone())
                 }
             }

src/utils/mod.rs

@@ -104,7 +104,7 @@ pub async fn user_auth_for_datasets(
                     break;
                 }
                 Permission::Resource(Action::Query, ParseableResourceType::Stream(stream)) => {
-                    if !PARSEABLE.check_or_load_stream(&stream).await {
+                    if !PARSEABLE.check_or_load_stream(stream).await {
                         return Err(actix_web::error::ErrorUnauthorized(format!(
                             "Stream not found: {stream}"
                         )));