quack_rs/copy_function/
mod.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2026 Tom F. <https://github.com/tomtom215/>
3// My way of giving something small back to the open source community
4// and encouraging more Rust development!
5
6//! Copy function registration (`DuckDB` 1.5.0+).
7//!
8//! Extensions can register custom `COPY TO` handlers that define how data is
9//! exported to a specific file format. The lifecycle consists of four phases:
10//!
11//! 1. **Bind** — inspect the output columns and configure the export.
12//! 2. **Global init** — set up global state (open file, allocate buffers).
13//! 3. **Sink** — receive data chunks to write.
14//! 4. **Finalize** — flush and close.
15//!
16//! # Example
17//!
18//! ```rust,no_run
19//! use quack_rs::copy_function::CopyFunctionBuilder;
20//!
21//! let builder = CopyFunctionBuilder::try_new("my_format")?;
22//! // .bind(my_bind_fn)
23//! // .global_init(my_init_fn)
24//! // .sink(my_sink_fn)
25//! // .finalize(my_finalize_fn)
26//! # Ok::<(), quack_rs::error::ExtensionError>(())
27//! ```
28
29pub mod info;
30
31pub use info::{CopyBindInfo, CopyFinalizeInfo, CopyGlobalInitInfo, CopySinkInfo};
32
33use std::ffi::CString;
34
35use libduckdb_sys::{
36    duckdb_connection, duckdb_copy_function_bind_info, duckdb_copy_function_finalize_info,
37    duckdb_copy_function_global_init_info, duckdb_copy_function_set_bind,
38    duckdb_copy_function_set_finalize, duckdb_copy_function_set_global_init,
39    duckdb_copy_function_set_name, duckdb_copy_function_set_sink, duckdb_copy_function_sink_info,
40    duckdb_create_copy_function, duckdb_data_chunk, duckdb_destroy_copy_function,
41    duckdb_register_copy_function, DuckDBSuccess,
42};
43
44use crate::error::ExtensionError;
45
46/// Callback type aliases for copy function phases.
47///
48/// Bind callback — called once to configure the export.
49pub type CopyBindFn = unsafe extern "C" fn(info: duckdb_copy_function_bind_info);
50
51/// Global init callback — called once to set up global state.
52pub type CopyGlobalInitFn = unsafe extern "C" fn(info: duckdb_copy_function_global_init_info);
53
54/// Sink callback — called once per data chunk to write data.
55pub type CopySinkFn =
56    unsafe extern "C" fn(info: duckdb_copy_function_sink_info, chunk: duckdb_data_chunk);
57
58/// Finalize callback — called once to flush and close.
59pub type CopyFinalizeFn = unsafe extern "C" fn(info: duckdb_copy_function_finalize_info);
60
61/// Builder for registering a custom `COPY TO` function.
62///
63/// All four lifecycle callbacks (bind, `global_init`, sink, finalize) should be
64/// set before calling [`register`][Self::register].
65#[must_use]
66pub struct CopyFunctionBuilder {
67    name: CString,
68    bind: Option<CopyBindFn>,
69    global_init: Option<CopyGlobalInitFn>,
70    sink: Option<CopySinkFn>,
71    finalize: Option<CopyFinalizeFn>,
72}
73
74impl CopyFunctionBuilder {
75    /// Creates a new copy function builder with the given format name.
76    ///
77    /// The name corresponds to the format identifier used in
78    /// `COPY table TO 'file' (FORMAT name)`.
79    ///
80    /// # Errors
81    ///
82    /// Returns `ExtensionError` if the name contains a null byte.
83    pub fn try_new(name: &str) -> Result<Self, ExtensionError> {
84        let c_name = CString::new(name)
85            .map_err(|_| ExtensionError::new("copy function name contains null byte"))?;
86        Ok(Self {
87            name: c_name,
88            bind: None,
89            global_init: None,
90            sink: None,
91            finalize: None,
92        })
93    }
94
95    /// Returns the name of this copy function.
96    #[must_use]
97    pub fn name(&self) -> &str {
98        self.name.to_str().unwrap_or("")
99    }
100
101    /// Sets the bind callback.
102    pub fn bind(mut self, f: CopyBindFn) -> Self {
103        self.bind = Some(f);
104        self
105    }
106
107    /// Sets the global init callback.
108    pub fn global_init(mut self, f: CopyGlobalInitFn) -> Self {
109        self.global_init = Some(f);
110        self
111    }
112
113    /// Sets the sink callback.
114    pub fn sink(mut self, f: CopySinkFn) -> Self {
115        self.sink = Some(f);
116        self
117    }
118
119    /// Sets the finalize callback.
120    pub fn finalize(mut self, f: CopyFinalizeFn) -> Self {
121        self.finalize = Some(f);
122        self
123    }
124
125    /// Registers the copy function on the given connection.
126    ///
127    /// # Errors
128    ///
129    /// Returns `ExtensionError` if required callbacks are missing or registration
130    /// fails.
131    ///
132    /// # Safety
133    ///
134    /// `con` must be a valid, open `duckdb_connection`.
135    pub unsafe fn register(self, con: duckdb_connection) -> Result<(), ExtensionError> {
136        let bind = self
137            .bind
138            .ok_or_else(|| ExtensionError::new("copy function bind callback not set"))?;
139        let sink = self
140            .sink
141            .ok_or_else(|| ExtensionError::new("copy function sink callback not set"))?;
142        let finalize = self
143            .finalize
144            .ok_or_else(|| ExtensionError::new("copy function finalize callback not set"))?;
145
146        // SAFETY: duckdb_create_copy_function allocates a new handle.
147        let func = unsafe { duckdb_create_copy_function() };
148
149        // SAFETY: func is a valid newly created handle.
150        unsafe {
151            duckdb_copy_function_set_name(func, self.name.as_ptr());
152            duckdb_copy_function_set_bind(func, Some(bind));
153            duckdb_copy_function_set_sink(func, Some(sink));
154            duckdb_copy_function_set_finalize(func, Some(finalize));
155        }
156
157        if let Some(global_init) = self.global_init {
158            // SAFETY: func is valid.
159            unsafe {
160                duckdb_copy_function_set_global_init(func, Some(global_init));
161            }
162        }
163
164        // SAFETY: con is valid, func is fully configured.
165        let result = unsafe { duckdb_register_copy_function(con, func) };
166
167        // SAFETY: func must be destroyed after registration.
168        let mut func_mut = func;
169        unsafe {
170            duckdb_destroy_copy_function(&raw mut func_mut);
171        }
172
173        if result == DuckDBSuccess {
174            Ok(())
175        } else {
176            Err(ExtensionError::new(format!(
177                "duckdb_register_copy_function failed for '{}'",
178                self.name.to_string_lossy()
179            )))
180        }
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn try_new_valid_name() {
190        let builder = CopyFunctionBuilder::try_new("parquet").unwrap();
191        assert_eq!(builder.name(), "parquet");
192    }
193
194    #[test]
195    fn try_new_null_byte_rejected() {
196        let result = CopyFunctionBuilder::try_new("bad\0name");
197        assert!(result.is_err());
198        let err = result.err().unwrap();
199        assert!(
200            err.to_string().contains("null byte"),
201            "error should mention null byte"
202        );
203    }
204
205    #[test]
206    fn builder_stores_bind_callback() {
207        unsafe extern "C" fn dummy_bind(_info: duckdb_copy_function_bind_info) {}
208        let builder = CopyFunctionBuilder::try_new("fmt")
209            .unwrap()
210            .bind(dummy_bind);
211        assert_eq!(builder.name(), "fmt");
212    }
213
214    #[test]
215    fn builder_stores_global_init_callback() {
216        unsafe extern "C" fn dummy_init(_info: duckdb_copy_function_global_init_info) {}
217        let builder = CopyFunctionBuilder::try_new("fmt")
218            .unwrap()
219            .global_init(dummy_init);
220        assert_eq!(builder.name(), "fmt");
221    }
222
223    #[test]
224    fn builder_stores_sink_callback() {
225        unsafe extern "C" fn dummy_sink(
226            _info: duckdb_copy_function_sink_info,
227            _chunk: duckdb_data_chunk,
228        ) {
229        }
230        let builder = CopyFunctionBuilder::try_new("fmt")
231            .unwrap()
232            .sink(dummy_sink);
233        assert_eq!(builder.name(), "fmt");
234    }
235
236    #[test]
237    fn builder_stores_finalize_callback() {
238        unsafe extern "C" fn dummy_finalize(_info: duckdb_copy_function_finalize_info) {}
239        let builder = CopyFunctionBuilder::try_new("fmt")
240            .unwrap()
241            .finalize(dummy_finalize);
242        assert_eq!(builder.name(), "fmt");
243    }
244
245    #[test]
246    fn full_builder_chain_compiles() {
247        unsafe extern "C" fn bind(_: duckdb_copy_function_bind_info) {}
248        unsafe extern "C" fn init(_: duckdb_copy_function_global_init_info) {}
249        unsafe extern "C" fn sink(_: duckdb_copy_function_sink_info, _: duckdb_data_chunk) {}
250        unsafe extern "C" fn finalize(_: duckdb_copy_function_finalize_info) {}
251
252        let builder = CopyFunctionBuilder::try_new("my_format")
253            .unwrap()
254            .bind(bind)
255            .global_init(init)
256            .sink(sink)
257            .finalize(finalize);
258        assert_eq!(builder.name(), "my_format");
259    }
260
261    #[test]
262    fn builder_stores_all_callbacks() {
263        unsafe extern "C" fn my_bind(_: duckdb_copy_function_bind_info) {}
264        unsafe extern "C" fn my_init(_: duckdb_copy_function_global_init_info) {}
265        unsafe extern "C" fn my_sink(_: duckdb_copy_function_sink_info, _: duckdb_data_chunk) {}
266        unsafe extern "C" fn my_finalize(_: duckdb_copy_function_finalize_info) {}
267
268        let b = CopyFunctionBuilder::try_new("f")
269            .unwrap()
270            .bind(my_bind)
271            .global_init(my_init)
272            .sink(my_sink)
273            .finalize(my_finalize);
274        assert!(b.bind.is_some());
275        assert!(b.global_init.is_some());
276        assert!(b.sink.is_some());
277        assert!(b.finalize.is_some());
278    }
279
280    #[test]
281    fn try_new_stores_name() {
282        let b = CopyFunctionBuilder::try_new("my_copy").unwrap();
283        assert_eq!(b.name(), "my_copy");
284    }
285
286    #[test]
287    fn callbacks_default_to_none() {
288        let b = CopyFunctionBuilder::try_new("fmt").unwrap();
289        assert!(b.bind.is_none());
290        assert!(b.global_init.is_none());
291        assert!(b.sink.is_none());
292        assert!(b.finalize.is_none());
293    }
294}
quack_rs/copy_function/mod.rs

quack_rs/copy_function/
mod.rs
