From ae5d9770a1c00326c446e7ec40bd3f13ca0ce8f4 Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Fri, 7 Feb 2025 19:25:14 +0400
Subject: [PATCH] feat(cubesql): Allow providing API type when getting load
 request meta (#9202)

---
 rust/cubesql/cubesql/src/compile/query_engine.rs |  8 ++++----
 rust/cubesql/cubesql/src/sql/postgres/shim.rs    | 14 +++++++-------
 rust/cubesql/cubesql/src/sql/session.rs          |  4 ++--
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/rust/cubesql/cubesql/src/compile/query_engine.rs b/rust/cubesql/cubesql/src/compile/query_engine.rs
index dac6c9cddc96a..2d1f79b75bd92 100644
--- a/rust/cubesql/cubesql/src/compile/query_engine.rs
+++ b/rust/cubesql/cubesql/src/compile/query_engine.rs
@@ -104,7 +104,7 @@ pub trait QueryEngine {
                     .log_load_state(
                         Some(span_id.clone()),
                         auth_context,
-                        state.get_load_request_meta(),
+                        state.get_load_request_meta("sql"),
                         "SQL API Query Planning".to_string(),
                         serde_json::json!({
                             "query": span_id.query_key.clone(),
@@ -286,7 +286,7 @@ pub trait QueryEngine {
                     .log_load_state(
                         Some(span_id.clone()),
                         auth_context,
-                        state.get_load_request_meta(),
+                        state.get_load_request_meta("sql"),
                         "SQL API Query Planning Success".to_string(),
                         serde_json::json!({
                             "query": span_id.query_key.clone(),
@@ -302,7 +302,7 @@ pub trait QueryEngine {
         // to catch all SQL generation errors during planning
         let rewrite_plan = Self::evaluate_wrapped_sql(
             self.transport_ref().clone(),
-            Arc::new(state.get_load_request_meta()),
+            Arc::new(state.get_load_request_meta("sql")),
             rewrite_plan,
         )
         .await?;
@@ -390,7 +390,7 @@ impl QueryEngine for SqlQueryEngine {
     ) -> Result<DFSessionContext, CompilationError> {
         let query_planner = Arc::new(CubeQueryPlanner::new(
             self.transport_ref().clone(),
-            state.get_load_request_meta(),
+            state.get_load_request_meta("sql"),
             self.config_ref().clone(),
         ));
         let mut ctx = DFSessionContext::with_state(
diff --git a/rust/cubesql/cubesql/src/sql/postgres/shim.rs b/rust/cubesql/cubesql/src/sql/postgres/shim.rs
index 5e99376d308ad..4247ce6b536f0 100644
--- a/rust/cubesql/cubesql/src/sql/postgres/shim.rs
+++ b/rust/cubesql/cubesql/src/sql/postgres/shim.rs
@@ -401,7 +401,7 @@ impl AsyncPostgresShim {
                             .log_load_state(
                                 span_id.clone(),
                                 auth_context,
-                                self.session.state.get_load_request_meta(),
+                                self.session.state.get_load_request_meta("sql"),
                                 "Load Request".to_string(),
                                 serde_json::json!({
                                     "query": span_id.as_ref().unwrap().query_key.clone(),
@@ -453,7 +453,7 @@ impl AsyncPostgresShim {
                                     .log_load_state(
                                         Some(span_id.clone()),
                                         auth_context,
-                                        self.session.state.get_load_request_meta(),
+                                        self.session.state.get_load_request_meta("sql"),
                                         "Data Query Status".to_string(),
                                         serde_json::json!({
                                             "isDataQuery": true
@@ -481,7 +481,7 @@ impl AsyncPostgresShim {
                                     .log_load_state(
                                         Some(span_id.clone()),
                                         auth_context.clone(),
-                                        self.session.state.get_load_request_meta(),
+                                        self.session.state.get_load_request_meta("sql"),
                                         "Data Query Status".to_string(),
                                         serde_json::json!({
                                             "isDataQuery": true,
@@ -496,7 +496,7 @@ impl AsyncPostgresShim {
                                     .log_load_state(
                                         Some(span_id.clone()),
                                         auth_context,
-                                        self.session.state.get_load_request_meta(),
+                                        self.session.state.get_load_request_meta("sql"),
                                         "Load Request Success".to_string(),
                                         serde_json::json!({
                                             "query": span_id.query_key.clone(),
@@ -602,7 +602,7 @@ impl AsyncPostgresShim {
                     .log_load_state(
                         Some(span_id.clone()),
                         auth_context,
-                        self.session.state.get_load_request_meta(),
+                        self.session.state.get_load_request_meta("sql"),
                         "SQL API Error".to_string(),
                         serde_json::json!({
                             "query": span_id.query_key.clone(),
@@ -1848,7 +1848,7 @@ impl AsyncPostgresShim {
                 .log_load_state(
                     span_id.clone(),
                     auth_context,
-                    self.session.state.get_load_request_meta(),
+                    self.session.state.get_load_request_meta("sql"),
                     "Load Request".to_string(),
                     serde_json::json!({
                         "query": {
@@ -1876,7 +1876,7 @@ impl AsyncPostgresShim {
                         .log_load_state(
                             Some(span_id.clone()),
                             auth_context,
-                            self.session.state.get_load_request_meta(),
+                            self.session.state.get_load_request_meta("sql"),
                             "Load Request Success".to_string(),
                             serde_json::json!({
                                 "query": {
diff --git a/rust/cubesql/cubesql/src/sql/session.rs b/rust/cubesql/cubesql/src/sql/session.rs
index 75a4541ffce85..b7a961af5f7c7 100644
--- a/rust/cubesql/cubesql/src/sql/session.rs
+++ b/rust/cubesql/cubesql/src/sql/session.rs
@@ -377,7 +377,7 @@ impl SessionState {
         Arc::clone(&self.temp_tables)
     }
 
-    pub fn get_load_request_meta(&self) -> LoadRequestMeta {
+    pub fn get_load_request_meta(&self, api_type: &str) -> LoadRequestMeta {
         let application_name = if let Some(var) = self.get_variable("application_name") {
             Some(var.value.to_string())
         } else {
@@ -386,7 +386,7 @@ impl SessionState {
 
         LoadRequestMeta::new(
             self.protocol.get_name().to_string(),
-            "sql".to_string(),
+            api_type.to_string(),
             application_name,
         )
     }