Merge branch 'trunk' into lukim/query-method

Author: lukekim

crates/adbc_client/src/lib.rs

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 pub mod databricks;
+pub mod spiceai;
 
 use std::collections::HashMap;
 

crates/adbc_client/src/spiceai.rs

@@ -0,0 +1,146 @@
+/*
+Copyright 2026 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+use std::collections::HashMap;
+
+use crate::{AdbcConnection, Result};
+
+/// Connect to a Spice.ai Flight SQL endpoint.
+///
+/// Uses [`AdbcConnection::create`] with the `"flightsql"` ADBC driver,
+/// passing `uri`, `username`, and `password` options.
+///
+/// # Arguments
+///
+/// - `flight_url`: The Flight SQL endpoint URI (e.g. `"grpc://localhost:50051"`).
+/// - `api_key`: Optional API key (passed as the `password` option; `username` is left empty).
+///
+/// # Example
+///
+/// ```rust,no_run
+/// use adbc_client::spiceai;
+///
+/// let mut conn = spiceai::connect("grpc://localhost:50051", Some("my-api-key"))
+///     .expect("Failed to connect");
+///
+/// let batches = conn.query("SELECT 1").expect("query failed");
+/// ```
+pub fn connect(flight_url: &str, api_key: Option<&str>) -> Result<AdbcConnection> {
+    let mut kwargs = HashMap::from([(
+        "uri".to_string(),
+        serde_json::Value::String(flight_url.to_string()),
+    )]);
+
+    if let Some(key) = api_key {
+        kwargs.insert(
+            "username".to_string(),
+            serde_json::Value::String(String::new()),
+        );
+        kwargs.insert(
+            "password".to_string(),
+            serde_json::Value::String(key.to_string()),
+        );
+    }
+
+    AdbcConnection::create("flightsql", kwargs)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Requires `SPICEAI_FLIGHT_URL` and `SPICEAI_API_KEY` env vars.
+    #[test]
+    #[ignore = "requires Spice.ai credentials in env"]
+    fn test_spiceai_query() {
+        let flight_url =
+            std::env::var("SPICEAI_FLIGHT_URL").expect("SPICEAI_FLIGHT_URL must be set");
+        let api_key = std::env::var("SPICEAI_API_KEY").ok();
+
+        let mut conn =
+            connect(&flight_url, api_key.as_deref()).expect("Failed to connect to Spice.ai");
+
+        let batches = conn
+            .query("SELECT 1 AS one")
+            .expect("Failed to execute query");
+
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert!(
+            total_rows > 0,
+            "Expected at least one row from Spice.ai query"
+        );
+    }
+
+    /// TPC-H Q22 – Global Sales Opportunity
+    ///
+    /// Requires `SPICEAI_FLIGHT_URL` and `SPICEAI_API_KEY` env vars.
+    /// SPICEAI_FLIGHT_URL="grpc://127.0.0.1:50051" cargo test -p adbc_client -- --ignored
+    #[test]
+    #[ignore = "requires Spice.ai credentials in env"]
+    fn test_tpch_q22() {
+        let flight_url =
+            std::env::var("SPICEAI_FLIGHT_URL").expect("SPICEAI_FLIGHT_URL must be set");
+        let api_key = std::env::var("SPICEAI_API_KEY").ok();
+
+        let mut conn =
+            connect(&flight_url, api_key.as_deref()).expect("Failed to connect to Spice.ai");
+
+        let batches = conn
+            .query(
+                "select \
+                    cntrycode, \
+                    count(*) as numcust, \
+                    sum(c_acctbal) as totacctbal \
+                from \
+                    ( \
+                        select \
+                            substring(c_phone from 1 for 2) as cntrycode, \
+                            c_acctbal \
+                        from \
+                            customer \
+                        where \
+                                substring(c_phone from 1 for 2) in \
+                                ('13', '31', '23', '29', '30', '18', '17') \
+                          and c_acctbal > ( \
+                            select \
+                                avg(c_acctbal) \
+                            from \
+                                customer \
+                            where \
+                                    c_acctbal > 0.00 \
+                              and substring(c_phone from 1 for 2) in \
+                                  ('13', '31', '23', '29', '30', '18', '17') \
+                        ) \
+                          and not exists ( \
+                                select \
+                                    * \
+                                from \
+                                    orders \
+                                where \
+                                        o_custkey = c_custkey \
+                            ) \
+                    ) as custsale \
+                group by \
+                    cntrycode \
+                order by \
+                    cntrycode",
+            )
+            .expect("Failed to execute TPC-H Q22");
+
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert!(total_rows > 0, "Expected at least one row from TPC-H Q22");
+    }
+}

crates/test-framework/src/queries/mod.rs

@@ -590,7 +590,7 @@ impl QueryOverrides {
 pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
     let queries = generate_tpch_queries!(
         q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12, q13, q14, q16, q17, q18, q19, q20, q21,
-        q22, simple_q1, simple_q2, simple_q3, simple_q4, simple_q5, simple_q6, simple_q7
+        q22
     );
 
     match overrides {
@@ -608,7 +608,7 @@ pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
             2, // Analysis error: [UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.UNSUPPORTED_CORRELATED_SCALAR_SUBQUERY] Unsupported subquery expression: Correlated scalar subqueries can only be used in filters, aggregations, projections, and UPDATE/MERGE/DELETE commands
             17 // Analysis error: [UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.UNSUPPORTED_CORRELATED_SCALAR_SUBQUERY] Unsupported subquery expression: Correlated scalar subqueries can only be used in filters, aggregations, projections, and UPDATE/MERGE/DELETE commands
         ),
-        Some(QueryOverrides::MySQL) => remove_tpch_query!(queries, simple_q7),
+        Some(QueryOverrides::MySQL) => queries,
         Some(QueryOverrides::DynamoDB) => remove_tpch_query!(
             queries, 6 // Unsupported Decimals
         ),
@@ -634,19 +634,11 @@ pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
             q19,
             q20,
             q21,
-            q22,
-            simple_q1,
-            simple_q2,
-            simple_q3,
-            simple_q4,
-            simple_q5,
-            simple_q6,
-            simple_q7
+            q22
         ),
         Some(QueryOverrides::Oracle) => generate_tpch_queries_override!(
             "oracle", q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12, q13, q14, q16, q17, q18,
-            q19, q20, q21, q22, simple_q1, simple_q2, simple_q3, simple_q4, simple_q5, simple_q6,
-            simple_q7
+            q19, q20, q21, q22
         ),
         Some(QueryOverrides::IcebergSF1) => generate_tpch_queries_override!(
             "iceberg_sf1",
@@ -670,19 +662,11 @@ pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
             q19,
             q20,
             q21,
-            q22,
-            simple_q1,
-            simple_q2,
-            simple_q3,
-            simple_q4,
-            simple_q5,
-            simple_q6,
-            simple_q7
+            q22
         ),
         Some(QueryOverrides::IcebergHadoop) => generate_tpch_queries_override!(
             "hadoop", q1, q2, q3, q4, q5, q6, q7, q8, q9, q10, q11, q12, q13, q14, q16, q17, q18,
-            q19, q20, q21, q22, simple_q1, simple_q2, simple_q3, simple_q4, simple_q5, simple_q6,
-            simple_q7
+            q19, q20, q21, q22
         ),
         Some(QueryOverrides::Spicecloud) => remove_tpch_query!(
             queries,
@@ -713,15 +697,7 @@ pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
             q18,
             q19,
             // q20, Binder Error; https://github.com/spiceai/spiceai/issues/7356
-            q21,
-            // q22, Binder Error; https://github.com/spiceai/spiceai/issues/7356
-            simple_q1,
-            simple_q2,
-            simple_q3,
-            simple_q4,
-            simple_q5,
-            simple_q6,
-            simple_q7
+            q21 // q22, Binder Error; https://github.com/spiceai/spiceai/issues/7356
         ),
         Some(QueryOverrides::GlueCatalog) => generate_tpch_queries_override!(
             "glue_catalog",
@@ -745,14 +721,7 @@ pub fn get_tpch_test_queries(overrides: Option<QueryOverrides>) -> Vec<Query> {
             q19,
             q20,
             q21,
-            q22,
-            simple_q1,
-            simple_q2,
-            simple_q3,
-            simple_q4,
-            simple_q5,
-            simple_q6,
-            simple_q7
+            q22
         ),
         Some(QueryOverrides::DuckDBPartitioned) => remove_tpch_query!(
             queries,