good_ormning/pg/query/
insert.rs

1use std::{
2    collections::{
3        HashMap,
4        HashSet,
5    },
6};
7use crate::{
8    pg::{
9        QueryResCount,
10        schema::{
11            field::Field,
12            table::Table,
13        },
14        types::SimpleSimpleType,
15    },
16    utils::Tokens,
17};
18use super::{
19    expr::{
20        Expr,
21        ExprType,
22        check_assignable,
23        ExprValName,
24    },
25    utils::{
26        QueryBody,
27        build_returning,
28        build_set,
29    },
30    select::Returning,
31};
32
33pub enum InsertConflict {
34    DoNothing,
35    DoUpdate {
36        conflict: Vec<Field>,
37        set: Vec<(Field, Expr)>,
38    },
39}
40
41pub struct Insert {
42    pub(crate) table: Table,
43    pub(crate) values: Vec<(Field, Expr)>,
44    pub(crate) on_conflict: Option<InsertConflict>,
45    pub(crate) returning: Vec<Returning>,
46}
47
48impl QueryBody for Insert {
49    fn build(
50        &self,
51        ctx: &mut super::utils::PgQueryCtx,
52        path: &rpds::Vector<String>,
53        res_count: QueryResCount,
54    ) -> (ExprType, Tokens) {
55        // Prep
56        let mut check_inserting_fields = HashSet::new();
57        for p in &self.values {
58            if p.0.type_.type_.opt {
59                continue;
60            }
61            if !check_inserting_fields.insert(p.0.clone()) {
62                ctx.errs.err(path, format!("Duplicate field {} in insert", p.0));
63            }
64        }
65        let mut scope = HashMap::new();
66        for (field, v) in match ctx.tables.get(&self.table) {
67            Some(t) => t,
68            None => {
69                ctx.errs.err(path, format!("Unknown table {} for insert", self.table));
70                return (ExprType(vec![]), Tokens::new());
71            },
72        } {
73            scope.insert(ExprValName::field(field), v.clone());
74            if !field.type_.type_.opt && field.type_.type_.type_.type_ != SimpleSimpleType::Auto &&
75                !check_inserting_fields.remove(field) {
76                ctx.errs.err(path, format!("{} is a non-optional field but is missing in insert", field));
77            }
78        }
79        drop(check_inserting_fields);
80
81        // Build query
82        let mut out = Tokens::new();
83        out.s("insert into").id(&self.table.id).s("(");
84        for (i, (field, _)) in self.values.iter().enumerate() {
85            if i > 0 {
86                out.s(",");
87            }
88            out.id(&field.id);
89        }
90        out.s(") values (");
91        for (i, (field, val)) in self.values.iter().enumerate() {
92            if i > 0 {
93                out.s(",");
94            }
95            let field_type = match ctx.tables.get(&field.table).and_then(|t| t.get(&field)) {
96                Some(t) => t,
97                None => {
98                    ctx.errs.err(path, format!("Insert destination value field {} is not known", field));
99                    continue;
100                },
101            };
102            let path = path.push_back(format!("Insert value {} ({})", i, field));
103            let res = val.build(ctx, &path, &scope);
104            check_assignable(&mut ctx.errs, &path, &field_type, &res.0);
105            out.s(&res.1.to_string());
106        }
107        out.s(")");
108        if let Some(conflict) = &self.on_conflict {
109            out.s("on conflict");
110            match conflict {
111                InsertConflict::DoNothing => {
112                    out.s("do nothing");
113                },
114                InsertConflict::DoUpdate { conflict, set } => {
115                    out.s("(");
116                    for (i, f) in conflict.iter().enumerate() {
117                        if i > 0 {
118                            out.s(",");
119                        }
120                        out.id(&f.id);
121                    }
122                    out.s(")");
123                    out.s("do update");
124                    build_set(ctx, path, &scope, &mut out, set);
125                },
126            }
127        }
128        match (&res_count, &self.on_conflict) {
129            (QueryResCount::MaybeOne, Some(InsertConflict::DoUpdate { .. })) => {
130                ctx.errs.err(path, format!("Insert with [on conflict update] will always return a row"));
131            },
132            (QueryResCount::One, Some(InsertConflict::DoNothing)) => {
133                ctx.errs.err(path, format!("Insert with [on conflict do nothing] may not return a row"));
134            },
135            (QueryResCount::Many, _) => {
136                ctx.errs.err(path, format!("Insert can at most return one row, but res count is many"));
137            },
138            (QueryResCount::None, _) | (QueryResCount::One, None) | (QueryResCount::MaybeOne, None) => {
139                // handled elsewhere, nop
140            },
141            (QueryResCount::One, Some(InsertConflict::DoUpdate { .. })) |
142            (QueryResCount::MaybeOne, Some(InsertConflict::DoNothing)) => {
143                // ok
144            },
145        }
146        let out_type = build_returning(ctx, path, &scope, &mut out, &self.returning, res_count);
147        (out_type, out)
148    }
149}