feat: API for date_bin query (#1093)

This PR adds a new API to get `date_bin` from a given stream.
This API can also be used to find the count(*) for the stream.
---------

Co-authored-by: Nikhil Sinha <[email protected]>
Co-authored-by: Devdutt Shenoi <[email protected]>

Author: 3 people

src/handlers/http/modal/query_server.rs

@@ -52,6 +52,7 @@ impl ParseableServer for QueryServer {
         config
             .service(
                 web::scope(&base_path())
+                    .service(Server::get_date_bin())
                     .service(Server::get_correlation_webscope())
                     .service(Server::get_query_factory())
                     .service(Server::get_liveness_factory())

src/handlers/http/modal/server.rs

@@ -80,6 +80,7 @@ impl ParseableServer for Server {
                     .service(Self::get_llm_webscope())
                     .service(Self::get_oauth_webscope(oidc_client))
                     .service(Self::get_user_role_webscope())
+                    .service(Self::get_date_bin())
                     .service(Self::get_metrics_webscope()),
             )
             .service(Self::get_ingest_otel_factory())
@@ -266,6 +267,10 @@ impl Server {
                     ),
             )
     }
+    pub fn get_date_bin() -> Resource {
+        web::resource("/datebin")
+            .route(web::post().to(query::get_date_bin).authorize(Action::Query))
+    }
 
     // get the query factory
     // POST "/query" ==> Get results of the SQL query passed in request body

src/handlers/http/query.rs

@@ -18,13 +18,15 @@
 
 use actix_web::http::header::ContentType;
 use actix_web::web::{self, Json};
-use actix_web::{FromRequest, HttpRequest, Responder};
+use actix_web::{FromRequest, HttpRequest, HttpResponse, Responder};
 use chrono::{DateTime, Utc};
 use datafusion::common::tree_node::TreeNode;
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::SessionState;
 use futures_util::Future;
 use http::StatusCode;
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
 use std::collections::HashMap;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -39,7 +41,7 @@ use crate::event::commit_schema;
 use crate::metrics::QUERY_EXECUTE_TIME;
 use crate::option::{Mode, CONFIG};
 use crate::query::error::ExecuteError;
-use crate::query::Query as LogicalQuery;
+use crate::query::{DateBinRequest, DateBinResponse, Query as LogicalQuery};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::rbac::Users;
 use crate::response::QueryResponse;
@@ -52,7 +54,7 @@ use crate::utils::user_auth_for_query;
 use super::modal::utils::logstream_utils::create_stream_and_schema_from_storage;
 
 /// Query Request through http endpoint.
-#[derive(Debug, serde::Deserialize, serde::Serialize, Clone)]
+#[derive(Debug, Deserialize, Serialize, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct Query {
     pub query: String,
@@ -66,7 +68,7 @@ pub struct Query {
     pub filter_tags: Option<Vec<String>>,
 }
 
-pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Responder, QueryError> {
+pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpResponse, QueryError> {
     let session_state = QUERY_SESSION.state();
     let raw_logical_plan = match session_state
         .create_logical_plan(&query_request.query)
@@ -81,11 +83,10 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
                 .await?
         }
     };
-
     let time_range =
         TimeRange::parse_human_time(&query_request.start_time, &query_request.end_time)?;
 
-    // create a visitor to extract the table names present in query
+    // Create a visitor to extract the table names present in query
     let mut visitor = TableScanVisitor::default();
     let _ = raw_logical_plan.visit(&mut visitor);
 
@@ -103,6 +104,31 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
     user_auth_for_query(&permissions, &tables)?;
 
     let time = Instant::now();
+    if let Some(column_name) = query.is_logical_plan_count_without_filters() {
+        let date_bin_request = DateBinRequest {
+            stream: table_name.clone(),
+            start_time: query_request.start_time.clone(),
+            end_time: query_request.end_time.clone(),
+            num_bins: 1,
+        };
+        let date_bin_records = date_bin_request.get_bin_density().await?;
+        let response = if query_request.fields {
+            json!({
+                "fields": vec![&column_name],
+                "records": vec![json!({column_name: date_bin_records[0].log_count})]
+            })
+        } else {
+            Value::Array(vec![json!({column_name: date_bin_records[0].log_count})])
+        };
+
+        let time = time.elapsed().as_secs_f64();
+
+        QUERY_EXECUTE_TIME
+            .with_label_values(&[&table_name])
+            .observe(time);
+
+        return Ok(HttpResponse::Ok().json(response));
+    }
     let (records, fields) = query.execute(table_name.clone()).await?;
 
     let response = QueryResponse {
@@ -122,6 +148,24 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<impl Respon
     Ok(response)
 }
 
+pub async fn get_date_bin(
+    req: HttpRequest,
+    date_bin: Json<DateBinRequest>,
+) -> Result<impl Responder, QueryError> {
+    let creds = extract_session_key_from_req(&req)?;
+    let permissions = Users.get_permissions(&creds);
+
+    // does user have access to table?
+    user_auth_for_query(&permissions, &[date_bin.stream.clone()])?;
+
+    let date_bin_records = date_bin.get_bin_density().await?;
+
+    Ok(web::Json(DateBinResponse {
+        fields: vec!["date_bin_timestamp".into(), "log_count".into()],
+        records: date_bin_records,
+    }))
+}
+
 pub async fn update_schema_when_distributed(tables: &Vec<String>) -> Result<(), QueryError> {
     if CONFIG.options.mode == Mode::Query {
         for table in tables {