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 pub fn column_logical_type(&self, idx: usize) -> LogicalTypeHandle {
160 self.stmt.column_logical_type(idx)
161 }
162}
163
164#[cfg(test)]
165mod test {
166 use crate::{Connection, Result};
167
168 #[test]
169 fn test_column_name_in_error() -> Result<()> {
170 use crate::{Error, types::Type};
171 let db = Connection::open_in_memory()?;
172 db.execute_batch(
173 "BEGIN;
174 CREATE TABLE foo(x INTEGER, y TEXT);
175 INSERT INTO foo VALUES(4, NULL);
176 END;",
177 )?;
178 let mut stmt = db.prepare("SELECT x as renamed, y FROM foo")?;
179 let mut rows = stmt.query([])?;
180 let row = rows.next()?.unwrap();
181 match row.get::<_, String>(0).unwrap_err() {
182 Error::InvalidColumnType(idx, name, ty) => {
183 assert_eq!(idx, 0);
184 assert_eq!(name, "renamed");
185 assert_eq!(ty, Type::Int);
186 }
187 e => {
188 panic!("Unexpected error type: {e:?}");
189 }
190 }
191 match row.get::<_, String>("y").unwrap_err() {
192 Error::InvalidColumnType(idx, name, ty) => {
193 assert_eq!(idx, 1);
194 assert_eq!(name, "y");
195 assert_eq!(ty, Type::Null);
196 }
197 e => {
198 panic!("Unexpected error type: {e:?}");
199 }
200 }
201 Ok(())
202 }
203
204 #[test]
205 fn test_column_index_duplicate_names() -> Result<()> {
206 let db = Connection::open_in_memory()?;
207 let mut stmt = db.prepare("SELECT 1 as a, 2 as a, 3 as a")?;
208 stmt.execute([])?;
209
210 assert_eq!(stmt.column_index("a")?, 0);
211 assert_eq!(stmt.column_index("A")?, 0);
212
213 Ok(())
214 }
215
216 #[test]
217 fn test_column_index_cache_invalidation() -> Result<()> {
218 let db = Connection::open_in_memory()?;
219 let mut stmt = db.prepare("SELECT 1 as col_a, 2 as col_b")?;
220 stmt.execute([])?;
221
222 assert_eq!(stmt.column_index("col_a")?, 0);
223 assert_eq!(stmt.column_index("col_b")?, 1);
224
225 stmt.execute([])?;
226
227 assert_eq!(stmt.column_index("col_a")?, 0);
228 assert_eq!(stmt.column_index("col_b")?, 1);
229
230 Ok(())
231 }
232
233 #[test]
234 fn test_column_index_nonexistent() -> Result<()> {
235 use crate::Error;
236 let db = Connection::open_in_memory()?;
237 let mut stmt = db.prepare("SELECT 1 as x, 2 as y")?;
238 stmt.execute([])?;
239
240 match stmt.column_index("nonexistent") {
241 Err(Error::InvalidColumnName(name)) => {
242 assert_eq!(name, "nonexistent");
243 }
244 other => panic!("Expected InvalidColumnName error, got: {other:?}"),
245 }
246
247 Ok(())
248 }
249
250 #[test]
251 #[should_panic(expected = "The statement was not executed yet")]
252 fn test_column_index_not_executed() {
253 let db = Connection::open_in_memory().unwrap();
254 let stmt = db.prepare("SELECT 1 as a").unwrap();
255 let _ = stmt.column_index("a");
256 }
257}