@@ -3,14 +3,25 @@ use crate::{
33 db::{Conn, Transaction},
44 schema::Schema,
55};
6−use anyhow::Context;
6+use anyhow::{bail, Context};
77use serde::{Deserialize, Serialize};
88
99#[derive(Serialize, Deserialize, Debug)]
1010pub struct AddColumn {
1111 pub table: String,
12− pub up: Option<String>,
1312 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+ },
1425}
1526
1627impl AddColumn {
@@ -32,6 +43,15 @@ impl AddColumn {
3243 )
3344 }
3445
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+
3555 fn not_null_constraint_name(&self, ctx: &MigrationContext) -> String {
3656 format!(
3757 "{}_add_column_not_null_{}_{}",
@@ -61,7 +81,7 @@ impl Action for AddColumn {
6181 let temp_column_name = self.temp_column_name(ctx);
6282
6383 let mut definition_parts = vec![
64− temp_column_name.to_string(),
84+ format!("\"{}\"", temp_column_name.to_string()),
6585 self.column.data_type.to_string(),
6686 ];
6787
@@ -86,53 +106,189 @@ impl Action for AddColumn {
86106 );
87107 db.run(&query).context("failed to add column")?;
88108
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+
89122 if let Some(up) = &self.up {
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− }
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")?;
131151
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")?;
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")?;
287+
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+ }
136292 }
137293
138294 // Add a temporary NOT NULL constraint if the column shouldn't be nullable.
@@ -142,10 +298,10 @@ impl Action for AddColumn {
142298 if !self.column.nullable {
143299 let query = format!(
144300 r#"
145− ALTER TABLE "{table}"
146− ADD CONSTRAINT "{constraint_name}"
147− CHECK ("{column}" IS NOT NULL) NOT VALID
148− "#,
301+ ALTER TABLE "{table}"
302+ ADD CONSTRAINT "{constraint_name}"
303+ CHECK ("{column}" IS NOT NULL) NOT VALID
304+ "#,
149305 table = self.table,
150306 constraint_name = self.not_null_constraint_name(ctx),
151307 column = temp_column_name,
@@ -167,11 +323,11 @@ impl Action for AddColumn {
167323 // Remove triggers and procedures
168324 let query = format!(
169325 r#"
170− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
171− DROP FUNCTION IF EXISTS "{trigger_name}";
326+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
327+ DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
172328 "#,
173− table = self.table,
174329 trigger_name = self.trigger_name(ctx),
330+ reverse_trigger_name = self.reverse_trigger_name(ctx),
175331 );
176332 transaction
177333 .run(&query)
@@ -262,11 +418,11 @@ impl Action for AddColumn {
262418 // Remove triggers and procedures
263419 let query = format!(
264420 r#"
265− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
266− DROP FUNCTION IF EXISTS "{trigger_name}";
421+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
422+ DROP FUNCTION IF EXISTS "{reverse_trigger_name}" CASCADE;
267423 "#,
268− table = self.table,
269424 trigger_name = self.trigger_name(ctx),
425+ reverse_trigger_name = self.reverse_trigger_name(ctx),
270426 );
271427 db.run(&query).context("failed to drop up trigger")?;
272428