wire duckdb

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

vortex-duckdb/src/lib.rs

@@ -6,12 +6,14 @@
 use std::ffi::CStr;
 use std::ffi::c_char;
 use std::sync::LazyLock;
+use std::sync::OnceLock;
 
 use vortex::VortexSessionDefault;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::CurrentThreadRuntime;
+use vortex::io::runtime::Handle;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
 
@@ -40,9 +42,22 @@ mod copy;
 mod e2e_test;
 
 // A global runtime for Vortex operations within DuckDB.
-static RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
-static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle()));
+static DEFAULT_RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
+static RUNTIME_OVERRIDE: OnceLock<Handle> = OnceLock::new();
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let handle = RUNTIME_OVERRIDE
+        .get()
+        .cloned()
+        .unwrap_or_else(|| DEFAULT_RUNTIME.handle());
+    VortexSession::default().with_handle(handle)
+});
+
+/// Configure the runtime handle used by the DuckDB extension.
+///
+/// Must be called before the session is first accessed (e.g., before registering table functions).
+pub fn configure_runtime(handle: Handle) {
+    let _ = RUNTIME_OVERRIDE.set(handle);
+}
 
 /// Register Vortex extension configuration options with DuckDB.
 /// This must be called before `register_table_functions` to take effect.