Skip to main content

shape_runtime/
query_builder.rs

1//! Query plan and SQL generation for data source pushdown.
2//!
3//! Provides a standalone query plan data structure and SQL generator.
4//! Used by DbTable and other Queryable implementations to build queries
5//! that can be pushed down to the data source.
6
7use shape_value::{FilterLiteral, FilterNode, FilterOp};
8
9/// A query plan that accumulates operations for SQL generation.
10#[derive(Debug, Clone, Default)]
11pub struct QueryPlan {
12    /// Table name to query
13    pub table: String,
14    /// Pushdown-able filter predicates
15    pub filters: Vec<FilterNode>,
16    /// Column projections (None = SELECT *)
17    pub projections: Option<Vec<String>>,
18    /// ORDER BY clauses: (column, descending)
19    pub order_by: Vec<(String, bool)>,
20    /// GROUP BY columns
21    pub group_by: Vec<String>,
22    /// LIMIT
23    pub limit: Option<usize>,
24    /// OFFSET
25    pub offset: Option<usize>,
26}
27
28impl QueryPlan {
29    /// Create a new query plan for a table
30    pub fn new(table: &str) -> Self {
31        Self {
32            table: table.to_string(),
33            ..Default::default()
34        }
35    }
36
37    /// Generate a SQL query string from this plan
38    pub fn to_sql(&self) -> String {
39        let mut sql = String::new();
40
41        // SELECT
42        sql.push_str("SELECT ");
43        match &self.projections {
44            Some(cols) if !cols.is_empty() => {
45                sql.push_str(
46                    &cols
47                        .iter()
48                        .map(|c| quote_ident(c))
49                        .collect::<Vec<_>>()
50                        .join(", "),
51                );
52            }
53            _ => sql.push('*'),
54        }
55
56        // FROM
57        sql.push_str(" FROM ");
58        sql.push_str(&quote_ident(&self.table));
59
60        // WHERE
61        if !self.filters.is_empty() {
62            sql.push_str(" WHERE ");
63            let filter_clauses: Vec<String> =
64                self.filters.iter().map(|f| filter_to_sql(f)).collect();
65            sql.push_str(&filter_clauses.join(" AND "));
66        }
67
68        // GROUP BY
69        if !self.group_by.is_empty() {
70            sql.push_str(" GROUP BY ");
71            sql.push_str(
72                &self
73                    .group_by
74                    .iter()
75                    .map(|c| quote_ident(c))
76                    .collect::<Vec<_>>()
77                    .join(", "),
78            );
79        }
80
81        // ORDER BY
82        if !self.order_by.is_empty() {
83            sql.push_str(" ORDER BY ");
84            let order_clauses: Vec<String> = self
85                .order_by
86                .iter()
87                .map(|(col, desc)| {
88                    if *desc {
89                        format!("{} DESC", quote_ident(col))
90                    } else {
91                        quote_ident(col)
92                    }
93                })
94                .collect();
95            sql.push_str(&order_clauses.join(", "));
96        }
97
98        // LIMIT
99        if let Some(limit) = self.limit {
100            sql.push_str(&format!(" LIMIT {}", limit));
101        }
102
103        // OFFSET
104        if let Some(offset) = self.offset {
105            sql.push_str(&format!(" OFFSET {}", offset));
106        }
107
108        sql
109    }
110}
111
112/// Generate SQL for a FilterNode
113pub fn filter_to_sql(node: &FilterNode) -> String {
114    match node {
115        FilterNode::Compare { column, op, value } => {
116            let op_str = match op {
117                FilterOp::Eq => "=",
118                FilterOp::Neq => "!=",
119                FilterOp::Gt => ">",
120                FilterOp::Gte => ">=",
121                FilterOp::Lt => "<",
122                FilterOp::Lte => "<=",
123            };
124            format!(
125                "{} {} {}",
126                quote_ident(column),
127                op_str,
128                literal_to_sql(value)
129            )
130        }
131        FilterNode::And(left, right) => {
132            format!("({} AND {})", filter_to_sql(left), filter_to_sql(right))
133        }
134        FilterNode::Or(left, right) => {
135            format!("({} OR {})", filter_to_sql(left), filter_to_sql(right))
136        }
137        FilterNode::Not(inner) => {
138            format!("NOT ({})", filter_to_sql(inner))
139        }
140    }
141}
142
143/// Generate SQL for a FilterLiteral
144pub fn literal_to_sql(lit: &FilterLiteral) -> String {
145    match lit {
146        FilterLiteral::Int(i) => i.to_string(),
147        FilterLiteral::Float(f) => format!("{}", f),
148        FilterLiteral::String(s) => format!("'{}'", s.replace('\'', "''")),
149        FilterLiteral::Bool(b) => {
150            if *b {
151                "TRUE".to_string()
152            } else {
153                "FALSE".to_string()
154            }
155        }
156        FilterLiteral::Null => "NULL".to_string(),
157    }
158}
159
160/// Quote a SQL identifier (simple quoting for safety)
161fn quote_ident(name: &str) -> String {
162    // Don't quote simple identifiers
163    if name.chars().all(|c| c.is_alphanumeric() || c == '_') && !name.is_empty() {
164        name.to_string()
165    } else {
166        format!("\"{}\"", name.replace('"', "\"\""))
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_simple_select() {
176        let plan = QueryPlan::new("users");
177        assert_eq!(plan.to_sql(), "SELECT * FROM users");
178    }
179
180    #[test]
181    fn test_select_with_projections() {
182        let mut plan = QueryPlan::new("users");
183        plan.projections = Some(vec!["name".to_string(), "age".to_string()]);
184        assert_eq!(plan.to_sql(), "SELECT name, age FROM users");
185    }
186
187    #[test]
188    fn test_where_clause() {
189        let mut plan = QueryPlan::new("users");
190        plan.filters.push(FilterNode::Compare {
191            column: "age".to_string(),
192            op: FilterOp::Gte,
193            value: FilterLiteral::Int(18),
194        });
195        assert_eq!(plan.to_sql(), "SELECT * FROM users WHERE age >= 18");
196    }
197
198    #[test]
199    fn test_compound_filters() {
200        let mut plan = QueryPlan::new("users");
201        plan.filters.push(FilterNode::And(
202            Box::new(FilterNode::Compare {
203                column: "age".to_string(),
204                op: FilterOp::Gte,
205                value: FilterLiteral::Int(18),
206            }),
207            Box::new(FilterNode::Compare {
208                column: "active".to_string(),
209                op: FilterOp::Eq,
210                value: FilterLiteral::Bool(true),
211            }),
212        ));
213        assert_eq!(
214            plan.to_sql(),
215            "SELECT * FROM users WHERE (age >= 18 AND active = TRUE)"
216        );
217    }
218
219    #[test]
220    fn test_or_filter() {
221        let filter = FilterNode::Or(
222            Box::new(FilterNode::Compare {
223                column: "role".to_string(),
224                op: FilterOp::Eq,
225                value: FilterLiteral::String("admin".to_string()),
226            }),
227            Box::new(FilterNode::Compare {
228                column: "role".to_string(),
229                op: FilterOp::Eq,
230                value: FilterLiteral::String("moderator".to_string()),
231            }),
232        );
233        assert_eq!(
234            filter_to_sql(&filter),
235            "(role = 'admin' OR role = 'moderator')"
236        );
237    }
238
239    #[test]
240    fn test_not_filter() {
241        let filter = FilterNode::Not(Box::new(FilterNode::Compare {
242            column: "deleted".to_string(),
243            op: FilterOp::Eq,
244            value: FilterLiteral::Bool(true),
245        }));
246        assert_eq!(filter_to_sql(&filter), "NOT (deleted = TRUE)");
247    }
248
249    #[test]
250    fn test_order_by() {
251        let mut plan = QueryPlan::new("users");
252        plan.order_by.push(("age".to_string(), true));
253        assert_eq!(plan.to_sql(), "SELECT * FROM users ORDER BY age DESC");
254    }
255
256    #[test]
257    fn test_order_by_asc() {
258        let mut plan = QueryPlan::new("users");
259        plan.order_by.push(("name".to_string(), false));
260        assert_eq!(plan.to_sql(), "SELECT * FROM users ORDER BY name");
261    }
262
263    #[test]
264    fn test_limit_offset() {
265        let mut plan = QueryPlan::new("users");
266        plan.limit = Some(100);
267        plan.offset = Some(50);
268        assert_eq!(plan.to_sql(), "SELECT * FROM users LIMIT 100 OFFSET 50");
269    }
270
271    #[test]
272    fn test_full_query() {
273        let mut plan = QueryPlan::new("users");
274        plan.projections = Some(vec!["name".to_string(), "age".to_string()]);
275        plan.filters.push(FilterNode::Compare {
276            column: "age".to_string(),
277            op: FilterOp::Gte,
278            value: FilterLiteral::Int(18),
279        });
280        plan.filters.push(FilterNode::Compare {
281            column: "active".to_string(),
282            op: FilterOp::Eq,
283            value: FilterLiteral::Bool(true),
284        });
285        plan.order_by.push(("age".to_string(), true));
286        plan.limit = Some(100);
287        assert_eq!(
288            plan.to_sql(),
289            "SELECT name, age FROM users WHERE age >= 18 AND active = TRUE ORDER BY age DESC LIMIT 100"
290        );
291    }
292
293    #[test]
294    fn test_string_literal_escaping() {
295        let lit = FilterLiteral::String("O'Brien".to_string());
296        assert_eq!(literal_to_sql(&lit), "'O''Brien'");
297    }
298
299    #[test]
300    fn test_null_comparison() {
301        let filter = FilterNode::Compare {
302            column: "email".to_string(),
303            op: FilterOp::Eq,
304            value: FilterLiteral::Null,
305        };
306        assert_eq!(filter_to_sql(&filter), "email = NULL");
307    }
308
309    #[test]
310    fn test_group_by() {
311        let mut plan = QueryPlan::new("orders");
312        plan.group_by.push("status".to_string());
313        assert_eq!(plan.to_sql(), "SELECT * FROM orders GROUP BY status");
314    }
315}