feat(odbc): introduce OdbcBufferSettings for performance tuning

lovasoa · lovasoa · commit 65829dbf27a3 · 2025-09-29T12:22:07.000+02:00
This commit adds OdbcBufferSettings to configure buffer settings for ODBC connections, allowing users to optimize memory usage and performance during data fetching. The new settings include batch size and maximum column size, with corresponding methods in OdbcConnectOptions for easy configuration.
diff --git a/sqlx-core/src/odbc/connection/mod.rs b/sqlx-core/src/odbc/connection/mod.rs
@@ -1,7 +1,8 @@
 use crate::connection::Connection;
 use crate::error::Error;
 use crate::odbc::{
-    Odbc, OdbcArguments, OdbcColumn, OdbcConnectOptions, OdbcQueryResult, OdbcRow, OdbcTypeInfo,
+    Odbc, OdbcArguments, OdbcBufferSettings, OdbcColumn, OdbcConnectOptions, OdbcQueryResult,
+    OdbcRow, OdbcTypeInfo,
 };
 use crate::transaction::Transaction;
 use either::Either;
@@ -67,12 +68,14 @@ pub(super) fn decode_column_name<T: ColumnNameDecode>(name: T, index: u16) -> St
 pub struct OdbcConnection {
     pub(crate) conn: SharedConnection<'static>,
     pub(crate) stmt_cache: HashMap<Arc<str>, SharedPreparedStatement>,
+    pub(crate) buffer_settings: OdbcBufferSettings,
 }
 
 impl std::fmt::Debug for OdbcConnection {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("OdbcConnection")
             .field("conn", &self.conn)
+            .field("buffer_settings", &self.buffer_settings)
             .finish()
     }
 }
@@ -108,6 +111,7 @@ impl OdbcConnection {
         Ok(Self {
             conn: shared_conn,
             stmt_cache: HashMap::new(),
+            buffer_settings: options.buffer_settings,
         })
     }
 
@@ -162,9 +166,10 @@ impl OdbcConnection {
         };
 
         let conn = Arc::clone(&self.conn);
+        let buffer_settings = self.buffer_settings;
         sqlx_rt::spawn(sqlx_rt::spawn_blocking(move || {
             let mut conn = conn.lock().expect("failed to lock connection");
-            if let Err(e) = execute_sql(&mut conn, maybe_prepared, args, &tx) {
+            if let Err(e) = execute_sql(&mut conn, maybe_prepared, args, &tx, buffer_settings) {
                 let _ = tx.send(Err(e));
             }
         }));
diff --git a/sqlx-core/src/odbc/connection/odbc_bridge.rs b/sqlx-core/src/odbc/connection/odbc_bridge.rs
@@ -1,34 +1,30 @@
 use super::decode_column_name;
 use crate::error::Error;
 use crate::odbc::{
-    connection::MaybePrepared, ColumnData, OdbcArgumentValue, OdbcArguments, OdbcBatch, OdbcColumn,
-    OdbcQueryResult, OdbcRow, OdbcTypeInfo,
+    connection::MaybePrepared, ColumnData, OdbcArgumentValue, OdbcArguments, OdbcBatch,
+    OdbcBufferSettings, OdbcColumn, OdbcQueryResult, OdbcRow, OdbcTypeInfo,
 };
 use either::Either;
 use flume::{SendError, Sender};
 use odbc_api::buffers::{AnySlice, BufferDesc, ColumnarAnyBuffer};
 use odbc_api::handles::{AsStatementRef, Nullability, Statement};
 use odbc_api::DataType;
 use odbc_api::{Cursor, IntoParameter, ResultSetMetadata};
-use std::cmp::min;
 use std::sync::Arc;
 
 // Bulk fetch implementation using columnar buffers instead of row-by-row fetching
 // This provides significant performance improvements by fetching rows in batches
 // and avoiding the slow `next_row()` method from odbc-api
-const BATCH_SIZE: usize = 128;
-const DEFAULT_TEXT_LEN: usize = 512;
-const DEFAULT_BINARY_LEN: usize = 1024;
-const DEFAULT_NUMERIC_TEXT_LEN: usize = 128;
-const MIN_TEXT_LEN: usize = 1024;
-const MAX_TEXT_LEN: usize = 4096;
 
 struct ColumnBinding {
     column: OdbcColumn,
     buffer_desc: BufferDesc,
 }
 
-fn build_bindings<C>(cursor: &mut C) -> Result<Vec<ColumnBinding>, Error>
+fn build_bindings<C>(
+    cursor: &mut C,
+    buffer_settings: &OdbcBufferSettings,
+) -> Result<Vec<ColumnBinding>, Error>
 where
     C: ResultSetMetadata,
 {
@@ -40,7 +36,13 @@ where
             .col_nullability(index as u16)
             .unwrap_or(Nullability::Unknown)
             .could_be_nullable();
-        let buffer_desc = map_buffer_desc(cursor, index as u16, &column.type_info, nullable)?;
+        let buffer_desc = map_buffer_desc(
+            cursor,
+            index as u16,
+            &column.type_info,
+            nullable,
+            buffer_settings,
+        )?;
         bindings.push(ColumnBinding {
             column,
             buffer_desc,
@@ -67,21 +69,22 @@ pub fn execute_sql(
     maybe_prepared: MaybePrepared,
     args: Option<OdbcArguments>,
     tx: &ExecuteSender,
+    buffer_settings: OdbcBufferSettings,
 ) -> Result<(), Error> {
     let params = prepare_parameters(args);
 
     let affected = match maybe_prepared {
         MaybePrepared::Prepared(prepared) => {
             let mut prepared = prepared.lock().expect("prepared statement lock");
             if let Some(cursor) = prepared.execute(&params[..])? {
-                handle_cursor(cursor, tx);
+                handle_cursor(cursor, tx, buffer_settings);
             }
             extract_rows_affected(&mut *prepared)
         }
         MaybePrepared::NotPrepared(sql) => {
             let mut preallocated = conn.preallocate().map_err(Error::from)?;
             if let Some(cursor) = preallocated.execute(&sql, &params[..])? {
-                handle_cursor(cursor, tx);
+                handle_cursor(cursor, tx, buffer_settings);
             }
             extract_rows_affected(&mut preallocated)
         }
@@ -127,19 +130,19 @@ fn to_param(arg: OdbcArgumentValue) -> Box<dyn odbc_api::parameter::InputParamet
     }
 }
 
-fn handle_cursor<C>(mut cursor: C, tx: &ExecuteSender)
+fn handle_cursor<C>(mut cursor: C, tx: &ExecuteSender, buffer_settings: OdbcBufferSettings)
 where
     C: Cursor + ResultSetMetadata,
 {
-    let bindings = match build_bindings(&mut cursor) {
+    let bindings = match build_bindings(&mut cursor, &buffer_settings) {
         Ok(b) => b,
         Err(e) => {
             send_error(tx, e);
             return;
         }
     };
 
-    match stream_rows(cursor, bindings, tx) {
+    match stream_rows(cursor, bindings, tx, buffer_settings) {
         Ok(true) => {
             let _ = send_done(tx, 0);
         }
@@ -181,6 +184,7 @@ fn map_buffer_desc<C>(
     _column_index: u16,
     type_info: &OdbcTypeInfo,
     nullable: bool,
+    buffer_settings: &OdbcBufferSettings,
 ) -> Result<BufferDesc, Error>
 where
     C: ResultSetMetadata,
@@ -200,9 +204,14 @@ where
         | DataType::Varbinary { length }
         | DataType::LongVarbinary { length } => BufferDesc::Binary {
             length: if let Some(length) = length {
-                length.get().clamp(MIN_TEXT_LEN, MAX_TEXT_LEN)
+                if length.get() < 255 {
+                    // Some drivers report 255 for max length
+                    length.get()
+                } else {
+                    buffer_settings.max_column_size
+                }
             } else {
-                MAX_TEXT_LEN
+                buffer_settings.max_column_size
             },
         },
         DataType::Char { length }
@@ -216,16 +225,20 @@ where
             ..
         } => BufferDesc::Text {
             max_str_len: if let Some(length) = length {
-                length.get().clamp(MIN_TEXT_LEN, MAX_TEXT_LEN)
+                if length.get() < 255 {
+                    length.get()
+                } else {
+                    buffer_settings.max_column_size
+                }
             } else {
-                MAX_TEXT_LEN
+                buffer_settings.max_column_size
             },
         },
         DataType::Unknown => BufferDesc::Text {
-            max_str_len: MAX_TEXT_LEN,
+            max_str_len: buffer_settings.max_column_size,
         },
         DataType::Decimal { .. } | DataType::Numeric { .. } => BufferDesc::Text {
-            max_str_len: min(DEFAULT_NUMERIC_TEXT_LEN, MAX_TEXT_LEN),
+            max_str_len: buffer_settings.max_column_size,
         },
     };
 
@@ -236,12 +249,13 @@ fn stream_rows<C>(
     cursor: C,
     bindings: Vec<ColumnBinding>,
     tx: &ExecuteSender,
+    buffer_settings: OdbcBufferSettings,
 ) -> Result<bool, Error>
 where
     C: Cursor + ResultSetMetadata,
 {
     let buffer_descriptions: Vec<_> = bindings.iter().map(|b| b.buffer_desc).collect();
-    let buffer = ColumnarAnyBuffer::from_descs(BATCH_SIZE, buffer_descriptions);
+    let buffer = ColumnarAnyBuffer::from_descs(buffer_settings.batch_size, buffer_descriptions);
     let mut row_set_cursor = cursor.bind_buffer(buffer)?;
 
     let mut receiver_open = true;
diff --git a/sqlx-core/src/odbc/mod.rs b/sqlx-core/src/odbc/mod.rs
@@ -20,6 +20,27 @@
 //! ```text
 //! odbc:DSN=MyDataSource
 //! ```
+//!
+//! ## Buffer Configuration
+//!
+//! You can configure buffer settings for performance tuning:
+//!
+//! ```rust
+//! use sqlx::odbc::{OdbcConnectOptions, OdbcBufferSettings};
+//!
+//! let mut opts = OdbcConnectOptions::from_str("DSN=MyDataSource")?;
+//!
+//! // Configure for high-throughput scenarios
+//! opts.buffer_settings(OdbcBufferSettings {
+//!     batch_size: 256,        // Fetch 256 rows at once
+//!     max_column_size: 2048,  // Limit text columns to 2048 chars
+//! });
+//!
+//! // Or configure individual settings
+//! opts.batch_size(512)
+//!     .max_column_size(1024);
+//! # Ok::<(), sqlx::Error>(())
+//! ```
 
 use crate::executor::Executor;
 
@@ -41,7 +62,7 @@ pub use arguments::{OdbcArgumentValue, OdbcArguments};
 pub use column::OdbcColumn;
 pub use connection::OdbcConnection;
 pub use database::Odbc;
-pub use options::OdbcConnectOptions;
+pub use options::{OdbcBufferSettings, OdbcConnectOptions};
 pub use query_result::OdbcQueryResult;
 pub use row::{OdbcBatch, OdbcRow};
 pub use statement::{OdbcStatement, OdbcStatementMetadata};
diff --git a/sqlx-core/src/odbc/options/mod.rs b/sqlx-core/src/odbc/options/mod.rs
@@ -8,22 +8,123 @@ use std::time::Duration;
 
 use crate::odbc::OdbcConnection;
 
+/// Configuration for ODBC buffer settings that control memory usage and performance characteristics.
+///
+/// These settings affect how SQLx fetches and processes data from ODBC data sources. Careful tuning
+/// of these parameters can significantly impact memory usage and query performance.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct OdbcBufferSettings {
+    /// The number of rows to fetch in each batch during bulk operations.
+    ///
+    /// **Performance Impact:**
+    /// - Higher values reduce the number of round-trips to the database but increase memory usage
+    /// - Lower values reduce memory usage but may increase latency due to more frequent fetches
+    /// - Typical range: 32-512 rows
+    ///
+    /// **Memory Impact:**
+    /// - Each batch allocates buffers for `batch_size * number_of_columns` cells
+    /// - For wide result sets, this can consume significant memory
+    ///
+    /// **Default:** 128 rows
+    pub batch_size: usize,
+
+    /// The maximum size (in characters) for text and binary columns when the database doesn't specify a length.
+    ///
+    /// **Performance Impact:**
+    /// - Higher values ensure large text fields are fully captured but increase memory allocation
+    /// - Lower values may truncate data but reduce memory pressure
+    /// - Affects VARCHAR, NVARCHAR, TEXT, and BLOB column types
+    ///
+    /// **Memory Impact:**
+    /// - Directly controls buffer size for variable-length columns
+    /// - Setting too high can waste memory; setting too low can cause data truncation
+    /// - Consider your data characteristics when tuning this value
+    pub max_column_size: usize,
+}
+
+impl Default for OdbcBufferSettings {
+    fn default() -> Self {
+        Self {
+            batch_size: 128,
+            max_column_size: 4096,
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct OdbcConnectOptions {
     pub(crate) conn_str: String,
     pub(crate) log_settings: LogSettings,
+    pub(crate) buffer_settings: OdbcBufferSettings,
 }
 
 impl OdbcConnectOptions {
     pub fn connection_string(&self) -> &str {
         &self.conn_str
     }
+
+    /// Sets the buffer configuration for this connection.
+    ///
+    /// The buffer settings control memory usage and performance characteristics
+    /// when fetching data from ODBC data sources.
+    ///
+    /// # Example
+    /// ```rust
+    /// use sqlx::odbc::{OdbcConnectOptions, OdbcBufferSettings};
+    ///
+    /// let mut opts = OdbcConnectOptions::from_str("DSN=MyDataSource")?;
+    /// opts.buffer_settings(OdbcBufferSettings {
+    ///     batch_size: 256,
+    ///     max_column_size: 2048,
+    /// });
+    /// # Ok::<(), sqlx::Error>(())
+    /// ```
+    pub fn buffer_settings(&mut self, settings: OdbcBufferSettings) -> &mut Self {
+        self.buffer_settings = settings;
+        self
+    }
+
+    /// Sets the batch size for bulk fetch operations.
+    ///
+    /// This controls how many rows are fetched at once during query execution.
+    /// Higher values can improve performance for large result sets but use more memory.
+    ///
+    /// # Panics
+    /// Panics if `batch_size` is 0.
+    pub fn batch_size(&mut self, batch_size: usize) -> &mut Self {
+        assert!(batch_size > 0, "batch_size must be greater than 0");
+        self.buffer_settings.batch_size = batch_size;
+        self
+    }
+
+    /// Sets the maximum column size for text and binary data.
+    ///
+    /// This controls the buffer size allocated for columns when the database
+    /// doesn't specify a maximum length. Larger values ensure complete data
+    /// capture but increase memory usage.
+    ///
+    /// # Panics
+    /// Panics if `max_column_size` is less than 1024 or greater than 4096.
+    pub fn max_column_size(&mut self, max_column_size: usize) -> &mut Self {
+        assert!(
+            (1024..=4096).contains(&max_column_size),
+            "max_column_size must be between 1024 and 4096"
+        );
+        self.buffer_settings.max_column_size = max_column_size;
+        self
+    }
+
+    /// Returns the current buffer settings for this connection.
+    pub fn buffer_settings_ref(&self) -> &OdbcBufferSettings {
+        &self.buffer_settings
+    }
 }
 
 impl Debug for OdbcConnectOptions {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("OdbcConnectOptions")
             .field("conn_str", &"<redacted>")
+            .field("buffer_settings", &self.buffer_settings)
             .finish()
     }
 }
@@ -51,6 +152,7 @@ impl FromStr for OdbcConnectOptions {
         Ok(Self {
             conn_str,
             log_settings: LogSettings::default(),
+            buffer_settings: OdbcBufferSettings::default(),
         })
     }
 }
