1use super::json::JsonBuilder;
6use crate::ast::Expr;
7
8pub trait ExprExt {
10 fn with_alias(self, alias: &str) -> Expr;
17
18 fn or_default(self, default: impl Into<Expr>) -> Expr;
25
26 fn json(self, key: &str) -> JsonBuilder;
33
34 fn path(self, dotted_path: &str) -> JsonBuilder;
41
42 fn cast(self, target_type: &str) -> Expr;
49
50 fn upper(self) -> Expr;
52
53 fn lower(self) -> Expr;
55
56 fn trim(self) -> Expr;
58
59 fn length(self) -> Expr;
61
62 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, }
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)], 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)) .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
223impl 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)); assert_eq!(path_segments[1], ("0".to_string(), false)); assert_eq!(path_segments[2], ("port".to_string(), true)); } 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}