Skip to main content

postrust_core/query/
mod.rs

1//! SQL query generation from execution plans.
2//!
3//! This module converts execution plans into parameterized SQL queries.
4
5mod builder;
6
7pub use builder::{QueryBuilder, INSERTED_COLUMN, PARENT_ROW_COLUMN};
8
9use crate::error::Result;
10use crate::plan::{ActionPlan, DbActionPlan};
11use postrust_sql::{SqlFragment, SqlParam};
12
13/// Build SQL from an action plan.
14pub fn build_query(plan: &ActionPlan, role: Option<&str>) -> Result<MainQuery> {
15    match plan {
16        ActionPlan::Db(db_plan) => build_db_query(db_plan, role),
17        ActionPlan::Info(_) => Ok(MainQuery::empty()),
18    }
19}
20
21/// The name the rows a mutation affected are read from.
22pub const MUTATION_RESULT: &str = "pgrst_mutation_result";
23
24/// Build SQL from a database action plan.
25fn build_db_query(plan: &DbActionPlan, role: Option<&str>) -> Result<MainQuery> {
26    let mut query = MainQuery::new();
27
28    // Add role switch if specified
29    if let Some(role) = role {
30        query.pre_statements.push(format!(
31            "SET LOCAL ROLE {}",
32            postrust_sql::escape_ident(role)
33        ));
34    }
35
36    match plan {
37        DbActionPlan::Read(read_tree) => {
38            query.main = QueryBuilder::build_read(read_tree)?;
39        }
40        DbActionPlan::MutateRead { mutate, read } => {
41            let mutation = QueryBuilder::build_mutate(mutate)?;
42
43            // The rows a mutation affected are a relation like any other, so
44            // `?select=` is answered by reading from them -- which is also
45            // what lets a mutation embed, alias, cast and compute exactly as a
46            // read does. Without the wrapper the `RETURNING` list was the
47            // whole answer, and none of that reached it.
48            query.main = match read {
49                Some(read_tree) => {
50                    let read_sql = QueryBuilder::build_read(read_tree)?;
51                    let mut frag = SqlFragment::new();
52                    frag.push("WITH ");
53                    frag.push(MUTATION_RESULT);
54                    frag.push(" AS (");
55                    frag.append(mutation);
56                    frag.push(") ");
57                    frag.append(read_sql);
58                    frag
59                }
60                None => mutation,
61            };
62        }
63        DbActionPlan::Call { call, read } => {
64            query.main = QueryBuilder::build_call(call, read.as_ref())?;
65        }
66    }
67
68    Ok(query)
69}
70
71/// A complete query with setup and main statement.
72#[derive(Clone, Debug, Default)]
73pub struct MainQuery {
74    /// Pre-query statements (SET commands)
75    pub pre_statements: Vec<String>,
76    /// Main query
77    pub main: SqlFragment,
78    /// Read query (for mutations with RETURNING)
79    pub read: Option<SqlFragment>,
80    /// Count query (for pagination)
81    pub count: Option<SqlFragment>,
82}
83
84impl MainQuery {
85    /// Create a new empty query.
86    pub fn new() -> Self {
87        Self::default()
88    }
89
90    /// Create an empty query (for info-only plans).
91    pub fn empty() -> Self {
92        Self::default()
93    }
94
95    /// Check if this query has a main statement.
96    pub fn has_main(&self) -> bool {
97        !self.main.is_empty()
98    }
99
100    /// Get the main SQL and parameters.
101    pub fn build_main(self) -> (String, Vec<SqlParam>) {
102        self.main.build()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::api_request::{QualifiedIdentifier, Range};
110    use crate::plan::{CoercibleField, CoercibleSelectField, MutatePlan, ReadPlan, ReadPlanTree};
111
112    /// The whole statement, parentheses included.
113    ///
114    /// A missing one here is a syntax error on every write with a body, and
115    /// nothing short of executing it against PostgreSQL -- or reading it --
116    /// will say so, since the fragments that build it are each well-formed.
117    #[test]
118    fn a_write_reads_its_body_and_is_read_from() {
119        let mutate = MutatePlan::Insert {
120            target: QualifiedIdentifier::new("test", "items"),
121            columns: vec![CoercibleField::simple("id", "int8")],
122            body: Some(bytes::Bytes::from(r#"{"id":9001}"#)),
123            on_conflict: None,
124            where_clauses: vec![],
125            returning: vec!["id".into()],
126            pk_cols: vec!["id".into()],
127            apply_defaults: true,
128            reports_inserted: false,
129        };
130        let read = ReadPlan {
131            select: vec![CoercibleSelectField::simple("id", "bigint")],
132            from: QualifiedIdentifier::unqualified(MUTATION_RESULT),
133            from_alias: None,
134            where_clauses: vec![],
135            order: vec![],
136            range: Range::default(),
137            rel_name: String::new(),
138            rel_to_parent: None,
139            rel_join_conds: vec![],
140            rel_join_type: None,
141            rel_select: vec![],
142            depth: 0,
143        };
144
145        let plan = ActionPlan::Db(DbActionPlan::MutateRead {
146            mutate,
147            read: Some(ReadPlanTree::leaf(read)),
148        });
149        let (sql, _) = build_query(&plan, None).unwrap().build_main();
150
151        assert_eq!(
152            sql,
153            "WITH pgrst_mutation_result AS (\
154             INSERT INTO \"test\".\"items\" (\"id\") \
155             SELECT pgrst_body.\"id\" \
156             FROM (SELECT $1::json AS json_data) pgrst_payload, \
157             LATERAL (SELECT \"id\" FROM json_to_record(pgrst_payload.json_data) \
158             AS _(\"id\" int8)) pgrst_body \
159             RETURNING \"items\".\"id\") \
160             SELECT \"id\" FROM \"pgrst_mutation_result\""
161        );
162    }
163}