refactor: query request flow

src/handlers/airplane.rs

@@ -21,7 +21,6 @@ use arrow_flight::flight_service_server::FlightServiceServer;
 use arrow_flight::PollInfo;
 use arrow_schema::ArrowError;
 
-use datafusion::common::tree_node::TreeNode;
 use serde_json::json;
 use std::net::SocketAddr;
 use std::time::Instant;
@@ -34,17 +33,13 @@ use tonic::transport::{Identity, Server, ServerTlsConfig};
 use tonic_web::GrpcWebLayer;
 
 use crate::handlers::http::cluster::get_ingestor_info;
-use crate::handlers::http::query::{into_query, update_schema_when_distributed};
 use crate::handlers::livetail::cross_origin_config;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::CONFIG;
-use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::utils::arrow::flight::{
     append_temporary_events, get_query_from_ticket, into_flight_data, run_do_get_rpc,
     send_to_ingester,
 };
-use crate::utils::time::TimeRange;
-use crate::utils::user_auth_for_query;
 use arrow_flight::{
     flight_service_server::FlightService, Action, ActionType, Criteria, Empty, FlightData,
     FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaAsIpc,
@@ -127,55 +122,41 @@ impl FlightService for AirServiceImpl {
     async fn do_get(&self, req: Request<Ticket>) -> Result<Response<Self::DoGetStream>, Status> {
         let key = extract_session_key(req.metadata())?;
 
-        let ticket = get_query_from_ticket(&req)?;
+        // try authorize
+        match Users.authorize(key.clone(), rbac::role::Action::Query, None, None) {
+            rbac::Response::Authorized => (),
+            rbac::Response::UnAuthorized => {
+                return Err(Status::permission_denied(
+                    "user is not authorized to access this resource",
+                ))
+            }
+            rbac::Response::ReloadRequired => {
+                return Err(Status::unauthenticated("reload required"))
+            }
+        }
 
-        info!("query requested to airplane: {:?}", ticket);
+        let query_request = get_query_from_ticket(&req)?;
+        info!("query requested to airplane: {:?}", query_request);
 
-        // get the query session_state
-        let session_state = QUERY_SESSION.state();
+        // map payload to query
+        let query = query_request.into_query(key).await.map_err(|e| {
+            error!("Error faced while constructing logical query: {e}");
+            Status::internal(format!("Failed to process query: {e}"))
+        })?;
 
-        // get the logical plan and extract the table name
-        let raw_logical_plan = session_state
-            .create_logical_plan(&ticket.query)
-            .await
-            .map_err(|err| {
-                error!("Datafusion Error: Failed to create logical plan: {}", err);
-                Status::internal("Failed to create logical plan")
-            })?;
-
-        let time_range = TimeRange::parse_human_time(&ticket.start_time, &ticket.end_time)
-            .map_err(|e| Status::internal(e.to_string()))?;
-        // create a visitor to extract the table name
-        let mut visitor = TableScanVisitor::default();
-        let _ = raw_logical_plan.visit(&mut visitor);
-
-        let streams = visitor.into_inner();
-
-        let stream_name = streams
-            .first()
-            .ok_or_else(|| Status::aborted("Malformed SQL Provided, Table Name Not Found"))?
+        let stream_name = query
+            .first_stream_name()
+            .ok_or_else(|| Status::internal("Failed to get stream name from query"))?
             .to_owned();
 
-        update_schema_when_distributed(&streams)
-            .await
-            .map_err(|err| Status::internal(err.to_string()))?;
-
-        // map payload to query
-        let query = into_query(&ticket, &session_state, time_range)
-            .await
-            .map_err(|_| Status::internal("Failed to parse query"))?;
-
         let event = if send_to_ingester(
             query.time_range.start.timestamp_millis(),
             query.time_range.end.timestamp_millis(),
         ) {
-            let sql = format!("select * from {}", &stream_name);
-            let start_time = ticket.start_time.clone();
-            let end_time = ticket.end_time.clone();
             let out_ticket = json!({
-                "query": sql,
-                "startTime": start_time,
-                "endTime": end_time
+                "query": format!("select * from {stream_name}"),
+                "startTime": query_request.start_time,
+                "endTime":  query_request.end_time,
             })
             .to_string();
 
@@ -190,33 +171,15 @@ impl FlightService for AirServiceImpl {
                 }
             }
             let mr = minute_result.iter().collect::<Vec<_>>();
-            let event = append_temporary_events(&stream_name, mr).await?;
+            let event = append_temporary_events(stream_name.as_str(), mr).await?;
             Some(event)
         } else {
             None
         };
 
-        // try authorize
-        match Users.authorize(key.clone(), rbac::role::Action::Query, None, None) {
-            rbac::Response::Authorized => (),
-            rbac::Response::UnAuthorized => {
-                return Err(Status::permission_denied(
-                    "user is not authorized to access this resource",
-                ))
-            }
-            rbac::Response::ReloadRequired => {
-                return Err(Status::unauthenticated("reload required"))
-            }
-        }
-
-        let permissions = Users.get_permissions(&key);
-
-        user_auth_for_query(&permissions, &streams).map_err(|_| {
-            Status::permission_denied("User Does not have permission to access this")
-        })?;
         let time = Instant::now();
         let (records, _) = query
-            .execute(stream_name.clone())
+            .execute()
             .await
             .map_err(|err| Status::internal(err.to_string()))?;
 
@@ -231,18 +194,17 @@ impl FlightService for AirServiceImpl {
             .collect::<Vec<_>>();
         let schema = Schema::try_merge(schemas).map_err(|err| Status::internal(err.to_string()))?;
          */
-        let out = into_flight_data(records);
 
         if let Some(event) = event {
             event.clear(&stream_name);
         }
 
         let time = time.elapsed().as_secs_f64();
         QUERY_EXECUTE_TIME
-            .with_label_values(&[&format!("flight-query-{}", stream_name)])
+            .with_label_values(&[&stream_name])
             .observe(time);
 
-        out
+        into_flight_data(records)
     }
 
     async fn do_put(

src/handlers/http/mod.rs

@@ -25,7 +25,7 @@ use serde_json::Value;
 
 use crate::{option::CONFIG, storage::STREAM_ROOT_DIRECTORY, HTTP_CLIENT};
 
-use self::{cluster::get_ingestor_info, query::Query};
+use self::{cluster::get_ingestor_info, query::QueryRequest};
 
 pub mod about;
 mod audit;
@@ -97,7 +97,7 @@ pub async fn fetch_schema(stream_name: &str) -> anyhow::Result<arrow_schema::Sch
 
 /// unused for now, might need it later
 #[allow(unused)]
-pub async fn send_query_request_to_ingestor(query: &Query) -> anyhow::Result<Vec<Value>> {
+pub async fn send_query_request_to_ingestor(query: &QueryRequest) -> anyhow::Result<Vec<Value>> {
     // send the query request to the ingestor
     let mut res = vec![];
     let ima = get_ingestor_info().await?;

src/handlers/http/modal/utils/logstream_utils.rs

@@ -32,7 +32,7 @@ use crate::{
     metadata::{self, SchemaVersion, STREAM_INFO},
     option::{Mode, CONFIG},
     static_schema::{convert_static_schema_to_arrow_schema, StaticSchema},
-    storage::{LogStream, ObjectStoreFormat, StreamType},
+    storage::{ObjectStoreFormat, StreamType},
     validator,
 };
 
@@ -443,9 +443,7 @@ pub async fn create_stream_and_schema_from_storage(stream_name: &str) -> Result<
     // Proceed to create log stream if it doesn't exist
     let storage = CONFIG.storage().get_object_store();
     let streams = storage.list_streams().await?;
-    if streams.contains(&LogStream {
-        name: stream_name.to_owned(),
-    }) {
+    if streams.iter().any(|stream| stream.name == stream_name) {
         let mut stream_metadata = ObjectStoreFormat::default();
         let stream_metadata_bytes = storage.create_stream_from_ingestor(stream_name).await?;
         if !stream_metadata_bytes.is_empty() {