propagate span_id

Author: vasilev-alex

packages/cubejs-backend-native/src/node_export.rs

@@ -212,6 +212,7 @@ async fn handle_sql_query(
         }
 
         let session_clone = Arc::clone(&session);
+        let span_id_clone = span_id.clone();
 
         let execute = || async move {
             // todo: can we use compiler_cache?
@@ -221,7 +222,8 @@ async fn handle_sql_query(
                 .map_err(|err| {
                     CubeError::internal(format!("Failed to get meta context: {}", err))
                 })?;
-            let query_plan = convert_sql_to_cube_query(sql_query, meta_context, session).await?;
+            let query_plan =
+                convert_sql_to_cube_query(sql_query, meta_context, session, span_id_clone).await?;
 
             let mut stream = get_df_batches(&query_plan).await?;
 

packages/cubejs-backend-native/src/sql4sql.rs

@@ -166,7 +166,8 @@ async fn handle_sql4sql_query(
             .await
             .map_err(|err| CubeError::internal(format!("Failed to get meta context: {err}")))?;
         let query_plan =
-            convert_sql_to_cube_query(sql_query, meta_context.clone(), session.clone()).await?;
+            convert_sql_to_cube_query(sql_query, meta_context.clone(), session.clone(), None)
+                .await?;
         let logical_plan = query_plan.try_as_logical_plan()?;
         get_sql(&session, meta_context, Arc::new(logical_plan.clone())).await
     })

packages/cubejs-backend-native/src/transport.rs

@@ -341,6 +341,8 @@ impl TransportService for NodeBridgeTransport {
     ) -> Result<Vec<RecordBatch>, CubeError> {
         trace!("[transport] Request ->");
 
+        println!("!!!! LOAD {:?}", span_id);
+
         let native_auth = ctx
             .as_any()
             .downcast_ref::<NativeAuthContext>()

rust/cubesql/cubesql/src/compile/router.rs

@@ -640,7 +640,8 @@ pub async fn convert_sql_to_cube_query(
     query: &str,
     meta: Arc<MetaContext>,
     session: Arc<Session>,
+    span_id: Option<Arc<SpanId>>,
 ) -> CompilationResult<QueryPlan> {
     let stmt = parse_sql_to_statement(&query, session.state.protocol.clone(), &mut None)?;
-    convert_statement_to_cube_query(stmt, meta, session, &mut None, None).await
+    convert_statement_to_cube_query(stmt, meta, session, &mut None, span_id).await
 }

rust/cubesql/cubesql/src/compile/service.rs

@@ -27,6 +27,6 @@ impl SqlService for SqlAuthDefaultImpl {
         meta: Arc<MetaContext>,
         session: Arc<Session>,
     ) -> CompilationResult<QueryPlan> {
-        convert_sql_to_cube_query(&query, meta, session).await
+        convert_sql_to_cube_query(&query, meta, session, None).await
     }
 }

rust/cubesql/cubesql/src/compile/test/mod.rs

@@ -1021,7 +1021,13 @@ impl TestContext {
 
     pub async fn convert_sql_to_cube_query(&self, query: &str) -> CompilationResult<QueryPlan> {
         // TODO push to_string() deeper
-        convert_sql_to_cube_query(&query.to_string(), self.meta.clone(), self.session.clone()).await
+        convert_sql_to_cube_query(
+            &query.to_string(),
+            self.meta.clone(),
+            self.session.clone(),
+            None,
+        )
+        .await
     }
 
     pub async fn execute_query_with_flags(
@@ -1131,6 +1137,7 @@ pub async fn convert_select_to_query_plan_with_meta(
         &query,
         meta_context.clone(),
         get_test_session(DatabaseProtocol::PostgreSQL, meta_context).await,
+        None,
     )
     .await;
 

rust/cubesql/cubesql/src/compile/test/test_cube_join.rs

@@ -503,6 +503,7 @@ async fn test_join_cubes_on_wrong_field_error() {
         .to_string(),
         meta.clone(),
         get_test_session(DatabaseProtocol::PostgreSQL, meta).await,
+        None,
     )
     .await;
 
@@ -527,6 +528,7 @@ async fn test_join_cubes_filter_from_wrong_side_error() {
             .to_string(),
         meta.clone(),
         get_test_session(DatabaseProtocol::PostgreSQL, meta).await,
+        None
     )
         .await;
 
@@ -554,6 +556,7 @@ async fn test_join_cubes_with_aggr_error() {
             .to_string(),
         meta.clone(),
         get_test_session(DatabaseProtocol::PostgreSQL, meta).await,
+        None
     )
         .await;
 

rust/cubesql/cubesql/src/compile/test/test_user_change.rs

@@ -163,7 +163,8 @@ async fn test_change_user_via_filter_or() {
         convert_sql_to_cube_query(
             &"SELECT COUNT(*) as cnt FROM KibanaSampleDataEcommerce WHERE __user = 'gopher' OR customer_gender = 'male'".to_string(),
             meta.clone(),
-            get_test_session(DatabaseProtocol::PostgreSQL, meta).await
+            get_test_session(DatabaseProtocol::PostgreSQL, meta).await,
+            None
         ).await;
 
     // TODO: We need to propagate error to result, to assert message