feat: add timeTaken to query response (#1270)

Author: parmesant

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;
@@ -122,22 +122,26 @@ 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!("{:?}", time.elapsed());
     let response = QueryResponse {
         records,
         fields,
         fill_null: query_request.send_null,
         with_fields: query_request.fields,
+        total_time,
     }
     .to_http()?;
 

src/response.rs

@@ -23,11 +23,14 @@ 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>,
     pub fill_null: bool,
     pub with_fields: bool,
+    pub total_time: String,
 }
 
 impl QueryResponse {
@@ -49,12 +52,14 @@ impl QueryResponse {
         let response = if self.with_fields {
             json!({
                 "fields": self.fields,
-                "records": values
+                "records": values,
             })
         } else {
             Value::Array(values)
         };
 
-        Ok(HttpResponse::Ok().json(response))
+        Ok(HttpResponse::Ok()
+            .insert_header((TIME_ELAPSED_HEADER, self.total_time.as_str()))
+            .json(response))
     }
 }