feat(table-providers): make sqlserver pool compile

Author: oktal

src-tauri/Cargo.lock

@@ -38,6 +38,7 @@ rand = "0.8.5"
 bytes = "1.7.1"
 arrow-schema = "52.2.0"
 indexmap = "2.4.0"
+table-providers = { version = "0.1.0", path = "table-providers" }
 
 [features]
 # this feature is used for production builds or when `devPath` points to the filesystem and the built-in dev server is disabled.

src-tauri/Cargo.toml

@@ -13,6 +13,7 @@ chrono = "0.4.38"
 datafusion = "41.0.0"
 datafusion-table-providers = "0.1.0"
 futures = "0.3.30"
+pin-project = "1.1.5"
 secrecy = "0.10.1"
 tiberius = { version = "0.12.3", features = ["chrono", "time", "tokio"] }
 tokio = { version = "1.40.0", features = ["net"] }

src-tauri/table-providers/Cargo.lock

@@ -1,25 +1,27 @@
 //! A connection pool for SQL Server
 
 use core::fmt;
+use std::borrow::Cow;
+use std::task::{ready, Poll};
 use std::{collections::HashMap, sync::Arc};
 
-use arrow::datatypes::{Field, Schema, SchemaRef};
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
 
 use async_stream::stream;
 use async_trait::async_trait;
+use datafusion::execution::RecordBatchStream;
 use datafusion::sql::TableReference;
-use datafusion::{
-    error::DataFusionError, execution::SendableRecordBatchStream,
-    physical_plan::stream::RecordBatchStreamAdapter,
-};
+use datafusion::{error::DataFusionError, execution::SendableRecordBatchStream};
 use datafusion_table_providers::sql::db_connection_pool;
 use datafusion_table_providers::sql::db_connection_pool::{
     dbconnection::{AsyncDbConnection, DbConnection},
     DbConnectionPool, JoinPushDown,
 };
-use futures::{stream, StreamExt};
+use futures::{Stream, StreamExt};
+use pin_project::pin_project;
 use secrecy::{ExposeSecret, SecretString};
-use tiberius::{AuthMethod, Config, EncryptionLevel, ToSql, ColumnData};
+use tiberius::{AuthMethod, ColumnData, Config, EncryptionLevel, IntoSql, Query, ToSql};
 use tokio::net::TcpStream;
 use tokio_util::compat::TokioAsyncWriteCompatExt;
 
@@ -84,6 +86,12 @@ impl fmt::Display for Error {
 
 impl std::error::Error for Error {}
 
+impl Into<DataFusionError> for Error {
+    fn into(self) -> DataFusionError {
+        DataFusionError::Execution(self.to_string())
+    }
+}
+
 pub struct SqlServerConnectionManager {
     config: tiberius::Config,
 }
@@ -210,11 +218,8 @@ impl DbConnectionPool<SqlServerPooledConnection, &'static dyn ToSql> for SqlServ
         Box<dyn DbConnection<SqlServerPooledConnection, &'static dyn ToSql>>,
         db_connection_pool::Error,
     > {
-        let pool = Arc::clone(&self.pool);
-        // let conn = pool.get().await?;
-        // Ok(Box::new(SqlServerConnection::new(conn)))
-
-        todo!();
+        let conn = self.pool.get_owned().await?;
+        Ok(Box::new(SqlServerConnection::new(conn)))
     }
 
     fn join_push_down(&self) -> JoinPushDown {
@@ -225,11 +230,97 @@ impl DbConnectionPool<SqlServerPooledConnection, &'static dyn ToSql> for SqlServ
 type ConnectionError = datafusion_table_providers::sql::db_connection_pool::dbconnection::Error;
 type GenericError = datafusion_table_providers::sql::db_connection_pool::dbconnection::GenericError;
 
-struct OwnedColumnData<'a>(ColumnData<'a>);
+#[pin_project]
+struct SqlRecordBatchStream<S> {
+    schema: Option<SchemaRef>,
+
+    #[pin]
+    stream: S,
+}
+
+impl<S> SqlRecordBatchStream<S> {
+    fn new(stream: S) -> Self {
+        Self {
+            schema: None,
+            stream,
+        }
+    }
+}
+
+impl<S, E> Stream for SqlRecordBatchStream<S>
+where
+    S: Stream<Item = Result<RecordBatch, E>>,
+    E: Into<DataFusionError>,
+{
+    type Item = Result<RecordBatch, DataFusionError>;
+
+    fn poll_next(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+
+        let batch = ready!(this.stream.poll_next(cx));
+        let Some(batch) = batch else {
+            return Poll::Ready(None);
+        };
+
+        if let Ok(batch) = &batch {
+            *this.schema = Some(batch.schema());
+        }
+
+        Poll::Ready(Some(batch.map_err(Into::into)))
+    }
+}
+
+impl<S, E> RecordBatchStream for SqlRecordBatchStream<S>
+where
+    S: Stream<Item = Result<RecordBatch, E>>,
+    E: Into<DataFusionError>,
+{
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone().unwrap_or(Arc::new(Schema::empty()))
+    }
+}
+
+fn to_owned<'a, T: ?Sized + ToOwned>(val: Cow<'a, T>) -> Cow<'static, T> {
+    match val {
+        Cow::Borrowed(val) => Cow::Owned(val.to_owned()),
+        Cow::Owned(val) => Cow::Owned(val),
+    }
+}
 
-impl<'a> ToSql for OwnedColumnData<'a> {
-    fn to_sql(&self) -> ColumnData<'_> {
-        self.0.clone()
+/// A [`ColumnData`] that owns the underlying data, meaning that it will
+/// transform Cow::Borrowed data to Cow::Owned when needed
+struct OwnedColumnData(ColumnData<'static>);
+impl<'a> From<ColumnData<'a>> for OwnedColumnData {
+    fn from(value: ColumnData<'a>) -> Self {
+        Self(match value {
+            ColumnData::U8(val) => ColumnData::U8(val),
+            ColumnData::I16(val) => ColumnData::I16(val),
+            ColumnData::I32(val) => ColumnData::I32(val),
+            ColumnData::I64(val) => ColumnData::I64(val),
+            ColumnData::F32(val) => ColumnData::F32(val),
+            ColumnData::F64(val) => ColumnData::F64(val),
+            ColumnData::Bit(val) => ColumnData::Bit(val),
+            ColumnData::Guid(val) => ColumnData::Guid(val),
+            ColumnData::Numeric(val) => ColumnData::Numeric(val),
+            ColumnData::DateTime(val) => ColumnData::DateTime(val),
+            ColumnData::SmallDateTime(val) => ColumnData::SmallDateTime(val),
+            ColumnData::Time(val) => ColumnData::Time(val),
+            ColumnData::Date(val) => ColumnData::Date(val),
+            ColumnData::DateTime2(val) => ColumnData::DateTime2(val),
+            ColumnData::DateTimeOffset(val) => ColumnData::DateTimeOffset(val),
+            ColumnData::String(val) => ColumnData::String(val.map(to_owned)),
+            ColumnData::Binary(val) => ColumnData::Binary(val.map(to_owned)),
+            ColumnData::Xml(val) => ColumnData::Xml(val.map(to_owned)),
+        })
+    }
+}
+
+impl IntoSql<'static> for OwnedColumnData {
+    fn into_sql(self) -> ColumnData<'static> {
+        self.0
     }
 }
 
@@ -267,19 +358,25 @@ impl<'a> AsyncDbConnection<SqlServerPooledConnection, &'a dyn ToSql> for SqlServ
         &self,
         sql: &str,
         params: &[&'a dyn ToSql],
-        projected_schema: Option<SchemaRef>,
+        _projected_schema: Option<SchemaRef>,
     ) -> Result<SendableRecordBatchStream, GenericError> {
         let conn = Arc::clone(&self.conn);
-        let params = params.iter().map(|p| OwnedColumnData(p.to_sql())).collect::<Vec<_>>();
 
-        let mut stream = Box::pin(stream! {
-            let mut conn = conn.lock().await;
+        let sql = sql.to_string();
+        let params = params
+            .iter()
+            .map(|p| OwnedColumnData::from(p.to_sql()))
+            .collect::<Vec<_>>();
 
-            let stream = conn
-                .query(sql, &params)
-                .await?;
+        let stream = stream! {
+            let mut conn = conn.lock().await;
+            let mut query = Query::new(sql.to_string());
+            for param in params {
+                query.bind(param);
+            }
+            let stream = query.query(&mut conn).await.map_err(Error::Sql)?;
 
-            let mut chunks = stream.chunks(8192).boxed().map(|rows| {
+            let mut chunks = stream.chunks(8192).map(|rows| {
                 let rows = rows.into_iter().collect::<Result<Vec<_>, _>>()?;
                 let rec = super::arrow::rows_to_arrow(rows)?;
 
@@ -289,32 +386,15 @@ impl<'a> AsyncDbConnection<SqlServerPooledConnection, &'a dyn ToSql> for SqlServ
             while let Some(chunk) = chunks.next().await {
                 yield chunk
             }
-        });
-
-        let Some(first_chunk) = stream.next().await else {
-            return Ok(Box::pin(RecordBatchStreamAdapter::new(
-                Arc::new(Schema::empty()),
-                stream::empty(),
-            )));
-        };
-
-        let first_chunk =
-            first_chunk.map_err(|e| ConnectionError::UnableToQueryArrow { source: e.into() })?;
-        let schema = first_chunk.schema();
-
-        Ok(Box::pin(RecordBatchStreamAdapter::new(schema, {
-            stream! {
-                yield Ok(first_chunk);
-
-                while let Some(chunk) = stream.next().await {
-                    yield chunk.map_err(|e| DataFusionError::Execution(format!("failed to fetch batch: {e}")))
-                }
+        }
+        .boxed();
 
-            }
-        })))
+        Ok(Box::pin(SqlRecordBatchStream::new(stream)))
     }
 
     async fn execute(&self, sql: &str, params: &[&'a dyn ToSql]) -> Result<u64, GenericError> {
-        todo!()
+        let mut conn = self.conn.lock().await;
+        let result = conn.execute(sql, params).await?;
+        Ok(result.into_iter().sum())
     }
 }