refactor(odbc): restructure OdbcStatement and enhance metadata handling

This commit refactors the OdbcStatement structure to encapsulate metadata, including columns and parameters, within a dedicated OdbcStatementMetadata struct. It also updates the OdbcConnection to cache prepared statement metadata, improving performance and reducing redundant metadata retrieval. Additionally, the prepare method is streamlined to utilize the new metadata structure.

Author: lovasoa

sqlx-core/src/any/connection/executor.rs

@@ -128,15 +128,7 @@ impl<'c> Executor<'c> for &'c mut AnyConnection {
                 AnyConnectionKind::Mssql(conn) => conn.prepare(sql).await.map(Into::into)?,
 
                 #[cfg(feature = "odbc")]
-                AnyConnectionKind::Odbc(conn) => {
-                    let (_, columns, parameters) = conn.prepare_metadata(sql).await?;
-                    crate::odbc::OdbcStatement {
-                        sql: sql.into(),
-                        columns,
-                        parameters,
-                    }
-                    .into()
-                }
+                AnyConnectionKind::Odbc(conn) => conn.prepare(sql).await.map(Into::into)?,
             })
         })
     }

sqlx-core/src/odbc/connection/executor.rs

@@ -6,7 +6,6 @@ use either::Either;
 use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
 use futures_util::TryStreamExt;
-use std::borrow::Cow;
 
 // run method removed; fetch_many implements streaming directly
 
@@ -59,14 +58,7 @@ impl<'c> Executor<'c> for &'c mut OdbcConnection {
     where
         'c: 'e,
     {
-        Box::pin(async move {
-            let (_, columns, parameters) = self.prepare_metadata(sql).await?;
-            Ok(OdbcStatement {
-                sql: Cow::Borrowed(sql),
-                columns,
-                parameters,
-            })
-        })
+        Box::pin(async move { self.prepare(sql).await })
     }
 
     #[doc(hidden)]

sqlx-core/src/odbc/connection/mod.rs

@@ -1,15 +1,48 @@
 use crate::connection::Connection;
 use crate::error::Error;
 use crate::odbc::blocking::run_blocking;
-use crate::odbc::{Odbc, OdbcArguments, OdbcColumn, OdbcConnectOptions, OdbcQueryResult, OdbcRow};
+use crate::odbc::{
+    Odbc, OdbcArguments, OdbcColumn, OdbcConnectOptions, OdbcQueryResult, OdbcRow, OdbcTypeInfo,
+};
 use crate::transaction::Transaction;
 use either::Either;
 mod odbc_bridge;
 use futures_core::future::BoxFuture;
 use futures_util::future;
-use odbc_bridge::{do_prepare, establish_connection, execute_sql, OdbcConn};
+use odbc_bridge::{establish_connection, execute_sql};
 // no direct spawn_blocking here; use run_blocking helper
-use std::sync::{Arc, Mutex};
+use crate::odbc::{OdbcStatement, OdbcStatementMetadata};
+use odbc_api::ResultSetMetadata;
+use std::borrow::Cow;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+fn collect_columns(
+    prepared: &mut odbc_api::Prepared<odbc_api::handles::StatementImpl<'_>>,
+) -> Vec<OdbcColumn> {
+    let count = prepared.num_result_cols().unwrap_or(0);
+    (1..=count)
+        .map(|i| create_column(prepared, i as u16))
+        .collect()
+}
+
+fn create_column(
+    stmt: &mut odbc_api::Prepared<odbc_api::handles::StatementImpl<'_>>,
+    index: u16,
+) -> OdbcColumn {
+    let mut cd = odbc_api::ColumnDescription::default();
+    let _ = stmt.describe_col(index, &mut cd);
+
+    OdbcColumn {
+        name: decode_column_name(cd.name, index),
+        type_info: OdbcTypeInfo::new(cd.data_type),
+        ordinal: usize::from(index.checked_sub(1).unwrap()),
+    }
+}
+
+fn decode_column_name(name_bytes: Vec<u8>, index: u16) -> String {
+    String::from_utf8(name_bytes).unwrap_or_else(|_| format!("col{}", index - 1))
+}
 
 mod executor;
 
@@ -19,86 +52,88 @@ mod executor;
 /// thread-pool via `spawn_blocking` and synchronize access with a mutex.
 #[derive(Debug)]
 pub struct OdbcConnection {
-    pub(crate) conn: Arc<Mutex<OdbcConn>>,
+    pub(crate) conn: odbc_api::SharedConnection<'static>,
+    pub(crate) stmt_cache: HashMap<u64, crate::odbc::statement::OdbcStatementMetadata>,
 }
 
 impl OdbcConnection {
-    #[inline]
-    async fn with_conn<T, F>(&self, f: F) -> Result<T, Error>
-    where
-        T: Send + 'static,
-        F: FnOnce(&mut OdbcConn) -> Result<T, Error> + Send + 'static,
-    {
-        let inner = self.conn.clone();
-        run_blocking(move || {
-            let mut conn = inner.lock().unwrap();
-            f(&mut conn)
-        })
-        .await
-    }
-
-    #[inline]
-    async fn with_conn_map<T, E, F>(&self, ctx: &'static str, f: F) -> Result<T, Error>
-    where
-        T: Send + 'static,
-        E: std::fmt::Display,
-        F: FnOnce(&mut OdbcConn) -> Result<T, E> + Send + 'static,
-    {
-        let inner = self.conn.clone();
-        run_blocking(move || {
-            let mut conn = inner.lock().unwrap();
-            f(&mut conn).map_err(|e| Error::Protocol(format!("{}: {}", ctx, e)))
-        })
-        .await
-    }
-
     pub(crate) async fn establish(options: &OdbcConnectOptions) -> Result<Self, Error> {
-        let conn = run_blocking({
+        let shared_conn = run_blocking({
             let options = options.clone();
-            move || establish_connection(&options)
+            move || {
+                let conn = establish_connection(&options)?;
+                let shared_conn = odbc_api::SharedConnection::new(std::sync::Mutex::new(conn));
+                Ok::<_, Error>(shared_conn)
+            }
         })
         .await?;
 
         Ok(Self {
-            conn: Arc::new(Mutex::new(conn)),
+            conn: shared_conn,
+            stmt_cache: HashMap::new(),
         })
     }
 
     /// Returns the name of the actual Database Management System (DBMS) this
     /// connection is talking to as reported by the ODBC driver.
     pub async fn dbms_name(&mut self) -> Result<String, Error> {
-        self.with_conn_map::<_, _, _>("Failed to get DBMS name", |conn| {
-            conn.conn.database_management_system_name()
+        let conn = Arc::clone(&self.conn);
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            conn_guard
+                .database_management_system_name()
+                .map_err(Error::from)
         })
         .await
     }
 
     pub(crate) async fn ping_blocking(&mut self) -> Result<(), Error> {
-        self.with_conn_map::<_, _, _>("Ping failed", |conn| {
-            conn.conn.execute("SELECT 1", (), None).map(|_| ())
+        let conn = Arc::clone(&self.conn);
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            conn_guard
+                .execute("SELECT 1", (), None)
+                .map_err(Error::from)
+                .map(|_| ())
         })
         .await
     }
 
     pub(crate) async fn begin_blocking(&mut self) -> Result<(), Error> {
-        self.with_conn_map::<_, _, _>("Failed to begin transaction", |conn| {
-            conn.conn.set_autocommit(false)
+        let conn = Arc::clone(&self.conn);
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            conn_guard.set_autocommit(false).map_err(Error::from)
         })
         .await
     }
 
     pub(crate) async fn commit_blocking(&mut self) -> Result<(), Error> {
-        self.with_conn_map::<_, _, _>("Failed to commit transaction", |conn| {
-            conn.conn.commit()?;
-            conn.conn.set_autocommit(true)
+        let conn = Arc::clone(&self.conn);
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            conn_guard.commit()?;
+            conn_guard.set_autocommit(true).map_err(Error::from)
         })
         .await
     }
 
     pub(crate) async fn rollback_blocking(&mut self) -> Result<(), Error> {
-        self.with_conn_map::<_, _, _>("Failed to rollback transaction", |conn| {
-            conn.conn.rollback()?;
-            conn.conn.set_autocommit(true)
+        let conn = Arc::clone(&self.conn);
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            conn_guard.rollback()?;
+            conn_guard.set_autocommit(true).map_err(Error::from)
         })
         .await
     }
@@ -111,23 +146,79 @@ impl OdbcConnection {
         let (tx, rx) = flume::bounded(64);
         let sql = sql.to_string();
         let args_move = args;
-        self.with_conn(move |conn| {
-            if let Err(e) = execute_sql(conn, &sql, args_move, &tx) {
+        let conn = Arc::clone(&self.conn);
+
+        run_blocking(move || {
+            let mut conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            if let Err(e) = execute_sql(&mut conn_guard, &sql, args_move, &tx) {
                 let _ = tx.send(Err(e));
             }
             Ok(())
         })
         .await?;
+
         Ok(rx)
     }
 
     pub(crate) async fn prepare_metadata(
         &mut self,
         sql: &str,
     ) -> Result<(u64, Vec<OdbcColumn>, usize), Error> {
+        use std::collections::hash_map::DefaultHasher;
+        use std::hash::{Hash, Hasher};
+
+        let mut hasher = DefaultHasher::new();
+        sql.hash(&mut hasher);
+        let key = hasher.finish();
+
+        // Check cache first
+        if let Some(metadata) = self.stmt_cache.get(&key) {
+            return Ok((key, metadata.columns.clone(), metadata.parameters));
+        }
+
+        // Create new prepared statement to get metadata
         let sql = sql.to_string();
-        self.with_conn(move |conn| do_prepare(conn, sql.into()))
-            .await
+        let conn = Arc::clone(&self.conn);
+
+        run_blocking(move || {
+            let conn_guard = conn
+                .lock()
+                .map_err(|_| Error::Protocol("Failed to lock connection".into()))?;
+            let mut prepared = conn_guard.prepare(&sql).map_err(Error::from)?;
+            let columns = collect_columns(&mut prepared);
+            let params = usize::from(prepared.num_params().unwrap_or(0));
+            Ok::<_, Error>((columns, params))
+        })
+        .await
+        .map(|(columns, params)| {
+            // Cache the metadata
+            let metadata = crate::odbc::statement::OdbcStatementMetadata {
+                columns: columns.clone(),
+                parameters: params,
+            };
+            self.stmt_cache.insert(key, metadata);
+            (key, columns, params)
+        })
+    }
+
+    pub(crate) async fn clear_cached_statements(&mut self) -> Result<(), Error> {
+        // Clear the statement metadata cache
+        self.stmt_cache.clear();
+        Ok(())
+    }
+
+    pub async fn prepare(&mut self, sql: &str) -> Result<OdbcStatement<'static>, Error> {
+        let (_, columns, parameters) = self.prepare_metadata(sql).await?;
+        let metadata = OdbcStatementMetadata {
+            columns,
+            parameters,
+        };
+        Ok(OdbcStatement {
+            sql: Cow::Owned(sql.to_string()),
+            metadata,
+        })
     }
 }
 
@@ -168,6 +259,10 @@ impl Connection for OdbcConnection {
     fn should_flush(&self) -> bool {
         false
     }
+
+    fn clear_cached_statements(&mut self) -> BoxFuture<'_, Result<(), Error>> {
+        Box::pin(self.clear_cached_statements())
+    }
 }
 
 // moved helpers to connection/inner.rs