From daa549384d4d38d4d5a798fd5c8ecca890236722 Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Mon, 25 Aug 2025 16:15:31 +0400
Subject: [PATCH] feat(cubesql): Report rewrite start/success events

Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
---
 .../cubesql/src/compile/query_engine.rs       | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/rust/cubesql/cubesql/src/compile/query_engine.rs b/rust/cubesql/cubesql/src/compile/query_engine.rs
index 705b001d484e7..0b28a4a4815b1 100644
--- a/rust/cubesql/cubesql/src/compile/query_engine.rs
+++ b/rust/cubesql/cubesql/src/compile/query_engine.rs
@@ -174,6 +174,22 @@ pub trait QueryEngine {
             )
             .map_err(|e| CompilationError::internal(e.to_string()))?;
 
+        let rewriting_start = SystemTime::now();
+        if let Some(span_id) = span_id.as_ref() {
+            if let Some(auth_context) = state.auth_context() {
+                self.transport_ref()
+                    .log_load_state(
+                        Some(span_id.clone()),
+                        auth_context,
+                        state.get_load_request_meta("sql"),
+                        "SQL API Plan Rewrite".to_string(),
+                        serde_json::json!({}),
+                    )
+                    .await
+                    .map_err(|e| CompilationError::internal(e.to_string()))?;
+            }
+        }
+
         let mut finalized_graph = self
             .compiler_cache_ref()
             .rewrite(
@@ -267,6 +283,23 @@ pub trait QueryEngine {
 
         let rewrite_plan = result?;
 
+        if let Some(span_id) = span_id.as_ref() {
+            if let Some(auth_context) = state.auth_context() {
+                self.transport_ref()
+                    .log_load_state(
+                        Some(span_id.clone()),
+                        auth_context,
+                        state.get_load_request_meta("sql"),
+                        "SQL API Plan Rewrite Success".to_string(),
+                        serde_json::json!({
+                            "duration": rewriting_start.elapsed().unwrap().as_millis() as u64,
+                        }),
+                    )
+                    .await
+                    .map_err(|e| CompilationError::internal(e.to_string()))?;
+            }
+        }
+
         // DF optimizes logical plan (second time) on physical plan creation
         // It's not safety to use all optimizers from DF for OLAP queries, because it will lead to errors
         // From another side, 99% optimizers cannot optimize anything