fix(core): fix count statement for BigQuery

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"}

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<char> {
         None
     }
+
+    fn col_alias_overrides(&self, _alias: &str) -> Result<Option<String>> {
+        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<Option<String>> {
+        // 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 {}

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<Option<String>> {
+        self.inner_dialect.col_alias_overrides(alias)
+    }
 }
 
 impl Default for WrenDialect {

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(())
     }