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