better_duck_core/raw/extracted.rs
1//! Parse a multi-statement SQL string into individually-preparable statements.
2//!
3//! DuckDB does the parsing: `duckdb_extract_statements` splits the query into N
4//! statements. There is deliberately **no** SQL-splitting heuristic on the Rust
5//! side — semicolons inside string literals, comments, or dollar-quoting make
6//! naive splitting wrong, and DuckDB's parser is the only correct source.
7
8use std::{
9 ffi::{CStr, CString},
10 ptr,
11 sync::Arc,
12};
13
14use crate::{
15 error::{EngineError, Error, Result},
16 ffi::{
17 duckdb_destroy_extracted, duckdb_extract_statements, duckdb_extract_statements_error,
18 duckdb_extracted_statements, duckdb_prepare_extracted_statement, duckdb_prepared_statement,
19 DuckDBError, Error as FFIError,
20 },
21 helpers::duck_result::result_from_duckdb_prepare,
22 raw::{
23 connection::{ConnectionInner, RawConnection},
24 statement::CachedStatement,
25 },
26};
27
28/// A parsed batch of SQL statements, each preparable on demand.
29///
30/// Produced by [`Connection::extract_statements`](crate::connection::Connection::extract_statements).
31/// Owns the
32/// `duckdb_extracted_statements` handle and destroys it in [`Drop`]; it retains
33/// the connection it parsed against (`Arc<ConnectionInner>`) so a statement
34/// prepared from it cannot outlive that connection — the same ownership rule as
35/// [`CachedStatement`].
36pub struct ExtractedStatements {
37 /// Keeps the parsing connection alive for at least as long as this batch and
38 /// any statement prepared from it.
39 connection: Arc<ConnectionInner>,
40 /// Owned handle; destroyed exactly once in `Drop`.
41 extracted: duckdb_extracted_statements,
42 /// Number of statements DuckDB parsed out of the query.
43 count: u64,
44}
45
46impl ExtractedStatements {
47 /// Extracts all statements from `sql` on `conn`.
48 ///
49 /// The SQL text is copied by DuckDB, so it need not outlive this call.
50 ///
51 /// # Errors
52 ///
53 /// [`Error::NulError`] if `sql` has an interior nul; otherwise, on a parse
54 /// failure, [`Error::Engine`] carrying DuckDB's extract-error message (an
55 /// [`EngineError::unavailable`] — `duckdb_extract_statements` exposes only a
56 /// message, no typed classification).
57 pub(crate) fn extract(
58 conn: &RawConnection,
59 sql: &str,
60 ) -> Result<ExtractedStatements> {
61 let c_sql = CString::new(sql)?;
62 let mut extracted: duckdb_extracted_statements = ptr::null_mut();
63 // SAFETY: `conn`'s handle is a valid open connection; `c_sql` is a valid
64 // null-terminated string that outlives the call; `&mut extracted` is a valid
65 // output pointer. DuckDB requires the handle to be destroyed regardless of
66 // the returned count, which `Drop` (success) or the error path (below) does.
67 let count =
68 unsafe { duckdb_extract_statements(conn.handle(), c_sql.as_ptr(), &mut extracted) };
69
70 if count == 0 {
71 // Parse failed: copy the error message out (it is freed by
72 // `duckdb_destroy_extracted`, so it must be copied first), then destroy.
73 // SAFETY: `extracted` is the (possibly-error) handle DuckDB just produced;
74 // `duckdb_extract_statements_error` returns a borrowed C string valid until
75 // the handle is destroyed, which we do immediately after copying.
76 let message = unsafe {
77 let raw = duckdb_extract_statements_error(extracted);
78 let msg =
79 (!raw.is_null()).then(|| CStr::from_ptr(raw).to_string_lossy().into_owned());
80 duckdb_destroy_extracted(&mut extracted);
81 msg
82 };
83 return Err(Error::Engine(EngineError::unavailable(Some(
84 message.unwrap_or_else(|| "failed to extract statements".to_owned()),
85 ))));
86 }
87
88 Ok(ExtractedStatements { connection: Arc::clone(conn.inner()), extracted, count })
89 }
90
91 /// The number of statements parsed from the query.
92 #[must_use]
93 pub fn len(&self) -> u64 {
94 self.count
95 }
96
97 /// Returns `true` if no statements were parsed.
98 ///
99 /// Always `false` in practice — a zero-statement extract is reported as an
100 /// error at extraction time — but provided so the
101 /// type satisfies the usual `len`/`is_empty` pairing.
102 #[must_use]
103 pub fn is_empty(&self) -> bool {
104 self.count == 0
105 }
106
107 /// Prepares the extracted statement at `index` (0-based), reusing the ordinary
108 /// [`CachedStatement`] wrapper.
109 ///
110 /// # Errors
111 ///
112 /// [`Error::DuckDBFailure`] if `index` is out of range, or if DuckDB fails to
113 /// prepare the statement.
114 pub fn prepare(
115 &self,
116 index: u64,
117 ) -> Result<CachedStatement> {
118 if index >= self.count {
119 return Err(Error::DuckDBFailure(
120 FFIError::new(DuckDBError),
121 Some(format!(
122 "extracted statement index {index} out of range (parsed {} statement(s))",
123 self.count
124 )),
125 ));
126 }
127
128 let mut stmt: duckdb_prepared_statement = ptr::null_mut();
129 // SAFETY: `self.connection`'s handle is valid and kept alive by the `Arc`;
130 // `self.extracted` is a live handle owned by `self`; `index` is in range
131 // (checked above); `&mut stmt` is a valid output pointer. DuckDB requires the
132 // prepared statement to be destroyed regardless of outcome — `CachedStatement`
133 // does that in its own `Drop`, and `result_from_duckdb_prepare` destroys it on
134 // the failure path.
135 let rc = unsafe {
136 duckdb_prepare_extracted_statement(
137 self.connection.handle(),
138 self.extracted,
139 index,
140 &mut stmt,
141 )
142 };
143 result_from_duckdb_prepare(rc, stmt)?;
144 // Extracted statements have no standalone per-statement SQL text; the cache
145 // key is unused for this path, so an empty label is fine.
146 Ok(CachedStatement::from_prepared(Arc::clone(&self.connection), stmt, Box::from("")))
147 }
148}
149
150impl std::fmt::Debug for ExtractedStatements {
151 fn fmt(
152 &self,
153 f: &mut std::fmt::Formatter<'_>,
154 ) -> std::fmt::Result {
155 f.debug_struct("ExtractedStatements").field("count", &self.count).finish_non_exhaustive()
156 }
157}
158
159impl Drop for ExtractedStatements {
160 fn drop(&mut self) {
161 if self.extracted.is_null() {
162 return;
163 }
164 // SAFETY: `self.extracted` is a valid, non-null handle owned exclusively by
165 // this value (null-guarded above); `duckdb_destroy_extracted` frees it and is
166 // called exactly once.
167 unsafe { duckdb_destroy_extracted(&mut self.extracted) };
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174 use crate::{config::Config, helpers::path::path_to_cstring, types::value::DuckValue};
175
176 fn conn() -> RawConnection {
177 let path = path_to_cstring(":memory:".as_ref()).unwrap();
178 let config = Config::default().with("duckdb_api", "rust").unwrap();
179 RawConnection::open_with_flags(&path, config).unwrap()
180 }
181
182 #[test]
183 fn extracts_and_prepares_each_statement_of_a_batch() {
184 let con = conn();
185 let batch =
186 ExtractedStatements::extract(&con, "SELECT 1 AS a; SELECT 'two' AS b; SELECT 3 AS c")
187 .unwrap();
188 assert_eq!(batch.len(), 3);
189 assert!(!batch.is_empty());
190
191 // Each parses/prepares independently and yields its own result.
192 let mut s0 = batch.prepare(0).unwrap();
193 let v = s0.execute().unwrap().next().unwrap().unwrap();
194 assert_eq!(v.get("a"), Some(&DuckValue::Int(1)));
195
196 let mut s2 = batch.prepare(2).unwrap();
197 let v = s2.execute().unwrap().next().unwrap().unwrap();
198 assert_eq!(v.get("c"), Some(&DuckValue::Int(3)));
199 }
200
201 #[test]
202 fn semicolon_inside_a_string_literal_is_not_a_split() {
203 // A naive splitter would see two statements; DuckDB's parser sees one.
204 let con = conn();
205 let batch = ExtractedStatements::extract(&con, "SELECT 'a;b' AS s").unwrap();
206 assert_eq!(batch.len(), 1);
207 let mut s = batch.prepare(0).unwrap();
208 let v = s.execute().unwrap().next().unwrap().unwrap();
209 assert_eq!(v.get("s"), Some(&DuckValue::Text("a;b".into())));
210 }
211
212 #[test]
213 fn out_of_range_index_is_rejected() {
214 let con = conn();
215 let batch = ExtractedStatements::extract(&con, "SELECT 1").unwrap();
216 assert!(matches!(
217 batch.prepare(5),
218 Err(Error::DuckDBFailure(_, Some(m))) if m.contains("out of range")
219 ));
220 }
221
222 #[test]
223 fn empty_query_reports_a_typed_engine_error() {
224 // An empty query extracts to zero statements — DuckDB returns count 0 with
225 // an extract error, which we surface as a typed engine error. (A *syntax*
226 // error like "SELCT 1" is deliberately not tested here: this DuckDB build
227 // throws a C++ parser exception across the C API on that path instead of
228 // returning 0, and a foreign exception crossing the FFI boundary aborts the
229 // process — an upstream limitation, not a wrapper bug.)
230 let con = conn();
231 let err = ExtractedStatements::extract(&con, "").unwrap_err();
232 assert!(matches!(err, Error::Engine(_)), "expected engine error, got {err:?}");
233 }
234
235 #[test]
236 fn interior_nul_query_is_rejected() {
237 let con = conn();
238 assert!(matches!(
239 ExtractedStatements::extract(&con, "SELECT 1\0; SELECT 2"),
240 Err(Error::NulError(_))
241 ));
242 }
243
244 #[test]
245 fn a_prepared_extracted_statement_outlives_the_batch() {
246 // The prepared statement retains the connection, not the batch, so dropping
247 // the batch first must not invalidate it.
248 let con = conn();
249 let batch = ExtractedStatements::extract(&con, "SELECT 42 AS v").unwrap();
250 let mut stmt = batch.prepare(0).unwrap();
251 drop(batch);
252 let v = stmt.execute().unwrap().next().unwrap().unwrap();
253 assert_eq!(v.get("v"), Some(&DuckValue::Int(42)));
254 }
255}