@@ -3,25 +3,14 @@ use crate::{
33 db::{Conn, Transaction},
44 schema::Schema,
55};
6−use anyhow::{bail, Context};
6+use anyhow::Context;
77use serde::{Deserialize, Serialize};
88
99#[derive(Serialize, Deserialize, Debug)]
1010pub struct AddColumn {
1111 pub table: String,
12+ pub up: Option<String>,
1213 pub column: Column,
13− pub up: Option<Transformation>,
14−}
15−
16−#[derive(Serialize, Deserialize, Debug)]
17−#[serde(untagged)]
18−pub enum Transformation {
19− Simple(String),
20− Update {
21− table: String,
22− value: String,
23− r#where: String,
24− },
2514}
2615
2716impl AddColumn {
@@ -43,15 +32,6 @@ impl AddColumn {
4332 )
4433 }
4534
46− fn reverse_trigger_name(&self, ctx: &MigrationContext) -> String {
47− format!(
48− "{}_add_column_{}_{}_rev",
49− ctx.prefix(),
50− self.table,
51− self.column.name
52− )
53− }
54−
5535 fn not_null_constraint_name(&self, ctx: &MigrationContext) -> String {
5636 format!(
5737 "{}_add_column_not_null_{}_{}",
@@ -81,7 +61,7 @@ impl Action for AddColumn {
8161 let temp_column_name = self.temp_column_name(ctx);
8262
8363 let mut definition_parts = vec![
84− format!("\"{}\"", temp_column_name.to_string()),
64+ temp_column_name.to_string(),
8565 self.column.data_type.to_string(),
8666 ];
8767
@@ -106,189 +86,53 @@ impl Action for AddColumn {
10686 );
10787 db.run(&query).context("failed to add column")?;
10888
109− let declarations: Vec<String> = table
110− .columns
111− .iter()
112− .map(|column| {
113− format!(
114− "\"{alias}\" public.{table}.{real_name}%TYPE := NEW.{real_name};",
115− table = table.real_name,
116− alias = column.name,
117− real_name = column.real_name,
118− )
119− })
120− .collect();
121−
12289 if let Some(up) = &self.up {
123− if let Transformation::Simple(up) = up {
124− // Add triggers to fill in values as they are inserted/updated
125− let query = format!(
126− r#"
127− CREATE OR REPLACE FUNCTION {trigger_name}()
128− RETURNS TRIGGER AS $$
129− BEGIN
130− IF NOT reshape.is_new_schema() THEN
131− DECLARE
132− {declarations}
133− BEGIN
134− NEW."{temp_column_name}" = {up};
135− END;
136− END IF;
137− RETURN NEW;
138− END
139− $$ language 'plpgsql';
140−
141− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
142− CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
143− "#,
144− temp_column_name = temp_column_name,
145− trigger_name = self.trigger_name(ctx),
146− up = up,
147− table = self.table,
148− declarations = declarations.join("\n"),
149− );
150− db.run(&query).context("failed to create up trigger")?;
151−
152− // Backfill values in batches
153− common::batch_touch_rows(db, &table.real_name, Some(&temp_column_name))
154− .context("failed to batch update existing rows")?;
155− }
156−
157− if let Transformation::Update {
158− table: from_table,
159− value,
160− r#where,
161− } = up
162− {
163− let existing_schema_name = match &ctx.existing_schema_name {
164− Some(name) => name,
165− None => bail!("can't use update without previous migration"),
166− };
167−
168− let from_table = schema.get_table(db, &from_table)?;
169−
170− let from_table_assignments: Vec<String> = from_table
171− .columns
172− .iter()
173− .map(|column| {
174− format!(
175− "{table}.{alias} = NEW.{real_name};",
176− table = from_table.name,
177− alias = column.name,
178− real_name = column.real_name,
179− )
180− })
181− .collect();
182−
183− // Add triggers to fill in values as they are inserted/updated
184− let query = format!(
185− r#"
186− CREATE OR REPLACE FUNCTION {trigger_name}()
187− RETURNS TRIGGER AS $$
188− #variable_conflict use_variable
189− BEGIN
190− IF NOT reshape.is_new_schema() THEN
191− DECLARE
192− {from_table} migration_{existing_schema_name}.{from_table}%ROWTYPE;
193− BEGIN
194− {assignments}
195−
196− -- Don't trigger reverse trigger when making this update
197− perform set_config('reshape.disable_triggers', 'TRUE', TRUE);
198−
199− UPDATE public."{changed_table_real}"
200− SET "{temp_column_name}" = {value}
201− WHERE {where};
202−
203− perform set_config('reshape.disable_triggers', '', TRUE);
204− END;
205− END IF;
206− RETURN NEW;
207− END
208− $$ language 'plpgsql';
209−
210− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{from_table_real}";
211− CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
212− "#,
213− assignments = from_table_assignments.join("\n"),
214− changed_table_real = table.real_name,
215− from_table = from_table.name,
216− from_table_real = from_table.real_name,
217− trigger_name = self.trigger_name(ctx),
218− // declarations = from_table_declarations.join("\n"),
219− temp_column_name = temp_column_name,
220− );
221− db.run(&query).context("failed to create up trigger")?;
222−
223− let from_table_columns = from_table
224− .columns
225− .iter()
226− .map(|column| format!("{} as {}", column.real_name, column.name))
227− .collect::<Vec<String>>()
228− .join(", ");
229−
230− let changed_table_assignments: Vec<String> = table
231− .columns
232− .iter()
233− .map(|column| {
234− format!(
235− "{table}.{alias} := NEW.{real_name};",
236− table = table.name,
237− alias = column.name,
238− real_name = column.real_name,
239− )
240− })
241− .collect();
242−
243− // Add triggers to fill in values as they are inserted/updated
244− let query = format!(
245− r#"
246− CREATE OR REPLACE FUNCTION {trigger_name}()
247− RETURNS TRIGGER AS $$
248− #variable_conflict use_variable
249− BEGIN
250− IF NOT reshape.is_new_schema() AND NOT current_setting('reshape.disable_triggers', TRUE) = 'TRUE' THEN
251− DECLARE
252− {changed_table} migration_{existing_schema_name}.{changed_table}%ROWTYPE;
253− __temp_row migration_{existing_schema_name}.{from_table}%ROWTYPE;
254− BEGIN
255− {changed_table_assignments}
256−
257− SELECT {from_table_columns}
258− INTO __temp_row
259− FROM migration_{existing_schema_name}.{from_table} {from_table}
260− WHERE {where};
261−
262− DECLARE
263− {from_table} migration_{existing_schema_name}.{from_table}%ROWTYPE;
264− BEGIN
265− {from_table} = __temp_row;
266− NEW.{temp_column_name} = {value};
267− END;
268− END;
269− END IF;
270− RETURN NEW;
271− END
272− $$ language 'plpgsql';
273−
274− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{changed_table_real}";
275− CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{changed_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
276− "#,
277− changed_table_assignments = changed_table_assignments.join("\n"),
278− changed_table_real = table.real_name,
279− changed_table = table.name,
280− from_table = from_table.name,
281− trigger_name = self.reverse_trigger_name(ctx),
282− temp_column_name = temp_column_name,
283− // declarations = declarations.join("\n"),
284− );
285− db.run(&query)
286− .context("failed to create reverse up trigger")?;
90+ let declarations: Vec<String> = table
91+ .columns
92+ .iter()
93+ .map(|column| {
94+ format!(
95+ "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
96+ table = table.real_name,
97+ alias = column.name,
98+ real_name = column.real_name,
99+ )
100+ })
101+ .collect();
102+
103+ // Add triggers to fill in values as they are inserted/updated
104+ let query = format!(
105+ r#"
106+ CREATE OR REPLACE FUNCTION {trigger_name}()
107+ RETURNS TRIGGER AS $$
108+ BEGIN
109+ IF NOT reshape.is_new_schema() THEN
110+ DECLARE
111+ {declarations}
112+ BEGIN
113+ NEW.{temp_column_name} = {up};
114+ END;
115+ END IF;
116+ RETURN NEW;
117+ END
118+ $$ language 'plpgsql';
119+
120+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
121+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
122+ "#,
123+ temp_column_name = temp_column_name,
124+ trigger_name = self.trigger_name(ctx),
125+ up = up,
126+ table = self.table,
127+ declarations = declarations.join("\n"),
128+ );
129+ db.run(&query).context("failed to create up trigger")?;
130+ }
287131
288− // Backfill values in batches by touching the from table
289− common::batch_touch_rows(db, &from_table.real_name, None)
290− .context("failed to batch update existing rows")?;
291− }
132+ // Backfill values in batches
133+ if self.up.is_some() {
134+ common::batch_touch_rows(db, &table.real_name, &temp_column_name)
135+ .context("failed to batch update existing rows")?;
292136 }
293137
294138 // Add a temporary NOT NULL constraint if the column shouldn't be nullable.
@@ -298,10 +142,10 @@ impl Action for AddColumn {
298142 if !self.column.nullable {
299143 let query = format!(
300144 r#"
301− ALTER TABLE "{table}"
302− ADD CONSTRAINT "{constraint_name}"
303− CHECK ("{column}" IS NOT NULL) NOT VALID
304− "#,
145+ ALTER TABLE "{table}"
146+ ADD CONSTRAINT "{constraint_name}"
147+ CHECK ("{column}" IS NOT NULL) NOT VALID
148+ "#,
305149 table = self.table,
306150 constraint_name = self.not_null_constraint_name(ctx),
307151 column = temp_column_name,
@@ -323,11 +167,11 @@ impl Action for AddColumn {
323167 // Remove triggers and procedures
324168 let query = format!(
325169 r#"
326− DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
327− DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
170+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
171+ DROP FUNCTION IF EXISTS "{trigger_name}";
328172 "#,
173+ table = self.table,
329174 trigger_name = self.trigger_name(ctx),
330− reverse_trigger_name = self.reverse_trigger_name(ctx),
331175 );
332176 transaction
333177 .run(&query)
@@ -418,11 +262,11 @@ impl Action for AddColumn {
418262 // Remove triggers and procedures
419263 let query = format!(
420264 r#"
421− DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
422− DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
265+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
266+ DROP FUNCTION IF EXISTS "{trigger_name}";
423267 "#,
268+ table = self.table,
424269 trigger_name = self.trigger_name(ctx),
425− reverse_trigger_name = self.reverse_trigger_name(ctx),
426270 );
427271 db.run(&query).context("failed to drop up trigger")?;
428272