Skip to main content

a3s_orm/compiler/
mod.rs

1mod dialect;
2mod validation;
3
4pub use dialect::{Dialect, MysqlDialect, PostgresDialect, SqliteDialect};
5
6use crate::ast::{
7    ConflictAction, ConflictValue, CteNode, DeleteNode, InsertNode, JoinKind, QueryNode,
8    SelectLockStrength, SelectLockWait, SelectNode, SetOperationKind, TableLockNode, TableNode,
9    UpdateNode,
10};
11use crate::error::{Error, Result};
12use crate::expression::{
13    BinaryOperator, Expression, OrderDirection, UnaryOperator, WindowBoundary, WindowFrameUnits,
14};
15use crate::query::PostgresTableLockMode;
16use crate::value::Value;
17use validation::{
18    validate_identifier, validate_sql_type, verify_insert_rows, verify_select_lock,
19    verify_update_assignments,
20};
21
22#[derive(Clone, Debug, PartialEq)]
23pub struct CompiledQuery {
24    pub sql: String,
25    pub parameters: Vec<Value>,
26}
27
28pub(crate) fn compile(query: QueryNode, dialect: &impl Dialect) -> Result<CompiledQuery> {
29    let mut compiler = Compiler {
30        dialect,
31        sql: String::new(),
32        parameters: Vec::new(),
33    };
34    match query {
35        QueryNode::Select(node) => compiler.select(*node)?,
36        QueryNode::Insert(node) => compiler.insert(node)?,
37        QueryNode::Update(node) => compiler.update(node)?,
38        QueryNode::Delete(node) => compiler.delete(node)?,
39        QueryNode::TableLock(node) => compiler.table_lock(node)?,
40    }
41    Ok(CompiledQuery {
42        sql: compiler.sql,
43        parameters: compiler.parameters,
44    })
45}
46
47struct Compiler<'a, D: Dialect> {
48    dialect: &'a D,
49    sql: String,
50    parameters: Vec<Value>,
51}
52
53impl<D: Dialect> Compiler<'_, D> {
54    fn select(&mut self, mut node: SelectNode) -> Result<()> {
55        if node.selections.is_empty() {
56            return Err(Error::EmptySelection);
57        }
58        verify_select_lock(&node)?;
59        self.ctes(std::mem::take(&mut node.ctes))?;
60        self.sql.push_str("select ");
61        if node.distinct {
62            self.sql.push_str("distinct ");
63        }
64        self.expression_list(&node.selections)?;
65        self.sql.push_str(" from ");
66        self.table(&node.from)?;
67        for join in node.joins {
68            self.sql.push(' ');
69            self.sql.push_str(match join.kind {
70                JoinKind::Inner => "inner join ",
71                JoinKind::Left => "left join ",
72                JoinKind::Right => "right join ",
73                JoinKind::Full => "full join ",
74            });
75            self.table(&join.table)?;
76            self.sql.push_str(" on ");
77            self.expression(&join.on)?;
78        }
79        self.filter(node.filter.as_ref())?;
80        if !node.group_by.is_empty() {
81            self.sql.push_str(" group by ");
82            self.expression_list(&node.group_by)?;
83        }
84        if let Some(having) = node.having.as_ref() {
85            self.sql.push_str(" having ");
86            self.expression(having)?;
87        }
88        for operation in node.set_operations {
89            if !operation.query.ctes.is_empty()
90                || !operation.query.order_by.is_empty()
91                || operation.query.limit.is_some()
92                || operation.query.offset.is_some()
93                || operation.query.lock.is_some()
94            {
95                return Err(Error::UnsupportedSetOperand);
96            }
97            self.sql.push_str(match operation.kind {
98                SetOperationKind::Union => " union ",
99                SetOperationKind::UnionAll => " union all ",
100                SetOperationKind::Intersect => " intersect ",
101                SetOperationKind::Except => " except ",
102            });
103            self.select(*operation.query)?;
104        }
105        if !node.order_by.is_empty() {
106            self.sql.push_str(" order by ");
107            for (index, (expression, direction)) in node.order_by.iter().enumerate() {
108                if index > 0 {
109                    self.sql.push_str(", ");
110                }
111                self.expression(expression)?;
112                self.sql.push_str(match direction {
113                    OrderDirection::Asc => " asc",
114                    OrderDirection::Desc => " desc",
115                });
116            }
117        }
118        if let Some(limit) = node.limit {
119            self.sql.push_str(" limit ");
120            self.parameter(Value::U64(limit));
121        }
122        if let Some(offset) = node.offset {
123            self.sql.push_str(" offset ");
124            self.parameter(Value::U64(offset));
125        }
126        if let Some(lock) = node.lock {
127            if !self.dialect.supports_select_row_locking() {
128                return Err(Error::Compilation(format!(
129                    "{} does not support select row locking",
130                    self.dialect.name()
131                )));
132            }
133            self.sql.push_str(match lock.strength {
134                SelectLockStrength::Update => " for update",
135                SelectLockStrength::NoKeyUpdate => " for no key update",
136                SelectLockStrength::Share => " for share",
137                SelectLockStrength::KeyShare => " for key share",
138            });
139            if !lock.tables.is_empty() {
140                self.sql.push_str(" of ");
141                for (index, table) in lock.tables.iter().enumerate() {
142                    if index > 0 {
143                        self.sql.push_str(", ");
144                    }
145                    self.identifier(table)?;
146                }
147            }
148            match lock.wait {
149                SelectLockWait::Block => {}
150                SelectLockWait::NoWait => self.sql.push_str(" nowait"),
151                SelectLockWait::SkipLocked => self.sql.push_str(" skip locked"),
152            }
153        }
154        Ok(())
155    }
156
157    fn insert(&mut self, node: InsertNode) -> Result<()> {
158        if node.rows.is_empty() || node.rows[0].is_empty() {
159            return Err(Error::EmptyInsert);
160        }
161        verify_insert_rows(&node)?;
162        self.sql.push_str("insert into ");
163        self.table(&node.table)?;
164        self.sql.push_str(" (");
165        for (index, assignment) in node.rows[0].iter().enumerate() {
166            if index > 0 {
167                self.sql.push_str(", ");
168            }
169            self.identifier(assignment.column)?;
170        }
171        self.sql.push_str(") values ");
172        for (row_index, row) in node.rows.into_iter().enumerate() {
173            if row_index > 0 {
174                self.sql.push_str(", ");
175            }
176            self.sql.push('(');
177            for (column_index, assignment) in row.into_iter().enumerate() {
178                if column_index > 0 {
179                    self.sql.push_str(", ");
180                }
181                self.parameter(assignment.value);
182            }
183            self.sql.push(')');
184        }
185        if let Some(conflict) = node.conflict {
186            if !self.dialect.supports_on_conflict() {
187                return Err(Error::Compilation(format!(
188                    "{} does not support on conflict clauses",
189                    self.dialect.name()
190                )));
191            }
192            self.sql.push_str(" on conflict (");
193            for (index, column) in conflict.target.iter().enumerate() {
194                if index > 0 {
195                    self.sql.push_str(", ");
196                }
197                self.identifier(column)?;
198            }
199            self.sql.push_str(") ");
200            match conflict.action {
201                Some(ConflictAction::DoNothing) => self.sql.push_str("do nothing"),
202                Some(ConflictAction::DoUpdate(assignments)) => {
203                    self.sql.push_str("do update set ");
204                    for (index, assignment) in assignments.into_iter().enumerate() {
205                        if index > 0 {
206                            self.sql.push_str(", ");
207                        }
208                        self.identifier(assignment.column)?;
209                        self.sql.push_str(" = ");
210                        match assignment.value {
211                            ConflictValue::Bound(value) => self.parameter(value),
212                            ConflictValue::Excluded { column, .. } => {
213                                self.sql.push_str("excluded.");
214                                self.identifier(column)?;
215                            }
216                        }
217                    }
218                }
219                None => return Err(Error::MissingConflictAction),
220            }
221        }
222        self.returning(&node.returning)
223    }
224
225    fn update(&mut self, mut node: UpdateNode) -> Result<()> {
226        if node.assignments.is_empty() {
227            return Err(Error::EmptyUpdate);
228        }
229        verify_update_assignments(&node.table, &node.assignments)?;
230        if !node.from.is_empty() && !self.dialect.supports_update_from() {
231            return Err(Error::Compilation(format!(
232                "{} does not support update from clauses",
233                self.dialect.name()
234            )));
235        }
236        self.ctes(std::mem::take(&mut node.ctes))?;
237        self.sql.push_str("update ");
238        self.table(&node.table)?;
239        self.sql.push_str(" set ");
240        for (index, assignment) in node.assignments.into_iter().enumerate() {
241            if index > 0 {
242                self.sql.push_str(", ");
243            }
244            self.identifier(assignment.column)?;
245            self.sql.push_str(" = ");
246            self.expression(&assignment.value)?;
247        }
248        if !node.from.is_empty() {
249            self.sql.push_str(" from ");
250            for (index, table) in node.from.iter().enumerate() {
251                if index > 0 {
252                    self.sql.push_str(", ");
253                }
254                self.table(table)?;
255            }
256        }
257        self.filter(node.filter.as_ref())?;
258        self.returning(&node.returning)
259    }
260
261    fn ctes(&mut self, ctes: Vec<CteNode>) -> Result<()> {
262        if ctes.is_empty() {
263            return Ok(());
264        }
265        let mut names = std::collections::HashSet::with_capacity(ctes.len());
266        for cte in &ctes {
267            if !names.insert(cte.name) {
268                return Err(Error::DuplicateCte(cte.name.to_owned()));
269            }
270        }
271        self.sql.push_str("with ");
272        for (index, cte) in ctes.into_iter().enumerate() {
273            if index > 0 {
274                self.sql.push_str(", ");
275            }
276            self.identifier(cte.name)?;
277            self.sql.push_str(" as (");
278            self.select(*cte.query)?;
279            self.sql.push(')');
280        }
281        self.sql.push(' ');
282        Ok(())
283    }
284
285    fn delete(&mut self, node: DeleteNode) -> Result<()> {
286        self.sql.push_str("delete from ");
287        self.table(&node.table)?;
288        self.filter(node.filter.as_ref())?;
289        self.returning(&node.returning)
290    }
291
292    fn table_lock(&mut self, node: TableLockNode) -> Result<()> {
293        if !self.dialect.supports_table_locking() {
294            return Err(Error::Compilation(format!(
295                "{} does not support table locking",
296                self.dialect.name()
297            )));
298        }
299        self.sql.push_str("lock table ");
300        self.table(&node.table)?;
301        self.sql.push_str(" in ");
302        self.sql.push_str(match node.mode {
303            PostgresTableLockMode::AccessShare => "access share",
304            PostgresTableLockMode::RowShare => "row share",
305            PostgresTableLockMode::RowExclusive => "row exclusive",
306            PostgresTableLockMode::ShareUpdateExclusive => "share update exclusive",
307            PostgresTableLockMode::Share => "share",
308            PostgresTableLockMode::ShareRowExclusive => "share row exclusive",
309            PostgresTableLockMode::Exclusive => "exclusive",
310            PostgresTableLockMode::AccessExclusive => "access exclusive",
311        });
312        self.sql.push_str(" mode");
313        if node.no_wait {
314            self.sql.push_str(" nowait");
315        }
316        Ok(())
317    }
318
319    fn returning(&mut self, expressions: &[Expression]) -> Result<()> {
320        if expressions.is_empty() {
321            return Ok(());
322        }
323        if !self.dialect.supports_returning() {
324            return Err(Error::Compilation(format!(
325                "{} does not support returning clauses",
326                self.dialect.name()
327            )));
328        }
329        self.sql.push_str(" returning ");
330        self.expression_list(expressions)
331    }
332
333    fn filter(&mut self, expression: Option<&Expression>) -> Result<()> {
334        if let Some(expression) = expression {
335            self.sql.push_str(" where ");
336            self.expression(expression)?;
337        }
338        Ok(())
339    }
340
341    fn expression_list(&mut self, expressions: &[Expression]) -> Result<()> {
342        for (index, expression) in expressions.iter().enumerate() {
343            if index > 0 {
344                self.sql.push_str(", ");
345            }
346            self.expression(expression)?;
347        }
348        Ok(())
349    }
350
351    fn expression(&mut self, expression: &Expression) -> Result<()> {
352        match expression {
353            Expression::Column { table, name } => {
354                self.identifier(table)?;
355                self.sql.push('.');
356                if *name == "*" {
357                    self.sql.push('*');
358                } else {
359                    self.identifier(name)?;
360                }
361            }
362            Expression::Value(value) => self.parameter(value.clone()),
363            Expression::Subquery(query) => self.subquery(&query.0, true)?,
364            Expression::Function { name, arguments } => {
365                self.identifier(name)?;
366                self.sql.push('(');
367                self.expression_list(arguments)?;
368                self.sql.push(')');
369            }
370            Expression::Coalesce(arguments) => {
371                if arguments.is_empty() {
372                    return Err(Error::Compilation(
373                        "coalesce requires at least one expression".to_owned(),
374                    ));
375                }
376                self.sql.push_str("coalesce(");
377                self.expression_list(arguments)?;
378                self.sql.push(')');
379            }
380            Expression::Least(arguments) => {
381                if arguments.is_empty() {
382                    return Err(Error::Compilation(
383                        "least requires at least one expression".to_owned(),
384                    ));
385                }
386                self.sql.push_str("least(");
387                self.expression_list(arguments)?;
388                self.sql.push(')');
389            }
390            Expression::Cast {
391                expression,
392                sql_type,
393            } => {
394                validate_sql_type(sql_type)?;
395                self.sql.push_str("cast(");
396                self.expression(expression)?;
397                self.sql.push_str(" as ");
398                self.identifier(sql_type)?;
399                self.sql.push(')');
400            }
401            Expression::Alias { expression, alias } => {
402                self.expression(expression)?;
403                self.sql.push_str(" as ");
404                self.identifier(alias)?;
405            }
406            Expression::Wildcard => self.sql.push('*'),
407            Expression::Window {
408                expression,
409                partition_by,
410                order_by,
411                frame,
412            } => {
413                self.expression(expression)?;
414                self.sql.push_str(" over (");
415                let mut needs_space = false;
416                if !partition_by.is_empty() {
417                    self.sql.push_str("partition by ");
418                    self.expression_list(partition_by)?;
419                    needs_space = true;
420                }
421                if !order_by.is_empty() {
422                    if needs_space {
423                        self.sql.push(' ');
424                    }
425                    self.sql.push_str("order by ");
426                    self.order_list(order_by)?;
427                    needs_space = true;
428                }
429                if let Some(frame) = frame {
430                    if matches!(frame.start, WindowBoundary::UnboundedFollowing)
431                        || matches!(frame.end, WindowBoundary::UnboundedPreceding)
432                    {
433                        return Err(Error::InvalidWindowFrame);
434                    }
435                    if needs_space {
436                        self.sql.push(' ');
437                    }
438                    self.sql.push_str(match frame.units {
439                        WindowFrameUnits::Rows => "rows",
440                        WindowFrameUnits::Range => "range",
441                        WindowFrameUnits::Groups => "groups",
442                    });
443                    self.sql.push_str(" between ");
444                    self.window_boundary(frame.start);
445                    self.sql.push_str(" and ");
446                    self.window_boundary(frame.end);
447                }
448                self.sql.push(')');
449            }
450            Expression::Binary {
451                left,
452                operator,
453                right,
454            } => {
455                self.sql.push('(');
456                self.expression(left)?;
457                self.sql.push_str(match operator {
458                    BinaryOperator::Eq => " = ",
459                    BinaryOperator::NotEq => " <> ",
460                    BinaryOperator::GreaterThan => " > ",
461                    BinaryOperator::GreaterThanOrEq => " >= ",
462                    BinaryOperator::LessThan => " < ",
463                    BinaryOperator::LessThanOrEq => " <= ",
464                    BinaryOperator::Like => " like ",
465                    BinaryOperator::In => " in ",
466                    BinaryOperator::Is => " is ",
467                    BinaryOperator::IsNot => " is not ",
468                    BinaryOperator::Add => " + ",
469                });
470                self.expression(right)?;
471                self.sql.push(')');
472            }
473            Expression::Unary {
474                operator,
475                expression,
476            } => match operator {
477                UnaryOperator::IsNull => {
478                    self.expression(expression)?;
479                    self.sql.push_str(" is null");
480                }
481                UnaryOperator::IsNotNull => {
482                    self.expression(expression)?;
483                    self.sql.push_str(" is not null");
484                }
485                UnaryOperator::Not => {
486                    self.sql.push_str("not (");
487                    self.expression(expression)?;
488                    self.sql.push(')');
489                }
490                UnaryOperator::Exists => {
491                    self.sql.push_str("exists ");
492                    match expression.as_ref() {
493                        Expression::Subquery(query) => self.subquery(&query.0, false)?,
494                        expression => self.expression(expression)?,
495                    }
496                }
497            },
498            Expression::And(expressions) => self.boolean_group(expressions, " and ")?,
499            Expression::Or(expressions) => self.boolean_group(expressions, " or ")?,
500        }
501        Ok(())
502    }
503
504    fn subquery(&mut self, query: &SelectNode, scalar: bool) -> Result<()> {
505        if scalar && query.selections.len() != 1 {
506            return Err(Error::InvalidScalarSubquery(query.selections.len()));
507        }
508        self.sql.push('(');
509        self.select(query.clone())?;
510        self.sql.push(')');
511        Ok(())
512    }
513
514    fn boolean_group(&mut self, expressions: &[Expression], separator: &str) -> Result<()> {
515        if expressions.is_empty() {
516            return Err(Error::Compilation(
517                "boolean expression group cannot be empty".to_string(),
518            ));
519        }
520        self.sql.push('(');
521        for (index, expression) in expressions.iter().enumerate() {
522            if index > 0 {
523                self.sql.push_str(separator);
524            }
525            self.expression(expression)?;
526        }
527        self.sql.push(')');
528        Ok(())
529    }
530
531    fn order_list(&mut self, order_by: &[(Expression, OrderDirection)]) -> Result<()> {
532        for (index, (expression, direction)) in order_by.iter().enumerate() {
533            if index > 0 {
534                self.sql.push_str(", ");
535            }
536            self.expression(expression)?;
537            self.sql.push_str(match direction {
538                OrderDirection::Asc => " asc",
539                OrderDirection::Desc => " desc",
540            });
541        }
542        Ok(())
543    }
544
545    fn window_boundary(&mut self, boundary: WindowBoundary) {
546        match boundary {
547            WindowBoundary::UnboundedPreceding => self.sql.push_str("unbounded preceding"),
548            WindowBoundary::Preceding(value) => {
549                self.sql.push_str(&value.to_string());
550                self.sql.push_str(" preceding");
551            }
552            WindowBoundary::CurrentRow => self.sql.push_str("current row"),
553            WindowBoundary::Following(value) => {
554                self.sql.push_str(&value.to_string());
555                self.sql.push_str(" following");
556            }
557            WindowBoundary::UnboundedFollowing => self.sql.push_str("unbounded following"),
558        }
559    }
560
561    fn table(&mut self, table: &TableNode) -> Result<()> {
562        self.identifier(table.name)?;
563        if let Some(alias) = table.alias {
564            self.sql.push_str(" as ");
565            self.identifier(alias)?;
566        }
567        Ok(())
568    }
569
570    fn identifier(&mut self, identifier: &str) -> Result<()> {
571        validate_identifier(identifier)?;
572        self.sql.push(self.dialect.identifier_quote());
573        for character in identifier.chars() {
574            if character == self.dialect.identifier_quote() {
575                self.sql.push(character);
576            }
577            self.sql.push(character);
578        }
579        self.sql.push(self.dialect.identifier_quote());
580        Ok(())
581    }
582
583    fn parameter(&mut self, value: Value) {
584        self.parameters.push(value);
585        self.sql
586            .push_str(&self.dialect.placeholder(self.parameters.len()));
587    }
588}