[flight] Enable TLS for tonic 0.12

Author: ccciudatu

Cargo.toml

@@ -68,7 +68,7 @@ trust-dns-resolver = "0.23.2"
 url = "2.5.1"
 pem = { version = "3.0.4", optional = true }
 tokio-rusqlite = { version = "0.5.1", optional = true }
-tonic = { version = "0.12.2", optional = true }
+tonic = { version = "0.12", optional = true, default-features = true, features = ["tls-native-roots", "tls-webpki-roots"] }
 itertools = "0.13.0"
 dyn-clone = { version = "1.0.17", optional = true }
 geo-types = "0.7.13"

src/flight.rs

@@ -20,6 +20,7 @@
 
 use std::any::Any;
 use std::collections::HashMap;
+use std::error::Error;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -35,7 +36,7 @@ use datafusion::datasource::TableProvider;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion_expr::{CreateExternalTable, Expr, TableType};
 use serde::{Deserialize, Serialize};
-use tonic::transport::Channel;
+use tonic::transport::{Channel, ClientTlsConfig};
 
 pub mod codec;
 mod exec;
@@ -107,16 +108,12 @@ impl FlightTableFactory {
         options: HashMap<String, String>,
     ) -> datafusion::common::Result<FlightTable> {
         let origin = entry_point.into();
-        let channel = Channel::from_shared(origin.clone())
-            .unwrap()
-            .connect()
-            .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let channel = flight_channel(&origin).await?;
         let metadata = self
             .driver
             .metadata(channel.clone(), &options)
             .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+            .map_err(to_df_err)?;
         let num_rows = precision(metadata.info.total_records);
         let total_byte_size = precision(metadata.info.total_bytes);
         let logical_schema = metadata.schema;
@@ -136,14 +133,6 @@ impl FlightTableFactory {
     }
 }
 
-fn precision(total: i64) -> Precision<usize> {
-    if total < 0 {
-        Precision::Absent
-    } else {
-        Precision::Exact(total as usize)
-    }
-}
-
 #[async_trait]
 impl TableProviderFactory for FlightTableFactory {
     async fn create(
@@ -292,7 +281,7 @@ impl TableProvider for FlightTable {
             .driver
             .metadata(self.channel.clone(), &self.options)
             .await
-            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+            .map_err(to_df_err)?;
         Ok(Arc::new(FlightExec::try_new(
             metadata,
             projection,
@@ -304,3 +293,26 @@ impl TableProvider for FlightTable {
         Some(self.stats.clone())
     }
 }
+
+fn to_df_err<E: Error + Send + Sync + 'static>(err: E) -> DataFusionError {
+    DataFusionError::External(Box::new(err))
+}
+
+async fn flight_channel(source: impl Into<String>) -> datafusion::common::Result<Channel> {
+    let tls_config = ClientTlsConfig::new().with_enabled_roots();
+    Channel::from_shared(source.into())
+        .map_err(to_df_err)?
+        .tls_config(tls_config)
+        .map_err(to_df_err)?
+        .connect()
+        .await
+        .map_err(to_df_err)
+}
+
+fn precision(total: i64) -> Precision<usize> {
+    if total < 0 {
+        Precision::Absent
+    } else {
+        Precision::Exact(total as usize)
+    }
+}

src/flight/codec.rs

@@ -20,6 +20,7 @@
 use std::sync::Arc;
 
 use crate::flight::exec::{FlightConfig, FlightExec};
+use crate::flight::to_df_err;
 use datafusion::common::DataFusionError;
 use datafusion_expr::registry::FunctionRegistry;
 use datafusion_physical_plan::ExecutionPlan;
@@ -37,8 +38,7 @@ impl PhysicalExtensionCodec for FlightPhysicalCodec {
         _registry: &dyn FunctionRegistry,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         if inputs.is_empty() {
-            let config: FlightConfig =
-                serde_json::from_slice(buf).map_err(|e| DataFusionError::External(Box::new(e)))?;
+            let config: FlightConfig = serde_json::from_slice(buf).map_err(to_df_err)?;
             Ok(Arc::from(FlightExec::from(config)))
         } else {
             Err(DataFusionError::Internal(
@@ -53,8 +53,7 @@ impl PhysicalExtensionCodec for FlightPhysicalCodec {
         buf: &mut Vec<u8>,
     ) -> datafusion::common::Result<()> {
         if let Some(flight) = node.as_any().downcast_ref::<FlightExec>() {
-            let mut bytes = serde_json::to_vec(flight.config())
-                .map_err(|e| DataFusionError::External(Box::new(e)))?;
+            let mut bytes = serde_json::to_vec(flight.config()).map_err(to_df_err)?;
             buf.append(&mut bytes);
             Ok(())
         } else {

src/flight/exec.rs

@@ -23,7 +23,7 @@ use std::fmt::{Debug, Formatter};
 use std::str::FromStr;
 use std::sync::Arc;
 
-use crate::flight::{FlightMetadata, FlightProperties, SizeLimits};
+use crate::flight::{flight_channel, to_df_err, FlightMetadata, FlightProperties, SizeLimits};
 use arrow_array::RecordBatch;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
@@ -41,7 +41,6 @@ use datafusion_physical_plan::{
 use futures::{StreamExt, TryStreamExt};
 use serde::{Deserialize, Serialize};
 use tonic::metadata::{AsciiMetadataKey, MetadataMap};
-use tonic::transport::Channel;
 
 /// Arrow Flight physical plan that maps flight endpoints to partitions
 #[derive(Clone, Debug)]
@@ -170,11 +169,7 @@ async fn flight_client(
     grpc_headers: &MetadataMap,
     size_limits: &SizeLimits,
 ) -> Result<FlightClient> {
-    let channel = Channel::from_shared(source.into())
-        .map_err(|e| DataFusionError::External(Box::new(e)))?
-        .connect()
-        .await
-        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+    let channel = flight_channel(source).await?;
     let inner_client = FlightServiceClient::new(channel)
         .max_encoding_message_size(size_limits.encoding)
         .max_decoding_message_size(size_limits.decoding);
@@ -212,10 +207,7 @@ async fn try_fetch_stream(
     schema: SchemaRef,
 ) -> arrow_flight::error::Result<SendableRecordBatchStream> {
     let ticket = Ticket::new(ticket.0.to_vec());
-    let stream = client
-        .do_get(ticket)
-        .await?
-        .map_err(|e| DataFusionError::External(Box::new(e)));
+    let stream = client.do_get(ticket).await?.map_err(to_df_err);
     Ok(Box::pin(RecordBatchStreamAdapter::new(
         schema.clone(),
         stream.map(move |item| item.and_then(|rb| enforce_schema(rb, &schema).map_err(Into::into))),