From b7d33e375fd19cee3322426f7eeb7fa24560bce1 Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 31 Aug 2026 12:31:36 +0800 Subject: [PATCH] fix(influxdb): merge compat query series and quote explicit columns --- crates/dbx-core/src/db/influxdb_driver.rs | 168 +++++++++++++++--- .../dbx-core/src/sql_dialect/table_select.rs | 27 ++- crates/dbx-core/src/sql_dialect/tests.rs | 2 +- .../app-tests/aiAgentEventSession.test.ts | 6 +- 4 files changed, 170 insertions(+), 33 deletions(-) diff --git a/crates/dbx-core/src/db/influxdb_driver.rs b/crates/dbx-core/src/db/influxdb_driver.rs index faf62fcdea..3ab2d90000 100644 --- a/crates/dbx-core/src/db/influxdb_driver.rs +++ b/crates/dbx-core/src/db/influxdb_driver.rs @@ -1,6 +1,7 @@ use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC}; use reqwest::{Certificate, Client as HttpClient, Response}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::fs; use std::time::{Duration, Instant}; @@ -21,6 +22,7 @@ pub struct InfluxdbClient { username: Option, password: Option, url_params: Option, + default_database: Option, version: InfluxdbApiVersion, org: Option, } @@ -40,6 +42,7 @@ impl InfluxdbClient { username, password, url_params, + default_database: None, version: InfluxdbApiVersion::V1, org: None, } @@ -60,6 +63,7 @@ impl InfluxdbClient { username, password, url_params, + default_database: None, version: InfluxdbApiVersion::V1, org: None, }) @@ -77,12 +81,15 @@ impl InfluxdbClient { InfluxdbApiVersion::V2 => (None, (!config.password.is_empty()).then_some(config.password.clone())), }; let http = build_http_client(Some(&config.ca_cert_path), timeout)?; + let default_database = + config.database.as_deref().map(str::trim).filter(|database| !database.is_empty()).map(str::to_string); Ok(Self { http, base_url: url.trim_end_matches('/').to_string(), username, password, url_params: config.url_params.clone(), + default_database, version, org, }) @@ -152,6 +159,7 @@ impl Clone for InfluxdbClient { username: self.username.clone(), password: self.password.clone(), url_params: self.url_params.clone(), + default_database: self.default_database.clone(), version: self.version, org: self.org.clone(), } @@ -192,6 +200,8 @@ struct InfluxSeries { name: String, columns: Vec, values: Vec>, + #[serde(default)] + tags: BTreeMap, } #[derive(Deserialize)] @@ -335,11 +345,20 @@ pub async fn test_connection(client: &InfluxdbClient, timeout: Duration) -> Resu req.send().await.map_err(|e| format!("InfluxDB connection failed: {e}")) }) .await?; - if !resp.status().is_success() { - let body = resp.text().await.unwrap_or_default(); - return Err(format!("InfluxDB error: {body}")); + if resp.status().is_success() { + return Ok(()); } - Ok(()) + let body = resp.text().await.unwrap_or_default(); + if let Some(database) = client.default_database.as_deref() { + // Some Influx-compatible gateways require db on every request and do + // not expose SHOW DATABASES. A configured database is enough to prove + // the connection in that mode. + let probe = influx_query_with_timeout(client, "SHOW MEASUREMENTS", Some(database), timeout).await; + if probe.is_ok() { + return Ok(()); + } + } + Err(format!("InfluxDB error: {body}")) } pub async fn list_databases(client: &InfluxdbClient) -> Result, String> { @@ -350,7 +369,18 @@ pub async fn list_databases(client: &InfluxdbClient) -> Result .map(|bucket| DatabaseInfo { name: bucket.name, ..Default::default() }) .collect()); } - let result = influx_query(client, "SHOW DATABASES", None).await?; + let result = match influx_query(client, "SHOW DATABASES", None).await { + Ok(result) => result, + Err(error) => { + if let Some(database) = client.default_database.as_deref() { + // Only hide the unsupported SHOW DATABASES response when the configured database is usable. + if influx_query(client, "SHOW MEASUREMENTS", Some(database)).await.is_ok() { + return Ok(vec![DatabaseInfo { name: database.to_string(), ..Default::default() }]); + } + } + return Err(error); + } + }; Ok(result .results .iter() @@ -470,39 +500,93 @@ pub async fn execute_query(client: &InfluxdbClient, database: &str, sql: &str) - } let start = Instant::now(); let json = influx_query(client, sql, Some(database)).await?; - let series = json.results.iter().flat_map(|r| &r.series).next(); - match series { - Some(s) => Ok(QueryResult { - columns: s.columns.clone(), + let series = json.results.iter().flat_map(|r| r.series.iter()).collect::>(); + if !series.is_empty() { + let (columns, rows) = merge_v1_series(&series); + return Ok(QueryResult { + column_sortables: columns.iter().map(|_| false).collect(), + columns, column_types: vec![], - column_sortables: s.columns.iter().map(|_| false).collect(), spatial_columns: vec![], spatial_values: vec![], - rows: s.values.clone(), - affected_rows: s.values.len() as u64, + affected_rows: rows.len() as u64, + rows, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), - }), - None => Ok(QueryResult { - columns: vec![], - column_types: vec![], - column_sortables: vec![], - spatial_columns: vec![], - spatial_values: vec![], - rows: vec![], - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, - elasticsearch_raw_body: None, - messages: Vec::new(), - }), + }); } + Ok(QueryResult { + columns: vec![], + column_types: vec![], + column_sortables: vec![], + spatial_columns: vec![], + spatial_values: vec![], + rows: vec![], + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + messages: Vec::new(), + }) +} + +fn merge_v1_series(series: &[&InfluxSeries]) -> (Vec, Vec>) { + let mut columns = Vec::new(); + let mut column_indexes = BTreeMap::new(); + let mut tag_keys = Vec::new(); + let mut tag_indexes = BTreeMap::new(); + for item in series { + for column in &item.columns { + if !column_indexes.contains_key(column) { + column_indexes.insert(column.clone(), columns.len()); + columns.push(column.clone()); + } + } + for tag in item.tags.keys() { + if !tag_indexes.contains_key(tag) { + tag_indexes.insert(tag.clone(), tag_keys.len()); + tag_keys.push(tag.clone()); + } + } + } + columns.extend(tag_keys.iter().cloned()); + let mut rows = Vec::new(); + for item in series { + let source_indexes = + item.columns.iter().enumerate().map(|(index, name)| (name, index)).collect::>(); + for source_row in &item.values { + let mut row = columns[..column_indexes.len()] + .iter() + .map(|name| { + source_indexes + .get(name) + .and_then(|index| source_row.get(*index)) + .cloned() + .unwrap_or(serde_json::Value::Null) + }) + .collect::>(); + row.extend(tag_keys.iter().map(|tag| { + item.tags.get(tag).cloned().map(serde_json::Value::String).unwrap_or(serde_json::Value::Null) + })); + rows.push(row); + } + } + (columns, rows) +} + +async fn influx_query_with_timeout( + client: &InfluxdbClient, + sql: &str, + database: Option<&str>, + timeout: Duration, +) -> Result { + with_connection_timeout("InfluxDB", timeout, influx_query(client, sql, database)).await } async fn get_columns_v2(client: &InfluxdbClient, bucket: &str, measurement: &str) -> Result, String> { @@ -707,6 +791,7 @@ mod tests { username: None, password: None, url_params: None, + default_database: None, version: InfluxdbApiVersion::V1, org: None, } @@ -720,6 +805,33 @@ mod tests { assert_eq!(url, "http://localhost:8086/query?db=sample&q=SHOW%20DATABASES"); } + #[test] + fn merges_v1_series_and_promotes_tags_to_columns() { + let first = InfluxSeries { + name: "dbx_smoke".to_string(), + columns: vec!["time".to_string(), "mean".to_string()], + values: vec![vec![json!("2026-08-31T00:00:00Z"), json!(10)]], + tags: BTreeMap::from([(String::from("host"), String::from("a"))]), + }; + let second = InfluxSeries { + name: "dbx_smoke".to_string(), + columns: vec!["time".to_string(), "mean".to_string()], + values: vec![vec![json!("2026-08-31T00:01:00Z"), json!(20)]], + tags: BTreeMap::from([(String::from("host"), String::from("b"))]), + }; + + let (columns, rows) = merge_v1_series(&[&first, &second]); + + assert_eq!(columns, vec!["time", "mean", "host"]); + assert_eq!( + rows, + vec![ + vec![json!("2026-08-31T00:00:00Z"), json!(10), json!("a")], + vec![json!("2026-08-31T00:01:00Z"), json!(20), json!("b")], + ] + ); + } + #[test] fn v2_config_uses_token_auth_and_org_urls() { let config: ConnectionConfig = serde_json::from_value(json!({ diff --git a/crates/dbx-core/src/sql_dialect/table_select.rs b/crates/dbx-core/src/sql_dialect/table_select.rs index a8155ea155..fcc2f68341 100644 --- a/crates/dbx-core/src/sql_dialect/table_select.rs +++ b/crates/dbx-core/src/sql_dialect/table_select.rs @@ -608,14 +608,21 @@ pub(super) fn build_select_columns( .collect::>() .join(", "); } - if !matches!(database_type, Some(DatabaseType::Hive | DatabaseType::Kyuubi | DatabaseType::Impala)) { + if !matches!( + database_type, + Some(DatabaseType::Hive | DatabaseType::Kyuubi | DatabaseType::Impala | DatabaseType::InfluxDb) + ) { return "*".to_string(); } columns .iter() .map(|column| { let ident = quote_table_identifier(database_type, column); - format!("{ident} AS {ident}") + if database_type == Some(DatabaseType::Hive) { + format!("{ident} AS {ident}") + } else { + ident + } }) .collect::>() .join(", ") @@ -770,6 +777,22 @@ mod tests { } } + #[test] + fn influxdb_table_select_quotes_explicit_columns() { + // InfluxQL needs double-quoted identifiers and only permits ORDER BY on + // the time column; explicit column lists must not fall back to "*". + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::InfluxDb), + database: Some("monitor".to_string()), + table_name: "cpu".to_string(), + columns: vec!["time".to_string(), "host".to_string(), "value".to_string()], + ..Default::default() + }), + "SELECT \"time\", \"host\", \"value\" FROM \"cpu\" ORDER BY time DESC LIMIT 100;" + ); + } + #[test] fn databricks_table_select_uses_backtick_identifiers() { assert_eq!( diff --git a/crates/dbx-core/src/sql_dialect/tests.rs b/crates/dbx-core/src/sql_dialect/tests.rs index 8a538cd2db..948fa8ae54 100644 --- a/crates/dbx-core/src/sql_dialect/tests.rs +++ b/crates/dbx-core/src/sql_dialect/tests.rs @@ -492,7 +492,7 @@ fn builds_table_data_where_and_schema_queries() { include_row_id: false, ..Default::default() }), - "SELECT `id` AS `id`, `name` AS `name` FROM `dbx_demo`.`connection_test` ORDER BY 1 LIMIT 2 OFFSET 1;" + "SELECT `id`, `name` FROM `dbx_demo`.`connection_test` ORDER BY 1 LIMIT 2 OFFSET 1;" ); assert_eq!( build_table_data_select_sql(TableDataSelectSqlOptions { diff --git a/packages/app-tests/aiAgentEventSession.test.ts b/packages/app-tests/aiAgentEventSession.test.ts index 30f6800670..3ad63e9763 100644 --- a/packages/app-tests/aiAgentEventSession.test.ts +++ b/packages/app-tests/aiAgentEventSession.test.ts @@ -9,8 +9,10 @@ test("Tauri AI agent stream events carry their session id", () => { assert.match(tauriAiCommandSource, /struct AiAgentEventPayload \{/); assert.match(tauriAiCommandSource, /session_id: String/); assert.match(tauriAiCommandSource, /#\[serde\(flatten\)\]\s*event: AgentEvent/); - assert.match(tauriAiCommandSource, /let event_session_id = session_id\.clone\(\);/); - assert.match(tauriAiCommandSource, /AiAgentEventPayload \{ session_id: event_session_id\.clone\(\), event \}/); + assert.match( + tauriAiCommandSource, + /AiAgentEventPayload \{ session_id: self\.session_id\.clone\(\), event \}/, + ); assert.match(tauriAiCommandSource, /app\.emit\("ai-agent-event", &payload\)/); });