git.iko.soy/reshape
on main help sign in
Add up option to create_table
6034dfefabianlindfors
src/migrations/create_table.rs+121−8
compare with main
unifiedsplit
@@ -1,6 +1,9 @@
1+use std::collections::HashMap;
2+
1use super::{common::ForeignKey, Action, Column, MigrationContext};
3use super::{common::ForeignKey, Action, Column, MigrationContext};
2use crate::{
4use crate::{
3 db::{Conn, Transaction},
5 db::{Conn, Transaction},
6+ migrations::common,
4 schema::Schema,
7 schema::Schema,
5};
8};
6use anyhow::Context;
9use anyhow::Context;
@@ -14,6 +17,21 @@ pub struct CreateTable {
14
17
15 #[serde(default)]
18 #[serde(default)]
16 pub foreign_keys: Vec<ForeignKey>,
19 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+ }
17}
35}
18
36
19#[typetag::serde(name = "create_table")]
37#[typetag::serde(name = "create_table")]
@@ -24,7 +42,7 @@ impl Action for CreateTable {
24
42
25 fn run(
43 fn run(
26 &self,
44 &self,
27− _ctx: &MigrationContext,
45+ ctx: &MigrationContext,
28 db: &mut dyn Conn,
46 db: &mut dyn Conn,
29 schema: &Schema,
47 schema: &Schema,
30 ) -> anyhow::Result<()> {
48 ) -> anyhow::Result<()> {
@@ -85,7 +103,7 @@ impl Action for CreateTable {
85 ));
103 ));
86 }
104 }
87
105
88− db.run(&format!(
106+ let query = &format!(
89 r#"
107 r#"
90 CREATE TABLE "{name}" (
108 CREATE TABLE "{name}" (
91 {definition}
109 {definition}
@@ -93,23 +111,118 @@ impl Action for CreateTable {
93 "#,
111 "#,
94 name = self.name,
112 name = self.name,
95 definition = definition_rows.join(",\n"),
113 definition = definition_rows.join(",\n"),
96− ))
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+ );
97− .context("failed to create table")?;
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+
98 Ok(())
194 Ok(())
99 }
195 }
100
196
101 fn complete<'a>(
197 fn complete<'a>(
102 &self,
198 &self,
103− _ctx: &MigrationContext,
199+ ctx: &MigrationContext,
104− _db: &'a mut dyn Conn,
200+ db: &'a mut dyn Conn,
105 ) -> anyhow::Result<Option<Transaction<'a>>> {
201 ) -> 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+
107 Ok(None)
211 Ok(None)
108 }
212 }
109
213
110 fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
214 fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
111
215
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+
113 db.run(&format!(
226 db.run(&format!(
114 r#"
227 r#"
115 DROP TABLE IF EXISTS "{name}"
228 DROP TABLE IF EXISTS "{name}"
src/migrations/remove_column.rs+3
@@ -122,6 +122,9 @@ impl Action for RemoveColumn {
122 })
122 })
123 .collect();
123 .collect();
124
124
125+ // TODO: If column is NOT NULL, remove the constraint and perform a NULL check in some other way.
126+ // Otherwise, it's not possible to update the NOT NULL column from another table even in the same transaction.
127+ // Either make the NULL check in here or maybe use a constraint trigger: https://www.postgresql.org/docs/9.0/sql-createconstraint.html.
125 let query = format!(
128 let query = format!(
126 r#"
129 r#"
127 CREATE OR REPLACE FUNCTION {trigger_name}()
130 CREATE OR REPLACE FUNCTION {trigger_name}()
tests/complex.rsnew file+210
@@ -0,0 +1,210 @@
1+mod common;
2+use common::Test;
3+
4+#[test]
5+fn extract_relation_into_new_table() {
6+ let mut test = Test::new("Extract relation into new table");
7+
8+ test.first_migration(
9+ r#"
10+ name = "create_tables"
11+
12+ [[actions]]
13+ type = "create_table"
14+ name = "accounts"
15+ primary_key = ["id"]
16+
17+ [[actions.columns]]
18+ name = "id"
19+ type = "INTEGER"
20+
21+ [[actions]]
22+ type = "create_table"
23+ name = "users"
24+ primary_key = ["id"]
25+
26+ [[actions.columns]]
27+ name = "id"
28+ type = "INTEGER"
29+
30+ [[actions.columns]]
31+ name = "account_id"
32+ type = "INTEGER"
33+
34+ [[actions.columns]]
35+ name = "account_role"
36+ type = "TEXT"
37+ "#,
38+ );
39+
40+ test.second_migration(
41+ r#"
42+ name = "add_account_user_connection"
43+
44+ [[actions]]
45+ type = "create_table"
46+ name = "user_account_connections"
47+ primary_key = ["account_id", "user_id"]
48+
49+ [[actions.columns]]
50+ name = "account_id"
51+ type = "INTEGER"
52+
53+ [[actions.columns]]
54+ name = "user_id"
55+ type = "INTEGER"
56+
57+ [[actions.columns]]
58+ name = "role"
59+ type = "TEXT"
60+
61+ [actions.up]
62+ table = "users"
63+ values = { user_id = "id", account_id = "account_id", role = "UPPER(account_role)" }
64+ where = "user_id = id"
65+
66+ [[actions]]
67+ type = "remove_column"
68+ table = "users"
69+ column = "account_id"
70+
71+ [actions.down]
72+ table = "user_account_connections"
73+ value = "account_id"
74+ where = "id = user_id"
75+
76+ [[actions]]
77+ type = "remove_column"
78+ table = "users"
79+ column = "account_role"
80+
81+ [actions.down]
82+ table = "user_account_connections"
83+ value = "LOWER(role)"
84+ where = "id = user_id"
85+ "#,
86+ );
87+
88+ test.after_first(|db| {
89+ db.simple_query("INSERT INTO accounts (id) VALUES (1)")
90+ .unwrap();
91+ db.simple_query("INSERT INTO users (id, account_id, account_role) VALUES (1, 1, 'admin')")
92+ .unwrap();
93+ });
94+
95+ test.intermediate(|old_db, new_db| {
96+ // Ensure connections was backfilled
97+ let rows: Vec<(i32, i32, String)> = new_db
98+ .query(
99+ "
100+ SELECT account_id, user_id, role
101+ FROM user_account_connections
102+ ",
103+ &[],
104+ )
105+ .unwrap()
106+ .iter()
107+ .map(|row| (row.get("account_id"), row.get("user_id"), row.get("role")))
108+ .collect();
109+ assert_eq!(1, rows.len());
110+
111+ let row = rows.first().unwrap();
112+ assert_eq!(1, row.0);
113+ assert_eq!(1, row.1);
114+ assert_eq!("ADMIN", row.2);
115+
116+ // Ensure inserted user in old schema creates a new connection
117+ old_db
118+ .simple_query(
119+ "INSERT INTO users (id, account_id, account_role) VALUES (2, 1, 'developer')",
120+ )
121+ .unwrap();
122+ assert!(
123+ new_db
124+ .query(
125+ "
126+ SELECT account_id, user_id, role
127+ FROM user_account_connections
128+ WHERE account_id = 1 AND user_id = 2 AND role = 'DEVELOPER'
129+ ",
130+ &[],
131+ )
132+ .unwrap()
133+ .len()
134+ == 1
135+ );
136+
137+ // Ensure updated user role in old schema updates connection in new schema
138+ old_db
139+ .simple_query("UPDATE users SET account_role = 'admin' WHERE id = 2")
140+ .unwrap();
141+ assert!(
142+ new_db
143+ .query(
144+ "
145+ SELECT account_id, user_id, role
146+ FROM user_account_connections
147+ WHERE account_id = 1 AND user_id = 2 AND role = 'ADMIN'
148+ ",
149+ &[],
150+ )
151+ .unwrap()
152+ .len()
153+ == 1
154+ );
155+
156+ // Ensure updated connection in new schema updates old schema user
157+ new_db
158+ .simple_query(
159+ "UPDATE user_account_connections SET role = 'DEVELOPER' WHERE account_id = 1 AND user_id = 2",
160+ )
161+ .unwrap();
162+ assert!(
163+ old_db
164+ .query(
165+ "
166+ SELECT id
167+ FROM users
168+ WHERE id = 2 AND account_id = 1 AND account_role = 'developer'
169+ ",
170+ &[],
171+ )
172+ .unwrap()
173+ .len()
174+ == 1
175+ );
176+
177+ // Ensure insert of user with connection through new schema updates user in old schema
178+ new_db
179+ .simple_query(
180+ r#"
181+ BEGIN;
182+ INSERT INTO users (id) VALUES (3);
183+ INSERT INTO user_account_connections (user_id, account_id, role) VALUES (3, 1, 'DEVELOPER');
184+ COMMIT;
185+ "#,
186+ )
187+ .unwrap();
188+ new_db
189+ .simple_query(
190+ "",
191+ )
192+ .unwrap();
193+ assert!(
194+ old_db
195+ .query(
196+ "
197+ SELECT id
198+ FROM users
199+ WHERE id = 3 AND account_id = 1 AND account_role = 'developer'
200+ ",
201+ &[],
202+ )
203+ .unwrap()
204+ .len()
205+ == 1
206+ );
207+ });
208+
209+ test.run();
210+}