Skip to main content

qail_core/ast/builders/
ext.rs

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