feature: get top datasets from CH

server/src/data/models.rs

@@ -4200,6 +4200,42 @@ impl SearchAnalyticsFilter {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize, ToSchema)]
+pub struct TopDatasetsResponse {
+    pub dataset_id: uuid::Uuid,
+    pub dataset_tracking_id: Option<String>,
+    pub total_queries: i64,
+}
+
+#[derive(Debug, Serialize, Deserialize, ToSchema, Row, Clone)]
+pub struct TopDatasetsResponseClickhouse {
+    #[serde(with = "clickhouse::serde::uuid")]
+    pub dataset_id: uuid::Uuid,
+    pub total_queries: i64,
+}
+
+impl From<TopDatasetsResponseClickhouse> for TopDatasetsResponse {
+    fn from(clickhouse_response: TopDatasetsResponseClickhouse) -> TopDatasetsResponse {
+        TopDatasetsResponse {
+            dataset_id: uuid::Uuid::from_bytes(*clickhouse_response.dataset_id.as_bytes()),
+            dataset_tracking_id: None,
+            total_queries: clickhouse_response.total_queries,
+        }
+    }
+}
+
+#[derive(Deserialize, Serialize, Clone, Debug, ToSchema, Display)]
+#[serde(rename_all = "snake_case")]
+pub enum TopDatasetsRequestTypes {
+    #[display(fmt = "search_queries")]
+    Search,
+    #[serde(rename = "rag")]
+    #[display(fmt = "rag_queries")]
+    RAG,
+    #[display(fmt = "recommendations")]
+    Recommendation,
+}
+
 #[derive(Debug, Serialize, Deserialize, ToSchema, Display, Clone)]
 #[serde(rename_all = "snake_case")]
 pub enum RagTypes {

server/src/handlers/analytics_handler.rs

@@ -2,9 +2,9 @@ use super::auth_handler::AdminOnly;
 use crate::{
     data::models::{
         CTRAnalytics, CTRAnalyticsResponse, CTRType, ClusterAnalytics, ClusterAnalyticsResponse,
-        DatasetAndOrgWithSubAndPlan, Pool, RAGAnalytics, RAGAnalyticsResponse,
+        DatasetAndOrgWithSubAndPlan, DateRange, Pool, RAGAnalytics, RAGAnalyticsResponse,
         RecommendationAnalytics, RecommendationAnalyticsResponse, SearchAnalytics,
-        SearchAnalyticsResponse,
+        SearchAnalyticsResponse, TopDatasetsRequestTypes,
     },
     errors::ServiceError,
     operators::analytics_operator::*,
@@ -563,3 +563,39 @@ pub async fn get_ctr_analytics(
 
     Ok(HttpResponse::Ok().json(response))
 }
+
+#[derive(Deserialize, Serialize, Clone, Debug, ToSchema)]
+pub struct GetTopDatasetsRequestBody {
+    pub r#type: TopDatasetsRequestTypes,
+    pub date_range: Option<DateRange>,
+}
+
+/// Get Top Datasets
+///
+/// This route allows you to view the top datasets for a given type.
+#[utoipa::path(
+    post,
+    path = "/analytics/top",
+    context_path = "/api",
+    tag = "Analytics",
+    request_body(content = GetTopDatasetsRequestBody, description = "JSON request payload to filter the top datasets", content_type = "application/json"),
+    responses(
+        (status = 200, description = "The top datasets for the request", body = Vec<TopDatasetsResponse>),
+
+        (status = 400, description = "Service error relating to getting top datasets", body = ErrorResponseBody),
+    ),
+    security(
+        ("ApiKey" = ["admin"]),
+    )
+)]
+pub async fn get_top_datasets(
+    _user: AdminOnly,
+    data: web::Json<GetTopDatasetsRequestBody>,
+    clickhouse_client: web::Data<clickhouse::Client>,
+    pool: web::Data<Pool>,
+) -> Result<HttpResponse, ServiceError> {
+    let top_datasets =
+        get_top_datasets_query(data.into_inner(), clickhouse_client.get_ref(), pool).await?;
+
+    Ok(HttpResponse::Ok().json(top_datasets))
+}

server/src/lib.rs

@@ -235,6 +235,7 @@ impl Modify for SecurityAddon {
         handlers::analytics_handler::send_ctr_data,
         handlers::analytics_handler::get_ctr_analytics,
         handlers::analytics_handler::set_query_rating,
+        handlers::analytics_handler::get_top_datasets,
         handlers::metrics_handler::get_metrics,
     ),
     components(
@@ -327,6 +328,7 @@ impl Modify for SecurityAddon {
             handlers::dataset_handler::UpdateDatasetRequest,
             handlers::dataset_handler::GetDatasetsPagination,
             data::models::DatasetConfigurationDTO,
+            handlers::analytics_handler::GetTopDatasetsRequestBody,
             operators::analytics_operator::HeadQueryResponse,
             operators::analytics_operator::LatencyGraphResponse,
             operators::analytics_operator::SearchUsageGraphResponse,
@@ -363,6 +365,8 @@ impl Modify for SecurityAddon {
             data::models::SearchResultType,
             data::models::RoleProxy,
             data::models::RAGUsageResponse,
+            data::models::TopDatasetsRequestTypes,
+            data::models::TopDatasetsResponse,
             data::models::RAGUsageGraphResponse,
             data::models::ClusterAnalytics,
             data::models::RAGAnalytics,
@@ -1121,6 +1125,10 @@ pub fn main() -> std::io::Result<()> {
                                 web::resource("/recommendations")
                                 .route(web::post().to(handlers::analytics_handler::get_recommendation_analytics)),
                             )
+                            .service(
+                                web::resource("/top")
+                                .route(web::post().to(handlers::analytics_handler::get_top_datasets)),
+                            )
                             .service(
                                 web::resource("/ctr")
                                 .route(web::put().to(handlers::analytics_handler::send_ctr_data))

server/src/operators/analytics_operator.rs

@@ -1,10 +1,4 @@
-use actix_web::web;
-use futures::future::join_all;
-use itertools::Itertools;
-use serde::{Deserialize, Serialize};
-use ureq::json;
-use utoipa::ToSchema;
-
+use super::chunk_operator::get_metadata_from_tracking_id_query;
 use crate::{
     data::models::{
         ClusterAnalyticsFilter, ClusterTopicsClickhouse, DatasetAnalytics, Granularity,
@@ -18,14 +12,22 @@ use crate::{
         SearchLatencyGraphClickhouse, SearchQueriesWithClicksCTRResponse,
         SearchQueriesWithClicksCTRResponseClickhouse, SearchQueriesWithoutClicksCTRResponse,
         SearchQueriesWithoutClicksCTRResponseClickhouse, SearchQueryEvent,
-        SearchQueryEventClickhouse, SearchSortBy, SearchTypeCount, SortOrder, UsageGraphPoint,
-        UsageGraphPointClickhouse,
+        SearchQueryEventClickhouse, SearchSortBy, SearchTypeCount, SortOrder, TopDatasetsResponse,
+        TopDatasetsResponseClickhouse, UsageGraphPoint, UsageGraphPointClickhouse,
     },
     errors::ServiceError,
-    handlers::analytics_handler::{CTRDataRequestBody, RateQueryRequest},
+    handlers::analytics_handler::{
+        CTRDataRequestBody, GetTopDatasetsRequestBody, RateQueryRequest,
+    },
 };
-
-use super::chunk_operator::get_metadata_from_tracking_id_query;
+use actix_web::web;
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
+use futures::future::join_all;
+use itertools::Itertools;
+use serde::{Deserialize, Serialize};
+use ureq::json;
+use utoipa::ToSchema;
 
 #[derive(Debug, Serialize, Deserialize, ToSchema)]
 pub struct SearchClusterResponse {
@@ -1353,3 +1355,85 @@ pub async fn set_query_rating_query(
 
     Ok(())
 }
+
+pub async fn get_top_datasets_query(
+    data: GetTopDatasetsRequestBody,
+    clickhouse_client: &clickhouse::Client,
+    pool: web::Data<Pool>,
+) -> Result<Vec<TopDatasetsResponse>, ServiceError> {
+    use crate::data::schema::datasets::dsl as datasets_columns;
+
+    let mut query_string = format!(
+        "SELECT 
+            dataset_id,
+            COUNT(*) as total_queries
+        FROM 
+            default.{}",
+        data.r#type
+    );
+
+    if let Some(date_range) = data.date_range {
+        if let Some(gt) = &date_range.gt {
+            query_string.push_str(&format!(" AND created_at > '{}'", gt));
+        }
+        if let Some(lt) = &date_range.lt {
+            query_string.push_str(&format!(" AND created_at < '{}'", lt));
+        }
+        if let Some(gte) = &date_range.gte {
+            query_string.push_str(&format!(" AND created_at >= '{}'", gte));
+        }
+        if let Some(lte) = &date_range.lte {
+            query_string.push_str(&format!(" AND created_at <= '{}'", lte));
+        }
+    }
+
+    query_string.push_str(
+        "
+        GROUP BY 
+            dataset_id
+        ORDER BY 
+            total_queries DESC
+        LIMIT 10",
+    );
+
+    let clickhouse_resp_data = clickhouse_client
+        .query(query_string.as_str())
+        .fetch_all::<TopDatasetsResponseClickhouse>()
+        .await
+        .map_err(|e| {
+            log::error!("Error fetching query: {:?}", e);
+            ServiceError::InternalServerError("Error fetching query".to_string())
+        })?;
+
+    let dataset_ids = clickhouse_resp_data
+        .iter()
+        .map(|x| x.dataset_id)
+        .collect::<Vec<_>>();
+    let mut conn = pool
+        .get()
+        .await
+        .map_err(|_| ServiceError::BadRequest("Could not get database connection".to_string()))?;
+    let dataset_id_and_tracking_ids = datasets_columns::datasets
+        .select((datasets_columns::id, datasets_columns::tracking_id))
+        .filter(datasets_columns::id.eq_any(dataset_ids))
+        .load::<(uuid::Uuid, Option<String>)>(&mut conn)
+        .await
+        .map_err(|e| {
+            log::error!("Error fetching dataset ids: {:?}", e);
+            ServiceError::InternalServerError("Error fetching dataset ids".to_string())
+        })?;
+
+    let response = clickhouse_resp_data
+        .into_iter()
+        .map(|x| {
+            let mut top_dataset_resps = TopDatasetsResponse::from(x.clone());
+            top_dataset_resps.dataset_tracking_id = dataset_id_and_tracking_ids
+                .iter()
+                .find(|(id, _)| id == &x.dataset_id)
+                .and_then(|(_, tracking_id)| tracking_id.clone());
+            top_dataset_resps
+        })
+        .collect::<Vec<_>>();
+
+    Ok(response)
+}

server/src/operators/message_operator.rs

@@ -69,15 +69,10 @@ pub async fn create_message_query(
 ) -> Result<(), ServiceError> {
     use crate::data::schema::messages::dsl::messages;
 
-    let mut conn = match pool.get().await {
-        Ok(conn) => conn,
-        Err(e) => {
-            log::error!("Error getting connection from pool: {:?}", e);
-            return Err(ServiceError::InternalServerError(
-                "Error getting postgres connection from pool".into(),
-            ));
-        }
-    };
+    let mut conn = pool
+        .get()
+        .await
+        .map_err(|_| ServiceError::BadRequest("Could not get database connection".to_string()))?;
 
     diesel::insert_into(messages)
         .values(&new_message)