Skip to main content

orbit_orm/schema/
blueprint.rs

1use super::Dialect;
2use super::column::{Column, ColumnType};
3use crate::query::validate;
4use crate::{Driver, Error};
5
6/// What a foreign key does to this table's rows when the referenced row is deleted or updated.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum Action {
9    Cascade,
10    Restrict,
11    SetNull,
12    SetDefault,
13    NoAction,
14}
15
16impl Action {
17    fn sql(self) -> &'static str {
18        match self {
19            Self::Cascade => "CASCADE",
20            Self::Restrict => "RESTRICT",
21            Self::SetNull => "SET NULL",
22            Self::SetDefault => "SET DEFAULT",
23            Self::NoAction => "NO ACTION",
24        }
25    }
26}
27
28/// A foreign key under construction. Chain `references`, `on_delete` and `on_update`.
29#[derive(Debug, Clone)]
30pub struct ForeignKey {
31    column: String,
32    target: Option<(String, String)>,
33    on_delete: Option<Action>,
34    on_update: Option<Action>,
35}
36
37impl ForeignKey {
38    /// The table and column that this column points at.
39    pub fn references(&mut self, table: &str, column: &str) -> &mut Self {
40        self.target = Some((table.to_owned(), column.to_owned()));
41        self
42    }
43
44    /// What happens when the referenced row is deleted.
45    pub fn on_delete(&mut self, action: Action) -> &mut Self {
46        self.on_delete = Some(action);
47        self
48    }
49
50    /// What happens when the referenced key is updated.
51    pub fn on_update(&mut self, action: Action) -> &mut Self {
52        self.on_update = Some(action);
53        self
54    }
55}
56
57#[derive(Debug, Clone)]
58struct Index {
59    columns: Vec<String>,
60    unique: bool,
61}
62
63/// The columns, indexes and foreign keys of a table, in the order they were added.
64#[derive(Debug, Clone, Default)]
65pub struct Blueprint {
66    columns: Vec<Column>,
67    indexes: Vec<Index>,
68    foreigns: Vec<ForeignKey>,
69}
70
71impl Blueprint {
72    pub(crate) fn new() -> Self {
73        Self::default()
74    }
75
76    fn add(&mut self, name: &str, ty: ColumnType) -> &mut Column {
77        self.columns.push(Column::new(name, ty));
78        self.columns.last_mut().expect("just pushed")
79    }
80
81    /// An auto-incrementing 64-bit primary key named `id`.
82    pub fn id(&mut self) -> &mut Column {
83        self.add("id", ColumnType::Id)
84    }
85
86    /// A `VARCHAR(255)`; `TEXT` on SQLite. Change the length with `.length(n)`.
87    pub fn string(&mut self, name: &str) -> &mut Column {
88        self.add(name, ColumnType::String(255))
89    }
90
91    /// A `TEXT` column.
92    pub fn text(&mut self, name: &str) -> &mut Column {
93        self.add(name, ColumnType::Text)
94    }
95
96    /// An `INTEGER` column.
97    pub fn integer(&mut self, name: &str) -> &mut Column {
98        self.add(name, ColumnType::Integer)
99    }
100
101    /// A `BIGINT` column; `INTEGER` on SQLite, whose integers are 64-bit already.
102    pub fn big_integer(&mut self, name: &str) -> &mut Column {
103        self.add(name, ColumnType::BigInteger)
104    }
105
106    /// A boolean: `BOOLEAN` on PostgreSQL, `TINYINT(1)` elsewhere.
107    pub fn boolean(&mut self, name: &str) -> &mut Column {
108        self.add(name, ColumnType::Boolean)
109    }
110
111    /// A date and time without a time zone: `DATETIME`, or `TIMESTAMP(0)` on PostgreSQL.
112    pub fn date_time(&mut self, name: &str) -> &mut Column {
113        self.add(name, ColumnType::DateTime)
114    }
115
116    /// Nullable `created_at` and `updated_at` columns.
117    pub fn timestamps(&mut self) {
118        self.date_time("created_at").nullable();
119        self.date_time("updated_at").nullable();
120    }
121
122    /// A unique constraint over `columns`, named `{table}_{columns}_unique`.
123    pub fn unique<I, S>(&mut self, columns: I)
124    where
125        I: IntoIterator<Item = S>,
126        S: AsRef<str>,
127    {
128        self.indexes.push(Index {
129            columns: collect(columns),
130            unique: true,
131        });
132    }
133
134    /// A plain index over `columns`, named `{table}_{columns}_index`.
135    pub fn index<I, S>(&mut self, columns: I)
136    where
137        I: IntoIterator<Item = S>,
138        S: AsRef<str>,
139    {
140        self.indexes.push(Index {
141            columns: collect(columns),
142            unique: false,
143        });
144    }
145
146    /// A foreign key on `column`, named `{table}_{column}_foreign`.
147    pub fn foreign(&mut self, column: &str) -> &mut ForeignKey {
148        self.foreigns.push(ForeignKey {
149            column: column.to_owned(),
150            target: None,
151            on_delete: None,
152            on_update: None,
153        });
154        self.foreigns.last_mut().expect("just pushed")
155    }
156
157    pub(crate) fn render(
158        &self,
159        driver: Driver,
160        dialect: Dialect,
161        table: &str,
162    ) -> Result<Vec<String>, Error> {
163        validate(table)?;
164        let q = |name: &str| driver.quote(name);
165        let list = |columns: &[String]| -> Result<String, Error> {
166            for c in columns {
167                validate(c)?;
168            }
169            Ok(columns.iter().map(|c| q(c)).collect::<Vec<_>>().join(", "))
170        };
171        let name_of =
172            |columns: &[String], suffix: &str| format!("{table}_{}_{suffix}", columns.join("_"));
173
174        let mut parts = Vec::new();
175        for column in &self.columns {
176            parts.push(column.render(driver, dialect)?);
177        }
178        let mut after = Vec::new();
179        for index in &self.indexes {
180            let columns = list(&index.columns)?;
181            if index.unique {
182                let name = name_of(&index.columns, "unique");
183                parts.push(format!("CONSTRAINT {} UNIQUE ({columns})", q(&name)));
184            } else {
185                let name = name_of(&index.columns, "index");
186                match dialect {
187                    Dialect::MySql => parts.push(format!("INDEX {} ({columns})", q(&name))),
188                    _ => after.push(format!(
189                        "CREATE INDEX {} ON {} ({columns})",
190                        q(&name),
191                        q(table)
192                    )),
193                }
194            }
195        }
196        for fk in &self.foreigns {
197            validate(&fk.column)?;
198            let Some((ref_table, ref_column)) = &fk.target else {
199                return Err(Error::Unsupported {
200                    operation: "a foreign key without references()",
201                    driver: driver.name(),
202                });
203            };
204            validate(ref_table)?;
205            validate(ref_column)?;
206            let name = name_of(std::slice::from_ref(&fk.column), "foreign");
207            let mut sql = format!(
208                "CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({})",
209                q(&name),
210                q(&fk.column),
211                q(ref_table),
212                q(ref_column)
213            );
214            if let Some(a) = fk.on_delete {
215                sql.push_str(&format!(" ON DELETE {}", a.sql()));
216            }
217            if let Some(a) = fk.on_update {
218                sql.push_str(&format!(" ON UPDATE {}", a.sql()));
219            }
220            parts.push(sql);
221        }
222        let mut out = vec![format!("CREATE TABLE {} ({})", q(table), parts.join(", "))];
223        out.extend(after);
224        Ok(out)
225    }
226}
227
228fn collect<I, S>(columns: I) -> Vec<String>
229where
230    I: IntoIterator<Item = S>,
231    S: AsRef<str>,
232{
233    columns.into_iter().map(|c| c.as_ref().to_owned()).collect()
234}