Allow to pass extra info to table functions

Via `register_table_function_with_extra_info`

Author: mlafeldt

crates/duckdb/src/vtab/function.rs

@@ -18,6 +18,11 @@ use std::{
     os::raw::c_char,
 };
 
+/// Callback function to drop extra info of type T
+unsafe extern "C" fn drop_extra_info<T>(ptr: *mut c_void) {
+    drop(unsafe { Box::from_raw(ptr.cast::<T>()) });
+}
+
 /// An interface to store and retrieve data during the function bind stage
 #[derive(Debug)]
 pub struct BindInfo {
@@ -36,6 +41,7 @@ impl BindInfo {
             duckdb_bind_add_result_column(self.ptr, c_str.as_ptr() as *const c_char, column_type.ptr);
         }
     }
+
     /// Report that an error has occurred while calling bind.
     ///
     /// # Arguments
@@ -51,16 +57,15 @@ impl BindInfo {
     /// # Arguments
     ///  * `extra_data`: The bind data object.
     ///  * `destroy`: The callback that will be called to destroy the bind data (if any)
-    ///
-    /// # Safety
-    ///
     pub unsafe fn set_bind_data(&self, data: *mut c_void, free_function: Option<unsafe extern "C" fn(*mut c_void)>) {
         duckdb_bind_set_bind_data(self.ptr, data, free_function);
     }
+
     /// Retrieves the number of regular (non-named) parameters to the function.
     pub fn get_parameter_count(&self) -> u64 {
         unsafe { duckdb_bind_get_parameter_count(self.ptr) }
     }
+
     /// Retrieves the parameter at the given index.
     ///
     /// # Arguments
@@ -107,10 +112,7 @@ impl BindInfo {
     pub fn set_cardinality(&self, cardinality: idx_t, is_exact: bool) {
         unsafe { duckdb_bind_set_cardinality(self.ptr, cardinality, is_exact) }
     }
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *const T {
         unsafe { duckdb_bind_get_extra_info(self.ptr).cast() }
     }
@@ -162,13 +164,11 @@ impl InitInfo {
         indices
     }
 
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *const T {
         unsafe { duckdb_init_get_extra_info(self.0).cast() }
     }
+
     /// Gets the bind data set by [`BindInfo::set_bind_data`] during the bind.
     ///
     /// Note that the bind data should be considered as read-only.
@@ -179,13 +179,15 @@ impl InitInfo {
     pub fn get_bind_data<T>(&self) -> *const T {
         unsafe { duckdb_init_get_bind_data(self.0).cast() }
     }
+
     /// Sets how many threads can process this table function in parallel (default: 1)
     ///
     /// # Arguments
     /// * `max_threads`: The maximum amount of threads that can process this table function
     pub fn set_max_threads(&self, max_threads: idx_t) {
         unsafe { duckdb_init_set_max_threads(self.0, max_threads) }
     }
+
     /// Report that an error has occurred while calling init.
     ///
     /// # Arguments
@@ -307,15 +309,35 @@ impl TableFunction {
 
     /// Assigns extra information to the table function that can be fetched during binding, etc.
     ///
+    /// For most use cases, prefer [`with_extra_info`](Self::with_extra_info) which handles memory management automatically.
+    ///
     /// # Arguments
     /// * `extra_info`: The extra information
     /// * `destroy`: The callback that will be called to destroy the bind data (if any)
     ///
     /// # Safety
+    /// The caller must ensure that `extra_info` is a valid pointer and that `destroy`
+    /// properly cleans up the data when called.
     pub unsafe fn set_extra_info(&self, extra_info: *mut c_void, destroy: duckdb_delete_callback_t) {
+        duckdb_table_function_set_extra_info(self.ptr, extra_info, destroy);
+    }
+
+    /// Assigns extra information to the table function that can be fetched during binding, init, and execution.
+    ///
+    /// This is a safe wrapper around [`set_extra_info`](Self::set_extra_info) that handles memory management automatically.
+    ///
+    /// # Arguments
+    /// * `info`: The extra information to store
+    pub fn with_extra_info<T>(&self, info: T) -> &Self
+    where
+        T: Send + Sync + 'static,
+    {
         unsafe {
-            duckdb_table_function_set_extra_info(self.ptr, extra_info, destroy);
+            let boxed = Box::new(info);
+            let ptr = Box::into_raw(boxed) as *mut c_void;
+            self.set_extra_info(ptr, Some(drop_extra_info::<T>));
         }
+        self
     }
 
     /// Sets the thread-local init function of the table function
@@ -383,13 +405,11 @@ impl<V: VTab> TableFunctionInfo<V> {
         }
     }
 
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *mut T {
         unsafe { duckdb_function_get_extra_info(self.ptr).cast() }
     }
+
     /// Gets the thread-local init data set by [`InitInfo::set_init_data`] during the local_init.
     ///
     /// # Arguments

crates/duckdb/src/vtab/mod.rs

@@ -150,6 +150,32 @@ impl Connection {
         }
         self.db.borrow_mut().register_table_function(table_function)
     }
+
+    /// Register the given TableFunction with custom extra info.
+    ///
+    /// This allows you to pass extra info that can be accessed during bind, init, and execution
+    /// via `BindInfo::get_extra_info`, `InitInfo::get_extra_info`, or `TableFunctionInfo::get_extra_info`.
+    #[inline]
+    pub fn register_table_function_with_extra_info<T: VTab, E>(&self, name: &str, extra_info: &E) -> Result<()>
+    where
+        E: Clone + Send + Sync + 'static,
+    {
+        let table_function = TableFunction::default();
+        table_function
+            .set_name(name)
+            .supports_pushdown(T::supports_pushdown())
+            .set_bind(Some(bind::<T>))
+            .set_init(Some(init::<T>))
+            .set_function(Some(func::<T>))
+            .with_extra_info(extra_info.clone());
+        for ty in T::parameters().unwrap_or_default() {
+            table_function.add_parameter(&ty);
+        }
+        for (name, ty) in T::named_parameters().unwrap_or_default() {
+            table_function.add_named_parameter(&name, &ty);
+        }
+        self.db.borrow_mut().register_table_function(table_function)
+    }
 }
 
 impl InnerConnection {
@@ -287,6 +313,57 @@ mod test {
         Ok(())
     }
 
+    // Test table function with extra info
+    struct PrefixVTab;
+
+    impl VTab for PrefixVTab {
+        type InitData = HelloInitData;
+        type BindData = HelloBindData;
+
+        fn bind(bind: &BindInfo) -> Result<Self::BindData, Box<dyn Error>> {
+            bind.add_result_column("column0", LogicalTypeHandle::from(LogicalTypeId::Varchar));
+            let name = bind.get_parameter(0).to_string();
+            Ok(HelloBindData { name })
+        }
+
+        fn init(_: &InitInfo) -> Result<Self::InitData, Box<dyn Error>> {
+            Ok(HelloInitData {
+                done: AtomicBool::new(false),
+            })
+        }
+
+        fn func(func: &TableFunctionInfo<Self>, output: &mut DataChunkHandle) -> Result<(), Box<dyn Error>> {
+            let init_data = func.get_init_data();
+            let bind_data = func.get_bind_data();
+            let prefix = unsafe { &*func.get_extra_info::<String>() };
+
+            if init_data.done.swap(true, Ordering::Relaxed) {
+                output.set_len(0);
+            } else {
+                let vector = output.flat_vector(0);
+                let result = CString::new(format!("{prefix} {}", bind_data.name))?;
+                vector.insert(0, result);
+                output.set_len(1);
+            }
+            Ok(())
+        }
+
+        fn parameters() -> Option<Vec<LogicalTypeHandle>> {
+            Some(vec![LogicalTypeHandle::from(LogicalTypeId::Varchar)])
+        }
+    }
+
+    #[test]
+    fn test_table_function_with_extra_info() -> Result<(), Box<dyn Error>> {
+        let conn = Connection::open_in_memory()?;
+        conn.register_table_function_with_extra_info::<PrefixVTab, _>("greet", &"Howdy".to_string())?;
+
+        let val = conn.query_row("select * from greet('partner')", [], |row| <(String,)>::try_from(row))?;
+        assert_eq!(val, ("Howdy partner".to_string(),));
+
+        Ok(())
+    }
+
     #[cfg(feature = "vtab-loadable")]
     use duckdb_loadable_macros::duckdb_entrypoint;