qail_core/ast/builders/
ext.rs

1//! Extension traits for Expr.
2//!
3//! Adds fluent methods to Expr for more ergonomic query building.
4
5use super::json::JsonBuilder;
6use crate::ast::Expr;
7
8/// Extension trait to add fluent methods to Expr
9pub trait ExprExt {
10    /// Add an alias to this expression.
11    ///
12    /// # Example
13    /// ```ignore
14    /// col("name").with_alias("user_name")
15    /// ```
16    fn with_alias(self, alias: &str) -> Expr;
17
18    /// COALESCE with a default value.
19    ///
20    /// # Example
21    /// ```ignore
22    /// col("name").or_default("Unknown")  // COALESCE(name, 'Unknown')
23    /// ```
24    fn or_default(self, default: impl Into<Expr>) -> Expr;
25
26    /// JSON text extraction (column->>'key').
27    ///
28    /// # Example
29    /// ```ignore
30    /// col("contact_info").json("phone")  // contact_info->>'phone'
31    /// ```
32    fn json(self, key: &str) -> JsonBuilder;
33
34    /// JSON path extraction with dot notation.
35    ///
36    /// # Example
37    /// ```ignore
38    /// col("metadata").path("vessel.0.port")  // metadata->'vessel'->0->>'port'
39    /// ```
40    fn path(self, dotted_path: &str) -> JsonBuilder;
41
42    /// Cast to a type: CAST(expr AS type)
43    ///
44    /// # Example
45    /// ```ignore
46    /// col("value").cast("int")  // CAST(value AS int)
47    /// ```
48    fn cast(self, target_type: &str) -> Expr;
49
50    /// UPPER(expr)
51    fn upper(self) -> Expr;
52
53    /// LOWER(expr)
54    fn lower(self) -> Expr;
55
56    /// TRIM(expr)
57    fn trim(self) -> Expr;
58
59    /// LENGTH(expr)
60    fn length(self) -> Expr;
61
62    /// ABS(expr)
63    fn abs(self) -> Expr;
64}
65
66impl ExprExt for Expr {
67    fn with_alias(self, alias: &str) -> Expr {
68        match self {
69            Expr::Named(name) => Expr::Aliased {
70                name,
71                alias: alias.to_string(),
72            },
73            Expr::Aggregate {
74                col,
75                func,
76                distinct,
77                filter,
78                ..
79            } => Expr::Aggregate {
80                col,
81                func,
82                distinct,
83                filter,
84                alias: Some(alias.to_string()),
85            },
86            Expr::Cast {
87                expr, target_type, ..
88            } => Expr::Cast {
89                expr,
90                target_type,
91                alias: Some(alias.to_string()),
92            },
93            Expr::Case {
94                when_clauses,
95                else_value,
96                ..
97            } => Expr::Case {
98                when_clauses,
99                else_value,
100                alias: Some(alias.to_string()),
101            },
102            Expr::FunctionCall { name, args, .. } => Expr::FunctionCall {
103                name,
104                args,
105                alias: Some(alias.to_string()),
106            },
107            Expr::Binary {
108                left, op, right, ..
109            } => Expr::Binary {
110                left,
111                op,
112                right,
113                alias: Some(alias.to_string()),
114            },
115            Expr::JsonAccess {
116                column,
117                path_segments,
118                ..
119            } => Expr::JsonAccess {
120                column,
121                path_segments,
122                alias: Some(alias.to_string()),
123            },
124            Expr::SpecialFunction { name, args, .. } => Expr::SpecialFunction {
125                name,
126                args,
127                alias: Some(alias.to_string()),
128            },
129            other => other, // Star, Aliased, Literal, etc. - return as-is
130        }
131    }
132
133    fn or_default(self, default: impl Into<Expr>) -> Expr {
134        Expr::FunctionCall {
135            name: "COALESCE".to_string(),
136            args: vec![self, default.into()],
137            alias: None,
138        }
139    }
140
141    fn json(self, key: &str) -> JsonBuilder {
142        let column = match self {
143            Expr::Named(name) => name,
144            _ => panic!("json() can only be called on column references"),
145        };
146        JsonBuilder {
147            column,
148            path_segments: vec![(key.to_string(), true)], // true = text extraction (->>)
149            alias: None,
150        }
151    }
152
153    fn path(self, dotted_path: &str) -> JsonBuilder {
154        let column = match self {
155            Expr::Named(name) => name,
156            _ => panic!("path() can only be called on column references"),
157        };
158
159        let segments: Vec<&str> = dotted_path.split('.').collect();
160        let len = segments.len();
161        let path_segments: Vec<(String, bool)> = segments
162            .into_iter()
163            .enumerate()
164            .map(|(i, segment)| (segment.to_string(), i == len - 1)) // Last segment as text
165            .collect();
166
167        JsonBuilder {
168            column,
169            path_segments,
170            alias: None,
171        }
172    }
173
174    fn cast(self, target_type: &str) -> Expr {
175        Expr::Cast {
176            expr: Box::new(self),
177            target_type: target_type.to_string(),
178            alias: None,
179        }
180    }
181
182    fn upper(self) -> Expr {
183        Expr::FunctionCall {
184            name: "UPPER".to_string(),
185            args: vec![self],
186            alias: None,
187        }
188    }
189
190    fn lower(self) -> Expr {
191        Expr::FunctionCall {
192            name: "LOWER".to_string(),
193            args: vec![self],
194            alias: None,
195        }
196    }
197
198    fn trim(self) -> Expr {
199        Expr::FunctionCall {
200            name: "TRIM".to_string(),
201            args: vec![self],
202            alias: None,
203        }
204    }
205
206    fn length(self) -> Expr {
207        Expr::FunctionCall {
208            name: "LENGTH".to_string(),
209            args: vec![self],
210            alias: None,
211        }
212    }
213
214    fn abs(self) -> Expr {
215        Expr::FunctionCall {
216            name: "ABS".to_string(),
217            args: vec![self],
218            alias: None,
219        }
220    }
221}
222
223// Implement ExprExt for &str to enable: "col_name".or_default("X")
224impl ExprExt for &str {
225    fn with_alias(self, alias: &str) -> Expr {
226        Expr::Aliased {
227            name: self.to_string(),
228            alias: alias.to_string(),
229        }
230    }
231
232    fn or_default(self, default: impl Into<Expr>) -> Expr {
233        Expr::FunctionCall {
234            name: "COALESCE".to_string(),
235            args: vec![Expr::Named(self.to_string()), default.into()],
236            alias: None,
237        }
238    }
239
240    fn json(self, key: &str) -> JsonBuilder {
241        JsonBuilder {
242            column: self.to_string(),
243            path_segments: vec![(key.to_string(), true)],
244            alias: None,
245        }
246    }
247
248    fn path(self, dotted_path: &str) -> JsonBuilder {
249        let segments: Vec<&str> = dotted_path.split('.').collect();
250        let len = segments.len();
251        let path_segments: Vec<(String, bool)> = segments
252            .into_iter()
253            .enumerate()
254            .map(|(i, segment)| (segment.to_string(), i == len - 1))
255            .collect();
256
257        JsonBuilder {
258            column: self.to_string(),
259            path_segments,
260            alias: None,
261        }
262    }
263
264    fn cast(self, target_type: &str) -> Expr {
265        Expr::Cast {
266            expr: Box::new(Expr::Named(self.to_string())),
267            target_type: target_type.to_string(),
268            alias: None,
269        }
270    }
271
272    fn upper(self) -> Expr {
273        Expr::FunctionCall {
274            name: "UPPER".to_string(),
275            args: vec![Expr::Named(self.to_string())],
276            alias: None,
277        }
278    }
279
280    fn lower(self) -> Expr {
281        Expr::FunctionCall {
282            name: "LOWER".to_string(),
283            args: vec![Expr::Named(self.to_string())],
284            alias: None,
285        }
286    }
287
288    fn trim(self) -> Expr {
289        Expr::FunctionCall {
290            name: "TRIM".to_string(),
291            args: vec![Expr::Named(self.to_string())],
292            alias: None,
293        }
294    }
295
296    fn length(self) -> Expr {
297        Expr::FunctionCall {
298            name: "LENGTH".to_string(),
299            args: vec![Expr::Named(self.to_string())],
300            alias: None,
301        }
302    }
303
304    fn abs(self) -> Expr {
305        Expr::FunctionCall {
306            name: "ABS".to_string(),
307            args: vec![Expr::Named(self.to_string())],
308            alias: None,
309        }
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::ast::builders::col;
317
318    #[test]
319    fn test_or_default() {
320        let expr = col("name").or_default("Unknown");
321        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "COALESCE"));
322    }
323
324    #[test]
325    fn test_json_fluent() {
326        let expr: Expr = col("info").json("phone").into();
327        assert!(matches!(expr, Expr::JsonAccess { column, .. } if column == "info"));
328    }
329
330    #[test]
331    fn test_path_fluent() {
332        let expr: Expr = col("metadata").path("vessel.0.port").into();
333        if let Expr::JsonAccess { path_segments, .. } = expr {
334            assert_eq!(path_segments.len(), 3);
335            assert_eq!(path_segments[0], ("vessel".to_string(), false)); // JSON
336            assert_eq!(path_segments[1], ("0".to_string(), false)); // JSON
337            assert_eq!(path_segments[2], ("port".to_string(), true)); // Text
338        } else {
339            panic!("Expected JsonAccess");
340        }
341    }
342
343    #[test]
344    fn test_str_or_default() {
345        let expr = "name".or_default("N/A");
346        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "COALESCE"));
347    }
348
349    #[test]
350    fn test_cast_fluent() {
351        let expr = col("value").cast("int4");
352        assert!(matches!(expr, Expr::Cast { target_type, .. } if target_type == "int4"));
353    }
354
355    #[test]
356    fn test_upper_fluent() {
357        let expr = col("name").upper();
358        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "UPPER"));
359    }
360
361    #[test]
362    fn test_lower_fluent() {
363        let expr = "email".lower();
364        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "LOWER"));
365    }
366
367    #[test]
368    fn test_trim_fluent() {
369        let expr = col("text").trim();
370        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "TRIM"));
371    }
372}