duckdb/column.rs
1use std::str;
2
3use arrow::datatypes::DataType;
4
5use crate::{Error, Result, Statement, core::LogicalTypeHandle};
6
7/// Information about a column of a DuckDB query.
8#[derive(Debug)]
9pub struct Column<'stmt> {
10 name: &'stmt str,
11 decl_type: Option<&'stmt str>,
12}
13
14impl Column<'_> {
15 /// Returns the name of the column.
16 #[inline]
17 pub fn name(&self) -> &str {
18 self.name
19 }
20
21 /// Returns the type of the column (`None` for expression).
22 #[inline]
23 pub fn decl_type(&self) -> Option<&str> {
24 self.decl_type
25 }
26}
27
28impl Statement<'_> {
29 /// Get all the column names in the result set of the prepared statement.
30 ///
31 /// If associated DB schema can be altered concurrently, you should make
32 /// sure that current statement has already been stepped once before
33 /// calling this method.
34 ///
35 /// # Caveats
36 /// Panics if the query has not been [`execute`](Statement::execute)d yet.
37 pub fn column_names(&self) -> Vec<String> {
38 self.stmt
39 .schema()
40 .fields()
41 .iter()
42 .map(|f| f.name().to_owned())
43 .collect()
44 }
45
46 /// Return the number of columns in the result set returned by the prepared
47 /// statement.
48 ///
49 /// If associated DB schema can be altered concurrently, you should make
50 /// sure that current statement has already been stepped once before
51 /// calling this method.
52 ///
53 /// # Example
54 ///
55 /// ```rust,no_run
56 /// # use duckdb::{Connection, Result};
57 /// fn get_column_count(conn: &Connection) -> Result<usize> {
58 /// let mut stmt = conn.prepare("SELECT id, name FROM people")?;
59 ///
60 /// // Option 1: Execute first, then get column count
61 /// stmt.execute([])?;
62 /// let count = stmt.column_count();
63 ///
64 /// // Option 2: Get column count from rows (avoids borrowing issues)
65 /// let mut stmt2 = conn.prepare("SELECT id, name FROM people")?;
66 /// let rows = stmt2.query([])?;
67 /// let count2 = rows.as_ref().unwrap().column_count();
68 ///
69 /// Ok(count)
70 /// }
71 /// ```
72 ///
73 /// # Caveats
74 /// Panics if the query has not been [`execute`](Statement::execute)d yet.
75 #[inline]
76 pub fn column_count(&self) -> usize {
77 self.stmt.column_count()
78 }
79
80 /// Check that column name reference lifetime is limited:
81 /// https://www.sqlite.org/c3ref/column_name.html
82 /// > The returned string pointer is valid...
83 ///
84 /// `column_name` reference can become invalid if `stmt` is reprepared
85 /// (because of schema change) when `query_row` is called. So we assert
86 /// that a compilation error happens if this reference is kept alive:
87 /// ```compile_fail
88 /// use duckdb::{Connection, Result};
89 /// fn main() -> Result<()> {
90 /// let db = Connection::open_in_memory()?;
91 /// let mut stmt = db.prepare("SELECT 1 as x")?;
92 /// let column_name = stmt.column_name(0)?;
93 /// let x = stmt.query_row([], |r| r.get::<_, i64>(0))?; // E0502
94 /// assert_eq!(1, x);
95 /// assert_eq!("x", column_name);
96 /// Ok(())
97 /// }
98 /// ```
99 #[inline]
100 pub(super) fn column_name_unwrap(&self, col: usize) -> &String {
101 // Just panic if the bounds are wrong for now, we never call this
102 // without checking first.
103 self.column_name(col).expect("Column out of bounds")
104 }
105
106 /// Returns the name assigned to a particular column in the result set
107 /// returned by the prepared statement.
108 ///
109 /// If associated DB schema can be altered concurrently, you should make
110 /// sure that current statement has already been stepped once before
111 /// calling this method.
112 ///
113 /// ## Failure
114 ///
115 /// Returns an `Error::InvalidColumnIndex` if `idx` is outside the valid
116 /// column range for this row.
117 ///
118 /// # Caveats
119 /// Panics if the query has not been [`execute`](Statement::execute)d yet
120 /// or when column name is not valid UTF-8.
121 #[inline]
122 pub fn column_name(&self, col: usize) -> Result<&String> {
123 self.stmt.column_name(col).ok_or(Error::InvalidColumnIndex(col))
124 }
125
126 /// Returns the column index in the result set for a given column name.
127 ///
128 /// If there is no AS clause then the name of the column is unspecified and
129 /// may change from one release of DuckDB to the next.
130 ///
131 /// If associated DB schema can be altered concurrently, you should make
132 /// sure that current statement has already been stepped once before
133 /// calling this method.
134 ///
135 /// # Failure
136 ///
137 /// Will return an `Error::InvalidColumnName` when there is no column with
138 /// the specified `name`.
139 ///
140 /// # Caveats
141 /// Panics if the query has not been [`execute`](Statement::execute)d yet.
142 #[inline]
143 pub fn column_index(&self, name: &str) -> Result<usize> {
144 self.stmt
145 .column_index(name)
146 .ok_or_else(|| Error::InvalidColumnName(String::from(name)))
147 }
148
149 /// Returns the declared data type of the column.
150 ///
151 /// # Caveats
152 /// Panics if the query has not been [`execute`](Statement::execute)d yet.
153 #[inline]
154 pub fn column_type(&self, idx: usize) -> DataType {
155 self.stmt.column_type(idx)
156 }
157
158 /// Returns the declared logical data type of the column.
159 ///
160 /// # Caveats
161 /// Panics if `idx` is outside the valid column range.
162 pub fn column_logical_type(&self, idx: usize) -> LogicalTypeHandle {
163 self.stmt.column_logical_type(idx)
164 }
165}
166
167#[cfg(test)]
168mod test {
169 use crate::{Connection, Result};
170
171 #[test]
172 fn test_column_name_in_error() -> Result<()> {
173 use crate::{Error, types::Type};
174 let db = Connection::open_in_memory()?;
175 db.execute_batch(
176 "BEGIN;
177 CREATE TABLE foo(x INTEGER, y TEXT);
178 INSERT INTO foo VALUES(4, NULL);
179 END;",
180 )?;
181 let mut stmt = db.prepare("SELECT x as renamed, y FROM foo")?;
182 let mut rows = stmt.query([])?;
183 let row = rows.next()?.unwrap();
184 match row.get::<_, String>(0).unwrap_err() {
185 Error::InvalidColumnType(idx, name, ty) => {
186 assert_eq!(idx, 0);
187 assert_eq!(name, "renamed");
188 assert_eq!(ty, Type::Int);
189 }
190 e => {
191 panic!("Unexpected error type: {e:?}");
192 }
193 }
194 match row.get::<_, String>("y").unwrap_err() {
195 Error::InvalidColumnType(idx, name, ty) => {
196 assert_eq!(idx, 1);
197 assert_eq!(name, "y");
198 assert_eq!(ty, Type::Null);
199 }
200 e => {
201 panic!("Unexpected error type: {e:?}");
202 }
203 }
204 Ok(())
205 }
206
207 #[test]
208 fn test_column_index_duplicate_names() -> Result<()> {
209 let db = Connection::open_in_memory()?;
210 let mut stmt = db.prepare("SELECT 1 as a, 2 as a, 3 as a")?;
211 stmt.execute([])?;
212
213 assert_eq!(stmt.column_index("a")?, 0);
214 assert_eq!(stmt.column_index("A")?, 0);
215
216 Ok(())
217 }
218
219 #[test]
220 fn test_column_index_cache_invalidation() -> Result<()> {
221 let db = Connection::open_in_memory()?;
222 let mut stmt = db.prepare("SELECT 1 as col_a, 2 as col_b")?;
223 stmt.execute([])?;
224
225 assert_eq!(stmt.column_index("col_a")?, 0);
226 assert_eq!(stmt.column_index("col_b")?, 1);
227
228 stmt.execute([])?;
229
230 assert_eq!(stmt.column_index("col_a")?, 0);
231 assert_eq!(stmt.column_index("col_b")?, 1);
232
233 Ok(())
234 }
235
236 #[test]
237 fn test_column_index_nonexistent() -> Result<()> {
238 use crate::Error;
239 let db = Connection::open_in_memory()?;
240 let mut stmt = db.prepare("SELECT 1 as x, 2 as y")?;
241 stmt.execute([])?;
242
243 match stmt.column_index("nonexistent") {
244 Err(Error::InvalidColumnName(name)) => {
245 assert_eq!(name, "nonexistent");
246 }
247 other => panic!("Expected InvalidColumnName error, got: {other:?}"),
248 }
249
250 Ok(())
251 }
252
253 #[test]
254 #[should_panic(expected = "The statement was not executed yet")]
255 fn test_column_index_not_executed() {
256 let db = Connection::open_in_memory().unwrap();
257 let stmt = db.prepare("SELECT 1 as a").unwrap();
258 let _ = stmt.column_index("a");
259 }
260}