fix(influxdb): merge compat query series and quote explicit columns

This commit is contained in:
zipg authored and GitHub committed 2026-08-31 12:31:36 +08:00
1 parent d1046d54f9
commit b7d33e375f
4 files changed
+170 -33

No files matched your search

+140 -28
View File
@@ -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<String>,
password: Option<String>,
url_params: Option<String>,
default_database: Option<String>,
version: InfluxdbApiVersion,
org: Option<String>,
}
@@ -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<String>,
values: Vec<Vec<serde_json::Value>>,
#[serde(default)]
tags: BTreeMap<String, String>,
}
#[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<Vec<DatabaseInfo>, String> {
@@ -350,7 +369,18 @@ pub async fn list_databases(client: &InfluxdbClient) -> Result<Vec<DatabaseInfo>
.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::<Vec<_>>();
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<String>, Vec<Vec<serde_json::Value>>) {
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::<BTreeMap<_, _>>();
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::<Vec<_>>();
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<InfluxJsonResult, String> {
with_connection_timeout("InfluxDB", timeout, influx_query(client, sql, database)).await
}
async fn get_columns_v2(client: &InfluxdbClient, bucket: &str, measurement: &str) -> Result<Vec<ColumnInfo>, 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!({
@@ -608,14 +608,21 @@ pub(super) fn build_select_columns(
.collect::<Vec<_>>()
.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::<Vec<_>>()
.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!(
+1 -1
View File
@@ -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 {