feature: add timer to usage request

Author: cdxker

server/src/handlers/organization_handler.rs

@@ -24,6 +24,7 @@ use actix_web::{web, HttpRequest, HttpResponse};
 use sanitize_html::rules;
 use sanitize_html::sanitize_str;
 use serde::{Deserialize, Serialize};
+use simple_server_timing_header::Timer;
 use utoipa::ToSchema;
 
 /// Get Organization
@@ -275,15 +276,23 @@ pub async fn get_organization_usage(
             }
         }))
     } else {
+        let mut usage_timer = Some(Timer::new());
+
         let extended_usage = get_extended_org_usage_by_id_query(
             org_id,
             data.date_range.clone(),
             clickhouse_client.get_ref(),
             pool,
+            &mut usage_timer,
         )
         .await?;
 
-        Ok(HttpResponse::Ok().json(extended_usage))
+        Ok(HttpResponse::Ok()
+            .insert_header((
+                Timer::header_key(),
+                usage_timer.expect("Timer should exist").header_value(),
+            ))
+            .json(extended_usage))
     }
 }
 

server/src/operators/organization_operator.rs

@@ -501,6 +501,7 @@ pub async fn get_extended_org_usage_by_id_query(
     date_range: Option<DateRange>,
     clickhouse_client: &clickhouse::Client,
     pool: web::Data<Pool>,
+    timer: &mut Option<simple_server_timing_header::Timer>,
 ) -> Result<ExtendedOrganizationUsageCount, ServiceError> {
     let usage = get_org_usage_by_id_query(organization_id, pool).await?;
 
@@ -521,6 +522,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching search queries {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched search_count");
+    }
+
     let message_tokens = clickhouse_client
         .query(&format_with_daterange(
             "
@@ -538,6 +543,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching message queries {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched message_count");
+    }
+
     let bytes_and_tokens_ingested = clickhouse_client
         .query(&format_with_daterange(
             "
@@ -555,6 +564,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching ingestion data {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched bytes_ingested and tokens_ingested");
+    }
+
     let ocr_pages = clickhouse_client
         .query(&format_with_daterange(
             "
@@ -571,6 +584,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching ingestion data {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched ocr_pages");
+    }
+
     let website_pages_scraped = clickhouse_client
         .query(&format_with_daterange(
             "
@@ -587,6 +604,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching ingestion data {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched website_pages_scraped");
+    }
+
     let events_ingested = clickhouse_client
         .query(&format_with_daterange(
             "
@@ -604,6 +625,10 @@ pub async fn get_extended_org_usage_by_id_query(
             ServiceError::InternalServerError(format!("Error fetching ingestion data {:?}", e))
         })?;
 
+    if let Some(timer) = timer {
+        timer.add("fetched events_ingested");
+    }
+
     Ok(ExtendedOrganizationUsageCount {
         dataset_count: usage.dataset_count,
         user_count: usage.user_count,