@@ -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 RemoveColumn {
1111 pub table: String,
1212 pub column: String,
13− pub down: Option<String>,
13+ pub down: 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 RemoveColumn {
@@ -39,10 +50,10 @@ impl Action for RemoveColumn {
3950 db: &mut dyn Conn,
4051 schema: &Schema,
4152 ) -> anyhow::Result<()> {
53+ let table = schema.get_table(db, &self.table)?;
54+
4255 // Add down trigger
4356 if let Some(down) = &self.down {
44− let table = schema.get_table(db, &self.table)?;
45−
4657 let declarations: Vec<String> = table
4758 .columns
4859 .iter()
@@ -56,32 +67,91 @@ impl Action for RemoveColumn {
5667 })
5768 .collect();
5869
59− let query = format!(
60− r#"
61− CREATE OR REPLACE FUNCTION {trigger_name}()
62− RETURNS TRIGGER AS $$
63− BEGIN
64− IF reshape.is_new_schema() THEN
65− DECLARE
66− {declarations}
67− BEGIN
68− NEW.{column_name} = {down};
69− END;
70− END IF;
71− RETURN NEW;
72− END
73− $$ language 'plpgsql';
74−
75− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
76− CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
77− "#,
78− column_name = self.column,
79− trigger_name = self.trigger_name(ctx),
80− down = down,
81− table = self.table,
82− declarations = declarations.join("\n"),
83− );
84− db.run(&query).context("failed to create down trigger")?;
70+ if let Transformation::Simple(down) = down {
71+ let query = format!(
72+ r#"
73+ CREATE OR REPLACE FUNCTION {trigger_name}()
74+ RETURNS TRIGGER AS $$
75+ BEGIN
76+ IF reshape.is_new_schema() THEN
77+ DECLARE
78+ {declarations}
79+ BEGIN
80+ NEW.{column_name} = {down};
81+ END;
82+ END IF;
83+ RETURN NEW;
84+ END
85+ $$ language 'plpgsql';
86+
87+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
88+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{table}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
89+ "#,
90+ column_name = self.column,
91+ trigger_name = self.trigger_name(ctx),
92+ down = down,
93+ table = self.table,
94+ declarations = declarations.join("\n"),
95+ );
96+ db.run(&query).context("failed to create down trigger")?;
97+ }
98+
99+ if let Transformation::Update {
100+ table: from_table,
101+ value,
102+ r#where,
103+ } = down
104+ {
105+ let existing_schema_name = match &ctx.existing_schema_name {
106+ Some(name) => name,
107+ None => bail!("can't use update without previous migration"),
108+ };
109+
110+ let from_table = schema.get_table(db, &from_table)?;
111+
112+ let declarations: Vec<String> = from_table
113+ .columns
114+ .iter()
115+ .map(|column| {
116+ format!(
117+ "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
118+ table = from_table.real_name,
119+ alias = column.name,
120+ real_name = column.real_name,
121+ )
122+ })
123+ .collect();
124+
125+ let query = format!(
126+ r#"
127+ CREATE OR REPLACE FUNCTION {trigger_name}()
128+ RETURNS TRIGGER AS $$
129+ #variable_conflict use_variable
130+ BEGIN
131+ IF reshape.is_new_schema() THEN
132+ DECLARE
133+ {declarations}
134+ BEGIN
135+ UPDATE "migration_{existing_schema_name}"."{changed_table}" "{changed_table}"
136+ SET "{column_name}" = {value}
137+ WHERE {where};
138+ END;
139+ END IF;
140+ RETURN NEW;
141+ END
142+ $$ language 'plpgsql';
143+
144+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{from_table_real}";
145+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
146+ "#,
147+ changed_table = self.table,
148+ from_table_real = from_table.real_name,
149+ column_name = self.column,
150+ trigger_name = self.trigger_name(ctx),
151+ declarations = declarations.join("\n"),
152+ );
153+ db.run(&query).context("failed to create down trigger")?;
154+ }
85155 }
86156
87157 Ok(())
@@ -106,12 +176,20 @@ impl Action for RemoveColumn {
106176 }
107177
108178 // Remove column, function and trigger
179+ let trigger_table = match &self.down {
180+ Some(Transformation::Update {
181+ table,
182+ value: _,
183+ r#where: _,
184+ }) => table,
185+ _ => &self.table,
186+ };
109187 let query = format!(
110188 r#"
111189 ALTER TABLE "{table}"
112190 DROP COLUMN IF EXISTS "{column}";
113191
114− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
192+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
115193 DROP FUNCTION IF EXISTS "{trigger_name}";
116194 "#,
117195 table = self.table,
@@ -134,12 +212,19 @@ impl Action for RemoveColumn {
134212
135213 fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
136214 // Remove function and trigger
215+ let trigger_table = match &self.down {
216+ Some(Transformation::Update {
217+ table,
218+ value: _,
219+ r#where: _,
220+ }) => table,
221+ _ => &self.table,
222+ };
137223 db.run(&format!(
138224 r#"
139− DROP TRIGGER IF EXISTS "{trigger_name}" ON "{table}";
225+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{trigger_table}";
140226 DROP FUNCTION IF EXISTS "{trigger_name}";
141227 "#,
142− table = self.table,
143228 trigger_name = self.trigger_name(ctx),
144229 ))
145230 .context("failed to drop down trigger")?;