Skip to main content

qail_pg/driver/
row.rs

1//! PostgreSQL Row Helpers
2//!
3//! Provides convenient methods to extract typed values from row data.
4//! Supports both text and binary result formats when column metadata is available.
5
6use super::PgRow;
7use crate::types::{FromPg, TypeError};
8
9/// Trait for types that can be constructed from a database row.
10///
11/// Implement this trait on your structs to enable typed fetching:
12/// ```ignore
13/// impl QailRow for User {
14///     fn columns() -> &'static [&'static str] {
15///         &["id", "name", "email"]
16///     }
17///     
18///     fn from_row(row: &PgRow) -> Self {
19///         User {
20///             // Metadata-aware decoders (OID + text/binary format).
21///             id: row
22///                 .try_get_by_name::<qail_pg::Uuid>("id")
23///                 .map(|v| v.0)
24///                 .unwrap_or_default(),
25///             name: row.try_get_by_name::<String>("name").unwrap_or_default(),
26///             email: row.try_get_by_name::<String>("email").unwrap_or_default(),
27///         }
28///     }
29/// }
30///
31/// // Then use:
32/// let users: Vec<User> = driver.fetch_typed::<User>(&query).await?;
33/// ```
34pub trait QailRow: Sized {
35    /// Return the column names this struct expects.
36    /// These are used to automatically build SELECT queries.
37    fn columns() -> &'static [&'static str];
38
39    /// Construct an instance from a PgRow.
40    /// Column indices match the order returned by `columns()`.
41    fn from_row(row: &PgRow) -> Self;
42}
43
44impl PgRow {
45    /// Decode a non-null column into any `FromPg` type using backend OID/format metadata.
46    ///
47    /// Returns:
48    /// - `TypeError::UnexpectedNull` if the column value is NULL
49    /// - `TypeError::InvalidData` if the column index is out of bounds or metadata is missing
50    /// - other `TypeError` variants from the target decoder
51    pub fn try_get<T: FromPg>(&self, idx: usize) -> Result<T, TypeError> {
52        let cell = self
53            .columns
54            .get(idx)
55            .ok_or_else(|| TypeError::InvalidData(format!("Column index {} out of bounds", idx)))?;
56
57        let bytes = cell.as_deref().ok_or(TypeError::UnexpectedNull)?;
58        let (oid, format) = self.column_type_meta(idx)?;
59        T::from_pg(bytes, oid, format)
60    }
61
62    /// Decode a possibly-null column into `Option<T>` using backend OID/format metadata.
63    ///
64    /// Returns `Ok(None)` when the column is SQL NULL.
65    pub fn try_get_opt<T: FromPg>(&self, idx: usize) -> Result<Option<T>, TypeError> {
66        let cell = self
67            .columns
68            .get(idx)
69            .ok_or_else(|| TypeError::InvalidData(format!("Column index {} out of bounds", idx)))?;
70
71        match cell {
72            None => Ok(None),
73            Some(bytes) => {
74                let (oid, format) = self.column_type_meta(idx)?;
75                Ok(Some(T::from_pg(bytes, oid, format)?))
76            }
77        }
78    }
79
80    /// Decode a non-null column by name into any `FromPg` type.
81    pub fn try_get_by_name<T: FromPg>(&self, name: &str) -> Result<T, TypeError> {
82        let idx = self
83            .column_index(name)
84            .ok_or_else(|| TypeError::InvalidData(format!("Unknown column name '{}'", name)))?;
85        self.try_get(idx)
86    }
87
88    /// Decode a possibly-null column by name into `Option<T>`.
89    pub fn try_get_opt_by_name<T: FromPg>(&self, name: &str) -> Result<Option<T>, TypeError> {
90        let idx = self
91            .column_index(name)
92            .ok_or_else(|| TypeError::InvalidData(format!("Unknown column name '{}'", name)))?;
93        self.try_get_opt(idx)
94    }
95
96    fn column_type_meta(&self, idx: usize) -> Result<(u32, i16), TypeError> {
97        let info = self.column_info.as_ref().ok_or_else(|| {
98            TypeError::InvalidData(
99                "Column metadata unavailable; use query APIs that preserve RowDescription"
100                    .to_string(),
101            )
102        })?;
103
104        let oid = info
105            .oids
106            .get(idx)
107            .copied()
108            .ok_or_else(|| TypeError::InvalidData(format!("Missing OID for column {}", idx)))?;
109        let format = info.formats.get(idx).copied().ok_or_else(|| {
110            TypeError::InvalidData(format!("Missing format code for column {}", idx))
111        })?;
112        Ok((oid, format))
113    }
114
115    /// Get a column value as String.
116    /// Returns None if the value is NULL or invalid UTF-8.
117    pub fn get_string(&self, idx: usize) -> Option<String> {
118        self.columns
119            .get(idx)?
120            .as_ref()
121            .and_then(|bytes| String::from_utf8(bytes.clone()).ok())
122    }
123
124    /// Get a column value as i32.
125    pub fn get_i32(&self, idx: usize) -> Option<i32> {
126        if self.column_info.is_some()
127            && let Ok(v) = self.try_get::<i32>(idx)
128        {
129            return Some(v);
130        }
131        let bytes = self.columns.get(idx)?.as_ref()?;
132        std::str::from_utf8(bytes).ok()?.parse().ok()
133    }
134
135    /// Get a column value as i64.
136    pub fn get_i64(&self, idx: usize) -> Option<i64> {
137        if self.column_info.is_some()
138            && let Ok(v) = self.try_get::<i64>(idx)
139        {
140            return Some(v);
141        }
142        let bytes = self.columns.get(idx)?.as_ref()?;
143        std::str::from_utf8(bytes).ok()?.parse().ok()
144    }
145
146    /// Get a column value as f64.
147    pub fn get_f64(&self, idx: usize) -> Option<f64> {
148        if self.column_info.is_some()
149            && let Ok(v) = self.try_get::<f64>(idx)
150        {
151            return Some(v);
152        }
153        let bytes = self.columns.get(idx)?.as_ref()?;
154        std::str::from_utf8(bytes).ok()?.parse().ok()
155    }
156
157    /// Get a column value as bool.
158    pub fn get_bool(&self, idx: usize) -> Option<bool> {
159        if self.column_info.is_some()
160            && let Ok(v) = self.try_get::<bool>(idx)
161        {
162            return Some(v);
163        }
164        let bytes = self.columns.get(idx)?.as_ref()?;
165        let s = std::str::from_utf8(bytes).ok()?;
166        match s {
167            "t" | "true" | "1" => Some(true),
168            "f" | "false" | "0" => Some(false),
169            _ => None,
170        }
171    }
172
173    /// Check if a column is NULL.
174    pub fn is_null(&self, idx: usize) -> bool {
175        self.columns.get(idx).map(|v| v.is_none()).unwrap_or(true)
176    }
177
178    /// Get raw bytes of a column.
179    pub fn get_bytes(&self, idx: usize) -> Option<&[u8]> {
180        self.columns.get(idx)?.as_ref().map(|v| v.as_slice())
181    }
182
183    /// Get number of columns in the row.
184    pub fn len(&self) -> usize {
185        self.columns.len()
186    }
187
188    /// Check if the row has no columns.
189    pub fn is_empty(&self) -> bool {
190        self.columns.is_empty()
191    }
192
193    /// Get a column value as UUID string.
194    /// Handles both text format (36-char string) and binary format (16 bytes).
195    pub fn get_uuid(&self, idx: usize) -> Option<String> {
196        let bytes = self.columns.get(idx)?.as_ref()?;
197
198        if bytes.len() == 16 {
199            // Binary format - decode 16 bytes
200            use crate::protocol::types::decode_uuid;
201            decode_uuid(bytes).ok()
202        } else {
203            // Text format - return as-is
204            String::from_utf8(bytes.clone()).ok()
205        }
206    }
207
208    /// Get a column value as JSON string.
209    /// Handles both JSON (text) and JSONB (version byte prefix) formats.
210    pub fn get_json(&self, idx: usize) -> Option<String> {
211        let bytes = self.columns.get(idx)?.as_ref()?;
212
213        if bytes.is_empty() {
214            return Some(String::new());
215        }
216
217        // JSONB has version byte (1) as first byte
218        if bytes[0] == 1 && bytes.len() > 1 {
219            String::from_utf8(bytes[1..].to_vec()).ok()
220        } else {
221            String::from_utf8(bytes.clone()).ok()
222        }
223    }
224
225    /// Get a column value as timestamp string (ISO 8601 format).
226    pub fn get_timestamp(&self, idx: usize) -> Option<String> {
227        let bytes = self.columns.get(idx)?.as_ref()?;
228        String::from_utf8(bytes.clone()).ok()
229    }
230
231    /// Get a column value as text array.
232    pub fn get_text_array(&self, idx: usize) -> Option<Vec<String>> {
233        let bytes = self.columns.get(idx)?.as_ref()?;
234        let s = std::str::from_utf8(bytes).ok()?;
235        Some(crate::protocol::types::decode_text_array(s))
236    }
237
238    /// Get a column value as integer array.
239    pub fn get_int_array(&self, idx: usize) -> Option<Vec<i64>> {
240        let bytes = self.columns.get(idx)?.as_ref()?;
241        let s = std::str::from_utf8(bytes).ok()?;
242        crate::protocol::types::decode_int_array(s).ok()
243    }
244
245    // ==================== ERGONOMIC SHORTCUTS ====================
246    // These methods reduce boilerplate by providing sensible defaults
247
248    /// Get string, defaulting to empty string if NULL.
249    /// Ergonomic shortcut: `row.text(0)` instead of `row.get_string(0).unwrap_or_default()`
250    pub fn text(&self, idx: usize) -> String {
251        self.get_string(idx).unwrap_or_default()
252    }
253
254    /// Get string with custom default if NULL.
255    /// Example: `row.text_or(1, "Unknown")`
256    pub fn text_or(&self, idx: usize, default: &str) -> String {
257        self.get_string(idx).unwrap_or_else(|| default.to_string())
258    }
259
260    /// Get i64, defaulting to 0 if NULL.
261    /// Ergonomic shortcut: `row.int(4)` instead of `row.get_i64(4).unwrap_or(0)`
262    pub fn int(&self, idx: usize) -> i64 {
263        self.get_i64(idx).unwrap_or(0)
264    }
265
266    /// Get f64, defaulting to 0.0 if NULL.
267    pub fn float(&self, idx: usize) -> f64 {
268        self.get_f64(idx).unwrap_or(0.0)
269    }
270
271    /// Get bool, defaulting to false if NULL.
272    pub fn boolean(&self, idx: usize) -> bool {
273        self.get_bool(idx).unwrap_or(false)
274    }
275
276    /// Parse timestamp as DateTime<Utc>.
277    /// Handles PostgreSQL timestamp formats automatically.
278    #[cfg(feature = "chrono")]
279    pub fn datetime(&self, idx: usize) -> Option<chrono::DateTime<chrono::Utc>> {
280        if let Ok(dt) = self.try_get::<chrono::DateTime<chrono::Utc>>(idx) {
281            return Some(dt);
282        }
283
284        let s = self.get_timestamp(idx)?;
285        // Try parsing various PostgreSQL timestamp formats
286        chrono::DateTime::parse_from_rfc3339(&s.replace(' ', "T"))
287            .ok()
288            .map(|dt| dt.with_timezone(&chrono::Utc))
289            .or_else(|| {
290                // Try PostgreSQL format: "2024-01-01 12:00:00.123456+00"
291                chrono::DateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S%.f%#z")
292                    .ok()
293                    .map(|dt| dt.with_timezone(&chrono::Utc))
294            })
295    }
296
297    /// Parse UUID column as uuid::Uuid type.
298    #[cfg(feature = "uuid")]
299    pub fn uuid_typed(&self, idx: usize) -> Option<uuid::Uuid> {
300        self.try_get::<uuid::Uuid>(idx).ok().or_else(|| {
301            self.get_uuid(idx)
302                .and_then(|s| uuid::Uuid::parse_str(&s).ok())
303        })
304    }
305
306    // ==================== GET BY COLUMN NAME ====================
307
308    /// Get column index by name.
309    pub fn column_index(&self, name: &str) -> Option<usize> {
310        self.column_info.as_ref()?.name_to_index.get(name).copied()
311    }
312
313    /// Get a String column by name.
314    pub fn get_string_by_name(&self, name: &str) -> Option<String> {
315        self.get_string(self.column_index(name)?)
316    }
317
318    /// Get an i32 column by name.
319    pub fn get_i32_by_name(&self, name: &str) -> Option<i32> {
320        self.get_i32(self.column_index(name)?)
321    }
322
323    /// Get an i64 column by name.
324    pub fn get_i64_by_name(&self, name: &str) -> Option<i64> {
325        self.get_i64(self.column_index(name)?)
326    }
327
328    /// Get a f64 column by name.
329    pub fn get_f64_by_name(&self, name: &str) -> Option<f64> {
330        self.get_f64(self.column_index(name)?)
331    }
332
333    /// Get a bool column by name.
334    pub fn get_bool_by_name(&self, name: &str) -> Option<bool> {
335        self.get_bool(self.column_index(name)?)
336    }
337
338    /// Get a UUID column by name.
339    pub fn get_uuid_by_name(&self, name: &str) -> Option<String> {
340        self.get_uuid(self.column_index(name)?)
341    }
342
343    /// Get a JSON column by name.
344    pub fn get_json_by_name(&self, name: &str) -> Option<String> {
345        self.get_json(self.column_index(name)?)
346    }
347
348    /// Check if a column is NULL by name.
349    pub fn is_null_by_name(&self, name: &str) -> bool {
350        self.column_index(name)
351            .map(|idx| self.is_null(idx))
352            .unwrap_or(true)
353    }
354
355    /// Get a timestamp column by name.
356    pub fn get_timestamp_by_name(&self, name: &str) -> Option<String> {
357        self.get_timestamp(self.column_index(name)?)
358    }
359
360    /// Get a text array column by name.
361    pub fn get_text_array_by_name(&self, name: &str) -> Option<Vec<String>> {
362        self.get_text_array(self.column_index(name)?)
363    }
364
365    /// Get an integer array column by name.
366    pub fn get_int_array_by_name(&self, name: &str) -> Option<Vec<i64>> {
367        self.get_int_array(self.column_index(name)?)
368    }
369
370    // ==================== ERGONOMIC BY-NAME SHORTCUTS ====================
371    // These mirror the positional shortcuts (text, boolean, int, etc.)
372    // but use column names — safe with RETURNING * regardless of column order.
373
374    /// Get string by column name, defaulting to empty string.
375    /// Example: `row.text_by_name("name")` instead of `row.get_string_by_name("name").unwrap_or_default()`
376    pub fn text_by_name(&self, name: &str) -> String {
377        self.get_string_by_name(name).unwrap_or_default()
378    }
379
380    /// Get bool by column name, defaulting to false.
381    pub fn boolean_by_name(&self, name: &str) -> bool {
382        self.get_bool_by_name(name).unwrap_or(false)
383    }
384
385    /// Get i64 by column name, defaulting to 0.
386    pub fn int_by_name(&self, name: &str) -> i64 {
387        self.get_i64_by_name(name).unwrap_or(0)
388    }
389
390    /// Get f64 by column name, defaulting to 0.0.
391    pub fn float_by_name(&self, name: &str) -> f64 {
392        self.get_f64_by_name(name).unwrap_or(0.0)
393    }
394
395    /// Parse timestamp by column name as DateTime<Utc>.
396    #[cfg(feature = "chrono")]
397    pub fn datetime_by_name(&self, name: &str) -> Option<chrono::DateTime<chrono::Utc>> {
398        self.datetime(self.column_index(name)?)
399    }
400
401    /// Parse UUID by column name as uuid::Uuid type.
402    #[cfg(feature = "uuid")]
403    pub fn uuid_typed_by_name(&self, name: &str) -> Option<uuid::Uuid> {
404        self.uuid_typed(self.column_index(name)?)
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use super::*;
411    use crate::protocol::types::oid;
412    use crate::types::{Json, Uuid};
413    use std::collections::HashMap;
414    use std::sync::Arc;
415
416    fn single_col_info(name: &str, oid: u32, format: i16) -> Arc<super::super::ColumnInfo> {
417        let mut name_to_index = HashMap::new();
418        name_to_index.insert(name.to_string(), 0);
419        Arc::new(super::super::ColumnInfo {
420            name_to_index,
421            oids: vec![oid],
422            formats: vec![format],
423        })
424    }
425
426    #[test]
427    fn test_get_string() {
428        let row = PgRow {
429            columns: vec![Some(b"hello".to_vec()), None, Some(b"world".to_vec())],
430            column_info: None,
431        };
432
433        assert_eq!(row.get_string(0), Some("hello".to_string()));
434        assert_eq!(row.get_string(1), None);
435        assert_eq!(row.get_string(2), Some("world".to_string()));
436    }
437
438    #[test]
439    fn test_get_i32() {
440        let row = PgRow {
441            columns: vec![
442                Some(b"42".to_vec()),
443                Some(b"-123".to_vec()),
444                Some(b"not_a_number".to_vec()),
445            ],
446            column_info: None,
447        };
448
449        assert_eq!(row.get_i32(0), Some(42));
450        assert_eq!(row.get_i32(1), Some(-123));
451        assert_eq!(row.get_i32(2), None);
452    }
453
454    #[test]
455    fn test_get_bool() {
456        let row = PgRow {
457            columns: vec![
458                Some(b"t".to_vec()),
459                Some(b"f".to_vec()),
460                Some(b"true".to_vec()),
461                Some(b"false".to_vec()),
462            ],
463            column_info: None,
464        };
465
466        assert_eq!(row.get_bool(0), Some(true));
467        assert_eq!(row.get_bool(1), Some(false));
468        assert_eq!(row.get_bool(2), Some(true));
469        assert_eq!(row.get_bool(3), Some(false));
470    }
471
472    #[test]
473    fn test_is_null() {
474        let row = PgRow {
475            columns: vec![Some(b"value".to_vec()), None],
476            column_info: None,
477        };
478
479        assert!(!row.is_null(0));
480        assert!(row.is_null(1));
481        assert!(row.is_null(99)); // Out of bounds
482    }
483
484    #[test]
485    fn test_try_get_i64_binary() {
486        let row = PgRow {
487            columns: vec![Some(42i64.to_be_bytes().to_vec())],
488            column_info: Some(single_col_info("count", oid::INT8, 1)),
489        };
490
491        let value: i64 = row.try_get(0).unwrap();
492        assert_eq!(value, 42);
493    }
494
495    #[test]
496    fn test_try_get_i64_text_by_name() {
497        let row = PgRow {
498            columns: vec![Some(b"123".to_vec())],
499            column_info: Some(single_col_info("total", oid::INT8, 0)),
500        };
501
502        let value: i64 = row.try_get_by_name("total").unwrap();
503        assert_eq!(value, 123);
504    }
505
506    #[test]
507    fn test_try_get_opt_null() {
508        let row = PgRow {
509            columns: vec![None],
510            column_info: Some(single_col_info("maybe_count", oid::INT8, 1)),
511        };
512
513        let value: Option<i64> = row.try_get_opt(0).unwrap();
514        assert_eq!(value, None);
515    }
516
517    #[test]
518    fn test_try_get_unexpected_null() {
519        let row = PgRow {
520            columns: vec![None],
521            column_info: Some(single_col_info("required_count", oid::INT8, 1)),
522        };
523
524        assert!(matches!(
525            row.try_get::<i64>(0),
526            Err(TypeError::UnexpectedNull)
527        ));
528    }
529
530    #[test]
531    fn test_try_get_uuid_binary() {
532        let uuid_bytes: [u8; 16] = [
533            0x55, 0x0e, 0x84, 0x00, 0xe2, 0x9b, 0x41, 0xd4, 0xa7, 0x16, 0x44, 0x66, 0x55, 0x44,
534            0x00, 0x00,
535        ];
536        let row = PgRow {
537            columns: vec![Some(uuid_bytes.to_vec())],
538            column_info: Some(single_col_info("id", oid::UUID, 1)),
539        };
540
541        let value: Uuid = row.try_get(0).unwrap();
542        assert_eq!(value.0, "550e8400-e29b-41d4-a716-446655440000");
543    }
544
545    #[test]
546    fn test_try_get_jsonb_binary() {
547        let mut bytes = vec![1u8];
548        bytes.extend_from_slice(br#"{"ok":true}"#);
549        let row = PgRow {
550            columns: vec![Some(bytes)],
551            column_info: Some(single_col_info("meta", oid::JSONB, 1)),
552        };
553
554        let value: Json = row.try_get(0).unwrap();
555        assert_eq!(value.0, r#"{"ok":true}"#);
556    }
557
558    #[test]
559    fn test_try_get_requires_column_metadata() {
560        let row = PgRow {
561            columns: vec![Some(b"42".to_vec())],
562            column_info: None,
563        };
564
565        assert!(matches!(
566            row.try_get::<i64>(0),
567            Err(TypeError::InvalidData(msg)) if msg.contains("metadata")
568        ));
569    }
570
571    #[test]
572    fn test_get_i64_uses_metadata_binary() {
573        let row = PgRow {
574            columns: vec![Some(777i64.to_be_bytes().to_vec())],
575            column_info: Some(single_col_info("v", oid::INT8, 1)),
576        };
577        assert_eq!(row.get_i64(0), Some(777));
578    }
579
580    #[test]
581    fn test_get_bool_uses_metadata_binary() {
582        let row = PgRow {
583            columns: vec![Some(vec![1u8])],
584            column_info: Some(single_col_info("flag", oid::BOOL, 1)),
585        };
586        assert_eq!(row.get_bool(0), Some(true));
587    }
588}