From 35eab8a8b1bccc46e19c59c4c0ca0c0a2256217e Mon Sep 17 00:00:00 2001 From: zipg Date: Fri, 14 Aug 2026 00:11:28 +0800 Subject: [PATCH] perf(query): optimize large-value result diagnostics --- .../__tests__/httpErrorParsing.spec.ts | 88 +++++++++- apps/desktop/src/lib/backend/http.ts | 68 +++++++- apps/desktop/src/lib/backend/tauri.ts | 21 ++- crates/dbx-core/src/db/mysql.rs | 153 ++++++++++++++++-- crates/dbx-core/src/query.rs | 92 ++++++++++- crates/dbx-web/src/routes/query.rs | 40 ++++- 6 files changed, 429 insertions(+), 33 deletions(-) diff --git a/apps/desktop/src/lib/backend/__tests__/httpErrorParsing.spec.ts b/apps/desktop/src/lib/backend/__tests__/httpErrorParsing.spec.ts index 8a5c2a0744..afedd055dd 100644 --- a/apps/desktop/src/lib/backend/__tests__/httpErrorParsing.spec.ts +++ b/apps/desktop/src/lib/backend/__tests__/httpErrorParsing.spec.ts @@ -1,5 +1,34 @@ import { describe, expect, test, vi } from "vitest"; -import { backendResponseError, importAgentDriver, installJdbcPluginLocal } from "@/lib/backend/http"; +import { backendResponseError, executeMulti, importAgentDriver, installJdbcPluginLocal } from "@/lib/backend/http"; +import { getDebugLogText } from "@/lib/backend/debugLog"; + +class MemoryStorage implements Storage { + private data = new Map(); + + get length() { + return this.data.size; + } + + clear() { + this.data.clear(); + } + + getItem(key: string) { + return this.data.get(key) ?? null; + } + + key(index: number) { + return [...this.data.keys()][index] ?? null; + } + + removeItem(key: string) { + this.data.delete(key); + } + + setItem(key: string, value: string) { + this.data.set(key, value); + } +} import { BackendErrorException } from "@/lib/backend/errorUtils"; const envelope = { @@ -70,3 +99,60 @@ describe("HTTP backend error parsing", () => { vi.unstubAllGlobals(); }); }); + +describe("HTTP query transport diagnostics", () => { + test("records timing and payload sizes without logging query contents", async () => { + vi.stubGlobal("localStorage", new MemoryStorage()); + const sql = "SELECT 'private-query-marker'"; + localStorage.setItem("dbx-debug-logging-enabled", "1"); + localStorage.removeItem("dbx-debug-log-entries"); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + new Response(JSON.stringify([{ columns: ["value"], rows: [["private-cell-marker"]] }]), { + status: 200, + headers: { + "content-type": "application/json", + "x-dbx-core-ms": "41", + "x-dbx-serialize-ms": "7", + }, + }), + ), + ); + + await executeMulti("private-connection-marker", "private-database-marker", sql, undefined, "trace-12345678"); + + const logs = getDebugLogText(); + expect(logs).toContain("[DBX][query-transport:http]"); + expect(logs).toContain('"backendCoreMs":"41"'); + expect(logs).toContain('"backendSerializeMs":"7"'); + expect(logs).toContain('"requestBytes":'); + expect(logs).toContain('"responseBytes":'); + expect(logs).toContain('"jsonParseMs":'); + expect(logs).not.toContain("private-query-marker"); + expect(logs).not.toContain("private-cell-marker"); + expect(logs).not.toContain("private-connection-marker"); + expect(logs).not.toContain("private-database-marker"); + + vi.unstubAllGlobals(); + }); + + test("preserves structured backend errors while logging only safe failure metadata", async () => { + vi.stubGlobal("localStorage", new MemoryStorage()); + localStorage.setItem("dbx-debug-logging-enabled", "1"); + localStorage.removeItem("dbx-debug-log-entries"); + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(new Response(JSON.stringify(envelope), { status: 400 }))); + + const error = await executeMulti("connection-1", "database-1", "SELECT 'failed-query-marker'", undefined, "trace-87654321").catch((value: unknown) => value); + + expect(error).toBeInstanceOf(BackendErrorException); + expect(error).toMatchObject({ backendError: expect.objectContaining({ code: envelope.code }) }); + const logs = getDebugLogText(); + expect(logs).toContain("[DBX][query-transport:http:error]"); + expect(logs).toContain('"status":400'); + expect(logs).not.toContain("failed-query-marker"); + expect(logs).not.toContain(envelope.detail); + + vi.unstubAllGlobals(); + }); +}); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index cf1b544dce..9dd794479d 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -203,6 +203,7 @@ import type { NacosSearchProgress, } from "@/types/nacos"; import { safeLocalStorageGet, safeLocalStorageSet } from "@/lib/backend/safeStorage"; +import { appendDebugLog, isDebugLoggingEnabled } from "@/lib/backend/debugLog"; import { normalizeConnectionTestResult } from "@/lib/connection/connectionDatabaseInfo"; import type { AnnotationFile, SchemaSnapshot } from "@/docs/types"; @@ -236,6 +237,55 @@ async function post(url: string, body: unknown): Promise { return res.json(); } +async function postQueryWithDiagnostics(url: string, body: unknown, traceId?: string): Promise { + if (!isDebugLoggingEnabled()) return post(url, body); + + const startedAt = performance.now(); + const serializedBody = JSON.stringify(body); + const response = await fetch(apiUrl(url), { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: serializedBody, + }); + const headersAt = performance.now(); + const responseText = await response.text(); + const bodyAt = performance.now(); + if (!response.ok) { + appendDebugLog("warn", "[DBX][query-transport:http:error]", { + traceId: traceId?.slice(0, 8), + status: response.status, + requestBytes: new TextEncoder().encode(serializedBody).byteLength, + responseBytes: new TextEncoder().encode(responseText).byteLength, + responseHeadersMs: Math.round(headersAt - startedAt), + responseBodyMs: Math.round(bodyAt - headersAt), + totalMs: Math.round(bodyAt - startedAt), + backendCoreMs: response.headers.get("x-dbx-core-ms"), + backendSerializeMs: response.headers.get("x-dbx-serialize-ms"), + }); + throw await backendResponseError( + new Response(responseText, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }), + ); + } + const result = JSON.parse(responseText) as T; + const parsedAt = performance.now(); + appendDebugLog("info", "[DBX][query-transport:http]", { + traceId: traceId?.slice(0, 8), + requestBytes: new TextEncoder().encode(serializedBody).byteLength, + responseBytes: new TextEncoder().encode(responseText).byteLength, + responseHeadersMs: Math.round(headersAt - startedAt), + responseBodyMs: Math.round(bodyAt - headersAt), + jsonParseMs: Math.round(parsedAt - bodyAt), + totalMs: Math.round(parsedAt - startedAt), + backendCoreMs: response.headers.get("x-dbx-core-ms"), + backendSerializeMs: response.headers.get("x-dbx-serialize-ms"), + }); + return result; +} + async function get(url: string): Promise { const res = await fetch(apiUrl(url)); if (!res.ok) throw await backendResponseError(res); @@ -940,14 +990,18 @@ export async function executeMulti( executionMode?: "simple"; }, ): Promise { - return post("/api/query/execute-multi", { - connectionId, - database, - sql, - schema, + return postQueryWithDiagnostics( + "/api/query/execute-multi", + { + connectionId, + database, + sql, + schema, + executionId, + ...options, + }, executionId, - ...options, - }); + ); } export interface ExecuteMultiProgress { diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 7ab7ca231c..b5d4598571 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -3,6 +3,7 @@ import { BackendErrorException, type BackendError } from "@/lib/backend/errorUti import { listen, type UnlistenFn } from "@tauri-apps/api/event"; import { normalizeRustMongoCommand, type MongoCommand } from "@/lib/mongo/mongoShellCommand"; import { ExternalSqlFileTooLargeError } from "@/lib/sql/sqlFileOpen"; +import { appendDebugLog, isDebugLoggingEnabled } from "@/lib/backend/debugLog"; /** Normalize Tauri rejections once at the public backend boundary. */ async function invokeBackend(command: string, args?: Record): Promise { @@ -1159,8 +1160,10 @@ export async function executeMulti( executionMode?: "simple"; }, ): Promise { + const diagnosticsEnabled = isDebugLoggingEnabled(); + const startedAt = diagnosticsEnabled ? performance.now() : 0; try { - return await invoke("execute_multi", { + const results = await invoke("execute_multi", { connectionId, database, sql, @@ -1168,7 +1171,23 @@ export async function executeMulti( executionId, ...options, }); + if (diagnosticsEnabled) { + appendDebugLog("info", "[DBX][query-transport:tauri]", { + traceId: executionId?.slice(0, 8), + totalMs: Math.round(performance.now() - startedAt), + resultCount: results.length, + rowCounts: results.map((result) => result.rows.length), + columnCounts: results.map((result) => result.columns.length), + }); + } + return results; } catch (error) { + if (diagnosticsEnabled) { + appendDebugLog("warn", "[DBX][query-transport:tauri:error]", { + traceId: executionId?.slice(0, 8), + totalMs: Math.round(performance.now() - startedAt), + }); + } throw new BackendErrorException(error); } } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 6d33c6971e..b377311a71 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -4127,9 +4127,13 @@ async fn execute_result_set_with_text_protocol_on_conn( max_rows: Option, max_result_bytes: Option, result_key_columns: &[String], + diagnostic_trace_id: Option<&str>, start: Instant, ) -> Result { + let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); + let dispatch_started_at = diagnostics_enabled.then(Instant::now); let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; + let dispatch_ms = dispatch_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()); if !advance_to_result_set_with_columns(&mut result).await? { let affected_rows = result.affected_rows(); let warnings = result.warnings(); @@ -4199,19 +4203,28 @@ async fn execute_result_set_with_text_protocol_on_conn( let mut spatial_values: Vec>> = Vec::new(); let mut large_value_cells = Vec::new(); let mut truncated = false; + let mut row_read_time = Duration::ZERO; + let mut row_convert_time = Duration::ZERO; let mut stream = result .stream::() .await .map_err(|e| e.to_string())? .ok_or_else(|| "Empty result set stream".to_string())?; - while let Some(row) = stream.next().await { + loop { + let read_started_at = diagnostics_enabled.then(Instant::now); + let next_row = stream.next().await; + if let Some(started_at) = read_started_at { + row_read_time += started_at.elapsed(); + } + let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; if result_rows.len() >= row_limit { truncated = true; break; } let row_index = result_rows.len(); + let convert_started_at = diagnostics_enabled.then(Instant::now); let (values, srids, mut row_large_values) = mysql_row_to_json_with_srids_and_previews( &row, &mut spatial_columns, @@ -4222,6 +4235,9 @@ async fn execute_result_set_with_text_protocol_on_conn( result_rows.push(values); spatial_values.push(srids); large_value_cells.append(&mut row_large_values); + if let Some(started_at) = convert_started_at { + row_convert_time += started_at.elapsed(); + } } drop(stream); @@ -4229,6 +4245,7 @@ async fn execute_result_set_with_text_protocol_on_conn( // terminator packet arrived, so `result.warnings()`/`result.info()` still // hold the previous statement's values; skip capture instead of reporting // stale messages. + let finalize_started_at = diagnostics_enabled.then(Instant::now); let messages = if truncated { drop(result); Vec::new() @@ -4238,6 +4255,20 @@ async fn execute_result_set_with_text_protocol_on_conn( drop(result); collect_mysql_server_messages(conn, warnings, &info).await }; + if diagnostics_enabled { + log::info!( + "[query][mysql-result] trace_id={} protocol=text dispatch_ms={} row_read_ms={} row_convert_ms={} finalize_ms={} rows={} columns={} truncated={} preview_cells={}", + diagnostic_trace_id.unwrap_or("none"), + dispatch_ms, + row_read_time.as_millis(), + row_convert_time.as_millis(), + finalize_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()), + result_rows.len(), + columns.len(), + truncated, + large_value_cells.len() + ); + } Ok(MySqlQueryResult { result: QueryResult { @@ -4266,6 +4297,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( max_rows: Option, max_result_bytes: Option, result_key_columns: &[String], + diagnostic_trace_id: Option<&str>, start: Instant, ) -> Result, String> { // Per-set warnings/info read after each `collect()` are only accurate when @@ -4279,9 +4311,14 @@ async fn execute_result_sets_with_text_protocol_on_conn( // per-set messages stay empty and only the final SHOW WARNINGS attachment // on the last result set reports warnings. let capture_per_set_messages = conn.opts().deprecate_eof(); + let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); + let dispatch_started_at = diagnostics_enabled.then(Instant::now); let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; + let dispatch_ms = dispatch_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()); let mut results: Vec = Vec::new(); let mut result_set_warnings: Vec = Vec::new(); + let mut row_read_time = Duration::ZERO; + let mut row_convert_time = Duration::ZERO; while advance_to_result_set_with_columns(&mut result).await? { let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); @@ -4295,16 +4332,26 @@ async fn execute_result_sets_with_text_protocol_on_conn( let protected_indexes = mysql_result_protected_column_indexes(result.columns_ref(), result_key_columns); let rows = if max_result_bytes.is_none() && should_collect_text_result_set(sql, row_limit, max_rows) { + let read_started_at = diagnostics_enabled.then(Instant::now); let rows: Vec = result.collect().await.map_err(|e| e.to_string())?; + if let Some(started_at) = read_started_at { + row_read_time += started_at.elapsed(); + } truncated = rows.len() > row_limit; - rows.iter() + let convert_started_at = diagnostics_enabled.then(Instant::now); + let converted_rows = 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() + .collect::>(); + if let Some(started_at) = convert_started_at { + row_convert_time += started_at.elapsed(); + } + converted_rows } else { let mut rows = Vec::new(); let mut stream = result @@ -4313,10 +4360,17 @@ async fn execute_result_sets_with_text_protocol_on_conn( .map_err(|e| e.to_string())? .ok_or_else(|| "Empty result set stream".to_string())?; - while let Some(row) = stream.next().await { + loop { + let read_started_at = diagnostics_enabled.then(Instant::now); + let next_row = stream.next().await; + if let Some(started_at) = read_started_at { + row_read_time += started_at.elapsed(); + } + let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; if rows.len() < row_limit { let row_index = rows.len(); + let convert_started_at = diagnostics_enabled.then(Instant::now); let (values, srids, mut row_large_values) = mysql_row_to_json_with_srids_and_previews( &row, &mut spatial_columns, @@ -4327,6 +4381,9 @@ async fn execute_result_sets_with_text_protocol_on_conn( rows.push(values); spatial_values.push(srids); large_value_cells.append(&mut row_large_values); + if let Some(started_at) = convert_started_at { + row_convert_time += started_at.elapsed(); + } } else { truncated = true; } @@ -4362,6 +4419,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( } if results.is_empty() { + let finalize_started_at = diagnostics_enabled.then(Instant::now); let affected_rows = result.affected_rows(); let warnings = result.warnings(); let info = result.info().into_owned(); @@ -4382,6 +4440,16 @@ async fn execute_result_sets_with_text_protocol_on_conn( elasticsearch_raw_body: None, messages, })); + if diagnostics_enabled { + log::info!( + "[query][mysql-result] trace_id={} protocol=text-multi dispatch_ms={} row_read_ms={} row_convert_ms={} finalize_ms={} result_count=1 row_counts=[0] column_counts=[0] truncated_sets=0 preview_cells=0", + diagnostic_trace_id.unwrap_or("none"), + dispatch_ms, + row_read_time.as_millis(), + row_convert_time.as_millis(), + finalize_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()) + ); + } return Ok(results); } // Without per-set capture (no CLIENT_DEPRECATE_EOF) the per-iteration @@ -4392,6 +4460,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( let last_index = result_set_warnings.len() - 1; result_set_warnings[last_index] = result.warnings(); } + let finalize_started_at = diagnostics_enabled.then(Instant::now); drop(result); // SHOW WARNINGS only reports the last executed statement, so detailed @@ -4410,6 +4479,22 @@ async fn execute_result_sets_with_text_protocol_on_conn( } } + if diagnostics_enabled { + log::info!( + "[query][mysql-result] trace_id={} protocol=text-multi dispatch_ms={} row_read_ms={} row_convert_ms={} finalize_ms={} result_count={} row_counts={:?} column_counts={:?} truncated_sets={} preview_cells={}", + diagnostic_trace_id.unwrap_or("none"), + dispatch_ms, + row_read_time.as_millis(), + row_convert_time.as_millis(), + finalize_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()), + results.len(), + results.iter().map(|result| result.result.rows.len()).collect::>(), + results.iter().map(|result| result.result.columns.len()).collect::>(), + results.iter().filter(|result| result.result.truncated).count(), + results.iter().map(|result| result.large_value_cells.len()).sum::() + ); + } + Ok(results) } @@ -4431,9 +4516,13 @@ async fn execute_result_set_with_prepared_protocol_on_conn( row_limit: usize, max_result_bytes: Option, result_key_columns: &[String], + diagnostic_trace_id: Option<&str>, start: Instant, ) -> Result { + let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); + let dispatch_started_at = diagnostics_enabled.then(Instant::now); let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?; + let dispatch_ms = dispatch_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()); let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); let mut spatial_columns = mysql_spatial_column_builder(result.columns_ref()); @@ -4445,19 +4534,28 @@ async fn execute_result_set_with_prepared_protocol_on_conn( let mut spatial_values: Vec>> = Vec::new(); let mut large_value_cells = Vec::new(); let mut truncated = false; + let mut row_read_time = Duration::ZERO; + let mut row_convert_time = Duration::ZERO; let mut stream = result .stream::() .await .map_err(|e| e.to_string())? .ok_or_else(|| "Empty result set stream".to_string())?; - while let Some(row) = stream.next().await { + loop { + let read_started_at = diagnostics_enabled.then(Instant::now); + let next_row = stream.next().await; + if let Some(started_at) = read_started_at { + row_read_time += started_at.elapsed(); + } + let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; if result_rows.len() >= row_limit { truncated = true; break; } let row_index = result_rows.len(); + let convert_started_at = diagnostics_enabled.then(Instant::now); let (values, srids, mut row_large_values) = mysql_row_to_json_with_srids_and_previews( &row, &mut spatial_columns, @@ -4468,6 +4566,9 @@ async fn execute_result_set_with_prepared_protocol_on_conn( result_rows.push(values); spatial_values.push(srids); large_value_cells.append(&mut row_large_values); + if let Some(started_at) = convert_started_at { + row_convert_time += started_at.elapsed(); + } } drop(stream); @@ -4475,6 +4576,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( // terminator packet arrived, so `result.warnings()`/`result.info()` still // hold the previous statement's values; skip capture instead of reporting // stale messages. + let finalize_started_at = diagnostics_enabled.then(Instant::now); let messages = if truncated { drop(result); Vec::new() @@ -4484,6 +4586,20 @@ async fn execute_result_set_with_prepared_protocol_on_conn( drop(result); collect_mysql_server_messages(conn, warnings, &info).await }; + if diagnostics_enabled { + log::info!( + "[query][mysql-result] trace_id={} protocol=prepared dispatch_ms={} row_read_ms={} row_convert_ms={} finalize_ms={} rows={} columns={} truncated={} preview_cells={}", + diagnostic_trace_id.unwrap_or("none"), + dispatch_ms, + row_read_time.as_millis(), + row_convert_time.as_millis(), + finalize_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()), + result_rows.len(), + columns.len(), + truncated, + large_value_cells.len() + ); + } Ok(MySqlQueryResult { result: QueryResult { @@ -4732,7 +4848,9 @@ pub async fn execute_query_on_conn_with_max_rows( max_rows: Option, dialect: MySqlQueryDialect, ) -> Result { - execute_query_on_conn_with_limits(conn, sql, bare, max_rows, None, &[], dialect).await.map(|result| result.result) + execute_query_on_conn_with_limits(conn, sql, bare, max_rows, None, &[], dialect, None) + .await + .map(|result| result.result) } pub async fn execute_query_on_conn_with_limits( @@ -4743,6 +4861,7 @@ pub async fn execute_query_on_conn_with_limits( max_result_bytes: Option, result_key_columns: &[String], dialect: MySqlQueryDialect, + diagnostic_trace_id: Option<&str>, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); @@ -4756,6 +4875,7 @@ pub async fn execute_query_on_conn_with_limits( max_rows, max_result_bytes, result_key_columns, + diagnostic_trace_id, start, ) .await @@ -4766,6 +4886,7 @@ pub async fn execute_query_on_conn_with_limits( row_limit, max_result_bytes, result_key_columns, + diagnostic_trace_id, start, ) .await @@ -4779,6 +4900,7 @@ pub async fn execute_query_on_conn_with_limits( max_rows, max_result_bytes, result_key_columns, + diagnostic_trace_id, start, ) .await @@ -4831,7 +4953,7 @@ pub async fn execute_query_results_on_conn_with_max_rows( max_rows: Option, dialect: MySqlQueryDialect, ) -> Result, String> { - execute_query_results_on_conn_with_limits(conn, sql, bare, max_rows, None, &[], dialect) + execute_query_results_on_conn_with_limits(conn, sql, bare, max_rows, None, &[], dialect, None) .await .map(|results| results.into_iter().map(|result| result.result).collect()) } @@ -4844,6 +4966,7 @@ pub async fn execute_query_results_on_conn_with_limits( max_result_bytes: Option, result_key_columns: &[String], dialect: MySqlQueryDialect, + diagnostic_trace_id: Option<&str>, ) -> Result, String> { if is_result_set_query(sql, dialect) && (bare || prefers_text_protocol_query(sql, dialect)) { let start = Instant::now(); @@ -4854,13 +4977,23 @@ pub async fn execute_query_results_on_conn_with_limits( max_rows, max_result_bytes, result_key_columns, + diagnostic_trace_id, start, ) .await } else { - execute_query_on_conn_with_limits(conn, sql, bare, max_rows, max_result_bytes, result_key_columns, dialect) - .await - .map(|result| vec![result]) + execute_query_on_conn_with_limits( + conn, + sql, + bare, + max_rows, + max_result_bytes, + result_key_columns, + dialect, + diagnostic_trace_id, + ) + .await + .map(|result| vec![result]) } } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 1b6d7c6d46..c03c5fd5bd 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -6,7 +6,7 @@ use sqlparser::ast::{ }; use sqlparser::dialect::{GenericDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::ops::ControlFlow; use std::sync::{ @@ -527,13 +527,23 @@ fn merge_large_value_cells( mut driver_cells: Vec, server_cells: Vec, ) -> Vec { + if driver_cells.is_empty() { + return server_cells; + } + if server_cells.is_empty() { + return driver_cells; + } + let mut driver_indexes = driver_cells + .iter() + .enumerate() + .map(|(index, cell)| ((cell.row_index, cell.column_index), index)) + .collect::>(); for server_cell in server_cells { - if let Some(driver_cell) = driver_cells - .iter_mut() - .find(|cell| cell.row_index == server_cell.row_index && cell.column_index == server_cell.column_index) - { - *driver_cell = server_cell; + let key = (server_cell.row_index, server_cell.column_index); + if let Some(index) = driver_indexes.get(&key).copied() { + driver_cells[index] = server_cell; } else { + driver_indexes.insert(key, driver_cells.len()); driver_cells.push(server_cell); } } @@ -1685,6 +1695,7 @@ async fn do_execute_typed( max_result_bytes, &options.result_key_columns, mysql_dialect, + options.execution_id.as_deref(), ), ) .await @@ -2732,6 +2743,7 @@ struct MysqlBatchConnection<'a> { result_key_columns: &'a [String], table_data_preview: bool, dialect: db::mysql::MySqlQueryDialect, + diagnostic_trace_id: Option<&'a str>, } impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> { @@ -2751,6 +2763,7 @@ impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> { self.max_result_bytes, self.result_key_columns, self.dialect, + self.diagnostic_trace_id, ), ) .await @@ -2948,6 +2961,8 @@ async fn execute_multi_mysql( options: QueryExecutionOptions, progress: Option<&ExecuteMultiProgressCallback>, ) -> Result, String> { + let trace_id = options.execution_id.as_deref().unwrap_or("none"); + let total_started_at = std::time::Instant::now(); let query_timeout = resolve_query_timeout(options.timeout_secs); let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let bare = mode == crate::connection::MysqlMode::Bare; @@ -2955,6 +2970,7 @@ async fn execute_multi_mysql( let max_result_bytes = options.max_result_bytes.filter(|value| *value > 0); let pipeline_non_result_statements = mysql_non_result_pipeline_enabled(statements.len(), options.continue_on_error, mode); + let checkout_started_at = std::time::Instant::now(); let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( pool, operation_budget.checkout_timeout, @@ -2973,13 +2989,16 @@ async fn execute_multi_mysql( return Ok(vec![ExecuteMultiResult::execution_error(error_query_result(err))]); } }; + let checkout_ms = checkout_started_at.elapsed().as_millis(); apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?; + let catalog_started_at = std::time::Instant::now(); wait_for_result_opt( cancel_token.clone(), query_timeout, db::mysql::apply_catalog_database_context(&mut conn, catalog_dialect, options.catalog.as_deref(), database), ) .await?; + let catalog_ms = catalog_started_at.elapsed().as_millis(); let pipeline_non_result_max_bytes = if pipeline_non_result_statements { db::mysql::max_allowed_packet_on_conn(&mut conn) .await @@ -3000,7 +3019,9 @@ async fn execute_multi_mysql( result_key_columns: &options.result_key_columns, table_data_preview: options.table_data_preview, dialect, + diagnostic_trace_id: options.execution_id.as_deref(), }; + let statements_started_at = std::time::Instant::now(); let (results, error_action) = execute_mysql_batch_statements( &mut executor, statements, @@ -3012,8 +3033,20 @@ async fn execute_multi_mysql( progress, ) .await; + let statements_ms = statements_started_at.elapsed().as_millis(); drop(executor); + log::info!( + "[query][mysql-batch] trace_id={} checkout_ms={} catalog_ms={} statements_ms={} total_ms={} result_count={} row_counts={:?}", + trace_id, + checkout_ms, + catalog_ms, + statements_ms, + total_started_at.elapsed().as_millis(), + results.len(), + results.iter().map(|result| result.result.rows.len()).collect::>() + ); + if matches!(error_action, Some(PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)) { state.remove_pool_by_key(pool_key).await; } @@ -6064,6 +6097,53 @@ for line in sys.stdin: ); } + #[test] + fn large_value_cell_merge_replaces_driver_entries_and_appends_server_entries() { + let driver_cells = vec![ + db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 10 }, + db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }, + ]; + let server_cells = vec![ + db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 100 }, + db::LargeValueCell { row_index: 5, column_index: 6, original_bytes: 200 }, + ]; + + assert_eq!( + merge_large_value_cells(driver_cells, server_cells), + vec![ + db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 100 }, + db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }, + db::LargeValueCell { row_index: 5, column_index: 6, original_bytes: 200 }, + ] + ); + } + + #[test] + fn large_value_cell_merge_preserves_single_source_inputs() { + let driver_cell = db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 10 }; + let server_cell = db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }; + + assert_eq!(merge_large_value_cells(vec![driver_cell.clone()], Vec::new()), vec![driver_cell]); + assert_eq!(merge_large_value_cells(Vec::new(), vec![server_cell.clone()]), vec![server_cell]); + } + + #[test] + fn large_value_cell_merge_handles_large_disjoint_inputs() { + const CELL_COUNT: usize = 100_000; + let driver_cells = (0..CELL_COUNT) + .map(|row_index| db::LargeValueCell { row_index, column_index: 0, original_bytes: 10 }) + .collect(); + let server_cells = (0..CELL_COUNT) + .map(|row_index| db::LargeValueCell { row_index, column_index: 1, original_bytes: 20 }) + .collect(); + + let merged = merge_large_value_cells(driver_cells, server_cells); + + assert_eq!(merged.len(), CELL_COUNT * 2); + assert_eq!(merged[CELL_COUNT].row_index, 0); + assert_eq!(merged[CELL_COUNT * 2 - 1].row_index, CELL_COUNT - 1); + } + #[test] fn sqlserver_batch_results_do_not_claim_statement_indexes() { assert_eq!(split_sql_batches("SELECT 1; SELECT 2;").len(), 1); diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 8c9f4a18c1..24f44c2b4b 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -1,7 +1,8 @@ use std::sync::Arc; use axum::extract::State; -use axum::http::HeaderMap; +use axum::http::{header, HeaderMap, HeaderValue}; +use axum::response::Response; use axum::Json; use serde::Deserialize; @@ -389,7 +390,7 @@ pub async fn execute_multi( State(state): State>, headers: HeaderMap, Json(req): Json, -) -> Result>, AppError> { +) -> Result { let allow_database_switch = req.client_session_id.as_deref().is_some_and(|id| !id.trim().is_empty()); super::mcp_policy::ensure_sql(&state, &headers, &req.connection_id, &req.database, &req.sql, allow_database_switch) .await?; @@ -403,6 +404,7 @@ pub async fn execute_multi( tracing::debug!(connection_id = %req.connection_id, "execute_multi"); + let core_started_at = std::time::Instant::now(); let result = dbx_core::query::execute_multi_core_with_options_for_client_typed( &state.app, &req.connection_id, @@ -430,15 +432,30 @@ pub async fn execute_multi( ) .await .map_err(|error| AppError::from(error.into_backend_error()))?; + let core_ms = core_started_at.elapsed().as_millis(); drop(registered); - Ok(execute_multi_response(result)) + execute_multi_response(result, core_ms) } fn execute_multi_response( result: Vec, -) -> Json> { - Json(result) + core_ms: u128, +) -> Result { + let serialize_started_at = std::time::Instant::now(); + let body = serde_json::to_vec(&result).map_err(|error| AppError::internal(error.to_string()))?; + let serialize_ms = serialize_started_at.elapsed().as_millis(); + let mut response = Response::new(axum::body::Body::from(body)); + response.headers_mut().insert(header::CONTENT_TYPE, HeaderValue::from_static("application/json")); + response.headers_mut().insert( + "x-dbx-core-ms", + HeaderValue::from_str(&core_ms.to_string()).map_err(|error| AppError::internal(error.to_string()))?, + ); + response.headers_mut().insert( + "x-dbx-serialize-ms", + HeaderValue::from_str(&serialize_ms.to_string()).map_err(|error| AppError::internal(error.to_string()))?, + ); + Ok(response) } pub async fn execute_batch( @@ -1054,8 +1071,8 @@ mod tests { assert_eq!(log.metadata["blocked"], "destructive_confirmation_required"); } - #[test] - fn execute_multi_response_preserves_nested_original_error_detail() { + #[tokio::test] + async fn execute_multi_response_preserves_nested_original_error_detail() { let result = dbx_core::query::ExecuteMultiResult { result: dbx_core::db::QueryResult { columns: vec!["Error".to_string()], @@ -1081,7 +1098,14 @@ mod tests { server_message: false, }; - let payload = serde_json::to_value(execute_multi_response(vec![result]).0).unwrap(); + let response = execute_multi_response(vec![result], 17).unwrap(); + assert_eq!(response.headers()["x-dbx-core-ms"], "17"); + assert!(response.headers()["x-dbx-serialize-ms"].to_str().unwrap().parse::().is_ok()); + assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json"); + let body = response.into_body(); + let payload = + serde_json::from_slice::(&axum::body::to_bytes(body, usize::MAX).await.unwrap()) + .unwrap(); assert_eq!(payload[0]["statement_index"], 1); assert_eq!(payload[0]["error"]["code"], "DBX-JDBC-4001");