Skip to main content

vibesql_executor/cache/
parameterized.rs

1//! Parameterized plan support for literal extraction and binding
2//!
3//! Allows queries with different literal values to share the same execution plan
4//! by extracting literals into parameters and binding them at execution time.
5
6use vibesql_ast::{Expression, Statement};
7use vibesql_types::SqlValue;
8
9/// Type-safe literal value representation
10#[derive(Clone, Debug, PartialEq)]
11pub enum LiteralValue {
12    Integer(i64),
13    Smallint(i16),
14    Bigint(i64),
15    Unsigned(u64),
16    Numeric(f64),
17    Float(f32),
18    Real(f64), // SQLite REAL is 8-byte IEEE float
19    Double(f64),
20    Character(String),
21    Varchar(String),
22    Boolean(bool),
23    Date(String),
24    Time(String),
25    Timestamp(String),
26    Blob(Vec<u8>),
27    Null,
28}
29
30impl LiteralValue {
31    /// Convert from SqlValue
32    pub fn from_sql_value(value: &SqlValue) -> Self {
33        match value {
34            SqlValue::Integer(n) => LiteralValue::Integer(*n),
35            SqlValue::Smallint(n) => LiteralValue::Smallint(*n),
36            SqlValue::Bigint(n) => LiteralValue::Bigint(*n),
37            SqlValue::Unsigned(n) => LiteralValue::Unsigned(*n),
38            SqlValue::Numeric(n) => LiteralValue::Numeric(*n),
39            SqlValue::Float(n) => LiteralValue::Float(*n),
40            SqlValue::Real(n) => LiteralValue::Real(*n),
41            SqlValue::Double(n) => LiteralValue::Double(*n),
42            SqlValue::Character(s) => LiteralValue::Character(s.to_string()),
43            SqlValue::Varchar(s) => LiteralValue::Varchar(s.to_string()),
44            SqlValue::Boolean(b) => LiteralValue::Boolean(*b),
45            SqlValue::Date(s) => LiteralValue::Date(s.to_string()),
46            SqlValue::Time(s) => LiteralValue::Time(s.to_string()),
47            SqlValue::Timestamp(s) => LiteralValue::Timestamp(s.to_string()),
48            SqlValue::Interval(s) => LiteralValue::Varchar(s.to_string()), /* Treat interval as */
49            // string for now
50            SqlValue::Vector(v) => {
51                // Treat vector as string representation for now
52                let formatted: Vec<String> = v.iter().map(|f| f.to_string()).collect();
53                LiteralValue::Varchar(format!("[{}]", formatted.join(", ")))
54            }
55            SqlValue::Blob(b) => LiteralValue::Blob(b.clone()),
56            SqlValue::Null => LiteralValue::Null,
57        }
58    }
59
60    /// Convert to SQL string representation
61    pub fn to_sql(&self) -> String {
62        match self {
63            LiteralValue::Integer(n) => n.to_string(),
64            LiteralValue::Smallint(n) => n.to_string(),
65            LiteralValue::Bigint(n) => n.to_string(),
66            LiteralValue::Unsigned(n) => n.to_string(),
67            LiteralValue::Numeric(n) => n.to_string(),
68            LiteralValue::Float(n) => n.to_string(),
69            LiteralValue::Real(n) => n.to_string(),
70            LiteralValue::Double(n) => n.to_string(),
71            LiteralValue::Character(s) | LiteralValue::Varchar(s) => {
72                format!("'{}'", s.replace("'", "''"))
73            }
74            LiteralValue::Boolean(b) => if *b { "true" } else { "false" }.to_string(),
75            LiteralValue::Date(s) => format!("DATE '{}'", s),
76            LiteralValue::Time(s) => format!("TIME '{}'", s),
77            LiteralValue::Timestamp(s) => format!("TIMESTAMP '{}'", s),
78            LiteralValue::Blob(b) => {
79                let hex: String = b.iter().map(|byte| format!("{:02X}", byte)).collect();
80                format!("x'{}'", hex)
81            }
82            LiteralValue::Null => "NULL".to_string(),
83        }
84    }
85}
86
87/// Position of a parameter in a query
88#[derive(Clone, Debug)]
89pub struct ParameterPosition {
90    pub position: usize,
91    pub context: String,
92}
93
94/// Query plan with parameter placeholders
95#[derive(Clone, Debug)]
96pub struct ParameterizedPlan {
97    pub normalized_query: String,
98    pub param_positions: Vec<ParameterPosition>,
99    pub literal_values: Vec<LiteralValue>,
100}
101
102impl ParameterizedPlan {
103    /// Create a new parameterized plan
104    pub fn new(
105        normalized_query: String,
106        param_positions: Vec<ParameterPosition>,
107        literal_values: Vec<LiteralValue>,
108    ) -> Self {
109        Self { normalized_query, param_positions, literal_values }
110    }
111
112    /// Bind literal values to create executable query
113    pub fn bind(&self, values: &[LiteralValue]) -> Result<String, String> {
114        if values.len() != self.param_positions.len() {
115            return Err(format!(
116                "Expected {} parameters, got {}",
117                self.param_positions.len(),
118                values.len()
119            ));
120        }
121
122        let mut result = self.normalized_query.clone();
123        let mut offset = 0;
124
125        for (i, value) in values.iter().enumerate() {
126            if let Some(_pos) = self.param_positions.get(i) {
127                let sql_value = value.to_sql();
128                let placeholder = "?";
129
130                if let Some(idx) = result[offset..].find(placeholder) {
131                    let insert_pos = offset + idx;
132                    result.replace_range(insert_pos..insert_pos + 1, &sql_value);
133                    offset = insert_pos + sql_value.len();
134                }
135            }
136        }
137
138        Ok(result)
139    }
140
141    /// Get a comparison key for the plan
142    pub fn comparison_key(&self) -> String {
143        self.normalized_query.clone()
144    }
145}
146
147/// Utility for extracting literals from AST
148pub struct LiteralExtractor;
149
150impl LiteralExtractor {
151    /// Extract literals from a statement AST in order of appearance
152    pub fn extract(stmt: &Statement) -> Vec<LiteralValue> {
153        let mut literals = Vec::new();
154        Self::extract_from_statement(stmt, &mut literals);
155        literals
156    }
157
158    fn extract_from_statement(stmt: &Statement, literals: &mut Vec<LiteralValue>) {
159        match stmt {
160            Statement::Select(select) => Self::extract_from_select(select, literals),
161            Statement::Insert(insert) => {
162                // Extract from VALUES or SELECT source
163                match &insert.source {
164                    vibesql_ast::InsertSource::Values(rows) => {
165                        for row in rows {
166                            for expr in row {
167                                Self::extract_from_expression(expr, literals);
168                            }
169                        }
170                    }
171                    vibesql_ast::InsertSource::Select(select) => {
172                        Self::extract_from_select(select, literals);
173                    }
174                    vibesql_ast::InsertSource::DefaultValues => {
175                        // No literals to extract for DEFAULT VALUES
176                    }
177                }
178            }
179            Statement::Update(update) => {
180                // Extract from assignments
181                for assignment in &update.assignments {
182                    Self::extract_from_expression(&assignment.value, literals);
183                }
184                // Extract from WHERE clause
185                if let Some(ref where_clause) = update.where_clause {
186                    match where_clause {
187                        vibesql_ast::WhereClause::Condition(expr) => {
188                            Self::extract_from_expression(expr, literals);
189                        }
190                        vibesql_ast::WhereClause::CurrentOf(_) => {
191                            // No literals in positioned update/delete
192                        }
193                    }
194                }
195            }
196            Statement::Delete(delete) => {
197                // Extract from WHERE clause
198                if let Some(ref where_clause) = delete.where_clause {
199                    match where_clause {
200                        vibesql_ast::WhereClause::Condition(expr) => {
201                            Self::extract_from_expression(expr, literals);
202                        }
203                        vibesql_ast::WhereClause::CurrentOf(_) => {
204                            // No literals in positioned update/delete
205                        }
206                    }
207                }
208            }
209            // Other statement types don't have literals we need to extract
210            _ => {}
211        }
212    }
213
214    fn extract_from_select(select: &vibesql_ast::SelectStmt, literals: &mut Vec<LiteralValue>) {
215        // Extract from SELECT items
216        for item in &select.select_list {
217            if let vibesql_ast::SelectItem::Expression { expr, .. } = item {
218                Self::extract_from_expression(expr, literals);
219            }
220        }
221
222        // Extract from FROM clause
223        if let Some(ref from) = select.from {
224            Self::extract_from_from_clause(from, literals);
225        }
226
227        // Extract from WHERE
228        if let Some(ref where_clause) = select.where_clause {
229            Self::extract_from_expression(where_clause, literals);
230        }
231
232        // Extract from GROUP BY
233        if let Some(ref group_by) = select.group_by {
234            Self::extract_from_group_by(group_by, literals);
235        }
236
237        // Extract from HAVING
238        if let Some(ref having) = select.having {
239            Self::extract_from_expression(having, literals);
240        }
241
242        // Extract from ORDER BY
243        if let Some(ref order_by) = select.order_by {
244            for item in order_by {
245                Self::extract_from_expression(&item.expr, literals);
246            }
247        }
248    }
249
250    fn extract_from_from_clause(from: &vibesql_ast::FromClause, literals: &mut Vec<LiteralValue>) {
251        match from {
252            vibesql_ast::FromClause::Join { left, right, condition, .. } => {
253                Self::extract_from_from_clause(left, literals);
254                Self::extract_from_from_clause(right, literals);
255                if let Some(expr) = condition {
256                    Self::extract_from_expression(expr, literals);
257                }
258            }
259            vibesql_ast::FromClause::Subquery { query, .. } => {
260                Self::extract_from_select(query, literals);
261            }
262            vibesql_ast::FromClause::Table { .. } => {
263                // No literals in table references
264            }
265            vibesql_ast::FromClause::Values { rows, .. } => {
266                // Extract literals from VALUES expressions
267                for row in rows {
268                    for expr in row {
269                        Self::extract_from_expression(expr, literals);
270                    }
271                }
272            }
273        }
274    }
275
276    fn extract_from_group_by(
277        group_by: &vibesql_ast::GroupByClause,
278        literals: &mut Vec<LiteralValue>,
279    ) {
280        match group_by {
281            vibesql_ast::GroupByClause::Simple(exprs) => {
282                for expr in exprs {
283                    Self::extract_from_expression(expr, literals);
284                }
285            }
286            vibesql_ast::GroupByClause::Rollup(elements)
287            | vibesql_ast::GroupByClause::Cube(elements) => {
288                Self::extract_from_grouping_elements(elements, literals);
289            }
290            vibesql_ast::GroupByClause::GroupingSets(sets) => {
291                Self::extract_from_grouping_sets(sets, literals);
292            }
293            vibesql_ast::GroupByClause::Mixed(items) => {
294                for item in items {
295                    match item {
296                        vibesql_ast::MixedGroupingItem::Simple(expr) => {
297                            Self::extract_from_expression(expr, literals);
298                        }
299                        vibesql_ast::MixedGroupingItem::Rollup(elements)
300                        | vibesql_ast::MixedGroupingItem::Cube(elements) => {
301                            Self::extract_from_grouping_elements(elements, literals);
302                        }
303                        vibesql_ast::MixedGroupingItem::GroupingSets(sets) => {
304                            Self::extract_from_grouping_sets(sets, literals);
305                        }
306                    }
307                }
308            }
309        }
310    }
311
312    fn extract_from_grouping_elements(
313        elements: &[vibesql_ast::GroupingElement],
314        literals: &mut Vec<LiteralValue>,
315    ) {
316        for element in elements {
317            match element {
318                vibesql_ast::GroupingElement::Single(expr) => {
319                    Self::extract_from_expression(expr, literals);
320                }
321                vibesql_ast::GroupingElement::Composite(exprs) => {
322                    for expr in exprs {
323                        Self::extract_from_expression(expr, literals);
324                    }
325                }
326            }
327        }
328    }
329
330    fn extract_from_grouping_sets(
331        sets: &[vibesql_ast::GroupingSet],
332        literals: &mut Vec<LiteralValue>,
333    ) {
334        for set in sets {
335            for expr in &set.columns {
336                Self::extract_from_expression(expr, literals);
337            }
338        }
339    }
340
341    fn extract_from_expression(expr: &Expression, literals: &mut Vec<LiteralValue>) {
342        match expr {
343            Expression::Literal(value) => {
344                literals.push(LiteralValue::from_sql_value(value));
345            }
346
347            Expression::BinaryOp { left, right, .. } => {
348                Self::extract_from_expression(left, literals);
349                Self::extract_from_expression(right, literals);
350            }
351
352            Expression::Conjunction(children)
353            | Expression::Disjunction(children)
354            | Expression::RowValueConstructor(children) => {
355                for child in children {
356                    Self::extract_from_expression(child, literals);
357                }
358            }
359
360            Expression::Collate { expr, .. } => {
361                Self::extract_from_expression(expr, literals);
362            }
363
364            Expression::Raise { error_message, .. } => {
365                if let Some(msg) = error_message {
366                    Self::extract_from_expression(msg, literals);
367                }
368            }
369
370            Expression::UnaryOp { expr, .. } => {
371                Self::extract_from_expression(expr, literals);
372            }
373
374            Expression::Function { args, .. } => {
375                for arg in args {
376                    Self::extract_from_expression(arg, literals);
377                }
378            }
379
380            Expression::AggregateFunction { args, .. } => {
381                for arg in args {
382                    Self::extract_from_expression(arg, literals);
383                }
384            }
385
386            Expression::IsNull { expr, .. } => {
387                Self::extract_from_expression(expr, literals);
388            }
389
390            Expression::IsDistinctFrom { left, right, .. } => {
391                Self::extract_from_expression(left, literals);
392                Self::extract_from_expression(right, literals);
393            }
394
395            Expression::IsTruthValue { expr, .. } => {
396                Self::extract_from_expression(expr, literals);
397            }
398
399            Expression::Case { operand, when_clauses, else_result } => {
400                if let Some(ref op) = operand {
401                    Self::extract_from_expression(op, literals);
402                }
403                for when in when_clauses {
404                    for cond in &when.conditions {
405                        Self::extract_from_expression(cond, literals);
406                    }
407                    Self::extract_from_expression(&when.result, literals);
408                }
409                if let Some(ref else_expr) = else_result {
410                    Self::extract_from_expression(else_expr, literals);
411                }
412            }
413
414            Expression::ScalarSubquery(subquery) => {
415                Self::extract_from_select(subquery, literals);
416            }
417
418            Expression::In { expr, subquery, .. } => {
419                Self::extract_from_expression(expr, literals);
420                Self::extract_from_select(subquery, literals);
421            }
422
423            Expression::InList { expr, values, .. } => {
424                Self::extract_from_expression(expr, literals);
425                for val in values {
426                    Self::extract_from_expression(val, literals);
427                }
428            }
429
430            Expression::Between { expr, low, high, .. } => {
431                Self::extract_from_expression(expr, literals);
432                Self::extract_from_expression(low, literals);
433                Self::extract_from_expression(high, literals);
434            }
435
436            Expression::Cast { expr, .. } => {
437                Self::extract_from_expression(expr, literals);
438            }
439
440            Expression::Position { substring, string, .. } => {
441                Self::extract_from_expression(substring, literals);
442                Self::extract_from_expression(string, literals);
443            }
444
445            Expression::Trim { removal_char, string, .. } => {
446                if let Some(ref ch) = removal_char {
447                    Self::extract_from_expression(ch, literals);
448                }
449                Self::extract_from_expression(string, literals);
450            }
451
452            Expression::Extract { expr, .. } => {
453                Self::extract_from_expression(expr, literals);
454            }
455
456            Expression::Like { expr, pattern, .. } => {
457                Self::extract_from_expression(expr, literals);
458                Self::extract_from_expression(pattern, literals);
459            }
460
461            Expression::Glob { expr, pattern, .. } => {
462                Self::extract_from_expression(expr, literals);
463                Self::extract_from_expression(pattern, literals);
464            }
465
466            Expression::Exists { subquery, .. } => {
467                Self::extract_from_select(subquery, literals);
468            }
469
470            Expression::QuantifiedComparison { expr, subquery, .. } => {
471                Self::extract_from_expression(expr, literals);
472                Self::extract_from_select(subquery, literals);
473            }
474
475            Expression::WindowFunction { function, over } => {
476                // Extract from function arguments
477                match function {
478                    vibesql_ast::WindowFunctionSpec::Aggregate { args, .. }
479                    | vibesql_ast::WindowFunctionSpec::Ranking { args, .. }
480                    | vibesql_ast::WindowFunctionSpec::Value { args, .. } => {
481                        for arg in args {
482                            Self::extract_from_expression(arg, literals);
483                        }
484                    }
485                }
486
487                // Extract from PARTITION BY
488                if let Some(ref partition_by) = over.partition_by {
489                    for expr in partition_by {
490                        Self::extract_from_expression(expr, literals);
491                    }
492                }
493
494                // Extract from ORDER BY
495                if let Some(ref order_by) = over.order_by {
496                    for item in order_by {
497                        Self::extract_from_expression(&item.expr, literals);
498                    }
499                }
500
501                // Extract from frame bounds
502                if let Some(ref frame) = over.frame {
503                    if let vibesql_ast::FrameBound::Preceding(expr)
504                    | vibesql_ast::FrameBound::Following(expr) = &frame.start
505                    {
506                        Self::extract_from_expression(expr, literals);
507                    }
508                    if let Some(
509                        vibesql_ast::FrameBound::Preceding(expr)
510                        | vibesql_ast::FrameBound::Following(expr),
511                    ) = &frame.end
512                    {
513                        Self::extract_from_expression(expr, literals);
514                    }
515                }
516            }
517
518            // These expressions don't contain literals
519            Expression::ColumnRef(_)
520            | Expression::Placeholder(_)
521            | Expression::NumberedPlaceholder(_)
522            | Expression::NamedPlaceholder(_)
523            | Expression::PseudoVariable { .. }
524            | Expression::Wildcard
525            | Expression::CurrentDate
526            | Expression::CurrentTime { .. }
527            | Expression::CurrentTimestamp { .. }
528            | Expression::Interval { .. }
529            | Expression::Default
530            | Expression::DuplicateKeyValue { .. }
531            | Expression::NextValue { .. }
532            | Expression::MatchAgainst { .. }
533            | Expression::SessionVariable { .. } => {}
534        }
535    }
536}
537
538#[cfg(test)]
539mod tests {
540    use super::*;
541
542    #[test]
543    fn test_literal_value_to_string() {
544        assert_eq!(LiteralValue::Integer(42).to_sql(), "42");
545        assert_eq!(LiteralValue::Varchar("hello".to_string()).to_sql(), "'hello'");
546        assert_eq!(LiteralValue::Boolean(true).to_sql(), "true");
547        assert_eq!(LiteralValue::Null.to_sql(), "NULL");
548    }
549
550    #[test]
551    fn test_literal_value_string_escape() {
552        assert_eq!(LiteralValue::Varchar("it's".to_string()).to_sql(), "'it''s'");
553    }
554
555    #[test]
556    fn test_literal_extraction_simple() {
557        use vibesql_ast::{
558            BinaryOperator, Expression, FromClause, SelectItem, SelectStmt, Statement,
559        };
560
561        // SELECT col0 FROM tab WHERE col1 > 25 AND col2 = 'John'
562        let stmt = Statement::Select(Box::new(SelectStmt {
563            with_clause: None,
564            distinct: false,
565            select_list: vec![SelectItem::Expression {
566                expr: Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple("col0", false)),
567                alias: None,
568                source_text: None,
569            }],
570            into_table: None,
571            into_variables: None,
572            from: Some(FromClause::Table {
573                index_hint: None,
574                name: "tab".to_string(),
575                alias: None,
576                column_aliases: None,
577                quoted: false,
578            }),
579            where_clause: Some(Expression::BinaryOp {
580                op: BinaryOperator::And,
581                left: Box::new(Expression::BinaryOp {
582                    op: BinaryOperator::GreaterThan,
583                    left: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
584                        "col1", false,
585                    ))),
586                    right: Box::new(Expression::Literal(SqlValue::Integer(25))),
587                }),
588                right: Box::new(Expression::BinaryOp {
589                    op: BinaryOperator::Equal,
590                    left: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
591                        "col2", false,
592                    ))),
593                    right: Box::new(Expression::Literal(SqlValue::Varchar(arcstr::ArcStr::from(
594                        "John",
595                    )))),
596                }),
597            }),
598            group_by: None,
599            having: None,
600            window_definitions: None,
601            order_by: None,
602            limit: None,
603            offset: None,
604            set_operation: None,
605            values: None,
606        }));
607
608        let literals = LiteralExtractor::extract(&stmt);
609
610        assert_eq!(literals.len(), 2);
611        assert_eq!(literals[0], LiteralValue::Integer(25));
612        assert_eq!(literals[1], LiteralValue::Varchar("John".to_string()));
613    }
614
615    #[test]
616    fn test_literal_extraction_in_list() {
617        use vibesql_ast::{Expression, FromClause, SelectItem, SelectStmt, Statement};
618
619        // SELECT * FROM tab WHERE id IN (1, 2, 3)
620        let stmt = Statement::Select(Box::new(SelectStmt {
621            with_clause: None,
622            distinct: false,
623            select_list: vec![SelectItem::Wildcard { alias: None }],
624            into_table: None,
625            into_variables: None,
626            from: Some(FromClause::Table {
627                index_hint: None,
628                name: "tab".to_string(),
629                alias: None,
630                column_aliases: None,
631                quoted: false,
632            }),
633            where_clause: Some(Expression::InList {
634                expr: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
635                    "id", false,
636                ))),
637                values: vec![
638                    Expression::Literal(SqlValue::Integer(1)),
639                    Expression::Literal(SqlValue::Integer(2)),
640                    Expression::Literal(SqlValue::Integer(3)),
641                ],
642                negated: false,
643            }),
644            group_by: None,
645            having: None,
646            window_definitions: None,
647            order_by: None,
648            limit: None,
649            offset: None,
650            set_operation: None,
651            values: None,
652        }));
653
654        let literals = LiteralExtractor::extract(&stmt);
655
656        assert_eq!(literals.len(), 3);
657        assert_eq!(literals[0], LiteralValue::Integer(1));
658        assert_eq!(literals[1], LiteralValue::Integer(2));
659        assert_eq!(literals[2], LiteralValue::Integer(3));
660    }
661
662    #[test]
663    fn test_parameterized_plan_bind() {
664        let plan = ParameterizedPlan::new(
665            "SELECT * FROM users WHERE age > ?".to_string(),
666            vec![ParameterPosition { position: 40, context: "age".to_string() }],
667            vec![LiteralValue::Integer(25)],
668        );
669
670        let result = plan.bind(&[LiteralValue::Integer(30)]).unwrap();
671        assert_eq!(result, "SELECT * FROM users WHERE age > 30");
672    }
673
674    #[test]
675    fn test_parameterized_plan_bind_string() {
676        let plan = ParameterizedPlan::new(
677            "SELECT * FROM users WHERE name = ?".to_string(),
678            vec![ParameterPosition { position: 40, context: "name".to_string() }],
679            vec![LiteralValue::Varchar("John".to_string())],
680        );
681
682        let result = plan.bind(&[LiteralValue::Varchar("Jane".to_string())]).unwrap();
683        assert_eq!(result, "SELECT * FROM users WHERE name = 'Jane'");
684    }
685
686    #[test]
687    fn test_parameterized_plan_bind_error() {
688        let plan = ParameterizedPlan::new(
689            "SELECT * FROM users WHERE age > ?".to_string(),
690            vec![ParameterPosition { position: 40, context: "age".to_string() }],
691            vec![LiteralValue::Integer(25)],
692        );
693
694        let result = plan.bind(&[LiteralValue::Integer(30), LiteralValue::Integer(40)]);
695        assert!(result.is_err());
696    }
697
698    #[test]
699    fn test_comparison_key() {
700        let plan = ParameterizedPlan::new("SELECT * FROM users".to_string(), vec![], vec![]);
701
702        assert_eq!(plan.comparison_key(), "SELECT * FROM users");
703    }
704}