feat(cubesql): Allow providing API type when getting load request meta (#9202)

Author: MazterQyou

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(

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": {

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,
         )
     }