feat: Implement streaming execution for ODBC connections

Co-authored-by: contact <[email protected]>

Author: cursoragent

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

@@ -21,8 +21,14 @@ impl<'c> Executor<'c> for &'c mut OdbcConnection {
         'c: 'e,
         E: Execute<'q, Self::Database> + 'q,
     {
-        let empty: Vec<Result<Either<OdbcQueryResult, OdbcRow>, Error>> = Vec::new();
-        Box::pin(futures_util::stream::iter(empty))
+        let sql = _query.sql().to_string();
+        Box::pin(try_stream! {
+            let rx = self.worker.execute_stream(&sql).await?;
+            while let Ok(item) = rx.recv_async().await {
+                r#yield!(item?);
+            }
+            Ok(())
+        })
     }
 
     fn fetch_optional<'e, 'q: 'e, E>(

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

@@ -5,7 +5,9 @@ use futures_channel::oneshot;
 use futures_intrusive::sync::Mutex;
 
 use crate::error::Error;
-use crate::odbc::OdbcConnectOptions;
+use crate::odbc::{OdbcColumn, OdbcConnectOptions, OdbcQueryResult, OdbcRow, OdbcTypeInfo};
+use either::Either;
+use odbc_api::Cursor;
 
 #[derive(Debug)]
 pub(crate) struct ConnectionWorker {
@@ -34,6 +36,10 @@ enum Command {
     Rollback {
         tx: oneshot::Sender<Result<(), Error>>,
     },
+    Execute {
+        sql: Box<str>,
+        tx: flume::Sender<Result<Either<OdbcQueryResult, OdbcRow>, Error>>,
+    },
 }
 
 impl ConnectionWorker {
@@ -119,6 +125,60 @@ impl ConnectionWorker {
                             let _ = tx.send(());
                             return;
                         }
+                        Command::Execute { sql, tx } => {
+                            // Helper closure to process using a given connection reference
+                            let process = |conn: &odbc_api::Connection<'static>| {
+                                match conn.execute(&sql, (), None) {
+                                    Ok(Some(mut cursor)) => {
+                                        use odbc_api::ResultSetMetadata;
+                                        let mut columns: Vec<OdbcColumn> = Vec::new();
+                                        if let Ok(count) = cursor.num_result_cols() {
+                                            for i in 1..=count {
+                                                let mut cd = odbc_api::ColumnDescription::default();
+                                                let _ = cursor.describe_col(i as u16, &mut cd);
+                                                let name = String::from_utf8(cd.name)
+                                                    .unwrap_or_else(|_| format!("col{}", i - 1));
+                                                columns.push(OdbcColumn {
+                                                    name,
+                                                    type_info: OdbcTypeInfo { name: format!("{:?}", cd.data_type), is_null: false },
+                                                    ordinal: (i - 1) as usize,
+                                                });
+                                            }
+                                        }
+
+                                        while let Ok(Some(mut row)) = cursor.next_row() {
+                                            let mut values: Vec<(OdbcTypeInfo, Option<Vec<u8>>)> = Vec::with_capacity(columns.len());
+                                            for i in 1..=columns.len() {
+                                                let mut buf = Vec::new();
+                                                match row.get_text(i as u16, &mut buf) {
+                                                    Ok(true) => values.push((OdbcTypeInfo { name: "TEXT".into(), is_null: false }, Some(buf))),
+                                                    Ok(false) => values.push((OdbcTypeInfo { name: "TEXT".into(), is_null: true }, None)),
+                                                    Err(e) => {
+                                                        let _ = tx.send(Err(Error::from(e)));
+                                                        return;
+                                                    }
+                                                }
+                                            }
+                                            let _ = tx.send(Ok(Either::Right(OdbcRow { columns: columns.clone(), values })));
+                                        }
+                                        let _ = tx.send(Ok(Either::Left(OdbcQueryResult { rows_affected: 0 })));
+                                    }
+                                    Ok(None) => {
+                                        let _ = tx.send(Ok(Either::Left(OdbcQueryResult { rows_affected: 0 })));
+                                    }
+                                    Err(e) => {
+                                        let _ = tx.send(Err(Error::from(e)));
+                                    }
+                                }
+                            };
+
+                            if let Some(conn) = shared.conn.try_lock() {
+                                process(&conn);
+                            } else {
+                                let guard = futures_executor::block_on(shared.conn.lock());
+                                process(&guard);
+                            }
+                        }
                     }
                 }
             })?;
@@ -173,4 +233,16 @@ impl ConnectionWorker {
         rx.await.map_err(|_| Error::WorkerCrashed)??;
         Ok(())
     }
+
+    pub(crate) async fn execute_stream(
+        &mut self,
+        sql: &str,
+    ) -> Result<flume::Receiver<Result<Either<OdbcQueryResult, OdbcRow>, Error>>, Error> {
+        let (tx, rx) = flume::bounded(64);
+        self.command_tx
+            .send_async(Command::Execute { sql: sql.into(), tx })
+            .await
+            .map_err(|_| Error::WorkerCrashed)?;
+        Ok(rx)
+    }
 }