Merge pull request #3 from spiceai/qianqian/v1

Update prices functions

Author: Jeadie

.gitignore

@@ -12,3 +12,6 @@ Cargo.lock
 
 # MSVC Windows builds of rustc generate these, which store debugging information
 *.pdb
+
+# Store local variables for test
+.env.local

Cargo.toml

@@ -4,18 +4,20 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-arrow-flight = { version = "47.0.0", features = ["flight-sql-experimental"] }
+arrow-flight = { version = "48.0.0", features = ["flight-sql-experimental"] }
 bytes = "1.5.0"
 prost = "0.12.1"
 prost-types = "0.12.1"
 rustls = "0.21.7"
 tokio = "1.32.0"
 rustls-native-certs = "0.6.3"
-tonic = { version = "0.10.0", default-features = false, features = ["transport", "tls"] }
+tonic = { version = "0.10.0", default-features = false, features = ["transport", "tls", "tls-roots"] }
 rustls-pemfile = "1.0.3"
 reqwest = { version = "0.11.21", features = ["json"] }
 serde = "1.0.188"
 serde_derive = "1.0.188"
 serde_json = "1.0.107"
 chrono = { version = "0.4.31", features = ["serde"] }
-
+dotenv = "0.15.0"
+arrow = "48.0.0"
+futures = "0.3.28"

src/client.rs

@@ -1,50 +1,45 @@
-
-use crate::{flight::{SqlFlightClient}, prices::PricesClient, tls::new_tls_flight_channel};
+use crate::{flight::SqlFlightClient, prices::PricesClient, tls::new_tls_flight_channel};
+use arrow_flight::decode::FlightRecordBatchStream;
 use std::error::Error;
-use arrow_flight::sql::client::FlightSqlServiceClient;
-use tonic::{transport::{Channel}};
-use tonic::{Streaming};
-use arrow_flight::{FlightData};
-
-
-pub async fn new_spice_client(api_key: String, http_addr: String, flight_addr: String, firecache_addr: String) -> Result<SpiceClient, Box<dyn Error>> {
-    let flight_chan = new_tls_flight_channel(flight_addr).await;
-    if flight_chan.is_err() {
-        return Err(flight_chan.err().expect("").into())
-    }
+use tonic::transport::Channel;
+
+pub async fn new_spice_client(api_key: String) -> Result<SpiceClient, Box<dyn Error>> {
+    return new_spice_client_with_address(
+        api_key.to_string(),
+        "https://data.spiceai.io".to_string(),
+        "https://flight.spiceai.io".to_string(),
+    )
+    .await;
+}
 
-    match new_tls_flight_channel(firecache_addr).await {
+pub async fn new_spice_client_with_address(
+    api_key: String,
+    http_addr: String,
+    flight_addr: String,
+) -> Result<SpiceClient, Box<dyn Error>> {
+    match new_tls_flight_channel(flight_addr).await {
         Err(e) => Err(e.into()),
-        Ok(firecache_chan) => Ok(SpiceClient::new(
-            http_addr,
-            api_key,
-            flight_chan.expect(""),
-            firecache_chan
-        ))
+        Ok(flight_chan) => Ok(SpiceClient::new(http_addr, api_key, flight_chan)),
     }
 }
 
 pub struct SpiceClient {
     flight: SqlFlightClient,
-    firecache: SqlFlightClient,
-    pub prices: PricesClient
+    pub prices: PricesClient,
 }
 
 impl SpiceClient {
-    pub fn new(http_addr: String, api_key: String, flight: Channel, firecache: Channel) -> Self {
+    pub fn new(http_addr: String, api_key: String, flight: Channel) -> Self {
         Self {
             flight: SqlFlightClient::new(flight, api_key.clone()),
-            firecache: SqlFlightClient::new(firecache, api_key.clone()),
-            prices: PricesClient::new(Some(http_addr), api_key)
+            prices: PricesClient::new(Some(http_addr), api_key),
         }
     }
-    pub async fn query(&mut self, query: String, timeout: Option<u32>) -> Result<Streaming<FlightData>, Box<dyn Error>> {
+    pub async fn query(
+        &mut self,
+        query: String,
+        timeout: Option<u32>,
+    ) -> Result<FlightRecordBatchStream, Box<dyn Error>> {
         self.flight.query(query, timeout).await
     }
-
-    pub async fn firecache_query(&mut self, query: String, timeout: Option<u32>) -> Result<Streaming<FlightData>, Box<dyn Error>> {
-        self.firecache.query(query, timeout).await
-    }
-
-    
-}
+}

src/flight.rs

@@ -1,10 +1,7 @@
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::sql::client::FlightSqlServiceClient;
 use std::error::Error;
-use prost::Message;
-use tonic::{transport::Channel, Streaming, Request};
-use tonic::transport::channel::{Endpoint, ClientTlsConfig,};
-use arrow_flight::sql::{client::FlightSqlServiceClient, CommandStatementQuery, ProstMessageExt};
-use arrow_flight::FlightData;
-use arrow_flight::FlightDescriptor;
+use tonic::transport::Channel;
 
 pub struct SqlFlightClient {
     client: FlightSqlServiceClient<Channel>,
@@ -15,55 +12,39 @@ impl SqlFlightClient {
     pub fn new(chan: Channel, api_key: String) -> Self {
         SqlFlightClient {
             api_key: api_key,
-            client: FlightSqlServiceClient::new(chan)
+            client: FlightSqlServiceClient::new(chan),
         }
     }
 
-    pub async fn authenticate(&mut self) -> Result<(), Box<dyn Error>> {
-        let parts: Vec<&str> = self.api_key.split("|").collect();
-        if parts.len() < 2 {
+    pub async fn authenticate(&mut self) -> std::result::Result<(), Box<dyn Error>> {
+        if self.api_key.split("|").collect::<String>().len() < 2 {
             return Err("Invalid API key format".into());
         }
-        match self.client.handshake(parts[0], parts[1]).await {
+        match self.client.handshake("", &self.api_key.clone()).await {
             Err(e) => Err(e.into()),
-            Ok(v) => {
-                self.client.set_token(String::from_utf8(v.to_vec()).expect("something"));
-                Ok(())
-            }
+            Ok(_) => Ok(()),
         }
     }
 
-    pub async fn query(&mut self, query: String, _timeout: Option<u32>) -> Result<Streaming<FlightData>, Box<dyn Error>> {
+    pub async fn query(
+        &mut self,
+        query: String,
+        _timeout: Option<u32>,
+    ) -> std::result::Result<FlightRecordBatchStream, Box<dyn Error>> {
         match self.authenticate().await {
-            Err(e) => {
-                return Err(e.into())
-            },
-            Ok(()) => {},
-        };
-        
-        let cmd = CommandStatementQuery {
-            query: query.clone(),
-            ..Default::default()
+            Err(e) => return Err(e.into()),
+            Ok(()) => {}
         };
-        let fd = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
-        let req = Request::new(fd);
-
-        match self.client.inner_mut().get_flight_info(req).await {
+        match self.client.execute(query, Option::None).await {
             Ok(resp) => {
-                let flight_info = resp.into_inner();
-                for ep in flight_info.endpoint {
+                for ep in resp.endpoint {
                     if let Some(tkt) = ep.ticket {
                         return self.client.do_get(tkt).await.map_err(|e| e.into());
                     }
                 }
                 Err("no tickets for flight endpoint".into())
-            },
-            // Err(e) => {
-            //     // Handle re-authentication similar to the Python client and then retry the request.
-            //     self.authenticate().await?;
-            //     self.query(query, _timeout)
-            // },
+            }
             Err(e) => Err(e.into()),
         }
     }
-}
+}

src/lib.rs

@@ -1,10 +1,8 @@
-
 mod client;
 mod flight;
-mod prices; 
+mod prices;
 mod tls;
 
-pub use client::{SpiceClient as Client};
-pub use flight::{SqlFlightClient};
+pub use client::{new_spice_client, SpiceClient as Client};
 
 // Further public exports and integrations