Add ClientBuilder and support for connecting to local spice runtime (#38)

* Add `ClientBuilder` and add support for local spice connection

* Update ci tests

* cleanup

* add docs comments

* update docs

* Update src/client.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* Update src/flight.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* Update src/client.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* Update src/client.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* Update src/client.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* Update src/client.rs

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

* update comments

* update builder api usage

* add builder method

* fix

---------

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

Author: ewgenius

.github/workflows/build.yml

@@ -48,6 +48,20 @@ jobs:
       - name: Build
         run: cargo build --verbose
 
+      - name: install Spice (https://install.spiceai.org)
+        run: |
+          curl https://install.spiceai.org | /bin/bash
+          echo "$HOME/.spice/bin" >> $GITHUB_PATH
+
+      - name: Init and start spice app
+        run: |
+          spice init spice_qs
+          cd spice_qs
+          spice add spiceai/quickstart
+          spice run &> spice.log &
+          # time to initialize added dataset
+          sleep 10
+
       - name: Run tests
         run: cargo test --verbose
         env:

README.md

@@ -14,14 +14,38 @@ cargo add spiceai
 
 <!-- NOTE: If you're changing the code examples below, make sure you update `tests/readme_test.rs`. -->
 
-### New client
+### Usage with locally running [spice runtime](https://github.com/spiceai/spiceai)
+
+Follow the [quiqstart guide](https://github.com/spiceai/spiceai?tab=readme-ov-file#%EF%B8%8F-quickstart-local-machine) to install and run spice locally
+
+```rust
+use spiceai::ClientBuilder;
+
+#[tokio::main]
+async fn main() {
+  let mut client = ClientBuilder::new()
+    .flight_url("http://localhost:50051")
+    .build()
+    .await
+    .unwrap();
+
+  let data = client.query("SELECT trip_distance, total_amount FROM taxi_trips ORDER BY trip_distance DESC LIMIT 10;").await;
+}
+```
+
+### New client with https://spice.ai cloud
 
 ```rust
-use spiceai::Client;
+use spiceai::ClientBuilder;
 
 #[tokio::main]
 async fn main() {
-  let mut client = Client::new("API_KEY").await.unwrap();
+  let mut client = ClientBuilder::new()
+    .api_key("API_KEY")
+    .use_spiceai_cloud()
+    .build()
+    .await
+    .unwrap();
 }
 ```
 
@@ -30,11 +54,17 @@ async fn main() {
 SQL Query
 
 ```rust
-use spiceai::Client;
+use spiceai::ClientBuilder;
 
 #[tokio::main]
 async fn main() {
-  let mut client = Client::new("API_KEY").await.unwrap();
+  let mut client = ClientBuilder::new()
+    .api_key("API_KEY")
+    .use_spiceai_cloud()
+    .build()
+    .await
+    .unwrap();
+
   let data = client.query("SELECT * FROM eth.recent_blocks LIMIT 10;").await;
 }
 
@@ -45,11 +75,17 @@ async fn main() {
 Firecache SQL Query
 
 ```rust
-use spiceai::Client;
+use spiceai::ClientBuilder;
 
 #[tokio::main]
 async fn main() {
-  let mut client = Client::new("API_KEY").await.unwrap();
+  let mut client = ClientBuilder::new()
+    .api_key("API_KEY")
+    .use_spiceai_cloud()
+    .build()
+    .await
+    .unwrap();
+
   let data = client.fire_query("SELECT * FROM eth.recent_blocks LIMIT 10;").await;
 }
 

src/client.rs

@@ -1,5 +1,5 @@
 use crate::{
-    config::{FIRECACHE_ADDR, FLIGHT_ADDR, HTTPS_ADDR},
+    config::{SPICE_CLOUD_FIRECACHE_ADDR, SPICE_CLOUD_FLIGHT_ADDR, SPICE_LOCAL_FLIGHT_ADDR},
     flight::SqlFlightClient,
     tls::new_tls_flight_channel,
 };
@@ -9,31 +9,25 @@ use std::error::Error;
 use tonic::transport::Channel;
 
 struct SpiceClientConfig {
-    https_addr: String,
     flight_channel: Channel,
     firecache_channel: Channel,
 }
 
 impl SpiceClientConfig {
-    fn new(https_addr: String, flight_channel: Channel, firecache_channel: Channel) -> Self {
+    fn new(flight_channel: Channel, firecache_channel: Channel) -> Self {
         SpiceClientConfig {
-            https_addr,
             flight_channel,
             firecache_channel,
         }
     }
 
     pub async fn load_from_default() -> Result<SpiceClientConfig, Box<dyn Error>> {
         let (flight_chan, firecache_chan) = try_join!(
-            new_tls_flight_channel(FLIGHT_ADDR),
-            new_tls_flight_channel(FIRECACHE_ADDR)
+            new_tls_flight_channel(SPICE_CLOUD_FLIGHT_ADDR),
+            new_tls_flight_channel(SPICE_CLOUD_FIRECACHE_ADDR)
         )?;
 
-        Ok(SpiceClientConfig::new(
-            HTTPS_ADDR.to_string(),
-            flight_chan,
-            firecache_chan,
-        ))
+        Ok(SpiceClientConfig::new(flight_chan, firecache_chan))
     }
 }
 
@@ -59,11 +53,15 @@ impl SpiceClient {
         let config = SpiceClientConfig::load_from_default().await?;
 
         Ok(Self {
-            flight: SqlFlightClient::new(config.flight_channel, api_key.to_string()),
-            firecache: SqlFlightClient::new(config.firecache_channel, api_key.to_string()),
+            flight: SqlFlightClient::new(config.flight_channel, Some(api_key.to_string())),
+            firecache: SqlFlightClient::new(config.firecache_channel, Some(api_key.to_string())),
         })
     }
 
+    pub fn builder() -> SpiceClientBuilder {
+        SpiceClientBuilder::new()
+    }
+
     /// Queries the Spice Flight endpoint with the given SQL query.
     /// ```
     /// # use spiceai::Client;
@@ -95,3 +93,100 @@ impl SpiceClient {
         self.firecache.query(query).await
     }
 }
+
+/// Builder for creating a `SpiceClient`.
+///
+/// By default the `SpiceClient` will use local spice runtime flight endpoint.
+/// Follow [spiceai quickstart](https://github.com/spiceai/spiceai?tab=readme-ov-file#%EF%B8%8F-quickstart-local-machine) to setup local spice runtime.
+/// ```
+/// # use spiceai::ClientBuilder;
+/// #
+/// # #[tokio::main]
+/// # async fn main() {
+/// #    let mut client = ClientBuilder::new()
+/// #      .build()
+/// #      .await
+/// #      .unwrap();
+/// # }
+/// ```
+/// To use default Spice.ai Cloud endpoints, you can use the `with_spiceai_cloud()` method.
+///
+/// ```
+/// # use spiceai::ClientBuilder;
+/// #
+/// # #[tokio::main]
+/// # async fn main() {
+/// #    let mut client = ClientBuilder::new()
+/// #      .api_key("API_KEY")
+/// #      .use_spiceai_cloud()
+/// #      .build()
+/// #      .await
+/// #      .unwrap();
+/// # }
+/// ```
+///
+pub struct SpiceClientBuilder {
+    api_key: Option<String>,
+    firecache_url: Option<String>,
+    flight_url: Option<String>,
+}
+
+impl Default for SpiceClientBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SpiceClientBuilder {
+    pub fn new() -> Self {
+        Self {
+            api_key: None,
+            firecache_url: None,
+            flight_url: None,
+        }
+    }
+
+    /// Configures the `SpiceClient` to use the given API key.
+    pub fn api_key(mut self, api_key: &str) -> Self {
+        self.api_key = Some(api_key.to_string());
+        self
+    }
+
+    /// Configures the `SpiceClient` to use the given Spice Firecache endpoint.
+    pub fn firecache_url(mut self, firecache_url: &str) -> Self {
+        self.firecache_url = Some(firecache_url.to_string());
+        self
+    }
+
+    /// Configures the `SpiceClient` to use the given Spice Flight endpoint.
+    pub fn flight_url(mut self, flight_url: &str) -> Self {
+        self.flight_url = Some(flight_url.to_string());
+        self
+    }
+
+    /// Configures the `SpiceClient` to use default Spice.ai Cloud endpoints.
+    /// Equivalent to calling `.firecache_url("https://firecache.spiceai.io")` and `.flight_url("https://flight.spiceai.io")`.
+    pub fn use_spiceai_cloud(mut self) -> Self {
+        self.flight_url = Some(SPICE_CLOUD_FLIGHT_ADDR.to_string());
+        self.firecache_url = Some(SPICE_CLOUD_FIRECACHE_ADDR.to_string());
+        self
+    }
+
+    /// Builds the `SpiceClient` with the specified configuration.
+    pub async fn build(self) -> Result<SpiceClient, Box<dyn Error>> {
+        let flight_channel = match self.flight_url {
+            Some(url) => new_tls_flight_channel(&url).await?,
+            None => new_tls_flight_channel(SPICE_LOCAL_FLIGHT_ADDR).await?,
+        };
+
+        let firecache_channel = match self.firecache_url {
+            Some(url) => new_tls_flight_channel(&url).await?,
+            None => new_tls_flight_channel(SPICE_CLOUD_FIRECACHE_ADDR).await?,
+        };
+
+        Ok(SpiceClient {
+            flight: SqlFlightClient::new(flight_channel, self.api_key.clone()),
+            firecache: SqlFlightClient::new(firecache_channel, self.api_key.clone()),
+        })
+    }
+}

src/config.rs

@@ -1,3 +1,5 @@
-pub const HTTPS_ADDR: &str = "https://data.spiceai.io";
-pub const FLIGHT_ADDR: &str = "https://flight.spiceai.io";
-pub const FIRECACHE_ADDR: &str = "https://firecache.spiceai.io";
+pub const SPICE_CLOUD_FLIGHT_ADDR: &str = "https://flight.spiceai.io";
+pub const SPICE_CLOUD_FIRECACHE_ADDR: &str = "https://firecache.spiceai.io";
+
+// default address for local spice runtime
+pub const SPICE_LOCAL_FLIGHT_ADDR: &str = "http://localhost:50051";

src/flight.rs

@@ -21,15 +21,15 @@ pub struct SqlFlightClient {
     token: Option<String>,
     headers: HashMap<String, String>,
     client: FlightServiceClient<Channel>,
-    api_key: String,
+    api_key: Option<String>,
 }
 
 fn status_to_arrow_error(status: tonic::Status) -> ArrowError {
     ArrowError::IpcError(format!("{status:?}"))
 }
 
 impl SqlFlightClient {
-    pub fn new(chan: Channel, api_key: String) -> Self {
+    pub fn new(chan: Channel, api_key: Option<String>) -> Self {
         SqlFlightClient {
             api_key,
             client: FlightServiceClient::new(chan),
@@ -81,11 +81,11 @@ impl SqlFlightClient {
         Ok(resp)
     }
 
-    async fn authenticate(&mut self) -> std::result::Result<(), Box<dyn Error>> {
-        if self.api_key.split('|').collect::<String>().len() < 2 {
+    async fn authenticate(&mut self, api_key: &str) -> std::result::Result<(), Box<dyn Error>> {
+        if api_key.split('|').collect::<String>().len() < 2 {
             return Err("Invalid API key format".into());
         }
-        self.handshake("", &self.api_key.to_string()).await?;
+        self.handshake("", api_key).await?;
         Ok(())
     }
 
@@ -115,7 +115,10 @@ impl SqlFlightClient {
         &mut self,
         query: &str,
     ) -> std::result::Result<FlightRecordBatchStream, Box<dyn Error>> {
-        self.authenticate().await?;
+        let api_key = self.api_key.clone();
+        if let Some(api_key) = api_key {
+            self.authenticate(&api_key).await?;
+        }
 
         let descriptor = FlightDescriptor::new_cmd(query.to_string());
         let req = self.set_request_headers(descriptor.into_request())?;

src/lib.rs

@@ -6,6 +6,7 @@ mod flight;
 mod tls;
 
 pub use client::SpiceClient as Client;
+pub use client::SpiceClientBuilder as ClientBuilder;
 
 // Further public exports and integrations
 pub use futures::StreamExt;

tests/client_test.rs

@@ -1,14 +1,17 @@
 #[cfg(test)]
 mod tests {
     use futures::stream::StreamExt;
-    use spiceai::Client;
+    use spiceai::{Client, ClientBuilder};
     use std::env;
     use std::path::Path;
 
     async fn new_client() -> Client {
         dotenv::from_path(Path::new(".env.local")).ok();
         let api_key = env::var("API_KEY").expect("API_KEY not found");
-        Client::new(&api_key)
+        ClientBuilder::new()
+            .api_key(&api_key)
+            .use_spiceai_cloud()
+            .build()
             .await
             .expect("Failed to create client")
     }
@@ -18,6 +21,11 @@ mod tests {
         new_client().await;
     }
 
+    #[tokio::test]
+    async fn test_new_client_builder() {
+        new_client().await;
+    }
+
     #[tokio::test]
     async fn test_query() {
         let mut spice_client = new_client().await;