Merge branch 'main' into rebase-spiceai-main-branches

Author: hozan23

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "datafusion-table-providers"
-version = "0.1.0"
+version = "0.2.0"
 readme = "README.md"
 edition = "2021"
 repository = "https://github.com/datafusion-contrib/datafusion-table-providers"
@@ -17,8 +17,6 @@ arrow-json = "53"
 async-stream = { version = "0.3.5", optional = true }
 async-trait = "0.1.80"
 num-bigint = "0.4.4"
-base64 = { version = "0.22.1", optional = true }
-bytes = { version = "1.7.1", optional = true }
 bigdecimal = "0.4.5"
 byteorder = "1.5.0"
 chrono = "0.4.38"
@@ -87,7 +85,8 @@ test-log = { version = "0.2.16", features = ["trace"] }
 rstest = "0.22.0"
 geozero = { version = "0.13.0", features = ["with-wkb"] }
 tokio-stream = { version = "0.1.15", features = ["net"] }
-arrow-schema = "52.2.0"
+arrow-schema = "53.1.0"
+prost = { version = "0.13"}
 
 [features]
 mysql = ["dep:mysql_async", "dep:async-stream"]
@@ -99,13 +98,10 @@ flight = [
   "dep:arrow-cast",
   "dep:arrow-flight",
   "dep:arrow-schema",
-  "dep:base64",
-  "dep:bytes",
   "dep:datafusion-expr",
   "dep:datafusion-physical-expr",
   "dep:datafusion-physical-plan",
   "dep:datafusion-proto",
-  "dep:prost",
   "dep:serde",
   "dep:tonic",
 ]

src/flight/sql.rs

@@ -18,22 +18,11 @@
 //! Default [FlightDriver] for Flight SQL
 
 use std::collections::HashMap;
-use std::str::FromStr;
 
 use arrow_flight::error::Result;
-use arrow_flight::flight_service_client::FlightServiceClient;
-use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt};
-use arrow_flight::{FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse};
-use arrow_schema::ArrowError;
+use arrow_flight::sql::client::FlightSqlServiceClient;
 use async_trait::async_trait;
-use base64::prelude::BASE64_STANDARD;
-use base64::Engine;
-use bytes::Bytes;
-use futures::{stream, TryStreamExt};
-use prost::Message;
-use tonic::metadata::AsciiMetadataKey;
 use tonic::transport::Channel;
-use tonic::IntoRequest;
 
 use crate::flight::{FlightDriver, FlightMetadata};
 
@@ -60,7 +49,7 @@ impl FlightDriver for FlightSqlDriver {
         channel: Channel,
         options: &HashMap<String, String>,
     ) -> Result<FlightMetadata> {
-        let mut client = FlightSqlClient::new(channel);
+        let mut client = FlightSqlServiceClient::new(channel);
         let headers = options.iter().filter_map(|(key, value)| {
             key.strip_prefix(HEADER_PREFIX)
                 .map(|header_name| (header_name, value))
@@ -75,147 +64,9 @@ impl FlightDriver for FlightSqlDriver {
         }
         let info = client.execute(options[QUERY].clone(), None).await?;
         let mut grpc_headers = HashMap::default();
-        if let Some(token) = client.token {
+        if let Some(token) = client.token() {
             grpc_headers.insert("authorization".into(), format!("Bearer {}", token));
         }
         FlightMetadata::try_new(info, grpc_headers)
     }
 }
-
-/////////////////////////////////////////////////////////////////////////
-// Shameless copy/paste from arrow-flight FlightSqlServiceClient
-// (only cherry-picked the functionality that we actually use).
-// This is only needed in order to access the bearer token received
-// during handshake, as the standard client does not expose this information.
-// The bearer token has to be passed to the clients that perform
-// the DoGet operation, since Dremio, Ballista and possibly others
-// expect the bearer token they produce with the handshake response
-// to be set on all subsequent requests, including DoGet.
-//
-// TODO: remove this and switch to the official client once
-//  https://github.com/apache/arrow-rs/pull/6254 is released,
-//  and remove a bunch of cargo dependencies, like base64 or bytes
-#[derive(Debug, Clone)]
-struct FlightSqlClient {
-    token: Option<String>,
-    headers: HashMap<String, String>,
-    flight_client: FlightServiceClient<Channel>,
-}
-
-impl FlightSqlClient {
-    /// Creates a new FlightSql client that connects to a server over an arbitrary tonic `Channel`
-    fn new(channel: Channel) -> Self {
-        Self {
-            token: None,
-            flight_client: FlightServiceClient::new(channel),
-            headers: HashMap::default(),
-        }
-    }
-
-    /// Perform a `handshake` with the server, passing credentials and establishing a session.
-    ///
-    /// If the server returns an "authorization" header, it is automatically parsed and set as
-    /// a token for future requests. Any other data returned by the server in the handshake
-    /// response is returned as a binary blob.
-    async fn handshake(
-        &mut self,
-        username: &str,
-        password: &str,
-    ) -> std::result::Result<Bytes, ArrowError> {
-        let cmd = HandshakeRequest {
-            protocol_version: 0,
-            payload: Default::default(),
-        };
-        let mut req = tonic::Request::new(stream::iter(vec![cmd]));
-        let val = BASE64_STANDARD.encode(format!("{username}:{password}"));
-        let val = format!("Basic {val}")
-            .parse()
-            .map_err(|_| ArrowError::ParseError("Cannot parse header".to_string()))?;
-        req.metadata_mut().insert("authorization", val);
-        let req = self.set_request_headers(req)?;
-        let resp = self
-            .flight_client
-            .handshake(req)
-            .await
-            .map_err(|e| ArrowError::IpcError(format!("Can't handshake {e}")))?;
-        if let Some(auth) = resp.metadata().get("authorization") {
-            let auth = auth
-                .to_str()
-                .map_err(|_| ArrowError::ParseError("Can't read auth header".to_string()))?;
-            let bearer = "Bearer ";
-            if !auth.starts_with(bearer) {
-                Err(ArrowError::ParseError("Invalid auth header!".to_string()))?;
-            }
-            let auth = auth[bearer.len()..].to_string();
-            self.token = Some(auth);
-        }
-        let responses: Vec<HandshakeResponse> = resp
-            .into_inner()
-            .try_collect()
-            .await
-            .map_err(|_| ArrowError::ParseError("Can't collect responses".to_string()))?;
-        let resp = match responses.as_slice() {
-            [resp] => resp.payload.clone(),
-            [] => Bytes::new(),
-            _ => Err(ArrowError::ParseError(
-                "Multiple handshake responses".to_string(),
-            ))?,
-        };
-        Ok(resp)
-    }
-
-    async fn execute(
-        &mut self,
-        query: String,
-        transaction_id: Option<Bytes>,
-    ) -> std::result::Result<FlightInfo, ArrowError> {
-        let cmd = CommandStatementQuery {
-            query,
-            transaction_id,
-        };
-        self.get_flight_info_for_command(cmd).await
-    }
-
-    async fn get_flight_info_for_command<M: ProstMessageExt>(
-        &mut self,
-        cmd: M,
-    ) -> std::result::Result<FlightInfo, ArrowError> {
-        let descriptor = FlightDescriptor::new_cmd(cmd.as_any().encode_to_vec());
-        let req = self.set_request_headers(descriptor.into_request())?;
-        let fi = self
-            .flight_client
-            .get_flight_info(req)
-            .await
-            .map_err(|status| ArrowError::IpcError(format!("{status:?}")))?
-            .into_inner();
-        Ok(fi)
-    }
-
-    fn set_header(&mut self, key: impl Into<String>, value: impl Into<String>) {
-        let key: String = key.into();
-        let value: String = value.into();
-        self.headers.insert(key, value);
-    }
-
-    fn set_request_headers<T>(
-        &self,
-        mut req: tonic::Request<T>,
-    ) -> std::result::Result<tonic::Request<T>, ArrowError> {
-        for (k, v) in &self.headers {
-            let k = AsciiMetadataKey::from_str(k.as_str()).map_err(|e| {
-                ArrowError::ParseError(format!("Cannot convert header key \"{k}\": {e}"))
-            })?;
-            let v = v.parse().map_err(|e| {
-                ArrowError::ParseError(format!("Cannot convert header value \"{v}\": {e}"))
-            })?;
-            req.metadata_mut().insert(k, v);
-        }
-        if let Some(token) = &self.token {
-            let val = format!("Bearer {token}").parse().map_err(|e| {
-                ArrowError::ParseError(format!("Cannot convert token to header value: {e}"))
-            })?;
-            req.metadata_mut().insert("authorization", val);
-        }
-        Ok(req)
-    }
-}