@@ -1,6 +1,9 @@
1+use std::collections::HashMap;
2+
13use super::{common::ForeignKey, Action, Column, MigrationContext};
24use crate::{
35 db::{Conn, Transaction},
6+ migrations::common,
47 schema::Schema,
58};
69use anyhow::Context;
@@ -14,6 +17,21 @@ pub struct CreateTable {
1417
1518 #[serde(default)]
1619 pub foreign_keys: Vec<ForeignKey>,
20+
21+ pub up: Option<Transformation>,
22+}
23+
24+#[derive(Serialize, Deserialize, Debug)]
25+pub struct Transformation {
26+ table: String,
27+ values: HashMap<String, String>,
28+ upsert_constraint: Option<String>,
29+}
30+
31+impl CreateTable {
32+ fn trigger_name(&self, ctx: &MigrationContext) -> String {
33+ format!("{}_create_table_{}", ctx.prefix(), self.name)
34+ }
1735}
1836
1937#[typetag::serde(name = "create_table")]
@@ -24,7 +42,7 @@ impl Action for CreateTable {
2442
2543 fn run(
2644 &self,
27− _ctx: &MigrationContext,
45+ ctx: &MigrationContext,
2846 db: &mut dyn Conn,
2947 schema: &Schema,
3048 ) -> anyhow::Result<()> {
@@ -85,7 +103,7 @@ impl Action for CreateTable {
85103 ));
86104 }
87105
88− db.run(&format!(
106+ let query = &format!(
89107 r#"
90108 CREATE TABLE "{name}" (
91109 {definition}
@@ -93,23 +111,118 @@ impl Action for CreateTable {
93111 "#,
94112 name = self.name,
95113 definition = definition_rows.join(",\n"),
96− ))
97− .context("failed to create table")?;
114+ );
115+ db.run(query).context("failed to create table")?;
116+
117+ if let Some(Transformation {
118+ table: from_table,
119+ values,
120+ upsert_constraint,
121+ }) = &self.up
122+ {
123+ let from_table = schema.get_table(db, &from_table)?;
124+
125+ let declarations: Vec<String> = from_table
126+ .columns
127+ .iter()
128+ .map(|column| {
129+ format!(
130+ "{alias} public.{table}.{real_name}%TYPE := NEW.{real_name};",
131+ table = from_table.real_name,
132+ alias = column.name,
133+ real_name = column.real_name,
134+ )
135+ })
136+ .collect();
137+
138+ let (insert_columns, insert_values): (Vec<&str>, Vec<&str>) = values
139+ .iter()
140+ .map(|(k, v)| -> (&str, &str) { (k, v) }) // Force &String to &str
141+ .unzip();
142+
143+ let update_set: Vec<String> = values
144+ .iter()
145+ .map(|(field, value)| format!("\"{field}\" = {value}"))
146+ .collect();
147+
148+ // Constraint to check for conflicts. Defaults to the primary key constraint.
149+ let conflict_constraint_name = match upsert_constraint {
150+ Some(custom_constraint) => custom_constraint.clone(),
151+ _ => format!("{table}_pkey", table = self.name),
152+ };
153+
154+ // Add triggers to fill in values as they are inserted/updated
155+ let query = format!(
156+ r#"
157+ CREATE OR REPLACE FUNCTION {trigger_name}()
158+ RETURNS TRIGGER AS $$
159+ #variable_conflict use_variable
160+ BEGIN
161+ IF NOT reshape.is_new_schema() THEN
162+ DECLARE
163+ {declarations}
164+ BEGIN
165+ INSERT INTO public."{changed_table_real}" ({columns})
166+ VALUES ({values})
167+ ON CONFLICT ON CONSTRAINT "{conflict_constraint_name}"
168+ DO UPDATE SET
169+ {updates};
170+ END;
171+ END IF;
172+ RETURN NEW;
173+ END
174+ $$ language 'plpgsql';
175+
176+ DROP TRIGGER IF EXISTS "{trigger_name}" ON "{from_table_real}";
177+ CREATE TRIGGER "{trigger_name}" BEFORE UPDATE OR INSERT ON "{from_table_real}" FOR EACH ROW EXECUTE PROCEDURE {trigger_name}();
178+ "#,
179+ changed_table_real = self.name,
180+ from_table_real = from_table.real_name,
181+ trigger_name = self.trigger_name(ctx),
182+ declarations = declarations.join("\n"),
183+ columns = insert_columns.join(", "),
184+ values = insert_values.join(", "),
185+ updates = update_set.join(",\n"),
186+ );
187+ db.run(&query).context("failed to create up trigger")?;
188+
189+ // Backfill values in batches by touching the from table
190+ common::batch_touch_rows(db, &from_table.real_name, None)
191+ .context("failed to batch update existing rows")?;
192+ }
193+
98194 Ok(())
99195 }
100196
101197 fn complete<'a>(
102198 &self,
103− _ctx: &MigrationContext,
104− _db: &'a mut dyn Conn,
199+ ctx: &MigrationContext,
200+ db: &'a mut dyn Conn,
105201 ) -> anyhow::Result<Option<Transaction<'a>>> {
106− // Do nothing
202+ // Remove triggers and procedures
203+ let query = format!(
204+ r#"
205+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
206+ "#,
207+ trigger_name = self.trigger_name(ctx),
208+ );
209+ db.run(&query).context("failed to drop up trigger")?;
210+
107211 Ok(None)
108212 }
109213
110214 fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
111215
112− fn abort(&self, _ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
216+ fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
217+ // Remove triggers and procedures
218+ let query = format!(
219+ r#"
220+ DROP FUNCTION IF EXISTS "{trigger_name}" CASCADE;
221+ "#,
222+ trigger_name = self.trigger_name(ctx),
223+ );
224+ db.run(&query).context("failed to drop up trigger")?;
225+
113226 db.run(&format!(
114227 r#"
115228 DROP TABLE IF EXISTS "{name}"