fix(rust/driver/datafusion): using datafusion driver in async runtime (#3712)

Closes #3711 

This is not an ideal solution, but we need to think about sync/async
ergonomics in the future.

At the moment, there is a problem: the inability to use this approach in
single-thread runtime.

---------

Signed-off-by: if0ne <pavel.agafonov.al@gmail.com>
Signed-off-by: Pavel Agafonov <pavel.agafonov.al@gmail.com>

Author: if0ne

rust/driver/datafusion/src/lib.rs

@@ -25,9 +25,9 @@ use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
 use datafusion_substrait::substrait::proto::Plan;
 use prost::Message;
 use std::fmt::Debug;
+use std::future::Future;
 use std::sync::Arc;
 use std::vec::IntoIter;
-use tokio::runtime::Runtime;
 
 use arrow_array::builder::{
     BooleanBuilder, Int32Builder, Int64Builder, ListBuilder, MapBuilder, MapFieldNames,
@@ -48,6 +48,31 @@ use adbc_core::{
     schemas, Connection, Database, Driver, Optionable, Statement,
 };
 
+pub enum Runtime {
+    Handle(tokio::runtime::Handle),
+    Tokio(tokio::runtime::Runtime),
+}
+
+impl Runtime {
+    pub fn new(handle: Option<tokio::runtime::Handle>) -> std::io::Result<Self> {
+        if let Some(handle) = handle {
+            Ok(Self::Handle(handle))
+        } else {
+            let runtime = tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()?;
+            Ok(Self::Tokio(runtime))
+        }
+    }
+
+    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        match self {
+            Runtime::Handle(handle) => tokio::task::block_in_place(|| handle.block_on(future)),
+            Runtime::Tokio(runtime) => runtime.block_on(future),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct SingleBatchReader {
     batch: Option<RecordBatch>,
@@ -109,13 +134,23 @@ impl RecordBatchReader for DataFusionReader {
 }
 
 #[derive(Default)]
-pub struct DataFusionDriver {}
+pub struct DataFusionDriver {
+    handle: Option<tokio::runtime::Handle>,
+}
+
+impl DataFusionDriver {
+    pub fn new(handle: Option<tokio::runtime::Handle>) -> Self {
+        Self { handle }
+    }
+}
 
 impl Driver for DataFusionDriver {
     type DatabaseType = DataFusionDatabase;
 
     fn new_database(&mut self) -> Result<Self::DatabaseType> {
-        Ok(Self::DatabaseType {})
+        Ok(Self::DatabaseType {
+            handle: self.handle.clone(),
+        })
     }
 
     fn new_database_with_opts(
@@ -127,15 +162,19 @@ impl Driver for DataFusionDriver {
             ),
         >,
     ) -> adbc_core::error::Result<Self::DatabaseType> {
-        let mut database = Self::DatabaseType {};
+        let mut database = Self::DatabaseType {
+            handle: self.handle.clone(),
+        };
         for (key, value) in opts {
             database.set_option(key, value)?;
         }
         Ok(database)
     }
 }
 
-pub struct DataFusionDatabase {}
+pub struct DataFusionDatabase {
+    handle: Option<tokio::runtime::Handle>,
+}
 
 impl Optionable for DataFusionDatabase {
     type Option = OptionDatabase;
@@ -186,10 +225,7 @@ impl Database for DataFusionDatabase {
     fn new_connection(&self) -> Result<Self::ConnectionType> {
         let ctx = SessionContext::new();
 
-        let runtime = tokio::runtime::Builder::new_multi_thread()
-            .enable_all()
-            .build()
-            .unwrap();
+        let runtime = Runtime::new(self.handle.clone()).unwrap();
 
         Ok(DataFusionConnection {
             runtime: Arc::new(runtime),
@@ -208,10 +244,7 @@ impl Database for DataFusionDatabase {
     ) -> adbc_core::error::Result<Self::ConnectionType> {
         let ctx = SessionContext::new();
 
-        let runtime = tokio::runtime::Builder::new_multi_thread()
-            .enable_all()
-            .build()
-            .unwrap();
+        let runtime = Runtime::new(self.handle.clone()).unwrap();
 
         let mut connection = DataFusionConnection {
             runtime: Arc::new(runtime),

rust/driver/datafusion/tests/test_datafusion.rs

@@ -26,8 +26,8 @@ use datafusion_substrait::logical_plan::producer::to_substrait_plan;
 use datafusion_substrait::substrait::proto::Plan;
 use prost::Message;
 
-fn get_connection() -> DataFusionConnection {
-    let mut driver = DataFusionDriver::default();
+fn get_connection(handle: Option<tokio::runtime::Handle>) -> DataFusionConnection {
+    let mut driver = DataFusionDriver::new(handle);
     let database = driver.new_database().unwrap();
     database.new_connection().unwrap()
 }
@@ -80,7 +80,7 @@ fn execute_substrait(connection: &mut DataFusionConnection, plan: Plan) -> Recor
 
 #[test]
 fn test_connection_options() {
-    let mut connection = get_connection();
+    let mut connection = get_connection(None);
 
     let current_catalog = connection
         .get_option_string(OptionConnection::CurrentCatalog)
@@ -119,7 +119,7 @@ fn test_connection_options() {
 
 #[test]
 fn test_get_objects_database() {
-    let mut connection = get_connection();
+    let mut connection = get_connection(None);
 
     let objects = get_objects(&connection);
 
@@ -134,7 +134,7 @@ fn test_get_objects_database() {
 
 #[test]
 fn test_execute_sql() {
-    let mut connection = get_connection();
+    let mut connection = get_connection(None);
 
     execute_update(&mut connection, "CREATE TABLE IF NOT EXISTS datafusion.public.example (c1 INT, c2 VARCHAR) AS VALUES(1,'HELLO'),(2,'DATAFUSION'),(3,'!')");
 
@@ -146,7 +146,7 @@ fn test_execute_sql() {
 
 #[test]
 fn test_ingest() {
-    let mut connection = get_connection();
+    let mut connection = get_connection(None);
 
     execute_update(&mut connection, "CREATE TABLE IF NOT EXISTS datafusion.public.example (c1 INT, c2 VARCHAR) AS VALUES(1,'HELLO'),(2,'DATAFUSION'),(3,'!')");
 
@@ -172,7 +172,7 @@ fn test_ingest() {
 
 #[test]
 fn test_execute_substrait() {
-    let mut connection = get_connection();
+    let mut connection = get_connection(None);
 
     execute_update(&mut connection, "CREATE TABLE IF NOT EXISTS datafusion.public.example (c1 INT, c2 VARCHAR) AS VALUES(1,'HELLO'),(2,'DATAFUSION'),(3,'!')");
 
@@ -198,3 +198,15 @@ fn test_execute_substrait() {
     assert_eq!(batch.num_rows(), 3);
     assert_eq!(batch.num_columns(), 2);
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_running_in_async() {
+    let mut connection = get_connection(Some(tokio::runtime::Handle::current()));
+
+    execute_update(&mut connection, "CREATE TABLE IF NOT EXISTS datafusion.public.example (c1 INT, c2 VARCHAR) AS VALUES(1,'HELLO'),(2,'DATAFUSION'),(3,'!')");
+
+    let batch = execute_sql_query(&mut connection, "SELECT * FROM datafusion.public.example");
+
+    assert_eq!(batch.num_rows(), 3);
+    assert_eq!(batch.num_columns(), 2);
+}