@@ -9,8 +9,19 @@ use serde::{Deserialize, Serialize};
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 {
@@ -87,52 +98,109 @@ impl Action for AddColumn {
8798 db.run(&query).context("failed to add column")?;
8899
89100 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− }
101+ if let Transformation::Simple(up) = up {
102+ let declarations: Vec<String> = table
103+ .columns
104+ .iter()
105+ .map(|column| {
106+ format!(
107+ "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
108+ table = table.real_name,
109+ alias = column.name,
110+ real_name = column.real_name,
111+ )
112+ })
113+ .collect();
131114
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")?;
115+ // Add triggers to fill in values as they are inserted/updated
116+ let query = format!(
117+ r#"
118+ CREATE OR REPLACE FUNCTION {trigger_name}()
119+ RETURNS TRIGGER AS $$
120+ BEGIN
121+ IF NOT reshape.is_new_schema() THEN
122+ DECLARE
123+ {declarations}
124+ BEGIN
125+ NEW.{temp_column_name} = {up};
126+ END;
127+ END IF;
128+ RETURN NEW;
129+ END
130+ $$ language 'plpgsql';
131+
132+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
133+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
134+ "#,
135+ temp_column_name = temp_column_name,
136+ trigger_name = self.trigger_name(ctx),
137+ up = up,
138+ table = self.table,
139+ declarations = declarations.join("\n"),
140+ );
141+ db.run(&query).context("failed to create up trigger")?;
142+
143+ // Backfill values in batches
144+ common::batch_touch_rows(db, &table.real_name, Some(&temp_column_name))
145+ .context("failed to batch update existing rows")?;
146+ }
147+
148+ if let Transformation::Update {
149+ table: from_table,
150+ value,
151+ r#where,
152+ } = up
153+ {
154+ let from_table = schema.get_table(db, &from_table)?;
155+
156+ let declarations: Vec<String> = from_table
157+ .columns
158+ .iter()
159+ .map(|column| {
160+ format!(
161+ "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
162+ table = from_table.real_name,
163+ alias = column.name,
164+ real_name = column.real_name,
165+ )
166+ })
167+ .collect();
168+
169+ // Add triggers to fill in values as they are inserted/updated
170+ let query = format!(
171+ r#"
172+ CREATE OR REPLACE FUNCTION {trigger_name}()
173+ RETURNS TRIGGER AS $$
174+ #variable_conflict use_variable
175+ BEGIN
176+ IF NOT reshape.is_new_schema() THEN
177+ DECLARE
178+ {declarations}
179+ BEGIN
180+ UPDATE public."{changed_table_real}"
181+ SET "{temp_column_name}" = {value}
182+ WHERE {where};
183+ END;
184+ END IF;
185+ RETURN NEW;
186+ END
187+ $$ language 'plpgsql';
188+
189+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{from_table_real}";
190+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
191+ "#,
192+ changed_table_real = table.real_name,
193+ from_table_real = from_table.real_name,
194+ trigger_name = self.trigger_name(ctx),
195+ declarations = declarations.join("\n"),
196+ temp_column_name = temp_column_name,
197+ );
198+ db.run(&query).context("failed to create up trigger")?;
199+
200+ // Backfill values in batches by touching the from table
201+ common::batch_touch_rows(db, &from_table.real_name, None)
202+ .context("failed to batch update existing rows")?;
203+ }
136204 }
137205
138206 // Add a temporary NOT NULL constraint if the column shouldn't be nullable.
@@ -167,10 +235,8 @@ impl Action for AddColumn {
167235 // Remove triggers and procedures
168236 let query = format!(
169237 r#"
170− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
171− DROP FUNCTION IF EXISTS "{trigger_name}";
238+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
172239 "#,
173− table = self.table,
174240 trigger_name = self.trigger_name(ctx),
175241 );
176242 transaction
@@ -262,10 +328,8 @@ impl Action for AddColumn {
262328 // Remove triggers and procedures
263329 let query = format!(
264330 r#"
265− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
266− DROP FUNCTION IF EXISTS "{trigger_name}";
331+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
267332 "#,
268− table = self.table,
269333 trigger_name = self.trigger_name(ctx),
270334 );
271335 db.run(&query).context("failed to drop up trigger")?;