qail_core/ast/builders/
ext.rs

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