Allow flightsql and spiceai connectors to override flight max message size (spiceai#5407)

* Allow `flightsql` and `spiceai` connectors to override default flight max message size

* Update

Author: sgrebnov

Cargo.lock

@@ -64,6 +64,7 @@ base64 = "0.22.1"
 bb8 = "0.8"
 bb8-postgres = "0.8"
 bytes = "1.10.0"
+byte-unit = "5.1.4"
 charset = "0.1.5"
 chrono = "0.4.38"
 clap = { version = "4.5.36", features = ["derive", "env"] }

Cargo.toml

@@ -12,7 +12,7 @@ version.workspace = true
 arrow.workspace = true
 async-stream.workspace = true
 async-trait.workspace = true
-byte-unit = "5.1.4"
+byte-unit.workspace = true
 datafusion.workspace = true
 fundu = { workspace = true }
 futures.workspace = true

crates/cache/Cargo.toml

@@ -246,6 +246,20 @@ impl FlightClient {
         self
     }
 
+    /// Overrides default maximum message size for encoding and decoding.
+    #[must_use]
+    pub fn with_max_message_size(
+        mut self,
+        max_encoding_message_size: usize,
+        max_decoding_message_size: usize,
+    ) -> Self {
+        self.flight_client = self
+            .flight_client
+            .max_encoding_message_size(max_encoding_message_size)
+            .max_decoding_message_size(max_decoding_message_size);
+        self
+    }
+
     /// Queries the flight service for the schema of the path.
     ///
     /// # Arguments

crates/flight_client/src/lib.rs

@@ -27,6 +27,7 @@ use crate::get_params_with_secrets;
 use crate::parameters::ParameterSpec;
 use crate::parameters::Parameters;
 use crate::secrets::Secrets;
+use app::App;
 use arrow_schema::SchemaRef;
 use arrow_tools::schema::schema_meta_get_computed_columns;
 use async_trait::async_trait;
@@ -579,6 +580,7 @@ pub struct ConnectorParams {
     pub(crate) parameters: Parameters,
     pub(crate) unsupported_type_action: Option<UnsupportedTypeAction>,
     pub(crate) component: ConnectorComponent,
+    pub(crate) app: Option<Arc<App>>,
 }
 
 pub struct ConnectorParamsBuilder {
@@ -601,7 +603,7 @@ impl ConnectorParamsBuilder {
     ) -> Result<ConnectorParams, Box<dyn std::error::Error + Send + Sync>> {
         let name = self.connector.to_string();
         let mut unsupported_type_action = None;
-        let (params, prefix, parameters) = match &self.component {
+        let (params, prefix, parameters, app) = match &self.component {
             ConnectorComponent::Catalog(catalog) => {
                 let guard = CATALOG_CONNECTOR_FACTORY_REGISTRY.lock().await;
                 let connector_factory = guard.get(&name);
@@ -616,6 +618,7 @@ impl ConnectorParamsBuilder {
                     get_params_with_secrets(Arc::clone(&secrets), &catalog.params).await,
                     factory.prefix(),
                     factory.parameters(),
+                    None,
                 )
             }
             ConnectorComponent::Dataset(dataset) => {
@@ -639,7 +642,12 @@ impl ConnectorParamsBuilder {
 
                 let params = get_params_with_secrets(Arc::clone(&secrets), &dataset.params).await;
 
-                (params, factory.prefix(), factory.parameters())
+                (
+                    params,
+                    factory.prefix(),
+                    factory.parameters(),
+                    Some(Arc::clone(&dataset.app)),
+                )
             }
         };
 
@@ -656,6 +664,7 @@ impl ConnectorParamsBuilder {
             parameters,
             unsupported_type_action: unsupported_type_action.map(UnsupportedTypeAction::from),
             component: self.component,
+            app,
         })
     }
 }

crates/runtime/src/dataconnector.rs

@@ -43,6 +43,14 @@ pub enum Error {
 
     #[snafu(display("Failed to connect to the Flight server.\n{source}"))]
     UnableToPerformHandshake { source: arrow::error::ArrowError },
+
+    #[snafu(display(
+        "Failed to apply parameter '{parameter}': {source}. Ensure the value is valid and retry.\nFor details, visit: https://spiceai.org/docs/components/data-connectors/flightsql#params"
+    ))]
+    InvalidParameterValue {
+        parameter: String,
+        source: Box<dyn std::error::Error + Send + Sync>,
+    },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -95,9 +103,24 @@ impl DataConnectorFactory for FlightSQLFactory {
                 .await
                 .context(UnableToConstructTlsChannelSnafu)?;
 
+            let max_message_size =
+                match params
+                    .app
+                    .as_ref()
+                    .and_then(|app| app.runtime.flight.as_ref())
+                {
+                    Some(flight) => flight.max_message_size_bytes().map_err(|err| {
+                        Error::InvalidParameterValue {
+                            parameter: "max_message_size".to_string(),
+                            source: err,
+                        }
+                    })?,
+                    None => None,
+                };
+
             let flight_client = FlightServiceClient::new(flight_channel)
-                .max_encoding_message_size(MAX_ENCODING_MESSAGE_SIZE)
-                .max_decoding_message_size(MAX_DECODING_MESSAGE_SIZE);
+                .max_encoding_message_size(max_message_size.unwrap_or(MAX_ENCODING_MESSAGE_SIZE))
+                .max_decoding_message_size(max_message_size.unwrap_or(MAX_DECODING_MESSAGE_SIZE));
 
             let mut client = FlightSqlServiceClient::new_from_inner(flight_client);
             let username = params.parameters.get("username").expose().ok();

crates/runtime/src/dataconnector/flightsql.rs

@@ -72,6 +72,14 @@ pub enum Error {
         value: Arc<str>,
         source: InvalidMetadataValue,
     },
+
+    #[snafu(display(
+        "Failed to apply parameter '{parameter}': {source}. Ensure the value is valid and retry.\nFor details, visit: https://spiceai.org/docs/components/data-connectors/spiceai#parameters"
+    ))]
+    InvalidParameterValue {
+        parameter: String,
+        source: Box<dyn std::error::Error + Send + Sync>,
+    },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -188,9 +196,12 @@ impl DataConnectorFactory for SpiceAIFactory {
                 .ok_or_else(|p| MissingRequiredParameterSnafu { parameter: p.0 }.build())?;
             let credentials = Credentials::new("", api_key.clone());
 
-            let flight_client = FlightClient::try_new(url, credentials, None)
+            let mut flight_client = FlightClient::try_new(url, credentials, None)
                 .await
                 .context(UnableToCreateFlightClientSnafu)?;
+
+            flight_client = configure_max_message_size(flight_client, &params)?;
+
             let flight_factory = FlightFactory::new(
                 "spice.ai",
                 flight_client,
@@ -211,6 +222,29 @@ impl DataConnectorFactory for SpiceAIFactory {
     }
 }
 
+/// Configures flight client's message size based on app parameters
+fn configure_max_message_size(
+    mut flight_client: FlightClient,
+    params: &ConnectorParams,
+) -> Result<FlightClient> {
+    if let Some(app) = params.app.as_ref() {
+        if let Some(flight) = app.runtime.flight.as_ref() {
+            if let Some(max_message_size) =
+                flight
+                    .max_message_size_bytes()
+                    .map_err(|err| Error::InvalidParameterValue {
+                        parameter: "max_message_size".to_string(),
+                        source: err,
+                    })?
+            {
+                flight_client =
+                    flight_client.with_max_message_size(max_message_size, max_message_size);
+            }
+        }
+    }
+    Ok(flight_client)
+}
+
 #[async_trait]
 impl DataConnector for SpiceAI {
     fn as_any(&self) -> &dyn Any {

crates/runtime/src/dataconnector/spiceai.rs

@@ -8,6 +8,7 @@ rust-version.workspace = true
 version.workspace = true
 
 [dependencies]
+byte-unit.workspace = true
 fundu.workspace = true
 regex.workspace = true
 schemars = { version = "0.8.22", optional = true }

crates/spicepod/Cargo.toml

@@ -56,6 +56,9 @@ pub struct Runtime {
     #[serde(default, skip_serializing_if = "is_default")]
     pub cors: CorsConfig,
 
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub flight: Option<Flight>,
+
     /// Configures where the runtime will store temporary files needed for operations like
     /// spilling to disk for queries & accelerations that are larger than memory.
     #[serde(default, skip_serializing_if = "Option::is_none")]
@@ -183,6 +186,30 @@ impl Default for TelemetryConfig {
     }
 }
 
+#[derive(Default, Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[cfg_attr(feature = "schemars", derive(JsonSchema))]
+pub struct Flight {
+    pub max_message_size: Option<String>,
+}
+
+impl Flight {
+    pub fn max_message_size_bytes(&self) -> Result<Option<usize>, Box<dyn Error + Send + Sync>> {
+        if let Some(size_str) = &self.max_message_size {
+            let size_in_bytes = usize::try_from(
+                byte_unit::Byte::parse_str(size_str, true)
+                    .map_err(|e| {
+                        format!("Failed to parse 'max_message_size' value '{size_str}': {e}")
+                    })?
+                    .as_u64(),
+            )
+            .unwrap_or_default();
+            Ok(Some(size_in_bytes))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[serde(deny_unknown_fields)]
 #[cfg_attr(feature = "schemars", derive(JsonSchema))]