From d208c9baa0249b196a2e2f4bff3c82891e1e742b Mon Sep 17 00:00:00 2001
From: zgxme <zhenggaoxiong@selectdb.com>
Date: Wed, 8 Jul 2026 21:45:54 +0800
Subject: [PATCH] fix: collect starrocks explain analyze plans

---
 benchmark.sh                | 30 +++++++++++++++++++++++++++++-
 engines/interface.sh        |  6 ++++++
 engines/starrocks_engine.sh | 10 ++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/benchmark.sh b/benchmark.sh
index 9bf22f6..6a5c924 100755
--- a/benchmark.sh
+++ b/benchmark.sh
@@ -563,6 +563,23 @@ min_query_duration() {
     printf '%s' "$min"
 }
 
+collect_query_analyze_plan() {
+    local query_name="$1"
+    local safe_query_name="$2"
+    local sql_content="$3"
+    local plan_dir="$4"
+    local analyze_plan_content
+    local analyze_plan_sql
+
+    analyze_plan_sql=$(printf '%s' "$sql_content" | sed -e '/^[[:space:]]*--/d' -e '/^[[:space:]]*#/d')
+    analyze_plan_content=$(engine_get_analyze_plan "${db}" "$analyze_plan_sql" 2>/dev/null || true)
+    if [ -n "$analyze_plan_content" ]; then
+        printf "%s\n" "$analyze_plan_content" > "$plan_dir/${safe_query_name}_analyze_plan.txt"
+    else
+        echo "Analyze plan collection returned empty for ${query_name}" >&2
+    fi
+}
+
 # Run benchmark queries
 run_query() {
     local query_dir="${QUERY_DIR:-query/}"
@@ -650,7 +667,11 @@ run_query() {
     local query_detail_csv="$RESULT_DIR/query_detail.csv"
     local plan_dir=""
     local profile_dir=""
-    if [[ "$plan" == "true" ]]; then
+    local collect_analyze_plan="false"
+    if [[ "${ENGINE_TYPE,,}" == "starrocks" ]]; then
+        collect_analyze_plan="true"
+    fi
+    if [[ "$plan" == "true" || "$collect_analyze_plan" == "true" ]]; then
         plan_dir="$RESULT_DIR/plan"
         mkdir -p "$plan_dir"
     fi
@@ -742,6 +763,10 @@ run_query() {
                 detail_result+=",$RUN_QUERY_DURATION"
             done
 
+            if [[ "$collect_analyze_plan" == "true" ]]; then
+                collect_query_analyze_plan "$query_name" "$safe_query_name" "$sql_content" "$plan_dir"
+            fi
+
             local hot_min="null"
             if (( hot_query_count > 0 )); then
                 hot_min=$(min_query_duration "${hot_values[@]}")
@@ -769,6 +794,9 @@ run_query() {
                 run_timed_query "$query_name" "$safe_query_name" "$t" "$sql_content"
                 times_result+=",$RUN_QUERY_DURATION"
             done
+            if [[ "$collect_analyze_plan" == "true" ]]; then
+                collect_query_analyze_plan "$query_name" "$safe_query_name" "$sql_content" "$plan_dir"
+            fi
             echo "$times_result" >> "$query_csv"
         fi
     done
diff --git a/engines/interface.sh b/engines/interface.sh
index 279b2e8..f2048ef 100644
--- a/engines/interface.sh
+++ b/engines/interface.sh
@@ -182,6 +182,12 @@ engine_get_plan() {
     return 1
 }
 
+# - engine_get_analyze_plan(db, sql): print execution plan with runtime stats
+engine_get_analyze_plan() {
+    echo "Analyze plan collection not supported by this engine, skipping..." >&2
+    return 1
+}
+
 engine_get_session_sql_content() {
     local apply_session="${1:-true}"
     local session_file="${SESSION_FILE:-session/session.sql}"
diff --git a/engines/starrocks_engine.sh b/engines/starrocks_engine.sh
index 8c1251a..4c50400 100644
--- a/engines/starrocks_engine.sh
+++ b/engines/starrocks_engine.sh
@@ -305,6 +305,16 @@ engine_get_plan() {
     mysql "${args[@]}" -e "EXPLAIN VERBOSE ${sql_statement}" 2>/dev/null || true
 }
 
+# Optional: fetch EXPLAIN ANALYZE output for a query.
+engine_get_analyze_plan() {
+    local db_name="$1"
+    local sql_statement="$2"
+    export MYSQL_PWD="${password:-}"
+    local args=(-h"$fe_host" -P"$fe_query_port" -u"$user" --batch --raw --skip-column-names)
+    [ -n "$db_name" ] && args+=(-D"$db_name")
+    mysql "${args[@]}" -e "EXPLAIN ANALYZE ${sql_statement}" 2>/dev/null || true
+}
+
 # 3. Generate JDBC DataSource XML configuration for Starrocks
 engine_get_jdbc_datasource() {
     # Escape any special characters in the password