diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py index 2541bceb9..0f3dddd2e 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py @@ -10,6 +10,7 @@ manifest = { "catalog": "wren", "schema": "public", + "dataSource": "bigquery", "models": [ { "name": "orders", @@ -359,7 +360,7 @@ async def test_order_by_nulls_last(client, manifest_str, connection_info): "sql": "SELECT letter FROM null_test ORDER BY id", }, headers={ - X_WREN_FALLBACK_DISABLE: "true", # Disable fallback to DuckDB + X_WREN_FALLBACK_DISABLE: "true", }, ) assert response.status_code == 200 @@ -377,7 +378,7 @@ async def test_order_by_nulls_last(client, manifest_str, connection_info): "sql": "SELECT letter FROM null_test ORDER BY id desc", }, headers={ - X_WREN_FALLBACK_DISABLE: "true", # Disable fallback to DuckDB + X_WREN_FALLBACK_DISABLE: "true", }, ) assert response.status_code == 200 @@ -386,3 +387,24 @@ async def test_order_by_nulls_last(client, manifest_str, connection_info): assert result["data"][0][0] == "two" assert result["data"][1][0] == "one" assert result["data"][2][0] == "three" + + +async def test_count(client, manifest_str, connection_info): + response = await client.post( + url=f"{base_url}/query", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT COUNT(*) FROM wren.public.orders", + }, + headers={ + X_WREN_FALLBACK_DISABLE: "true", + }, + ) + assert response.status_code == 200 + result = response.json() + assert len(result["data"]) == 1 + assert ( + result["data"][0][0] == 15000 + ) # Adjust based on actual data count in orders table + assert result["dtypes"] == {"count_40_42_41": "int64"} diff --git a/wren-core/core/src/mdl/dialect/inner_dialect.rs b/wren-core/core/src/mdl/dialect/inner_dialect.rs index a6a281f07..be203e527 100644 --- a/wren-core/core/src/mdl/dialect/inner_dialect.rs +++ b/wren-core/core/src/mdl/dialect/inner_dialect.rs @@ -48,6 +48,10 @@ pub trait InnerDialect: Send + Sync { fn identifier_quote_style(&self, _identifier: &str) -> Option { None } + + fn col_alias_overrides(&self, _alias: &str) -> Result> { + Ok(None) + } } /// [get_inner_dialect] returns the suitable InnerDialect for the given data source. @@ -89,6 +93,29 @@ impl InnerDialect for BigQueryDialect { fn unnest_as_table_factor(&self) -> bool { true } + + fn col_alias_overrides(&self, alias: &str) -> Result> { + // Check if alias contains any special characters not supported by BigQuery col names + // https://cloud.google.com/bigquery/docs/schemas#flexible-column-names + let special_chars: [char; 20] = [ + '!', '"', '$', '(', ')', '*', ',', '.', '/', ';', '?', '@', '[', '\\', ']', + '^', '`', '{', '}', '~', + ]; + + if alias.chars().any(|c| special_chars.contains(&c)) { + let mut encoded_name = String::new(); + for c in alias.chars() { + if special_chars.contains(&c) { + encoded_name.push_str(&format!("_{}", c as u32)); + } else { + encoded_name.push(c); + } + } + Ok(Some(encoded_name)) + } else { + Ok(Some(alias.to_string())) + } + } } pub struct OracleDialect {} diff --git a/wren-core/core/src/mdl/dialect/wren_dialect.rs b/wren-core/core/src/mdl/dialect/wren_dialect.rs index 56dfa55dd..d54772195 100644 --- a/wren-core/core/src/mdl/dialect/wren_dialect.rs +++ b/wren-core/core/src/mdl/dialect/wren_dialect.rs @@ -91,6 +91,10 @@ impl Dialect for WrenDialect { fn unnest_as_table_factor(&self) -> bool { self.inner_dialect.unnest_as_table_factor() } + + fn col_alias_overrides(&self, alias: &str) -> Result> { + self.inner_dialect.col_alias_overrides(alias) + } } impl Default for WrenDialect { diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs index 06f19722d..d983ef9eb 100644 --- a/wren-core/core/src/mdl/mod.rs +++ b/wren-core/core/src/mdl/mod.rs @@ -1189,7 +1189,7 @@ mod test { sql, ) .await?; - assert_snapshot!(actual, @"SELECT \"UNNEST(make_array(Int64(1),Int64(2),Int64(3)))\" FROM UNNEST([1, 2, 3])"); + assert_snapshot!(actual, @r#"SELECT "UNNEST_40make_array_40Int64_401_41_44Int64_402_41_44Int64_403_41_41_41" FROM UNNEST([1, 2, 3])"#); Ok(()) }