fix(mysql): preserve stored procedure result sets

crates/dbx-core/src/db/mysql.rs

@@ -3676,6 +3676,91 @@ async fn execute_result_set_with_text_protocol_on_conn(
     })
 }
 
+async fn execute_result_sets_with_text_protocol_on_conn(
+    conn: &mut mysql_async::Conn,
+    sql: &str,
+    row_limit: usize,
+    max_rows: Option<usize>,
+    start: Instant,
+) -> Result<Vec<QueryResult>, String> {
+    let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
+    let mut results = Vec::new();
+
+    while advance_to_result_set_with_columns(&mut result).await? {
+        let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
+        let column_types: Vec<String> = result.columns_ref().iter().map(mysql_column_type_name).collect();
+        let mut spatial_columns = mysql_spatial_column_builder(result.columns_ref());
+        let mut spatial_values = Vec::new();
+        let mut truncated = false;
+
+        let rows = if should_collect_text_result_set(sql, row_limit, max_rows) {
+            let rows: Vec<mysql_async::Row> = result.collect().await.map_err(|e| e.to_string())?;
+            truncated = rows.len() > row_limit;
+            rows.iter()
+                .take(row_limit)
+                .map(|row| {
+                    let (values, srids) = mysql_row_to_json_with_srids(row, &mut spatial_columns);
+                    spatial_values.push(srids);
+                    values
+                })
+                .collect()
+        } else {
+            let mut rows = Vec::new();
+            let mut stream = result
+                .stream::<mysql_async::Row>()
+                .await
+                .map_err(|e| e.to_string())?
+                .ok_or_else(|| "Empty result set stream".to_string())?;
+
+            while let Some(row) = stream.next().await {
+                let row = row.map_err(|e| e.to_string())?;
+                if rows.len() < row_limit {
+                    let (values, srids) = mysql_row_to_json_with_srids(&row, &mut spatial_columns);
+                    rows.push(values);
+                    spatial_values.push(srids);
+                } else {
+                    truncated = true;
+                }
+            }
+            rows
+        };
+
+        results.push(QueryResult {
+            columns,
+            column_types,
+            column_sortables: vec![],
+            spatial_columns: spatial_columns.finish(),
+            spatial_values,
+            rows,
+            affected_rows: 0,
+            execution_time_ms: start.elapsed().as_millis(),
+            truncated,
+            session_id: None,
+            has_more: false,
+            elasticsearch_raw_body: None,
+        });
+    }
+
+    if results.is_empty() {
+        results.push(QueryResult {
+            columns: vec![],
+            column_types: Vec::new(),
+            column_sortables: vec![],
+            spatial_columns: vec![],
+            spatial_values: vec![],
+            rows: vec![],
+            affected_rows: result.affected_rows(),
+            execution_time_ms: start.elapsed().as_millis(),
+            truncated: false,
+            session_id: None,
+            has_more: false,
+            elasticsearch_raw_body: None,
+        });
+    }
+
+    Ok(results)
+}
+
 async fn advance_to_result_set_with_columns(
     result: &mut mysql_async::QueryResult<'_, '_, mysql_async::TextProtocol>,
 ) -> Result<bool, String> {
@@ -3997,6 +4082,22 @@ pub async fn execute_query_on_conn_with_max_rows(
     }
 }
 
+pub async fn execute_query_results_on_conn_with_max_rows(
+    conn: &mut mysql_async::Conn,
+    sql: &str,
+    bare: bool,
+    max_rows: Option<usize>,
+    dialect: MySqlQueryDialect,
+) -> Result<Vec<QueryResult>, String> {
+    if is_result_set_query(sql, dialect) && (bare || prefers_text_protocol_query(sql, dialect)) {
+        let start = Instant::now();
+        execute_result_sets_with_text_protocol_on_conn(conn, sql, query_result_row_limit(max_rows), max_rows, start)
+            .await
+    } else {
+        execute_query_on_conn_with_max_rows(conn, sql, bare, max_rows, dialect).await.map(|result| vec![result])
+    }
+}
+
 fn prefers_text_protocol_query(sql: &str, dialect: MySqlQueryDialect) -> bool {
     // User-entered result-set queries are not parameterized in DBX. Text protocol
     // avoids binary result decoding bugs in MySQL-compatible servers and proxies.