Fix formatting and address comments

Author: crystalxyz

core/src/flight/exec.rs

@@ -193,13 +193,12 @@ async fn flight_stream(
     for loc in partition.locations.iter() {
         let get_client = || async { flight_client(loc, grpc_headers.as_ref(), &size_limits).await };
         let client = run_async_with_tokio(get_client).await?;
-        // let client = flight_client(loc, grpc_headers.as_ref(), &size_limits).await?;
         match try_fetch_stream(client, &partition.ticket, schema.clone()).await {
             Ok(stream) => return Ok(stream),
             Err(e) => errors.push(Box::new(e)),
         }
     }
-    let err = errors.into_iter().last().unwrap_or_else(|| {
+    let err = errors.into_iter().next_back().unwrap_or_else(|| {
         Box::new(FlightError::ProtocolError(format!(
             "No available location for endpoint {:?}",
             partition.locations

core/src/sql/db_connection_pool/runtime.rs

@@ -3,11 +3,22 @@
 // tokio runtime, so we need to start a new one.
 use std::{future::Future, sync::OnceLock};
 
-use tokio::runtime::{Handle, Runtime};
+use tokio::runtime::Handle;
 
-pub fn get_tokio_runtime() -> &'static Runtime {
-    static RUNTIME: OnceLock<Runtime> = OnceLock::new();
-    RUNTIME.get_or_init(|| Runtime::new().expect("Failed to create Tokio runtime"))
+pub(crate) struct TokioRuntime(tokio::runtime::Runtime);
+
+#[inline]
+pub(crate) fn get_tokio_runtime() -> &'static TokioRuntime {
+    static RUNTIME: OnceLock<TokioRuntime> = OnceLock::new();
+    RUNTIME.get_or_init(|| TokioRuntime(tokio::runtime::Runtime::new().unwrap()))
+}
+
+pub fn execute_in_tokio<F, Fut, T>(f: F) -> T
+where
+    F: FnOnce() -> Fut,
+    Fut: Future<Output = T>,
+{
+    get_tokio_runtime().0.block_on(f())
 }
 
 pub async fn run_async_with_tokio<F, Fut, T, E>(f: F) -> Result<T, E>
@@ -17,13 +28,14 @@ where
 {
     match Handle::try_current() {
         Ok(_) => f().await,
-        Err(_) => get_tokio_runtime().block_on(f()),
+        // Err(_) => get_tokio_runtime().0.block_on(f()),
+        Err(_) => execute_in_tokio(f),
     }
 }
 
 pub fn run_sync_with_tokio<T, E>(f: impl FnOnce() -> Result<T, E>) -> Result<T, E> {
     match Handle::try_current() {
         Ok(_) => f(),
-        Err(_) => get_tokio_runtime().block_on(async { f() }),
+        Err(_) => execute_in_tokio(|| async { f() }),
     }
 }

python/src/utils.rs

@@ -1,25 +1,17 @@
 use pyo3::{exceptions::PyException, prelude::*};
-use std::{future::Future, sync::OnceLock};
+use std::future::Future;
 
+use datafusion_table_providers::sql::db_connection_pool::runtime::execute_in_tokio;
 use pyo3::types::PyDict;
 use std::collections::HashMap;
 
-pub(crate) struct TokioRuntime(tokio::runtime::Runtime);
-
-#[inline]
-pub(crate) fn get_tokio_runtime() -> &'static TokioRuntime {
-    static RUNTIME: OnceLock<TokioRuntime> = OnceLock::new();
-    RUNTIME.get_or_init(|| TokioRuntime(tokio::runtime::Runtime::new().unwrap()))
-}
-
 /// Utility to collect rust futures with GIL released
 pub fn wait_for_future<F>(py: Python, f: F) -> F::Output
 where
     F: Future + Send,
     F::Output: Send,
 {
-    let runtime: &tokio::runtime::Runtime = &get_tokio_runtime().0;
-    py.allow_threads(|| runtime.block_on(f))
+    py.allow_threads(|| execute_in_tokio(|| f))
 }
 
 pub fn to_pyerr<T: ToString>(err: T) -> PyErr {