on main help sign in
Refactor Schema module
40ba58bfabianlindfors
src/migrations/alter_column.rs+10−3
@@ -303,14 +303,21 @@ impl Action for AlterColumn {
303 // Instead, we rename the schema column but point it to the old column
303 // Instead, we rename the schema column but point it to the old column
304 if self.can_short_circuit() {
304 if self.can_short_circuit() {
305 if let Some(new_name) = &self.changes.name {
305 if let Some(new_name) = &self.changes.name {
306− schema.set_column_alias(&self.table, &self.column, new_name);
306+ schema.change_table(&self.table, |table_changes| {
307− println!("Schema is now: {:?}", schema);
307+ table_changes.change_column(&self.column, |column_changes| {
308+ column_changes.set_name(new_name);
309+ });
310+ });
308 }
311 }
309
312
310 return Ok(());
313 return Ok(());
311 }
314 }
312
315
316+ schema.change_table(&self.table, |table_changes| {
317+ table_changes.change_column(&self.column, |column_changes| {
313− schema.set_column_alias(&self.table, &self.temporary_column_name(ctx), &self.column);
318+ column_changes.set_column(&self.temporary_column_name(ctx));
319+ });
320+ });
314
321
315 Ok(())
322 Ok(())
316 }
323 }
src/migrations/remove_column.rs+5−1
@@ -98,7 +98,11 @@ impl Action for RemoveColumn {
98 }
98 }
99
99
100 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
100 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
101− schema.set_column_hidden(&self.table, &self.column);
101+ schema.change_table(&self.table, |table_changes| {
102+ table_changes.change_column(&self.column, |column_changes| {
103+ column_changes.set_removed();
104+ })
105+ });
102
106
103 Ok(())
107 Ok(())
104 }
108 }
src/migrations/remove_table.rs+3−1
@@ -31,7 +31,9 @@ impl Action for RemoveTable {
31 }
31 }
32
32
33 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
33 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
34− schema.set_table_hidden(&self.table);
34+ schema.change_table(&self.table, |table_changes| {
35+ table_changes.set_removed();
36+ });
35
37
36 Ok(())
38 Ok(())
37 }
39 }
src/migrations/rename_table.rs+3−1
@@ -34,7 +34,9 @@ impl Action for RenameTable {
34 }
34 }
35
35
36 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
36 fn update_schema(&self, _ctx: &Context, schema: &mut Schema) -> anyhow::Result<()> {
37− schema.set_table_alias(&self.table, &self.new_name);
37+ schema.change_table(&self.table, |table_changes| {
38+ table_changes.set_name(&self.new_name);
39+ });
38
40
39 Ok(())
41 Ok(())
40 }
42 }
src/schema.rs+164−140
@@ -1,144 +1,136 @@
1use crate::db::Conn;
1use crate::db::Conn;
2−use anyhow::anyhow;
3−use bimap::BiMap;
4−use std::collections::{hash_map::Entry, HashMap, HashSet};
2+use std::collections::{HashMap, HashSet};
5−
3+
4+// Schema tracks changes made to tables and columns during a migration.
5+// These changes are not applied until the migration is completed but
6+// need to be taken into consideration when creating views for a migration
7+// and when a user references a table or column in a migration.
8+//
9+// The changes to a table are tracked by a `TableChanges` struct. The possible
10+// changes are:
11+// - Changing the name which updates `current_name`.
12+// - Removing which sets the `removed` flag.
13+//
14+// Changes to a column are tracked by a `ColumnChanges` struct which reside in
15+// the corresponding `TableChanges`. The possible changes are:
16+// - Changing the name which updates `current_name`.
17+// - Changing the backing column which will add the new column to the end of
18+// `intermediate_columns`. This is used when temporary columns are
19+// introduced which will eventually replace the current column.
20+// - Removing which sets the `removed` flag.
21+//
22+// Schema provides some schema introspection methods, `get_tables` and `get_table`,
23+// which will retrieve the current schema from the database and apply the changes.
6#[derive(Debug)]
24#[derive(Debug)]
7pub struct Schema {
25pub struct Schema {
8− table_alias_to_name: BiMap<String, String>,
26+ table_changes: Vec<TableChanges>,
9− hidden_tables: HashSet<String>,
10− table_schemas_by_name: HashMap<String, TableSchema>,
11}
27}
12
28
13impl Schema {
29impl Schema {
14 pub fn new() -> Schema {
30 pub fn new() -> Schema {
15 Schema {
31 Schema {
16− table_alias_to_name: BiMap::new(),
32+ table_changes: Vec::new(),
17− hidden_tables: HashSet::new(),
18− table_schemas_by_name: HashMap::new(),
19 }
33 }
20 }
34 }
21
35
22− pub fn set_table_alias(&mut self, current_name: &str, alias: &str) {
36+ pub fn change_table<F>(&mut self, current_name: &str, f: F)
23− self.table_alias_to_name
37+ where
24− .insert(alias.to_string(), current_name.to_string());
38+ F: FnOnce(&mut TableChanges),
25− }
39+ {
26−
40+ let table_change_index = self
27− pub fn set_table_hidden(&mut self, table: &str) {
41+ .table_changes
28− let real_name = self
42+ .iter()
29− .get_table_real_name(table)
43+ .position(|table| table.current_name == current_name)
30− .unwrap_or_else(|| table.to_string());
31− self.hidden_tables.insert(real_name);
32− }
33−
34− pub fn set_column_alias(&mut self, table: &str, current_name: &str, alias: &str) {
35− let table_real_name = self
36− .get_table_real_name(table)
37− .unwrap_or_else(|| table.to_string());
44+ .unwrap_or_else(|| {
38−
39− let column = match self.table_schemas_by_name.entry(table_real_name) {
40− Entry::Occupied(o) => o.into_mut(),
41− Entry::Vacant(v) => v.insert(TableSchema::new()),
42− };
43−
44− column
45− .column_alias_to_name
46− .insert(alias.to_string(), current_name.to_string());
45+ let new_changes = TableChanges::new(current_name.to_string());
47− }
46+ self.table_changes.push(new_changes);
48−
47+ self.table_changes.len() - 1
49− pub fn set_column_hidden(&mut self, table: &str, name: &str) {
50− let table_real_name = self
51− .get_table_real_name(table)
52− .unwrap_or_else(|| table.to_string());
53− let column = match self.table_schemas_by_name.entry(table_real_name) {
54− Entry::Occupied(o) => o.into_mut(),
55− Entry::Vacant(v) => v.insert(TableSchema::new()),
56− };
48+ });
57−
58− column.hidden_columns.insert(name.to_string());
59− }
60
49
61− fn get_table_real_name(&self, name: &str) -> Option<String> {
50+ let table_changes = &mut self.table_changes[table_change_index];
62− self.table_alias_to_name
51+ f(table_changes)
63− .get_by_left(name)
64− .map(|real_name| real_name.to_string())
65 }
52 }
53+}
66
54
67− fn get_table_alias_from_real_name<'a>(&'a self, real_name: &'a str) -> Option<String> {
55+#[derive(Debug)]
68− self.table_alias_to_name
56+pub struct TableChanges {
57+ current_name: String,
69− .get_by_right(real_name)
58+ real_name: String,
70− .map(|name| name.to_string())
59+ column_changes: Vec<ColumnChanges>,
60+ removed: bool,
71− }
61+}
72
62
73− fn is_table_hidden(&self, table: &str) -> bool {
63+impl TableChanges {
74− let table_real_name = self
64+ fn new(name: String) -> Self {
75− .get_table_real_name(table)
65+ Self {
76− .unwrap_or_else(|| table.to_string());
66+ current_name: name.to_string(),
77− self.hidden_tables.contains(&table_real_name)
67+ real_name: name.to_string(),
68+ column_changes: Vec::new(),
69+ removed: false,
70+ }
78 }
71 }
79
72
80− fn get_column_real_name<'a>(&'a self, table: &'a str, name: &'a str) -> Option<&'a str> {
73+ pub fn set_name(&mut self, name: &str) {
81− let table_real_name = self
74+ self.current_name = name.to_string();
82− .get_table_real_name(table)
83− .unwrap_or_else(|| table.to_string());
84−
85− self.table_schemas_by_name
86− .get(&table_real_name)
87− .map(|column| column.get_column_real_name(name))
88 }
75 }
89
76
90− fn get_column_alias_from_real_name(&self, table: &str, real_name: &str) -> Option<String> {
77+ pub fn change_column<F>(&mut self, current_name: &str, f: F)
91− let table_real_name = self
78+ where
92− .get_table_real_name(table)
79+ F: FnOnce(&mut ColumnChanges),
80+ {
81+ let column_change_index = self
82+ .column_changes
83+ .iter()
84+ .position(|column| column.current_name == current_name)
93− .unwrap_or_else(|| table.to_string());
85+ .unwrap_or_else(|| {
86+ let new_changes = ColumnChanges::new(current_name.to_string());
87+ self.column_changes.push(new_changes);
88+ self.column_changes.len() - 1
89+ });
94
90
95− self.table_schemas_by_name
91+ let column_changes = &mut self.column_changes[column_change_index];
96− .get(&table_real_name)
92+ f(column_changes)
97− .and_then(|column| column.get_column_alias_from_real_name(real_name))
98 }
93 }
99
94
100− fn is_column_hidden(&self, table: &str, name: &str) -> bool {
95+ pub fn set_removed(&mut self) {
101− let table_real_name = self
96+ self.removed = true;
102− .get_table_real_name(table)
103− .unwrap_or_else(|| table.to_string());
104− self.table_schemas_by_name
105− .get(&table_real_name)
106− .map(|column| {
107− let alias = column
108− .get_column_alias_from_real_name(name)
109− .unwrap_or_else(|| name.to_string());
110− column.hidden_columns.contains(&alias)
111− })
112− .unwrap_or(false)
113 }
97 }
114}
98}
115
99
116#[derive(Debug)]
100#[derive(Debug)]
117−struct TableSchema {
101+pub struct ColumnChanges {
118− column_alias_to_name: BiMap<String, String>,
102+ current_name: String,
119− hidden_columns: HashSet<String>,
103+ original_name: String,
104+ intermediate_names: Vec<String>,
105+ removed: bool,
120}
106}
121
107
122−impl TableSchema {
108+impl ColumnChanges {
123− fn new() -> TableSchema {
109+ fn new(name: String) -> Self {
124− TableSchema {
110+ Self {
125− column_alias_to_name: BiMap::new(),
111+ current_name: name.to_string(),
126− hidden_columns: HashSet::new(),
112+ original_name: name.to_string(),
113+ intermediate_names: Vec::new(),
114+ removed: false,
127 }
115 }
128 }
116 }
129
117
130− fn get_column_real_name<'a>(&'a self, name: &'a str) -> &'a str {
118+ pub fn set_name(&mut self, name: &str) {
131− if let Some(real_name) = self.column_alias_to_name.get_by_left(name) {
119+ self.current_name = name.to_string();
132− real_name
133− } else {
134− name
135− }
136 }
120 }
137
121
138− fn get_column_alias_from_real_name(&self, real_name: &str) -> Option<String> {
122+ pub fn set_column(&mut self, column_name: &str) {
139− self.column_alias_to_name
123+ self.intermediate_names.push(column_name.to_string())
140− .get_by_right(real_name)
124+ }
141− .map(|name| name.to_string())
125+
126+ pub fn set_removed(&mut self) {
127+ self.removed = true;
128+ }
129+
130+ fn real_name(&self) -> &str {
131+ self.intermediate_names
132+ .last()
133+ .unwrap_or_else(|| &self.original_name)
142 }
134 }
143}
135}
144
136
@@ -168,26 +160,46 @@ impl Schema {
168 )?
160 )?
169 .iter()
161 .iter()
170 .map(|row| row.get::<'_, _, String>("table_name"))
162 .map(|row| row.get::<'_, _, String>("table_name"))
171− .filter(|real_name| !self.is_table_hidden(real_name))
163+ .filter_map(|real_name| {
164+ let table_changes = self
165+ .table_changes
166+ .iter()
167+ .find(|changes| changes.real_name == real_name);
168+
169+ // Skip table if it has been removed
170+ if let Some(changes) = table_changes {
171+ if changes.removed {
172+ return None;
173+ }
174+ }
175+
172− .map(|real_name| self.get_table_by_real_name(db, &real_name))
176+ Some(self.get_table_by_real_name(db, &real_name))
177+ })
173 .collect()
178 .collect()
174 }
179 }
175
180
176 pub fn get_table(&self, db: &mut dyn Conn, table_name: &str) -> anyhow::Result<Table> {
181 pub fn get_table(&self, db: &mut dyn Conn, table_name: &str) -> anyhow::Result<Table> {
177− let real_table_name = self
182+ let table_changes = self
178− .get_table_real_name(table_name)
183+ .table_changes
184+ .iter()
185+ .find(|changes| changes.current_name == table_name);
186+
187+ let real_table_name = table_changes
188+ .map(|changes| changes.real_name.to_string())
179 .unwrap_or_else(|| table_name.to_string());
189 .unwrap_or_else(|| table_name.to_string());
190+
180 self.get_table_by_real_name(db, &real_table_name)
191 self.get_table_by_real_name(db, &real_table_name)
181 }
192 }
182
193
183− pub fn get_table_by_real_name(
194+ fn get_table_by_real_name(
184 &self,
195 &self,
185 db: &mut dyn Conn,
196 db: &mut dyn Conn,
186 real_table_name: &str,
197 real_table_name: &str,
187 ) -> anyhow::Result<Table> {
198 ) -> anyhow::Result<Table> {
188− if self.is_table_hidden(real_table_name) {
199+ let table_changes = self
189− return Err(anyhow!("no table named {}", real_table_name));
200+ .table_changes
190− }
201+ .iter()
202+ .find(|changes| changes.real_name == real_table_name);
191
203
192 let real_columns: Vec<(String, String, bool)> = db
204 let real_columns: Vec<(String, String, bool)> = db
193 .query(&format!(
205 .query(&format!(
@@ -209,27 +221,43 @@ impl Schema {
209 })
221 })
210 .collect();
222 .collect();
211
223
212− let mut columns: Vec<Column> = Vec::new();
224+ let mut ignore_columns: HashSet<String> = HashSet::new();
213−
225+ let mut aliases: HashMap<String, &str> = HashMap::new();
214− for (column_name, data_type, nullable) in real_columns {
226+
215− if self.is_column_hidden(real_table_name, &column_name) {
227+ if let Some(changes) = table_changes {
216− continue;
228+ for column_changes in &changes.column_changes {
229+ if column_changes.removed {
230+ ignore_columns.insert(column_changes.real_name().to_string());
231+ } else {
232+ aliases.insert(
233+ column_changes.real_name().to_string(),
234+ &column_changes.current_name,
235+ );
236+ }
237+
238+ if !column_changes.intermediate_names.is_empty() {
239+ ignore_columns.insert(column_changes.original_name.to_string());
240+
241+ let non_used_intermediate = &column_changes.intermediate_names
242+ [..column_changes.intermediate_names.len() - 1];
243+ for intermediate in non_used_intermediate {
244+ ignore_columns.insert(intermediate.to_string());
245+ }
246+ }
217 }
247 }
248+ }
249+
250+ let mut columns: Vec<Column> = Vec::new();
218
251
219− if is_column_temporary(&column_name) {
252+ for (real_name, data_type, nullable) in real_columns {
253+ if ignore_columns.contains(&*real_name) {
220 continue;
254 continue;
221 }
255 }
222
256
223− let (name, real_name) = if let Some(alias) =
257+ let name = aliases
224− self.get_column_alias_from_real_name(real_table_name, &column_name)
258+ .get(&real_name)
225− {
259+ .map(|alias| alias.to_string())
226− (alias, column_name)
227− } else {
228− let real_name = self
229− .get_column_real_name(real_table_name, &column_name)
230− .unwrap_or(&column_name);
231− (column_name.to_string(), real_name.to_string())
260+ .unwrap_or_else(|| real_name.to_string());
232− };
233
261
234 columns.push(Column {
262 columns.push(Column {
235 name,
263 name,
@@ -239,12 +267,12 @@ impl Schema {
239 });
267 });
240 }
268 }
241
269
242− let table_name = self
270+ let current_table_name = table_changes
243− .get_table_alias_from_real_name(real_table_name)
271+ .map(|changes| changes.current_name.as_ref())
244− .unwrap_or_else(|| real_table_name.to_string());
272+ .unwrap_or_else(|| real_table_name);
245
273
246 let table = Table {
274 let table = Table {
247− name: table_name,
275+ name: current_table_name.to_string(),
248 real_name: real_table_name.to_string(),
276 real_name: real_table_name.to_string(),
249 columns: columns,
277 columns: columns,
250 };
278 };
@@ -252,7 +280,3 @@ impl Schema {
252 Ok(table)
280 Ok(table)
253 }
281 }
254}
282}
255−
256−fn is_column_temporary(name: &str) -> bool {
257− name.starts_with("__reshape_")
258−}