feat(query): show database server messages in query results

This commit is contained in:
Eddy Lei authored and GitHub committed 2026-08-07 15:18:32 +08:00
1 parent 4c7ddd266b
commit a5e46fd659
45 files changed
+1401 -51

No files matched your search

+58 -3
View File
@@ -10,7 +10,7 @@ use crate::query::QueryExecutionOptions;
use crate::query_execution_sql::{build_explain_sql, supports_explain_plan, supports_sql_query, ExplainSqlOptions};
use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions};
use crate::sql_risk::SqlRisk;
use crate::types::QueryResult;
use crate::types::{QueryMessage, QueryResult};
/// Maximum number of tables returned by list_tables tool.
const LIST_TABLES_LIMIT: usize = 200;
@@ -734,7 +734,10 @@ fn format_query_result_as_text(result: &QueryResult, limit: usize) -> Result<Str
// A result without columns is a command result, not an empty result set.
// This is how drivers represent DML that does not use RETURNING.
if result.columns.is_empty() {
return Ok(format!("Query executed. {} row(s) affected.", result.affected_rows));
return Ok(append_server_messages(
format!("Query executed. {} row(s) affected.", result.affected_rows),
&result.messages,
));
}
let mut lines = Vec::new();
@@ -774,7 +777,20 @@ fn format_query_result_as_text(result: &QueryResult, limit: usize) -> Result<Str
// Stats line
lines.push(format!("({} rows, {}ms)", result.rows.len(), result.execution_time_ms));
Ok(lines.join("\n"))
Ok(append_server_messages(lines.join("\n"), &result.messages))
}
/// Append server messages in the same style as the MCP `format_query_result`
/// renderer: a `Server messages:` section with `- SEVERITY: message` lines.
fn append_server_messages(mut output: String, messages: &[QueryMessage]) -> String {
if messages.is_empty() {
return output;
}
output.push_str("\n\nServer messages:");
for message in messages {
output.push_str(&format!("\n- {}", message.format_line()));
}
output
}
/// Get sample data from a table via the get_sample_data tool.
@@ -1363,6 +1379,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -1393,6 +1410,44 @@ for line in sys.stdin:
);
}
#[test]
fn query_result_formatter_appends_server_messages() {
let mut dml = query_result(vec![], vec![], 2);
dml.messages = vec![
QueryMessage {
severity: "notice".to_string(),
message: "hello world".to_string(),
code: Some("00000".to_string()),
detail: None,
hint: Some("use a table".to_string()),
},
QueryMessage {
severity: "WARNING".to_string(),
message: "careful".to_string(),
code: None,
detail: None,
hint: None,
},
];
assert_eq!(
format_query_result_as_text(&dml, 50).unwrap(),
"Query executed. 2 row(s) affected.\n\nServer messages:\n- NOTICE: hello world (code: 00000, hint: use a table)\n- WARNING: careful"
);
let mut result = query_result(vec!["id"], vec![vec![serde_json::json!(1)]], 0);
result.messages = vec![QueryMessage {
severity: "INFO".to_string(),
message: "print output".to_string(),
code: None,
detail: None,
hint: None,
}];
assert_eq!(
format_query_result_as_text(&result, 50).unwrap(),
"| id |\n|---|\n| 1 |\n(1 rows, 1ms)\n\nServer messages:\n- INFO: print output"
);
}
#[test]
fn sample_data_sql_uses_database_identifier_and_limit_syntax() {
assert_eq!(
+1
View File
@@ -5215,6 +5215,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
assert_eq!(gaussdb_identifier_quote_from_query_result(&result).as_deref(), Some("`"));
@@ -567,6 +567,7 @@ fn limited_query_result(result: ChJsonResult, execution_time_ms: u128, max_rows:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -738,6 +739,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -358,6 +358,7 @@ fn query_result(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -1298,6 +1298,7 @@ fn parse_elasticsearch_response_with_sql_parser(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
} else {
Ok(json_response_result(status, &body, start))
@@ -1323,6 +1324,7 @@ fn parse_elasticsearch_response_with_sql_parser(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
} else {
Ok(json_response_result(status, &body, start))
@@ -1494,6 +1496,7 @@ fn raw_json_response_result(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -1560,6 +1563,7 @@ fn parse_elasticsearch_rest_response_with_sql_parser(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -1996,6 +2000,7 @@ pub(crate) fn parse_tabular_sql_response(
session_id: body.get("cursor").and_then(|cursor| cursor.as_str()).map(str::to_string),
has_more: body.get("cursor").and_then(|cursor| cursor.as_str()).is_some(),
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -485,6 +485,7 @@ pub async fn execute_query(client: &InfluxdbClient, database: &str, sql: &str) -
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}),
None => Ok(QueryResult {
columns: vec![],
@@ -499,6 +500,7 @@ pub async fn execute_query(client: &InfluxdbClient, database: &str, sql: &str) -
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}),
}
}
@@ -645,6 +647,7 @@ fn parse_flux_csv(text: &str, start: Instant) -> Result<QueryResult, String> {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
+259 -4
View File
@@ -21,7 +21,8 @@ use crate::sql::{starts_with_executable_sql_keyword, starts_with_executable_sql_
use crate::types::{
ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode,
CompletionAssistantObjectKind, CompletionAssistantRequest, CompletionAssistantResponse, DatabaseInfo,
ForeignKeyInfo, IndexInfo, ObjectInfo, ObjectStatistics, QueryResult, SpatialColumnBuilder, TableInfo, TriggerInfo,
ForeignKeyInfo, IndexInfo, ObjectInfo, ObjectStatistics, QueryMessage, QueryResult, SpatialColumnBuilder,
TableInfo, TriggerInfo,
};
use super::file_validator::validate_file_path;
@@ -3772,6 +3773,80 @@ async fn ping_conn_with_timeout_and_cancel(
}
}
/// Maps `SHOW WARNINGS` rows (`Level`, `Code`, `Message`) to query messages.
fn mysql_warning_rows_to_messages(rows: Vec<(String, u16, String)>) -> Vec<QueryMessage> {
rows.into_iter()
.map(|(level, code, message)| QueryMessage {
severity: level,
code: Some(code.to_string()),
message,
detail: None,
hint: None,
})
.collect()
}
/// Maps a non-empty OK-packet info string (e.g. `Records: 3 Duplicates: 0
/// Warnings: 1`) to an INFO query message.
fn mysql_info_message(info: &str) -> Option<QueryMessage> {
let info = info.trim();
if info.is_empty() {
return None;
}
Some(QueryMessage { severity: "INFO".to_string(), code: None, message: info.to_string(), detail: None, hint: None })
}
fn mysql_warnings_fallback_message(warnings: u16) -> QueryMessage {
QueryMessage {
severity: "Warning".to_string(),
code: None,
message: format!("{warnings} warning(s)"),
detail: None,
hint: None,
}
}
/// Builds the message list from an OK-packet info string plus the outcome of a
/// `SHOW WARNINGS` query: `None` when the query failed (MySQL-compatible
/// proxies such as Doris/StarRocks may not support it), `Some(rows)` with its
/// rows otherwise. An empty successful result still falls back to the
/// count-only message so a nonzero warning count is never silently dropped.
fn mysql_server_messages_from_warnings(
info: &str,
warnings: u16,
warning_rows: Option<Vec<(String, u16, String)>>,
) -> Vec<QueryMessage> {
let mut messages: Vec<QueryMessage> = mysql_info_message(info).into_iter().collect();
if warnings == 0 {
return messages;
}
match warning_rows {
Some(rows) if !rows.is_empty() => messages.extend(mysql_warning_rows_to_messages(rows)),
_ => messages.push(mysql_warnings_fallback_message(warnings)),
}
messages
}
/// Best-effort collection of server messages for a finished statement: the
/// OK-packet info string plus `SHOW WARNINGS` output when the server reported
/// warnings. `SHOW WARNINGS` runs on the same connection; all errors are
/// swallowed (MySQL-compatible proxies such as Doris/StarRocks may not support
/// it) and fall back to a count-only message.
async fn collect_mysql_server_messages(conn: &mut mysql_async::Conn, warnings: u16, info: &str) -> Vec<QueryMessage> {
let warning_rows = if warnings == 0 {
None
} else {
match conn.query_iter("SHOW WARNINGS").await {
Ok(result) => match result.try_collect_and_drop::<(String, u16, String)>().await {
Ok(rows) => rows.into_iter().collect::<Result<Vec<_>, _>>().ok(),
Err(_) => None,
},
Err(_) => None,
}
};
mysql_server_messages_from_warnings(info, warnings, warning_rows)
}
async fn execute_result_set_with_text_protocol_on_conn(
conn: &mut mysql_async::Conn,
sql: &str,
@@ -3781,6 +3856,11 @@ async fn execute_result_set_with_text_protocol_on_conn(
) -> Result<QueryResult, String> {
let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
if !advance_to_result_set_with_columns(&mut result).await? {
let affected_rows = result.affected_rows();
let warnings = result.warnings();
let info = result.info().into_owned();
drop(result);
let messages = collect_mysql_server_messages(conn, warnings, &info).await;
return Ok(QueryResult {
columns: vec![],
column_types: Vec::new(),
@@ -3788,12 +3868,13 @@ async fn execute_result_set_with_text_protocol_on_conn(
spatial_columns: vec![],
spatial_values: vec![],
rows: vec![],
affected_rows: result.affected_rows(),
affected_rows,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
});
}
let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
@@ -3802,6 +3883,11 @@ async fn execute_result_set_with_text_protocol_on_conn(
if should_collect_text_result_set(sql, row_limit, max_rows) {
let rows: Vec<mysql_async::Row> = result.collect_and_drop().await.map_err(|e| e.to_string())?;
// collect_and_drop consumed the result; warnings/info now reflect the
// trailing EOF/OK packet of the finished result set.
let warnings = conn.get_warnings();
let info = conn.info().into_owned();
let messages = collect_mysql_server_messages(conn, warnings, &info).await;
let truncated = rows.len() > row_limit;
let mut spatial_values = Vec::new();
let result_rows = rows
@@ -3827,6 +3913,7 @@ async fn execute_result_set_with_text_protocol_on_conn(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
});
}
@@ -3849,6 +3936,21 @@ async fn execute_result_set_with_text_protocol_on_conn(
result_rows.push(values);
spatial_values.push(srids);
}
drop(stream);
// A truncated stream broke out of the row loop before this result set's
// terminator packet arrived, so `result.warnings()`/`result.info()` still
// hold the previous statement's values; skip capture instead of reporting
// stale messages.
let messages = if truncated {
drop(result);
Vec::new()
} else {
let warnings = result.warnings();
let info = result.info().into_owned();
drop(result);
collect_mysql_server_messages(conn, warnings, &info).await
};
Ok(QueryResult {
columns,
@@ -3863,6 +3965,7 @@ async fn execute_result_set_with_text_protocol_on_conn(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
})
}
@@ -3873,8 +3976,20 @@ async fn execute_result_sets_with_text_protocol_on_conn(
max_rows: Option<usize>,
start: Instant,
) -> Result<Vec<QueryResult>, String> {
// Per-set warnings/info read after each `collect()` are only accurate when
// the connection negotiated CLIENT_DEPRECATE_EOF: with legacy EOF packets
// the next result set's column-definition EOF clobbers the previous OK
// state, and a following column-less statement's OK packet would be
// misattributed to the wrong set. mysql_async does not expose the
// negotiated capability publicly, so gate on the client-side
// `deprecate_eof` option requested when the pool was built (DBX disables
// it automatically when a proxy only speaks legacy EOF). When disabled,
// 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 mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
let mut results = Vec::new();
let mut results: Vec<QueryResult> = Vec::new();
let mut result_set_warnings: Vec<u16> = 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();
@@ -3915,6 +4030,13 @@ async fn execute_result_sets_with_text_protocol_on_conn(
rows
};
let warnings = if capture_per_set_messages { result.warnings() } else { 0 };
let messages: Vec<QueryMessage> = if capture_per_set_messages {
mysql_info_message(&result.info()).into_iter().collect()
} else {
Vec::new()
};
result_set_warnings.push(warnings);
results.push(QueryResult {
columns,
column_types,
@@ -3928,10 +4050,16 @@ async fn execute_result_sets_with_text_protocol_on_conn(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
});
}
if results.is_empty() {
let affected_rows = result.affected_rows();
let warnings = result.warnings();
let info = result.info().into_owned();
drop(result);
let messages = collect_mysql_server_messages(conn, warnings, &info).await;
results.push(QueryResult {
columns: vec![],
column_types: Vec::new(),
@@ -3939,13 +4067,40 @@ async fn execute_result_sets_with_text_protocol_on_conn(
spatial_columns: vec![],
spatial_values: vec![],
rows: vec![],
affected_rows: result.affected_rows(),
affected_rows,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
});
return Ok(results);
}
// Without per-set capture (no CLIENT_DEPRECATE_EOF) the per-iteration
// counts were skipped above; the connection state after the loop still
// reflects the last statement's terminator, so use its warnings count for
// the final SHOW WARNINGS attachment on the last result set.
if !capture_per_set_messages {
let last_index = result_set_warnings.len() - 1;
result_set_warnings[last_index] = result.warnings();
}
drop(result);
// SHOW WARNINGS only reports the last executed statement, so detailed
// warning rows can only be attached to the last result set; earlier sets
// with warnings get a count-only fallback message.
let last_index = results.len() - 1;
for (index, warnings) in result_set_warnings.iter().copied().enumerate() {
if warnings == 0 {
continue;
}
if index == last_index {
let mut messages = collect_mysql_server_messages(conn, warnings, "").await;
results[index].messages.append(&mut messages);
} else {
results[index].messages.push(mysql_warnings_fallback_message(warnings));
}
}
Ok(results)
@@ -3993,6 +4148,21 @@ async fn execute_result_set_with_prepared_protocol_on_conn(
result_rows.push(values);
spatial_values.push(srids);
}
drop(stream);
// A truncated stream broke out of the row loop before this result set's
// terminator packet arrived, so `result.warnings()`/`result.info()` still
// hold the previous statement's values; skip capture instead of reporting
// stale messages.
let messages = if truncated {
drop(result);
Vec::new()
} else {
let warnings = result.warnings();
let info = result.info().into_owned();
drop(result);
collect_mysql_server_messages(conn, warnings, &info).await
};
Ok(QueryResult {
columns,
@@ -4007,6 +4177,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
})
}
@@ -4258,7 +4429,13 @@ pub async fn execute_query_on_conn_with_max_rows(
}
};
let affected_rows = result.affected_rows();
let warnings = result.warnings();
let info = result.info().into_owned();
let drop_result = result.drop_result().await;
// Collect server messages before restoring the timestamp defaults: the
// restore issues `SET SESSION explicit_defaults_for_timestamp`, which
// clears the diagnostics area and would leave SHOW WARNINGS empty.
let messages = collect_mysql_server_messages(conn, warnings, &info).await;
restore_explicit_timestamp_defaults_for_query(conn, previous_explicit_timestamp_defaults).await;
drop_result.map_err(|e| e.to_string())?;
@@ -4275,6 +4452,7 @@ pub async fn execute_query_on_conn_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages,
})
}
}
@@ -5011,6 +5189,83 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res
mod tests {
use super::*;
#[test]
fn mysql_info_message_maps_non_empty_info_strings() {
let message = mysql_info_message("Records: 3 Duplicates: 0 Warnings: 1").unwrap();
assert_eq!(message.severity, "INFO");
assert_eq!(message.message, "Records: 3 Duplicates: 0 Warnings: 1");
assert_eq!(message.code, None);
assert_eq!(message.detail, None);
assert_eq!(message.hint, None);
assert!(mysql_info_message("").is_none());
assert!(mysql_info_message(" ").is_none());
}
#[test]
fn mysql_warning_rows_to_messages_maps_levels_and_codes() {
let messages = mysql_warning_rows_to_messages(vec![
("Warning".to_string(), 1265, "Data truncated for column 'a' at row 1".to_string()),
("Note".to_string(), 1051, "Unknown table 't'".to_string()),
]);
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].severity, "Warning");
assert_eq!(messages[0].code.as_deref(), Some("1265"));
assert_eq!(messages[0].message, "Data truncated for column 'a' at row 1");
assert_eq!(messages[0].detail, None);
assert_eq!(messages[0].hint, None);
assert_eq!(messages[1].severity, "Note");
assert_eq!(messages[1].code.as_deref(), Some("1051"));
assert!(mysql_warning_rows_to_messages(Vec::new()).is_empty());
}
#[test]
fn mysql_warnings_fallback_message_reports_count() {
let message = mysql_warnings_fallback_message(3);
assert_eq!(message.severity, "Warning");
assert_eq!(message.message, "3 warning(s)");
assert_eq!(message.code, None);
}
#[test]
fn mysql_server_messages_from_warnings_maps_info_and_warning_rows() {
let rows = vec![("Warning".to_string(), 1265, "Data truncated for column 'a' at row 1".to_string())];
let messages = mysql_server_messages_from_warnings("Records: 1 Duplicates: 0 Warnings: 1", 1, Some(rows));
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].severity, "INFO");
assert_eq!(messages[1].severity, "Warning");
assert_eq!(messages[1].code.as_deref(), Some("1265"));
}
#[test]
fn mysql_server_messages_from_warnings_falls_back_when_show_warnings_returns_no_rows() {
let messages = mysql_server_messages_from_warnings("", 2, Some(Vec::new()));
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].severity, "Warning");
assert_eq!(messages[0].message, "2 warning(s)");
}
#[test]
fn mysql_server_messages_from_warnings_falls_back_when_show_warnings_fails() {
let messages = mysql_server_messages_from_warnings("", 2, None);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].message, "2 warning(s)");
}
#[test]
fn mysql_server_messages_from_warnings_skips_warnings_when_count_is_zero() {
let messages = mysql_server_messages_from_warnings("Query OK", 0, None);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].severity, "INFO");
assert_eq!(messages[0].message, "Query OK");
}
#[test]
fn mysql_sql_statement_limit_reserves_packet_headroom() {
let packet_bytes = 64 * 1024 * 1024;
+303 -9
View File
@@ -15,13 +15,16 @@ use std::collections::{BTreeSet, HashMap};
use std::fs::File;
use std::future::Future;
use std::io::BufReader;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex, OnceLock, Weak};
use std::time::{Duration, Instant};
use tokio::task::JoinHandle;
use tokio_postgres::config::SslMode;
use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
use tokio_postgres::types::{FromSql, Kind, Type};
use tokio_postgres::{NoTls, Row, SimpleQueryMessage};
use tokio_postgres::{AsyncMessage, NoTls, Row, SimpleQueryMessage, Socket};
use tokio_util::sync::CancellationToken;
use super::file_validator::validate_file_path;
@@ -31,7 +34,8 @@ use crate::types::{
ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode,
CompletionAssistantObjectKind, CompletionAssistantRequest, CompletionAssistantResponse, DatabaseInfo,
DatabaseStorageInfo, ExtensionInfo, ForeignKeyInfo, FunctionInfo, IndexInfo, ObjectInfo, ObjectStatistics,
OwnerInfo, QueryResult, RuleInfo, SchemaInfo, SequenceInfo, SpatialColumnBuilder, TableInfo, TriggerInfo,
OwnerInfo, QueryMessage, QueryResult, RuleInfo, SchemaInfo, SequenceInfo, SpatialColumnBuilder, TableInfo,
TriggerInfo,
};
pub(crate) const GAUSSDB_COMPATIBILITY_SQL: &str =
@@ -1056,6 +1060,7 @@ async fn execute_select_prepared(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}))
}
@@ -1150,6 +1155,7 @@ async fn execute_select_text(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -1484,6 +1490,188 @@ async fn connect_with_local_timezone(url: &str, fallback_timeout: Duration, time
connect_with_optional_local_timezone(url, fallback_timeout, Some(timezone)).await
}
/// Identity of one physical backend connection for notice attribution:
/// (server address, server port, backend PID). The PID alone is not unique
/// across different servers, so the server's own address/port disambiguate
/// (`inet_server_addr()` is NULL for Unix sockets, hence the fallback).
type PostgresConnectionKey = (String, i32, i32);
const POSTGRES_CONNECTION_IDENTITY_SQL: &str = "SELECT pg_backend_pid(), \
COALESCE(host(inet_server_addr()), 'unix'), \
COALESCE(inet_server_port(), current_setting('port')::integer)";
/// Notice buffers for live connections, keyed by connection identity. Entries
/// are weak so they disappear once the pooled connection (and its driver
/// task) is dropped.
fn postgres_notice_buffers() -> &'static Mutex<HashMap<PostgresConnectionKey, Weak<Mutex<Vec<QueryMessage>>>>> {
static BUFFERS: OnceLock<Mutex<HashMap<PostgresConnectionKey, Weak<Mutex<Vec<QueryMessage>>>>>> = OnceLock::new();
BUFFERS.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Connection identity cache, keyed by the per-physical-connection statement
/// cache pointer (same identity pattern as `postgres_single_schema_clients`).
/// The `Weak` guards against pointer reuse: a new connection whose
/// `StatementCache` lands on a freed entry's address fails the `ptr_eq` check
/// and is treated as a miss. The value is `None` when the identity query
/// failed, so it is issued at most once per physical connection — it must
/// never be retried from `drain_postgres_notices`, which can run inside the
/// read-only transaction used for EXPLAIN, where a failing query would abort
/// the user's statement.
fn postgres_client_keys(
) -> &'static Mutex<HashMap<usize, (Weak<deadpool_postgres::StatementCache>, Option<PostgresConnectionKey>)>> {
static KEYS: OnceLock<
Mutex<HashMap<usize, (Weak<deadpool_postgres::StatementCache>, Option<PostgresConnectionKey>)>>,
> = OnceLock::new();
KEYS.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Establishes connections like deadpool's `ConfigConnectImpl`, but drives
/// each connection with a task that captures `NoticeResponse` messages
/// (`RAISE NOTICE`/`WARNING`, etc.) into a per-backend buffer instead of
/// discarding them. Query execution drains the buffer so notices are
/// attached to the `QueryResult` of the statement that raised them.
struct NoticeCapturingConnect<T>
where
T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
T::Stream: Sync + Send,
T::TlsConnect: Sync + Send,
<T::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
tls: T,
}
impl<T> deadpool_postgres::Connect for NoticeCapturingConnect<T>
where
T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
T::Stream: Sync + Send,
T::TlsConnect: Sync + Send,
<T::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
fn connect(
&self,
pg_config: &tokio_postgres::Config,
) -> Pin<
Box<dyn Future<Output = Result<(tokio_postgres::Client, JoinHandle<()>), tokio_postgres::Error>> + Send + '_>,
> {
let tls = self.tls.clone();
let pg_config = pg_config.clone();
Box::pin(async move {
let (client, mut connection) = pg_config.connect(tls).await?;
// No query can complete before the connection is being driven, so
// the notice buffer is handed to the driver task through a slot
// that is filled once the backend PID is known.
let notice_buffer = Arc::new(Mutex::new(None::<Arc<Mutex<Vec<QueryMessage>>>>));
let task_buffer = Arc::clone(&notice_buffer);
let conn_task = tokio::spawn(async move {
loop {
match std::future::poll_fn(|cx| connection.poll_message(cx)).await {
Some(Ok(AsyncMessage::Notice(error))) => {
let message = QueryMessage {
severity: error.severity().to_string(),
message: error.message().to_string(),
code: Some(error.code().code().to_string()),
detail: error.detail().map(str::to_string),
hint: error.hint().map(str::to_string),
};
let buffer = task_buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).clone();
match buffer {
Some(buffer) => {
buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).push(message);
}
None => {
log::info!("[postgres][notice] {}: {}", message.severity, message.message);
}
}
}
// LISTEN/NOTIFY messages are not surfaced anywhere.
Some(Ok(_)) => {}
Some(Err(err)) => {
log::warn!("[postgres] connection driver error: {err}");
break;
}
None => break,
}
}
});
// Best-effort: without the connection identity, notices cannot be
// attributed to query results on this connection and are logged
// by the driver task instead. Never fail the connection over this.
if let Ok(row) = client.query_one(POSTGRES_CONNECTION_IDENTITY_SQL, &[]).await {
let key: PostgresConnectionKey = (row.get(1), row.get(2), row.get(0));
let buffer = Arc::new(Mutex::new(Vec::new()));
let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
buffers.retain(|_, weak| weak.strong_count() > 0);
buffers.insert(key, Arc::downgrade(&buffer));
drop(buffers);
*notice_buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(buffer);
}
Ok((client, conn_task))
})
}
}
/// Cached identity lookup. `Some(Some(key))` = resolved, `Some(None)` =
/// identity query failed before (do not retry), `None` = never seen (or the
/// entry belonged to a dropped connection whose address was reused).
fn cached_postgres_client_key(client: &deadpool_postgres::Client) -> Option<Option<PostgresConnectionKey>> {
let cache_key = Arc::as_ptr(&client.statement_cache) as usize;
let mut keys = postgres_client_keys().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
match keys.get(&cache_key) {
Some((cached, key)) if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, &client.statement_cache)) => {
Some(key.clone())
}
Some(_) => {
keys.remove(&cache_key);
None
}
None => None,
}
}
/// Best-effort identity resolution. Failures are cached as `None` so the
/// identity query runs at most once per physical connection.
async fn resolve_postgres_client_key(client: &deadpool_postgres::Client) -> Option<PostgresConnectionKey> {
if let Some(key) = cached_postgres_client_key(client) {
return key;
}
let key: Option<PostgresConnectionKey> = client
.query_one(POSTGRES_CONNECTION_IDENTITY_SQL, &[])
.await
.ok()
.map(|row| (row.get(1), row.get(2), row.get(0)));
let mut keys = postgres_client_keys().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
keys.retain(|_, (cached, _)| cached.strong_count() > 0);
let cache_key = Arc::as_ptr(&client.statement_cache) as usize;
keys.insert(cache_key, (Arc::downgrade(&client.statement_cache), key.clone()));
key
}
async fn postgres_client_key(client: &deadpool_postgres::Client) -> Option<PostgresConnectionKey> {
resolve_postgres_client_key(client).await
}
fn take_notices_for_key(key: &PostgresConnectionKey) -> Vec<QueryMessage> {
let buffer = {
let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
buffers.retain(|_, weak| weak.strong_count() > 0);
buffers.get(key).and_then(Weak::upgrade)
};
let Some(buffer) = buffer else {
return Vec::new();
};
let notices = std::mem::take(&mut *buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()));
notices
}
async fn drain_postgres_notices(client: &deadpool_postgres::Client) -> Vec<QueryMessage> {
match postgres_client_key(client).await {
Some(key) => take_notices_for_key(&key),
None => Vec::new(),
}
}
async fn connect_with_optional_local_timezone(
url: &str,
fallback_timeout: Duration,
@@ -1512,9 +1700,9 @@ async fn connect_with_optional_local_timezone(
postgres_url.accepts_invalid_certs,
postgres_url.verifies_hostname,
)?;
let mgr = deadpool_postgres::Manager::from_config(
let mgr = deadpool_postgres::Manager::from_connect(
pg_config.clone(),
tokio_postgres_rustls::MakeRustlsConnect::new(tls_config),
NoticeCapturingConnect { tls: tokio_postgres_rustls::MakeRustlsConnect::new(tls_config) },
mgr_config,
);
let pool = Pool::builder(mgr)
@@ -3341,6 +3529,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -3890,12 +4079,18 @@ pub async fn checkout_postgres_client(
},
None => get_future.await,
};
if result.is_ok() {
if let Ok(client) = &result {
log::debug!(
"[db:pool.checkout:done] elapsed_ms={} timeout_ms={}",
start.elapsed().as_millis(),
checkout_timeout.as_millis()
);
// Resolve the notice-attribution identity here, outside of any
// transaction: a lazy first lookup from `drain_postgres_notices`
// could run inside the read-only EXPLAIN transaction, where a
// failing identity query would abort the user's statement. Cached
// after the first checkout of each physical connection.
let _ = resolve_postgres_client_key(client).await;
}
result
}
@@ -3960,16 +4155,19 @@ async fn execute_query_with_max_rows_inner(
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if postgres_statement_returns_rows(sql) {
// Discard stale notices from infrastructure statements (e.g. the timezone
// SET issued at connect time) so only messages raised by this statement
// are attached to its result.
let _ = drain_postgres_notices(client).await;
let result = if postgres_statement_returns_rows(sql) {
if prefer_text_protocol {
execute_select_text(client, sql, start, row_limit, None, progress_clock.as_deref()).await
} else {
execute_select_query_with_progress(client, sql, start, row_limit, progress_clock.as_deref()).await
}
} else {
let affected = client.execute(sql, &[]).await.map_err(pg_error_to_string)?;
Ok(QueryResult {
client.execute(sql, &[]).await.map_err(pg_error_to_string).map(|affected| QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: Vec::new(),
@@ -3982,7 +4180,21 @@ async fn execute_query_with_max_rows_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
};
match result {
Ok(mut result) => {
result.messages = drain_postgres_notices(client).await;
Ok(result)
}
Err(error) => {
// Drop notices so an errored statement's messages cannot leak
// into the next query on this pooled connection.
let _ = drain_postgres_notices(client).await;
Err(error)
}
}
}
@@ -4734,6 +4946,87 @@ mod tests {
assert_eq!(gaussdb_identifier_quote_for_compatibility_mode(""), None);
}
fn test_query_message(message: &str) -> QueryMessage {
QueryMessage {
severity: "NOTICE".to_string(),
message: message.to_string(),
code: Some("00000".to_string()),
detail: None,
hint: None,
}
}
#[test]
fn take_notices_for_key_returns_buffered_notices_and_empties_buffer() {
let key = ("test-host".to_string(), 9_000_001, 9_000_001);
let buffer = Arc::new(Mutex::new(vec![test_query_message("first"), test_query_message("second")]));
postgres_notice_buffers()
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.insert(key.clone(), Arc::downgrade(&buffer));
let notices = take_notices_for_key(&key);
assert_eq!(notices.len(), 2);
assert_eq!(notices[0].message, "first");
assert_eq!(notices[1].message, "second");
assert_eq!(notices[0].severity, "NOTICE");
assert_eq!(notices[0].code.as_deref(), Some("00000"));
// The buffer was drained but stays registered while the connection lives.
assert!(take_notices_for_key(&key).is_empty());
buffer.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).push(test_query_message("third"));
let notices = take_notices_for_key(&key);
assert_eq!(notices.len(), 1);
assert_eq!(notices[0].message, "third");
}
#[test]
fn take_notices_for_key_prunes_dead_buffers_and_misses_return_empty() {
let live_key = ("test-host".to_string(), 9_000_002, 9_000_002);
let dead_key = ("test-host".to_string(), 9_000_003, 9_000_003);
let live = Arc::new(Mutex::new(vec![test_query_message("live")]));
let dead = Arc::new(Mutex::new(vec![test_query_message("dead")]));
{
let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
buffers.insert(live_key.clone(), Arc::downgrade(&live));
buffers.insert(dead_key.clone(), Arc::downgrade(&dead));
}
drop(dead);
assert!(take_notices_for_key(&("test-host".to_string(), 9_000_004, 9_000_004)).is_empty());
assert!(take_notices_for_key(&dead_key).is_empty());
let buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
assert!(!buffers.contains_key(&dead_key));
assert!(buffers.contains_key(&live_key));
drop(buffers);
let notices = take_notices_for_key(&live_key);
assert_eq!(notices.len(), 1);
assert_eq!(notices[0].message, "live");
}
#[test]
fn take_notices_for_key_distinguishes_same_pid_on_different_servers() {
// Backend PIDs collide across servers; the (address, port, pid) key
// keeps notice attribution separate.
let key_a = ("server-a".to_string(), 5432, 42);
let key_b = ("server-b".to_string(), 5432, 42);
let buffer_a = Arc::new(Mutex::new(vec![test_query_message("from-a")]));
let buffer_b = Arc::new(Mutex::new(vec![test_query_message("from-b")]));
{
let mut buffers = postgres_notice_buffers().lock().unwrap_or_else(|poisoned| poisoned.into_inner());
buffers.insert(key_a.clone(), Arc::downgrade(&buffer_a));
buffers.insert(key_b.clone(), Arc::downgrade(&buffer_b));
}
let notices_a = take_notices_for_key(&key_a);
let notices_b = take_notices_for_key(&key_b);
assert_eq!(notices_a.len(), 1);
assert_eq!(notices_a[0].message, "from-a");
assert_eq!(notices_b.len(), 1);
assert_eq!(notices_b[0].message, "from-b");
}
#[test]
fn postgres_json_arrays_decode_elements_without_jsonb_version_bytes() {
let json_raw = pg_array_binary(
@@ -6586,6 +6879,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let columns = redshift_columns_from_query_result(result);
+2
View File
@@ -306,6 +306,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -358,6 +359,7 @@ fn query_result_from_rqlite_result(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
+2
View File
@@ -2530,6 +2530,7 @@ fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option<usize
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
} else {
conn.execute_batch(sql).map_err(|e| e.to_string())?;
@@ -2546,6 +2547,7 @@ fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option<usize
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
})
+81 -5
View File
@@ -1,7 +1,7 @@
use crate::query::MAX_ROWS;
use crate::sql::starts_with_executable_sql_keyword;
use crate::types::{
ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, LinkedServerInfo, ObjectStatistics, QueryResult,
ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, LinkedServerInfo, ObjectStatistics, QueryMessage, QueryResult,
SpatialColumnBuilder, TableInfo, TriggerInfo,
};
use futures::{FutureExt, TryStreamExt};
@@ -486,15 +486,41 @@ where
(output, messages)
}
/// Map captured tiberius INFO-token texts to generic query messages. Tiberius
/// exposes only the message text in its tracing event, so severity is always
/// `INFO` and no code/detail/hint is available.
fn sqlserver_query_messages(messages: &[String]) -> Vec<QueryMessage> {
messages
.iter()
.map(|message| QueryMessage {
severity: "INFO".to_string(),
message: message.clone(),
code: None,
detail: None,
hint: None,
})
.collect()
}
fn query_result_with_server_messages(result: QueryResult, messages: Vec<String>) -> QueryResult {
query_result_with_server_messages_metadata(result, messages).result
}
fn query_result_with_server_messages_metadata(mut result: QueryResult, messages: Vec<String>) -> SqlServerBatchResult {
if messages.is_empty() || !result.columns.is_empty() || !result.rows.is_empty() {
if messages.is_empty() {
return SqlServerBatchResult { result, server_message: false };
}
if !result.columns.is_empty() || !result.rows.is_empty() {
// Tabular results keep their shape and additionally carry the
// messages; only the empty-result synthesis below leaves
// `messages` empty so consumers do not render the same text twice.
result.messages = sqlserver_query_messages(&messages);
return SqlServerBatchResult { result, server_message: false };
}
// Empty results synthesize the legacy single-"Message"-column grid; the
// text lives only there, so `messages` stays empty to avoid consumers
// rendering the same text twice.
result.columns = vec![SQLSERVER_MESSAGE_COLUMN.to_string()];
result.column_types = vec!["nvarchar".to_string()];
result.rows = messages.into_iter().map(|message| vec![serde_json::Value::String(message)]).collect();
@@ -520,6 +546,7 @@ fn server_messages_query_result(messages: Vec<String>, start: Instant) -> Option
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
messages,
))
@@ -585,6 +612,7 @@ async fn collect_first_result_limited(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -1399,6 +1427,7 @@ fn push_sqlserver_result_set(results: &mut Vec<QueryResult>, result: Option<SqlS
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
});
}
}
@@ -1456,6 +1485,7 @@ fn push_sqlserver_ordered_events(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
server_message: false,
});
@@ -2785,6 +2815,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
messages,
))
@@ -2829,6 +2860,7 @@ pub(crate) async fn execute_batch_with_max_rows_metadata(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
messages,
)]);
@@ -2913,6 +2945,7 @@ pub(crate) async fn execute_simple_batch_with_max_rows_metadata(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
server_message: false,
});
@@ -3182,12 +3215,12 @@ mod tests {
is_sqlserver_legacy_duplicate_probe_error, is_sqlserver_spatial_column, is_sqlserver_variant_column,
push_sqlserver_ordered_events, query_result_with_server_messages, query_result_with_server_messages_metadata,
requires_simple_query_batch, restore_sqlserver_legacy_probe_output_names,
restore_sqlserver_spatial_column_types, sqlserver_batch_can_use_execute, sqlserver_bulk_token_row,
sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql,
restore_sqlserver_spatial_column_types, server_messages_query_result, sqlserver_batch_can_use_execute,
sqlserver_bulk_token_row, sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql,
sqlserver_dml_output_returns_rows, sqlserver_done_trace_event, sqlserver_filter_definition_error,
sqlserver_hidden_schema_names, sqlserver_indexes_sql, sqlserver_legacy_indexes_sql, sqlserver_legacy_probe,
sqlserver_legacy_probe_with_nonce, sqlserver_legacy_wildcard_metadata_query, sqlserver_list_objects_sql,
sqlserver_list_schemas_sql, sqlserver_list_tables_sql, sqlserver_probe_explicit_alias,
sqlserver_list_schemas_sql, sqlserver_list_tables_sql, sqlserver_probe_explicit_alias, sqlserver_query_messages,
sqlserver_schema_name_predicate, sqlserver_spatial_marker, sqlserver_supports_session_database_switch,
sqlserver_table_comment_sql, sqlserver_triggers_sql, sqlserver_visible_object_predicate,
strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, SqlServerProbeOutputNameOverride,
@@ -3237,10 +3270,15 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let result = query_result_with_server_messages(empty, vec!["DBCC execution completed".to_string()]);
assert_eq!(result.columns, vec!["Message"]);
assert_eq!(result.column_types, vec!["nvarchar"]);
assert_eq!(result.rows, vec![vec![serde_json::json!("DBCC execution completed")]]);
// The synthesized grid already carries the text; `messages` stays
// empty so consumers do not render the same text twice.
assert!(result.messages.is_empty());
let message = query_result_with_server_messages_metadata(
QueryResult {
@@ -3256,6 +3294,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
},
vec!["PRINT output".to_string()],
);
@@ -3274,11 +3313,16 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let result = query_result_with_server_messages_metadata(select, vec!["informational".to_string()]);
assert!(!result.server_message);
assert_eq!(result.result.columns, vec!["id"]);
assert_eq!(result.result.rows, vec![vec![serde_json::json!(1)]]);
// Non-empty results keep their shape but still carry the server messages.
assert_eq!(result.result.messages.len(), 1);
assert_eq!(result.result.messages[0].severity, "INFO");
assert_eq!(result.result.messages[0].message, "informational");
}
#[test]
@@ -3421,6 +3465,36 @@ mod tests {
assert_eq!(results[1].result.affected_rows, 2);
}
#[test]
fn sqlserver_query_messages_map_info_severity() {
let messages = sqlserver_query_messages(&["first".to_string(), "second".to_string()]);
assert_eq!(messages.len(), 2);
for (message, expected) in messages.iter().zip(["first", "second"]) {
assert_eq!(message.severity, "INFO");
assert_eq!(message.message, expected);
assert!(message.code.is_none());
assert!(message.detail.is_none());
assert!(message.hint.is_none());
}
assert!(sqlserver_query_messages(&[]).is_empty());
}
#[test]
fn sqlserver_server_messages_query_result_synthesizes_grid() {
assert!(server_messages_query_result(vec![], Instant::now()).is_none());
let result = server_messages_query_result(vec!["print output".to_string()], Instant::now())
.expect("non-empty messages yield a result");
assert!(result.server_message);
assert_eq!(result.result.columns, vec!["Message"]);
assert_eq!(result.result.column_types, vec!["nvarchar"]);
assert_eq!(result.result.rows, vec![vec![serde_json::json!("print output")]]);
// The synthesized grid already carries the text; `messages` stays
// empty so consumers do not render the same text twice.
assert!(result.result.messages.is_empty());
}
#[test]
fn sqlserver_endpoint_splits_named_instance_hosts() {
assert_eq!(
@@ -3939,6 +4013,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
.unwrap();
@@ -4281,6 +4356,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
strip_dbx_sqlserver_row_number_column(&mut result, sql);
+4
View File
@@ -355,6 +355,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
});
}
@@ -376,6 +377,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
} else {
// Batch multiple statements into a single pipeline for transactional integrity
@@ -394,6 +396,7 @@ pub async fn execute_query_with_max_rows(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -574,6 +577,7 @@ fn query_result_from_turso_result(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
+2
View File
@@ -832,6 +832,7 @@ fn json_to_query_result(status: u16, body: Value, start: Instant) -> QueryResult
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -865,6 +866,7 @@ fn values_to_query_result(items: Vec<Value>, start: Instant) -> QueryResult {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -339,6 +339,7 @@ fn series_result_to_query_result(series: Vec<SeriesResult>, start: Instant) -> Q
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -375,6 +376,7 @@ fn simple_result(rows: Vec<Vec<Value>>, value_type: &str, start: Instant) -> Que
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
+1
View File
@@ -889,6 +889,7 @@ fn query_result(columns: Vec<String>, rows: Vec<Vec<serde_json::Value>>, affecte
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
+17
View File
@@ -2484,6 +2484,7 @@ fn error_query_result(message: String) -> db::QueryResult {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -2501,6 +2502,7 @@ fn empty_query_result(execution_time_ms: u128) -> db::QueryResult {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
@@ -2731,6 +2733,7 @@ pub async fn execute_statements(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -3232,6 +3235,7 @@ async fn exec_tx_pg_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}),
(Err(e), Ok(_)) => Err(e),
(Ok(_), Err(reset_err)) => Err(reset_err),
@@ -3324,6 +3328,7 @@ async fn exec_tx_mysql_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -3390,6 +3395,7 @@ async fn exec_tx_sqlite_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
})
})
@@ -3464,6 +3470,7 @@ async fn exec_tx_explicit_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -3531,6 +3538,7 @@ async fn exec_tx_none_inner(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -4061,6 +4069,7 @@ async fn execute_manual_txn_postgres_statement(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -4103,6 +4112,7 @@ async fn execute_manual_txn_mysql_statement(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
} else {
let result = conn.query_iter(sql).await.map_err(|e| format!("Query failed: {e}"))?;
@@ -4121,6 +4131,7 @@ async fn execute_manual_txn_mysql_statement(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
}
@@ -4156,6 +4167,7 @@ pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) -
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -4183,6 +4195,7 @@ pub async fn rollback_manual_transaction(state: &AppState, txn_session_id: &str)
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
}
@@ -4972,6 +4985,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let mut executor = FakeMysqlBatchExecutor {
outcomes: std::collections::VecDeque::from([
@@ -5499,6 +5513,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
})
.await;
@@ -5523,6 +5538,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
})
})
.await;
@@ -6215,6 +6231,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let normalized = normalize_query_result_for_js(result);
+9
View File
@@ -2848,6 +2848,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
assert_eq!(mysql_external_driver_ddl_from_query_result(result).unwrap(), "CREATE TABLE `users` (`id` bigint);");
@@ -2868,6 +2869,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
assert_eq!(
@@ -3572,6 +3574,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let tables = presto_like_tables_from_query_result(&result);
@@ -3619,6 +3622,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let columns = presto_like_columns_from_query_result(&result);
@@ -3790,6 +3794,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
assert_eq!(oracle_table_comment_from_query_result(result).unwrap().as_deref(), Some("Customer table"));
@@ -3807,6 +3812,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
assert_eq!(oracle_table_comment_from_query_result(empty).unwrap(), None);
@@ -3842,6 +3848,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let comments = table_comments_from_query_result(result);
@@ -3964,6 +3971,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let columns = oracle_columns_from_query_result(result);
@@ -4068,6 +4076,7 @@ for line in sys.stdin:
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let stats = oracle_object_statistics_from_query_result(result);
+1
View File
@@ -209,6 +209,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
};
let extensions = extension_infos_from_query_result(result, true);
+1
View File
@@ -530,6 +530,7 @@ async fn fetch_table_export_batch(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
});
}
@@ -688,6 +688,7 @@ fn execute_change_transaction(
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}),
}
}
+1
View File
@@ -6772,6 +6772,7 @@ mod tests {
session_id: None,
has_more: false,
elasticsearch_raw_body: None,
messages: Vec::new(),
}
}
+122 -1
View File
@@ -285,6 +285,42 @@ impl SpatialColumnBuilder {
}
}
/// A message emitted by the database server while executing a statement:
/// PostgreSQL `RAISE NOTICE`/`WARNING`, MySQL warnings and OK-packet info
/// strings, SQL Server `PRINT`/`RAISERROR` info messages, and similar.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryMessage {
/// Severity/level as reported by the server (e.g. `NOTICE`, `WARNING`,
/// `INFO`, `ERROR`, MySQL's `Note`/`Warning`).
pub severity: String,
pub message: String,
/// Server error/condition code (PostgreSQL SQLSTATE, MySQL error code).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub code: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub detail: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub hint: Option<String>,
}
impl QueryMessage {
/// One-line text rendering shared by the CLI, MCP, and agent tool output:
/// `SEVERITY: message` plus inline `(code: …, detail: …, hint: …)` extras.
pub fn format_line(&self) -> String {
let mut line = format!("{}: {}", self.severity.to_uppercase(), self.message);
let extras = [
self.code.as_ref().map(|value| format!("code: {value}")),
self.detail.as_ref().map(|value| format!("detail: {value}")),
self.hint.as_ref().map(|value| format!("hint: {value}")),
];
let extras: Vec<_> = extras.into_iter().flatten().collect();
if !extras.is_empty() {
line.push_str(&format!(" ({})", extras.join(", ")));
}
line
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
pub columns: Vec<String>,
@@ -321,6 +357,11 @@ pub struct QueryResult {
/// between the tabular view and the original JSON.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub elasticsearch_raw_body: Option<String>,
/// Messages emitted by the database server while executing the statement
/// (notices, warnings, info messages). Empty for drivers that do not
/// capture server messages.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub messages: Vec<QueryMessage>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -453,7 +494,87 @@ pub struct OwnerInfo {
#[cfg(test)]
mod tests {
use super::{ObjectInfo, ObjectSourceKind, SpatialColumn, SpatialColumnBuilder};
use super::{ObjectInfo, ObjectSourceKind, QueryMessage, SpatialColumn, SpatialColumnBuilder};
#[test]
fn query_message_format_line_uppercases_severity() {
let message = QueryMessage {
severity: "notice".to_string(),
message: "hello world".to_string(),
code: None,
detail: None,
hint: None,
};
assert_eq!(message.format_line(), "NOTICE: hello world");
}
#[test]
fn query_message_format_line_appends_code_detail_hint_extras() {
let message = QueryMessage {
severity: "WARNING".to_string(),
message: "careful".to_string(),
code: Some("01000".to_string()),
detail: Some("column truncated".to_string()),
hint: Some("widen the column".to_string()),
};
assert_eq!(
message.format_line(),
"WARNING: careful (code: 01000, detail: column truncated, hint: widen the column)"
);
}
#[test]
fn query_message_format_line_skips_missing_extras() {
let message = QueryMessage {
severity: "INFO".to_string(),
message: "Records: 3".to_string(),
code: None,
detail: None,
hint: Some("use a table".to_string()),
};
assert_eq!(message.format_line(), "INFO: Records: 3 (hint: use a table)");
}
#[test]
fn query_message_omits_empty_optional_fields_in_json() {
let minimal = QueryMessage {
severity: "NOTICE".to_string(),
message: "hello".to_string(),
code: None,
detail: None,
hint: None,
};
assert_eq!(
serde_json::to_value(&minimal).unwrap(),
serde_json::json!({ "severity": "NOTICE", "message": "hello" })
);
let full = QueryMessage {
severity: "NOTICE".to_string(),
message: "hello".to_string(),
code: Some("00000".to_string()),
detail: Some("d".to_string()),
hint: Some("h".to_string()),
};
assert_eq!(
serde_json::to_value(&full).unwrap(),
serde_json::json!({ "severity": "NOTICE", "message": "hello", "code": "00000", "detail": "d", "hint": "h" })
);
}
#[test]
fn query_message_deserializes_without_optional_fields() {
let message: QueryMessage = serde_json::from_str(r#"{"severity":"Note","message":"Records: 1"}"#).unwrap();
assert_eq!(message.severity, "Note");
assert_eq!(message.message, "Records: 1");
assert_eq!(message.code, None);
assert_eq!(message.detail, None);
assert_eq!(message.hint, None);
}
#[test]
fn list_objects_payload_preserves_optional_validity() {