Updates for counts API

src/alerts/alerts_utils.rs

@@ -339,6 +339,33 @@ fn get_filter_expr(where_clause: &Conditions) -> Expr {
     }
 }
 
+pub fn get_filter_string(where_clause: &Conditions) -> Result<String, String> {
+    match &where_clause.operator {
+        Some(op) => match op {
+            LogicalOperator::And => {
+                let mut exprs = vec![];
+                for condition in &where_clause.condition_config {
+                    let value = NumberOrString::from_string(condition.value.to_string());
+                    match value {
+                        NumberOrString::Number(val) => exprs.push(format!(
+                            "{} {} {}",
+                            condition.column, condition.operator, val
+                        )),
+                        NumberOrString::String(val) => exprs.push(format!(
+                            "{} {} '{}'",
+                            condition.column, condition.operator, val
+                        )),
+                    }
+                }
+
+                Ok(exprs.join(" AND "))
+            }
+            _ => Err(String::from("Invalid option 'or'")),
+        },
+        _ => Err(String::from("Invalid option 'null'")),
+    }
+}
+
 fn match_alert_operator(expr: &ConditionConfig) -> Expr {
     // the form accepts value as a string
     // if it can be parsed as a number, then parse it

src/handlers/http/query.rs

@@ -18,6 +18,7 @@
 
 use crate::event::error::EventError;
 use crate::handlers::http::fetch_schema;
+use crate::utils::arrow::record_batches_to_json;
 use actix_web::http::header::ContentType;
 use actix_web::web::{self, Json};
 use actix_web::{Either, FromRequest, HttpRequest, HttpResponse, Responder};
@@ -31,8 +32,9 @@ use futures::stream::once;
 use futures::{future, Stream, StreamExt};
 use futures_util::Future;
 use http::StatusCode;
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
-use serde_json::json;
+use serde_json::{json, Value};
 use std::collections::HashMap;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -44,7 +46,7 @@ use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::Mode;
 use crate::parseable::{StreamNotFound, PARSEABLE};
 use crate::query::error::ExecuteError;
-use crate::query::{execute, CountsRequest, CountsResponse, Query as LogicalQuery};
+use crate::query::{execute, CountsRequest, Query as LogicalQuery};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::rbac::Users;
 use crate::response::QueryResponse;
@@ -193,6 +195,7 @@ async fn handle_count_query(
         start_time: query_request.start_time.clone(),
         end_time: query_request.end_time.clone(),
         num_bins: 1,
+        conditions: None,
     };
     let count_records = counts_req.get_bin_density().await?;
     let count = count_records[0].count;
@@ -358,22 +361,59 @@ fn create_batch_processor(
         Err(e) => Err(actix_web::error::ErrorInternalServerError(e)),
     }
 }
+
 pub async fn get_counts(
     req: HttpRequest,
     counts_request: Json<CountsRequest>,
 ) -> Result<impl Responder, QueryError> {
     let creds = extract_session_key_from_req(&req)?;
     let permissions = Users.get_permissions(&creds);
 
+    let body = counts_request.into_inner();
+
     // does user have access to table?
-    user_auth_for_datasets(&permissions, &[counts_request.stream.clone()])?;
+    user_auth_for_datasets(&permissions, &[body.stream.clone()])?;
+
+    // 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 query_request = Query {
+            query: sql,
+            start_time: body.start_time,
+            end_time: body.end_time,
+            send_null: true,
+            fields: true,
+            streaming: false,
+            filter_tags: None,
+        };
 
-    let records = counts_request.get_bin_density().await?;
+        let (records, _) = get_records_and_fields(&query_request, &req).await?;
 
-    Ok(web::Json(CountsResponse {
-        fields: vec!["start_time".into(), "end_time".into(), "count".into()],
-        records,
-    }))
+        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", "endTime", "count"],
+                "records": records,
+            });
+
+            return Ok(web::Json(res));
+        } else {
+            return Err(QueryError::CustomError(
+                "No data returned for counts SQL".into(),
+            ));
+        }
+    }
+
+    let records = body.get_bin_density().await?;
+    let res = json!({
+        "fields": vec!["start_time", "endTime", "count"],
+        "records": records,
+    });
+    Ok(web::Json(res))
 }
 
 pub async fn update_schema_when_distributed(tables: &Vec<String>) -> Result<(), EventError> {

src/prism/logstream/mod.rs

@@ -349,6 +349,7 @@ impl PrismDatasetRequest {
             start_time: "1h".to_owned(),
             end_time: "now".to_owned(),
             num_bins: 10,
+            conditions: None,
         };
 
         let records = count_request.get_bin_density().await?;

src/query/mod.rs

@@ -47,6 +47,8 @@ use tokio::runtime::Runtime;
 use self::error::ExecuteError;
 use self::stream_schema_provider::GlobalSchemaProvider;
 pub use self::stream_schema_provider::PartialTimeFilter;
+use crate::alerts::alerts_utils::get_filter_string;
+use crate::alerts::Conditions;
 use crate::catalog::column::{Int64Type, TypedStatistics};
 use crate::catalog::manifest::Manifest;
 use crate::catalog::snapshot::Snapshot;
@@ -297,7 +299,7 @@ impl Query {
 }
 
 /// Record of counts for a given time bin.
-#[derive(Debug, Serialize, Clone)]
+#[derive(Debug, Serialize, Clone, Deserialize)]
 pub struct CountsRecord {
     /// Start time of the bin
     pub start_time: String,
@@ -312,6 +314,15 @@ struct TimeBounds {
     end: DateTime<Utc>,
 }
 
+#[derive(Debug, Deserialize, Clone)]
+#[serde(rename_all = "camelCase")]
+pub struct CountConditions {
+    /// Optional conditions for filters
+    pub conditions: Option<Conditions>,
+    /// GroupBy columns
+    pub group_by: Option<Vec<String>>,
+}
+
 /// Request for counts, received from API/SQL query.
 #[derive(Debug, Deserialize, Clone)]
 #[serde(rename_all = "camelCase")]
@@ -324,6 +335,8 @@ pub struct CountsRequest {
     pub end_time: String,
     /// Number of bins to divide the time range into
     pub num_bins: u64,
+    /// Conditions
+    pub conditions: Option<CountConditions>,
 }
 
 impl CountsRequest {
@@ -429,6 +442,34 @@ impl CountsRequest {
 
         bounds
     }
+
+    /// This function will get executed only if self.conditions is some
+    pub async fn get_df_sql(&self) -> Result<String, QueryError> {
+        let count_conditions = self.conditions.as_ref().unwrap();
+
+        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', {}.p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') AS TEXT) as start_time, DATE_BIN('1 minute', p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') + INTERVAL '1 minute' as end_time", self.stream)
+        } else if dur.num_minutes() > 60 * 10 && dur.num_minutes() < 60 * 240 {
+            // date_bin 1 hour
+            String::from("CAST(DATE_BIN('1 hour', p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') AS TEXT) as start_time, DATE_BIN('1 hour', p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') + INTERVAL '1 hour' as end_time")
+        } else {
+            // date_bin 1 day
+            String::from("CAST(DATE_BIN('1 day', p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') AS TEXT) as start_time, DATE_BIN('1 day', p_timestamp, TIMESTAMP '2025-01-01T00:00:00.000Z') + INTERVAL '1 day' as end_time")
+        };
+
+        let query = if let Some(conditions) = &count_conditions.conditions {
+            let f = get_filter_string(conditions).map_err(QueryError::CustomError)?;
+            format!("SELECT {date_bin}, COUNT(*) as count FROM {} WHERE {} GROUP BY end_time,start_time ORDER BY end_time",self.stream,f)
+        } else {
+            format!("SELECT {date_bin}, COUNT(*) as count FROM {} GROUP BY end_time,start_time ORDER BY end_time",self.stream)
+        };
+        Ok(query)
+    }
 }
 
 /// Response for the counts API