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