From 947e0490bd46162a0f0985363c8534e0ff17aebe Mon Sep 17 00:00:00 2001
From: anant <anant.v09@proton.me>
Date: Tue, 25 Mar 2025 12:44:39 +0530
Subject: [PATCH 1/2] add `timeTaken` to query response

---
 src/handlers/http/query.rs | 4 +++-
 src/response.rs            | 9 +++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/handlers/http/query.rs b/src/handlers/http/query.rs
index 7d9c33a45..cca589839 100644
--- a/src/handlers/http/query.rs
+++ b/src/handlers/http/query.rs
@@ -67,6 +67,7 @@ pub struct Query {
 }
 
 pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpResponse, QueryError> {
+    let start = Instant::now();
     let session_state = QUERY_SESSION.state();
     let raw_logical_plan = match session_state
         .create_logical_plan(&query_request.query)
@@ -132,12 +133,13 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
     }
 
     let (records, fields) = execute(query, &table_name).await?;
-
+    let total_time = format!("{:?}", start.elapsed());
     let response = QueryResponse {
         records,
         fields,
         fill_null: query_request.send_null,
         with_fields: query_request.fields,
+        total_time,
     }
     .to_http()?;
 
diff --git a/src/response.rs b/src/response.rs
index fe3159f71..78f2b7898 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -28,6 +28,7 @@ pub struct QueryResponse {
     pub fields: Vec<String>,
     pub fill_null: bool,
     pub with_fields: bool,
+    pub total_time: String,
 }
 
 impl QueryResponse {
@@ -49,10 +50,14 @@ impl QueryResponse {
         let response = if self.with_fields {
             json!({
                 "fields": self.fields,
-                "records": values
+                "records": values,
+                "totalTime": self.total_time
             })
         } else {
-            Value::Array(values)
+            json!({
+                "records": values,
+                "totalTime": self.total_time
+            })
         };
 
         Ok(HttpResponse::Ok().json(response))

From ccc540b22868cd00941e71b608e1c33d5d614185 Mon Sep 17 00:00:00 2001
From: anant <anant.v09@proton.me>
Date: Tue, 25 Mar 2025 14:09:23 +0530
Subject: [PATCH 2/2] modified time taken to be sent as response header

---
 src/handlers/http/query.rs | 10 ++++++----
 src/response.rs            | 12 ++++++------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/handlers/http/query.rs b/src/handlers/http/query.rs
index cca589839..fce7164c2 100644
--- a/src/handlers/http/query.rs
+++ b/src/handlers/http/query.rs
@@ -44,7 +44,7 @@ use crate::query::error::ExecuteError;
 use crate::query::{execute, CountsRequest, CountsResponse, Query as LogicalQuery};
 use crate::query::{TableScanVisitor, QUERY_SESSION};
 use crate::rbac::Users;
-use crate::response::QueryResponse;
+use crate::response::{QueryResponse, TIME_ELAPSED_HEADER};
 use crate::storage::object_storage::commit_schema_to_storage;
 use crate::storage::ObjectStorageError;
 use crate::utils::actix::extract_session_key_from_req;
@@ -67,7 +67,6 @@ pub struct Query {
 }
 
 pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpResponse, QueryError> {
-    let start = Instant::now();
     let session_state = QUERY_SESSION.state();
     let raw_logical_plan = match session_state
         .create_logical_plan(&query_request.query)
@@ -123,17 +122,20 @@ pub async fn query(req: HttpRequest, query_request: Query) -> Result<HttpRespons
             Value::Array(vec![json!({column_name: count})])
         };
 
+        let total_time = format!("{:?}", time.elapsed());
         let time = time.elapsed().as_secs_f64();
 
         QUERY_EXECUTE_TIME
             .with_label_values(&[&table_name])
             .observe(time);
 
-        return Ok(HttpResponse::Ok().json(response));
+        return Ok(HttpResponse::Ok()
+            .insert_header((TIME_ELAPSED_HEADER, total_time.as_str()))
+            .json(response));
     }
 
     let (records, fields) = execute(query, &table_name).await?;
-    let total_time = format!("{:?}", start.elapsed());
+    let total_time = format!("{:?}", time.elapsed());
     let response = QueryResponse {
         records,
         fields,
diff --git a/src/response.rs b/src/response.rs
index 78f2b7898..7aca5c44f 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -23,6 +23,8 @@ use itertools::Itertools;
 use serde_json::{json, Value};
 use tracing::info;
 
+pub const TIME_ELAPSED_HEADER: &str = "p-time-elapsed";
+
 pub struct QueryResponse {
     pub records: Vec<RecordBatch>,
     pub fields: Vec<String>,
@@ -51,15 +53,13 @@ impl QueryResponse {
             json!({
                 "fields": self.fields,
                 "records": values,
-                "totalTime": self.total_time
             })
         } else {
-            json!({
-                "records": values,
-                "totalTime": self.total_time
-            })
+            Value::Array(values)
         };
 
-        Ok(HttpResponse::Ok().json(response))
+        Ok(HttpResponse::Ok()
+            .insert_header((TIME_ELAPSED_HEADER, self.total_time.as_str()))
+            .json(response))
     }
 }