Skip to content

Commit

Permalink
feat(cubesql): Allow providing API type when getting load request meta (
Browse files Browse the repository at this point in the history 
#9202)
  • Loading branch information
@MazterQyou
MazterQyou authored Feb 7, 2025
1 parent f8e523b commit ae5d977
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 13 deletions.
8 changes: 4 additions & 4 deletions rust/cubesql/cubesql/src/compile/query_engine.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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(),
Expand All @@ -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?;
Expand Down Expand Up @@ -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(
Expand Down
14 changes: 7 additions & 7 deletions rust/cubesql/cubesql/src/sql/postgres/shim.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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(),
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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": {
Expand Down
4 changes: 2 additions & 2 deletions rust/cubesql/cubesql/src/sql/session.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -386,7 +386,7 @@ impl SessionState {

LoadRequestMeta::new(
self.protocol.get_name().to_string(),
"sql".to_string(),
api_type.to_string(),
application_name,
)
}
Expand Down

0 comments on commit ae5d977

Please sign in to comment.