Skip to main content

orbit_orm/
query.rs

1use std::marker::PhantomData;
2
3use crate::backend;
4use crate::driver::rows::DriverRows;
5use crate::transaction::Executor;
6use crate::{Driver, Error, Model, Value};
7
8/// A comparison for `where_op`.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Operator {
11    Lt,
12    Le,
13    Gt,
14    Ge,
15    Ne,
16    Like,
17}
18
19impl Operator {
20    fn sql(self) -> &'static str {
21        match self {
22            Self::Lt => "<",
23            Self::Le => "<=",
24            Self::Gt => ">",
25            Self::Ge => ">=",
26            Self::Ne => "!=",
27            Self::Like => "LIKE",
28        }
29    }
30}
31
32/// The sort direction for `order_by`.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum Direction {
35    Asc,
36    Desc,
37}
38
39/// How a condition joins the one before it. The first condition has no joiner.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41enum Joiner {
42    And,
43    Or,
44}
45
46/// A condition, with its names as given. They are quoted when the SQL is built, for the driver
47/// that runs it.
48#[derive(Debug, Clone)]
49enum Condition {
50    Compare {
51        column: String,
52        op: &'static str,
53        value: Value,
54    },
55    In {
56        column: String,
57        values: Vec<Value>,
58    },
59    Null {
60        column: String,
61        negated: bool,
62    },
63    /// A parenthesised group, built by `where_group`.
64    Group(Vec<Clause>),
65}
66
67#[derive(Debug, Clone)]
68struct Clause {
69    joiner: Joiner,
70    condition: Condition,
71}
72
73/// A query on one model's table. Start it with `Model::query()` or `Query::table()`, then
74/// finish it with a terminal method such as `get`, `first` or `insert`.
75pub struct Query<M: Model> {
76    table: String,
77    conditions: Vec<Clause>,
78    orders: Vec<(String, Direction)>,
79    columns: Option<Vec<String>>,
80    limit: Option<u64>,
81    offset: Option<u64>,
82    lock: bool,
83    error: Option<Error>,
84    model: PhantomData<M>,
85}
86
87/// The values bound to one statement, in placeholder order. Each call to `bind` adds one and
88/// returns the placeholder for it, in the driver's syntax.
89pub(crate) struct Params {
90    driver: Driver,
91    pub(crate) values: Vec<Value>,
92}
93
94impl Params {
95    pub(crate) fn new(driver: Driver) -> Self {
96        Self {
97            driver,
98            values: Vec::new(),
99        }
100    }
101
102    /// Adds `value` and returns its placeholder. A `NULL` is written as `NULL` on drivers that
103    /// inline it, and is not bound.
104    pub(crate) fn bind(&mut self, value: Value) -> String {
105        if matches!(value, Value::Null) && self.driver.inlines_null() {
106            return "NULL".to_owned();
107        }
108        self.values.push(value);
109        self.driver.placeholder(self.values.len())
110    }
111}
112
113/// Checks a table or column name. It must not be empty, and must not contain a backtick or a
114/// double quote, which are the quote characters of the supported databases.
115fn validate(name: &str) -> Result<(), Error> {
116    if name.is_empty() || name.contains(['`', '"']) {
117        return Err(Error::InvalidIdentifier(name.to_owned()));
118    }
119    Ok(())
120}
121
122/// Checks a name, or records the error and returns the name for the terminal to report.
123fn checked(name: &str, error: &mut Option<Error>) -> String {
124    if let Err(e) = validate(name) {
125        error.get_or_insert(e);
126    }
127    name.to_owned()
128}
129
130/// Checks a column that `insert` or `update` writes against `FILLABLE`.
131fn fillable<M: Model>(column: &str) -> Result<(), Error> {
132    if !M::FILLABLE.contains(&column) {
133        return Err(Error::NotFillable {
134            table: M::TABLE,
135            column: column.to_owned(),
136        });
137    }
138    validate(column)
139}
140
141impl<M: Model> Query<M> {
142    /// Starts a query on the model's table.
143    pub fn table() -> Self {
144        let mut error = None;
145        let table = checked(M::TABLE, &mut error);
146        Self {
147            table,
148            error,
149            ..Self::blank()
150        }
151    }
152
153    /// A query with no table, for the conditions of a group.
154    fn blank() -> Self {
155        Self {
156            table: String::new(),
157            conditions: Vec::new(),
158            orders: Vec::new(),
159            columns: None,
160            limit: None,
161            offset: None,
162            lock: false,
163            error: None,
164            model: PhantomData,
165        }
166    }
167
168    /// Adds `column = value`. A `NULL` value becomes `IS NULL`, as in Laravel.
169    pub fn where_eq(self, column: &str, value: impl Into<Value>) -> Self {
170        self.compare(Joiner::And, column, "=", value.into())
171    }
172
173    /// Adds `OR column = value`.
174    pub fn or_where_eq(self, column: &str, value: impl Into<Value>) -> Self {
175        self.compare(Joiner::Or, column, "=", value.into())
176    }
177
178    /// Adds `column <op> value`. `!=` with a `NULL` value becomes `IS NOT NULL`.
179    pub fn where_op(self, column: &str, op: Operator, value: impl Into<Value>) -> Self {
180        self.compare(Joiner::And, column, op.sql(), value.into())
181    }
182
183    /// Adds `OR column <op> value`. `!=` with a `NULL` value becomes `IS NOT NULL`.
184    pub fn or_where_op(self, column: &str, op: Operator, value: impl Into<Value>) -> Self {
185        self.compare(Joiner::Or, column, op.sql(), value.into())
186    }
187
188    fn compare(self, joiner: Joiner, column: &str, op: &'static str, value: Value) -> Self {
189        let mut this = self;
190        let column = checked(column, &mut this.error);
191        let condition = match (value, op) {
192            (Value::Null, "=") => Condition::Null {
193                column,
194                negated: false,
195            },
196            (Value::Null, "!=") => Condition::Null {
197                column,
198                negated: true,
199            },
200            (value, op) => Condition::Compare { column, op, value },
201        };
202        this.push(joiner, condition)
203    }
204
205    /// Adds `column IN (values)`. An empty list matches no rows, as in Laravel.
206    pub fn where_in<V: Into<Value>>(
207        self,
208        column: &str,
209        values: impl IntoIterator<Item = V>,
210    ) -> Self {
211        self.in_list(Joiner::And, column, values)
212    }
213
214    /// Adds `OR column IN (values)`. An empty list matches no rows.
215    pub fn or_where_in<V: Into<Value>>(
216        self,
217        column: &str,
218        values: impl IntoIterator<Item = V>,
219    ) -> Self {
220        self.in_list(Joiner::Or, column, values)
221    }
222
223    fn in_list<V: Into<Value>>(
224        mut self,
225        joiner: Joiner,
226        column: &str,
227        values: impl IntoIterator<Item = V>,
228    ) -> Self {
229        let column = checked(column, &mut self.error);
230        let values = values.into_iter().map(Into::into).collect();
231        self.push(joiner, Condition::In { column, values })
232    }
233
234    /// Adds `column IS NULL`.
235    pub fn where_null(self, column: &str) -> Self {
236        self.null(Joiner::And, column, false)
237    }
238
239    /// Adds `OR column IS NULL`.
240    pub fn or_where_null(self, column: &str) -> Self {
241        self.null(Joiner::Or, column, false)
242    }
243
244    /// Adds `column IS NOT NULL`.
245    pub fn where_not_null(self, column: &str) -> Self {
246        self.null(Joiner::And, column, true)
247    }
248
249    /// Adds `OR column IS NOT NULL`.
250    pub fn or_where_not_null(self, column: &str) -> Self {
251        self.null(Joiner::Or, column, true)
252    }
253
254    fn null(mut self, joiner: Joiner, column: &str, negated: bool) -> Self {
255        let column = checked(column, &mut self.error);
256        self.push(joiner, Condition::Null { column, negated })
257    }
258
259    /// Adds a group of conditions joined by `AND`, in parentheses. The closure gets an empty
260    /// query and adds its conditions to it, as Laravel's `where(fn)` does:
261    ///
262    /// ```ignore
263    /// .where_group(|q| q.where_eq("title", "a").or_where_eq("title", "b"))
264    /// ```
265    ///
266    /// Only where conditions are kept from the group; `order_by`, `limit` and the other
267    /// settings are ignored. An empty group adds nothing.
268    pub fn where_group(self, f: impl FnOnce(Self) -> Self) -> Self {
269        self.group(Joiner::And, f)
270    }
271
272    /// Adds a group of conditions joined by `OR`, in parentheses. See `where_group`.
273    pub fn or_where_group(self, f: impl FnOnce(Self) -> Self) -> Self {
274        self.group(Joiner::Or, f)
275    }
276
277    fn group(mut self, joiner: Joiner, f: impl FnOnce(Self) -> Self) -> Self {
278        let inner = f(Self::blank());
279        if let Some(error) = inner.error {
280            self.error.get_or_insert(error);
281        }
282        if inner.conditions.is_empty() {
283            return self;
284        }
285        self.push(joiner, Condition::Group(inner.conditions))
286    }
287
288    fn push(mut self, joiner: Joiner, condition: Condition) -> Self {
289        self.conditions.push(Clause { joiner, condition });
290        self
291    }
292
293    /// Adds a sort. Later calls sort within the earlier ones.
294    pub fn order_by(mut self, column: &str, direction: Direction) -> Self {
295        let column = checked(column, &mut self.error);
296        self.orders.push((column, direction));
297        self
298    }
299
300    /// Chooses the columns that `get` and `first` select. Without it, all columns are selected.
301    pub fn select(mut self, columns: &[&str]) -> Self {
302        if columns.is_empty() {
303            return self;
304        }
305        let names = columns
306            .iter()
307            .map(|column| checked(column, &mut self.error))
308            .collect();
309        self.columns = Some(names);
310        self
311    }
312
313    /// Returns at most `limit` rows.
314    pub fn limit(mut self, limit: u64) -> Self {
315        self.limit = Some(limit);
316        self
317    }
318
319    /// Skips the first `offset` rows.
320    pub fn offset(mut self, offset: u64) -> Self {
321        self.offset = Some(offset);
322        self
323    }
324
325    /// Locks the matching rows (`SELECT ... FOR UPDATE`). It applies to `get`, `first`, `find`,
326    /// `count` and `exists`. In a transaction the lock lasts until it ends. On the pool it lasts
327    /// only for the statement, in autocommit, as Laravel's `lockForUpdate()` does. SQLite has no
328    /// row locks, so the clause is left out there, as Laravel's SQLite grammar does.
329    pub fn for_update(mut self) -> Self {
330        self.lock = true;
331        self
332    }
333
334    /// Returns every row that matches.
335    pub async fn get<'a>(mut self, db: impl Executor<'a>) -> Result<Vec<M>, Error>
336    where
337        M: DriverRows,
338    {
339        self.take_error()?;
340        self.rows(db, self.limit).await
341    }
342
343    /// Returns every row that matches, decoded as `T` instead of the model. Use it with `select`
344    /// to read some columns into a tuple or a struct that implements `FromRow`. The clauses are
345    /// the same as `get`.
346    ///
347    /// ```ignore
348    /// let titles = Post::query()
349    ///     .select(&["id", "title"])
350    ///     .get_as::<(i64, String)>(&db)
351    ///     .await?;
352    /// ```
353    pub async fn get_as<'a, T>(mut self, db: impl Executor<'a>) -> Result<Vec<T>, Error>
354    where
355        T: DriverRows + Send + Unpin,
356    {
357        self.take_error()?;
358        self.rows::<T>(db, self.limit).await
359    }
360
361    /// Returns the first row decoded as `T`, in the order given by `order_by`. See `get_as`.
362    pub async fn first_as<'a, T>(mut self, db: impl Executor<'a>) -> Result<Option<T>, Error>
363    where
364        T: DriverRows + Send + Unpin,
365    {
366        self.take_error()?;
367        let rows = self.rows::<T>(db, Some(1)).await?;
368        Ok(rows.into_iter().next())
369    }
370
371    /// Returns the first row in the order given by `order_by`. Without `order_by` there is no
372    /// order, as with Laravel's `first`, so call `order_by` when the row matters.
373    pub async fn first<'a>(mut self, db: impl Executor<'a>) -> Result<Option<M>, Error>
374    where
375        M: DriverRows,
376    {
377        self.take_error()?;
378        let rows = self.rows::<M>(db, Some(1)).await?;
379        Ok(rows.into_iter().next())
380    }
381
382    /// Returns the first row, or an error when none matches.
383    pub async fn first_or_fail<'a>(self, db: impl Executor<'a>) -> Result<M, Error>
384    where
385        M: DriverRows,
386    {
387        self.first(db)
388            .await?
389            .ok_or(Error::NotFound { table: M::TABLE })
390    }
391
392    /// Returns the row with the given primary key, or `None` when there is none.
393    pub async fn find<'a>(
394        self,
395        db: impl Executor<'a>,
396        id: impl Into<Value>,
397    ) -> Result<Option<M>, Error>
398    where
399        M: DriverRows,
400    {
401        self.where_eq(M::PRIMARY_KEY, id).first(db).await
402    }
403
404    /// Returns the row with the given primary key, or an error when there is none.
405    pub async fn find_or_fail<'a>(
406        self,
407        db: impl Executor<'a>,
408        id: impl Into<Value>,
409    ) -> Result<M, Error>
410    where
411        M: DriverRows,
412    {
413        self.find(db, id)
414            .await?
415            .ok_or(Error::NotFound { table: M::TABLE })
416    }
417
418    /// Counts the rows that match.
419    pub async fn count<'a>(mut self, db: impl Executor<'a>) -> Result<i64, Error> {
420        self.take_error()?;
421        let target = db.target();
422        let driver = target.driver();
423        let mut params = Params::new(driver);
424        let sql = format!(
425            "SELECT COUNT(*) FROM {}{}{}",
426            self.table_sql(driver),
427            self.where_sql(driver, &mut params),
428            self.lock_sql(driver)
429        );
430        Ok(backend::scalar::<i64>(target, &sql, &params.values).await?)
431    }
432
433    /// Returns whether any row matches.
434    pub async fn exists<'a>(mut self, db: impl Executor<'a>) -> Result<bool, Error> {
435        self.take_error()?;
436        let target = db.target();
437        let driver = target.driver();
438        let mut params = Params::new(driver);
439        // The flag is cast to an integer so that every driver returns the same type.
440        let sql = format!(
441            "SELECT CAST(CASE WHEN EXISTS (SELECT 1 FROM {}{}{}) THEN 1 ELSE 0 END AS {})",
442            self.table_sql(driver),
443            self.where_sql(driver, &mut params),
444            self.lock_sql(driver),
445            driver.bigint()
446        );
447        let found = backend::scalar::<i64>(target, &sql, &params.values).await?;
448        Ok(found != 0)
449    }
450
451    /// Inserts one row and returns its auto-increment id. The columns must be in `FILLABLE`.
452    /// On PostgreSQL the id comes from `RETURNING` on the primary key, so the primary key must be
453    /// an integer column.
454    pub async fn insert<'a>(
455        mut self,
456        db: impl Executor<'a>,
457        values: &[(&str, Value)],
458    ) -> Result<u64, Error> {
459        self.take_error()?;
460        let target = db.target();
461        let driver = target.driver();
462        let mut params = Params::new(driver);
463        let returning = driver.insert_returns_id().then(|| {
464            format!(
465                " RETURNING {}::{}",
466                driver.quote(M::PRIMARY_KEY),
467                driver.bigint()
468            )
469        });
470        let sql = self.insert_sql(driver, "INSERT", &mut params, values, returning.as_deref())?;
471        if returning.is_some() {
472            let id: i64 = backend::scalar(target, &sql, &params.values).await?;
473            return u64::try_from(id).map_err(|_| Error::Database(sqlx::Error::RowNotFound));
474        }
475        let outcome = backend::execute(target, &sql, &params.values).await?;
476        outcome.insert_id.ok_or(Error::Unsupported {
477            operation: "reading the insert id",
478            driver: driver.name(),
479        })
480    }
481
482    /// Replaces the row that has the same unique key or primary key as `values`, or inserts the
483    /// row when none has. Returns the rows affected, as the database reports them: MySQL gives 2
484    /// when a row was replaced, since it deletes the old row and inserts the new one, and SQLite
485    /// gives 1.
486    ///
487    /// This is Diesel's `replace_into`, which issues `REPLACE INTO` on MySQL and SQLite. It is not
488    /// Laravel's `upsert`, which is `INSERT ... ON DUPLICATE KEY UPDATE` or `ON CONFLICT`. The old
489    /// row is deleted, so its auto-increment id changes and columns not in `values` take their
490    /// defaults. The columns must be in `FILLABLE`.
491    ///
492    /// PostgreSQL has no `REPLACE`, and `INSERT ... ON CONFLICT` would keep the old row. So on
493    /// PostgreSQL this returns `Error::Unsupported`.
494    pub async fn replace<'a>(
495        mut self,
496        db: impl Executor<'a>,
497        values: &[(&str, Value)],
498    ) -> Result<u64, Error> {
499        self.take_error()?;
500        let target = db.target();
501        let driver = target.driver();
502        if !driver.supports_replace() {
503            return Err(Error::Unsupported {
504                operation: "replace",
505                driver: driver.name(),
506            });
507        }
508        let mut params = Params::new(driver);
509        let sql = self.insert_sql(driver, "REPLACE", &mut params, values, None)?;
510        let outcome = backend::execute(target, &sql, &params.values).await?;
511        Ok(outcome.rows_affected)
512    }
513
514    /// Builds `INSERT INTO table (columns) VALUES (...)`, with `returning` appended. The columns
515    /// must be in `FILLABLE`.
516    fn insert_sql(
517        &self,
518        driver: Driver,
519        keyword: &str,
520        params: &mut Params,
521        values: &[(&str, Value)],
522        returning: Option<&str>,
523    ) -> Result<String, Error> {
524        let mut columns = Vec::with_capacity(values.len());
525        let mut placeholders = Vec::with_capacity(values.len());
526        for (column, value) in values {
527            fillable::<M>(column)?;
528            columns.push(driver.quote(column));
529            placeholders.push(params.bind(value.clone()));
530        }
531        let table = self.table_sql(driver);
532        let returning = returning.unwrap_or("");
533        Ok(if columns.is_empty() {
534            format!(
535                "{keyword} INTO {table} {}{returning}",
536                driver.empty_insert()
537            )
538        } else {
539            format!(
540                "{keyword} INTO {table} ({}) VALUES ({}){returning}",
541                columns.join(", "),
542                placeholders.join(", ")
543            )
544        })
545    }
546
547    /// Sets the columns on every matching row and returns how many rows matched. The columns
548    /// must be in `FILLABLE`. Order and limit are ignored.
549    pub async fn update<'a>(
550        mut self,
551        db: impl Executor<'a>,
552        values: &[(&str, Value)],
553    ) -> Result<u64, Error> {
554        self.take_error()?;
555        if values.is_empty() {
556            return Ok(0);
557        }
558        let target = db.target();
559        let driver = target.driver();
560        let mut params = Params::new(driver);
561        let mut sets = Vec::with_capacity(values.len());
562        for (column, value) in values {
563            fillable::<M>(column)?;
564            let placeholder = params.bind(value.clone());
565            sets.push(format!("{} = {placeholder}", driver.quote(column)));
566        }
567        let mut sql = format!("UPDATE {} SET {}", self.table_sql(driver), sets.join(", "));
568        sql.push_str(&self.where_sql(driver, &mut params));
569        let outcome = backend::execute(target, &sql, &params.values).await?;
570        Ok(outcome.rows_affected)
571    }
572
573    /// Deletes every matching row and returns how many rows were deleted. Order and limit are
574    /// ignored.
575    pub async fn delete<'a>(mut self, db: impl Executor<'a>) -> Result<u64, Error> {
576        self.take_error()?;
577        let target = db.target();
578        let driver = target.driver();
579        let mut params = Params::new(driver);
580        let sql = format!(
581            "DELETE FROM {}{}",
582            self.table_sql(driver),
583            self.where_sql(driver, &mut params)
584        );
585        let outcome = backend::execute(target, &sql, &params.values).await?;
586        Ok(outcome.rows_affected)
587    }
588
589    /// Returns the first error recorded while building the query.
590    fn take_error(&mut self) -> Result<(), Error> {
591        match self.error.take() {
592            Some(e) => Err(e),
593            None => Ok(()),
594        }
595    }
596
597    /// The lock clause when `for_update` was asked for, or nothing.
598    fn lock_sql(&self, driver: Driver) -> &'static str {
599        if self.lock { driver.lock_clause() } else { "" }
600    }
601
602    /// The quoted table name.
603    fn table_sql(&self, driver: Driver) -> String {
604        driver.quote(&self.table)
605    }
606
607    async fn rows<'a, T>(&self, db: impl Executor<'a>, limit: Option<u64>) -> Result<Vec<T>, Error>
608    where
609        T: DriverRows + Send + Unpin,
610    {
611        let target = db.target();
612        let driver = target.driver();
613        let columns = match &self.columns {
614            Some(names) => names
615                .iter()
616                .map(|name| driver.quote(name))
617                .collect::<Vec<_>>()
618                .join(", "),
619            None => "*".to_owned(),
620        };
621        let mut params = Params::new(driver);
622        let sql = self.select_sql(driver, &columns, limit, &mut params);
623        Ok(backend::fetch_all::<T>(target, &sql, &params.values).await?)
624    }
625
626    /// Builds `SELECT` with the where clauses, order, limit, offset and lock.
627    fn select_sql(
628        &self,
629        driver: Driver,
630        columns: &str,
631        limit: Option<u64>,
632        params: &mut Params,
633    ) -> String {
634        let mut sql = format!("SELECT {columns} FROM {}", self.table_sql(driver));
635        sql.push_str(&self.where_sql(driver, params));
636        if !self.orders.is_empty() {
637            let terms: Vec<String> = self
638                .orders
639                .iter()
640                .map(|(column, direction)| match direction {
641                    Direction::Asc => format!("{} ASC", driver.quote(column)),
642                    Direction::Desc => format!("{} DESC", driver.quote(column)),
643                })
644                .collect();
645            sql.push_str(&format!(" ORDER BY {}", terms.join(", ")));
646        }
647        if let Some(limit) = limit {
648            sql.push_str(&format!(" LIMIT {}", params.bind(Value::UInt(limit))));
649        } else if let (true, Some(all)) = (self.offset.is_some(), driver.offset_only_limit()) {
650            // MySQL and SQLite need a LIMIT before OFFSET, so they get their largest value.
651            sql.push_str(&format!(" LIMIT {all}"));
652        }
653        if let Some(offset) = self.offset {
654            sql.push_str(&format!(" OFFSET {}", params.bind(Value::UInt(offset))));
655        }
656        sql.push_str(self.lock_sql(driver));
657        sql
658    }
659
660    /// Builds ` WHERE ...` from the conditions, or nothing.
661    fn where_sql(&self, driver: Driver, params: &mut Params) -> String {
662        if self.conditions.is_empty() {
663            return String::new();
664        }
665        format!(" WHERE {}", render(driver, &self.conditions, params))
666    }
667}
668
669/// Joins the clauses with their `AND` or `OR`. Values are bound in the order they appear.
670fn render(driver: Driver, clauses: &[Clause], params: &mut Params) -> String {
671    let mut sql = String::new();
672    for (index, clause) in clauses.iter().enumerate() {
673        if index > 0 {
674            sql.push_str(match clause.joiner {
675                Joiner::And => " AND ",
676                Joiner::Or => " OR ",
677            });
678        }
679        match &clause.condition {
680            Condition::Compare { column, op, value } => {
681                let placeholder = params.bind(value.clone());
682                sql.push_str(&format!("{} {op} {placeholder}", driver.quote(column)));
683            }
684            Condition::In { values, .. } if values.is_empty() => sql.push_str("0 = 1"),
685            Condition::In { column, values } => {
686                let placeholders: Vec<String> =
687                    values.iter().map(|v| params.bind(v.clone())).collect();
688                sql.push_str(&format!(
689                    "{} IN ({})",
690                    driver.quote(column),
691                    placeholders.join(", ")
692                ));
693            }
694            Condition::Null { column, negated } => {
695                let not = if *negated { "NOT " } else { "" };
696                sql.push_str(&format!("{} IS {not}NULL", driver.quote(column)));
697            }
698            Condition::Group(inner) => {
699                sql.push_str(&format!("({})", render(driver, inner, params)));
700            }
701        }
702    }
703    sql
704}