updates

- handled counts request

Author: parmesant

src/handlers/http/cluster/mod.rs

@@ -1114,7 +1114,7 @@ struct QuerierStatus {
     last_used: Option<Instant>,
 }
 
-async fn get_available_querier() -> Result<QuerierMetadata, QueryError> {
+pub async fn get_available_querier() -> Result<QuerierMetadata, QueryError> {
     // Get all querier metadata
     let querier_metadata: Vec<NodeMetadata> = get_node_info(NodeType::Querier).await?;
 

src/handlers/http/query.rs

@@ -19,6 +19,7 @@
 use crate::event::error::EventError;
 use crate::handlers::http::fetch_schema;
 use crate::option::Mode;
+use crate::rbac::map::SessionKey;
 use crate::utils::arrow::record_batches_to_json;
 use actix_web::http::header::ContentType;
 use actix_web::web::{self, Json};
@@ -43,7 +44,7 @@ use std::time::Instant;
 use tokio::task::JoinSet;
 use tracing::{error, warn};
 
-use crate::event::commit_schema;
+use crate::event::{DEFAULT_TIMESTAMP_KEY, commit_schema};
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::parseable::{PARSEABLE, StreamNotFound};
 use crate::query::error::ExecuteError;
@@ -79,7 +80,7 @@ pub struct Query {
 /// TODO: Improve this function and make this a part of the query API
 pub async fn get_records_and_fields(
     query_request: &Query,
-    req: &HttpRequest,
+    creds: &SessionKey,
 ) -> Result<(Option<Vec<RecordBatch>>, Option<Vec<String>>), QueryError> {
     let session_state = QUERY_SESSION.state();
     let time_range =
@@ -89,8 +90,8 @@ pub async fn get_records_and_fields(
     create_streams_for_distributed(tables.clone()).await?;
 
     let query: LogicalQuery = into_query(query_request, &session_state, time_range).await?;
-    let creds = extract_session_key_from_req(req)?;
-    let permissions = Users.get_permissions(&creds);
+
+    let permissions = Users.get_permissions(creds);
 
     user_auth_for_datasets(&permissions, &tables).await?;
 
@@ -350,7 +351,12 @@ pub async fn get_counts(
     // if the user has given a sql query (counts call with filters applied), then use this flow
     // this could include filters or group by
     if body.conditions.is_some() {
-        let sql = body.get_df_sql().await?;
+        let time_partition = PARSEABLE
+            .get_stream(&body.stream)?
+            .get_time_partition()
+            .unwrap_or_else(|| DEFAULT_TIMESTAMP_KEY.into());
+
+        let sql = body.get_df_sql(time_partition).await?;
 
         let query_request = Query {
             query: sql,
@@ -362,7 +368,9 @@ pub async fn get_counts(
             filter_tags: None,
         };
 
-        let (records, _) = get_records_and_fields(&query_request, &req).await?;
+        let creds = extract_session_key_from_req(&req)?;
+
+        let (records, _) = get_records_and_fields(&query_request, &creds).await?;
 
         if let Some(records) = records {
             let json_records = record_batches_to_json(&records)?;

src/prism/logstream/mod.rs

@@ -20,28 +20,35 @@ use std::sync::Arc;
 
 use actix_web::http::header::ContentType;
 use arrow_schema::Schema;
-use chrono::Utc;
+use chrono::{TimeDelta, Utc};
 use http::StatusCode;
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
+use serde_json::{Value, json};
 use tracing::warn;
 
 use crate::{
     LOCK_EXPECT,
+    alerts::alert_structs::{ConditionConfig, Conditions},
+    event::DEFAULT_TIMESTAMP_KEY,
     handlers::http::{
         cluster::{
             fetch_stats_from_ingestors,
             utils::{IngestionStats, QueriedStats, StorageStats, merge_queried_stats},
         },
         logstream::error::StreamError,
-        query::{QueryError, update_schema_when_distributed},
+        query::{Query, QueryError, get_records_and_fields, update_schema_when_distributed},
     },
     hottier::{HotTierError, HotTierManager, StreamHotTier},
     parseable::{PARSEABLE, StreamNotFound},
-    query::{CountsRequest, CountsResponse, error::ExecuteError},
+    query::{CountConditions, CountsRequest, CountsResponse, error::ExecuteError},
     rbac::{Users, map::SessionKey, role::Action},
     stats,
     storage::{StreamInfo, StreamType, retention::Retention},
-    utils::time::TimeParseError,
+    utils::{
+        arrow::record_batches_to_json,
+        time::{TimeParseError, truncate_to_minute},
+    },
     validator::error::HotTierValidationError,
 };
 
@@ -218,7 +225,7 @@ pub struct PrismDatasetResponse {
 
 /// Request parameters for retrieving Prism dataset information.
 /// Defines which streams to query
-#[derive(Deserialize, Default)]
+#[derive(Deserialize, Default, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct PrismDatasetRequest {
     /// List of stream names to query
@@ -292,7 +299,7 @@ impl PrismDatasetRequest {
 
         // Process stream data
         match get_prism_logstream_info(&stream).await {
-            Ok(info) => Ok(Some(self.build_dataset_response(stream, info).await?)),
+            Ok(info) => Ok(Some(self.build_dataset_response(stream, info, &key).await?)),
             Err(err) => Err(err),
         }
     }
@@ -312,12 +319,13 @@ impl PrismDatasetRequest {
         &self,
         stream: String,
         info: PrismLogstreamInfo,
+        key: &SessionKey,
     ) -> Result<PrismDatasetResponse, PrismLogstreamError> {
         // Get hot tier info
         let hottier = self.get_hot_tier_info(&stream).await?;
 
         // Get counts
-        let counts = self.get_counts(&stream).await?;
+        let counts = self.get_counts(&stream, key).await?;
 
         Ok(PrismDatasetResponse {
             stream,
@@ -346,20 +354,84 @@ impl PrismDatasetRequest {
         }
     }
 
-    async fn get_counts(&self, stream: &str) -> Result<CountsResponse, PrismLogstreamError> {
+    async fn get_counts(
+        &self,
+        stream: &str,
+        key: &SessionKey,
+    ) -> Result<CountsResponse, PrismLogstreamError> {
+        let end = truncate_to_minute(Utc::now());
+        let start = end - TimeDelta::hours(1);
+
+        let conditions = if PARSEABLE.get_stream(stream)?.get_time_partition().is_some() {
+            Some(CountConditions {
+                conditions: Some(Conditions {
+                    operator: Some(crate::alerts::LogicalOperator::And),
+                    condition_config: vec![
+                        ConditionConfig {
+                            column: DEFAULT_TIMESTAMP_KEY.into(),
+                            operator: crate::alerts::WhereConfigOperator::GreaterThanOrEqual,
+                            value: Some(start.to_rfc3339()),
+                        },
+                        ConditionConfig {
+                            column: DEFAULT_TIMESTAMP_KEY.into(),
+                            operator: crate::alerts::WhereConfigOperator::LessThan,
+                            value: Some(end.to_rfc3339()),
+                        },
+                    ],
+                }),
+                group_by: None,
+            })
+        } else {
+            None
+        };
+
         let count_request = CountsRequest {
             stream: stream.to_owned(),
-            start_time: "1h".to_owned(),
-            end_time: "now".to_owned(),
+            start_time: start.to_rfc3339(),
+            end_time: end.to_rfc3339(),
             num_bins: 10,
-            conditions: None,
+            conditions,
         };
 
-        let records = count_request.get_bin_density().await?;
-        Ok(CountsResponse {
-            fields: vec!["start_time".into(), "end_time".into(), "count".into()],
-            records,
-        })
+        if count_request.conditions.is_some() {
+            // forward request to querier
+            let query = count_request
+                .get_df_sql(DEFAULT_TIMESTAMP_KEY.into())
+                .await?;
+
+            let query_request = Query {
+                query,
+                start_time: start.to_rfc3339(),
+                end_time: end.to_rfc3339(),
+                send_null: true,
+                fields: true,
+                streaming: false,
+                filter_tags: None,
+            };
+
+            let (records, _) = get_records_and_fields(&query_request, key).await?;
+            if let Some(records) = records {
+                let json_records = record_batches_to_json(&records)?;
+                let records = json_records.into_iter().map(Value::Object).collect_vec();
+
+                let res = json!({
+                    "fields": vec!["start_time", "end_time", "count"],
+                    "records": records,
+                });
+
+                Ok(serde_json::from_value(res)?)
+            } else {
+                Err(PrismLogstreamError::Anyhow(anyhow::Error::msg(
+                    "No data returned for counts SQL",
+                )))
+            }
+        } else {
+            let records = count_request.get_bin_density().await?;
+            Ok(CountsResponse {
+                fields: vec!["start_time".into(), "end_time".into(), "count".into()],
+                records,
+            })
+        }
     }
 }
 
@@ -381,6 +453,10 @@ pub enum PrismLogstreamError {
     Execute(#[from] ExecuteError),
     #[error("Auth: {0}")]
     Auth(#[from] actix_web::Error),
+    #[error("SerdeError: {0}")]
+    SerdeError(#[from] serde_json::Error),
+    #[error("ReqwestError: {0}")]
+    ReqwestError(#[from] reqwest::Error),
 }
 
 impl actix_web::ResponseError for PrismLogstreamError {
@@ -393,6 +469,8 @@ impl actix_web::ResponseError for PrismLogstreamError {
             PrismLogstreamError::Query(_) => StatusCode::INTERNAL_SERVER_ERROR,
             PrismLogstreamError::TimeParse(_) => StatusCode::NOT_FOUND,
             PrismLogstreamError::Execute(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            PrismLogstreamError::SerdeError(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            PrismLogstreamError::ReqwestError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             PrismLogstreamError::Auth(_) => StatusCode::UNAUTHORIZED,
         }
     }

src/query/mod.rs

@@ -441,37 +441,31 @@ impl CountsRequest {
     }
 
     /// This function will get executed only if self.conditions is some
-    pub async fn get_df_sql(&self) -> Result<String, QueryError> {
+    pub async fn get_df_sql(&self, time_column: String) -> Result<String, QueryError> {
         // unwrap because we have asserted that it is some
         let count_conditions = self.conditions.as_ref().unwrap();
 
-        // get time partition column
-        let time_partition = PARSEABLE
-            .get_stream(&self.stream)?
-            .get_time_partition()
-            .unwrap_or_else(|| DEFAULT_TIMESTAMP_KEY.into());
-
         let time_range = TimeRange::parse_human_time(&self.start_time, &self.end_time)?;
 
         let dur = time_range.end.signed_duration_since(time_range.start);
 
         let date_bin = if dur.num_minutes() <= 60 * 10 {
             // date_bin 1 minute
             format!(
-                "CAST(DATE_BIN('1 minute', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 minute', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 minute' as end_time",
-                self.stream, self.stream
+                "CAST(DATE_BIN('1 minute', \"{}\".\"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 minute', \"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 minute' as end_time",
+                self.stream
             )
         } else if dur.num_minutes() > 60 * 10 && dur.num_minutes() < 60 * 240 {
             // date_bin 1 hour
             format!(
-                "CAST(DATE_BIN('1 hour', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 hour', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 hour' as end_time",
-                self.stream, self.stream
+                "CAST(DATE_BIN('1 hour', \"{}\".\"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 hour', \"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 hour' as end_time",
+                self.stream
             )
         } else {
             // date_bin 1 day
             format!(
-                "CAST(DATE_BIN('1 day', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 day', \"{}\".\"{time_partition}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 day' as end_time",
-                self.stream, self.stream
+                "CAST(DATE_BIN('1 day', \"{}\".\"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') AS TEXT) as start_time, DATE_BIN('1 day', \"{time_column}\", TIMESTAMP '1970-01-01 00:00:00+00') + INTERVAL '1 day' as end_time",
+                self.stream
             )
         };
 
@@ -492,7 +486,7 @@ impl CountsRequest {
 }
 
 /// Response for the counts API
-#[derive(Debug, Serialize, Clone)]
+#[derive(Debug, Serialize, Clone, Deserialize)]
 pub struct CountsResponse {
     /// Fields in the log stream
     pub fields: Vec<String>,