Skip to content

Commit 3879a3c

Browse files
Julius de Bruijnspacekookie
Julius de Bruijn
authored and
spacekookie
committed
Allow adding foreign keys
1 parent 31cf039 commit 3879a3c

File tree

8 files changed

+193
-8
lines changed

8 files changed

+193
-8
lines changed

‎src/backend/mod.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,4 +96,12 @@ pub trait SqlGenerator {
9696

9797
/// Drop a multi-column index
9898
fn drop_index(name: &str) -> String;
99+
100+
/// Add a foreign key
101+
fn add_foreign_key(
102+
columns: &[String],
103+
table: &str,
104+
relation_columns: &[String],
105+
schema: Option<&str>,
106+
) -> String;
99107
}

‎src/backend/mssql.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,27 @@ impl SqlGenerator for MsSql {
152152
fn drop_index(name: &str) -> String {
153153
format!("DROP INDEX [{}]", name)
154154
}
155+
156+
fn add_foreign_key(
157+
columns: &[String],
158+
table: &str,
159+
relation_columns: &[String],
160+
schema: Option<&str>,
161+
) -> String {
162+
let columns: Vec<_> = columns.into_iter().map(|c| format!("[{}]", c)).collect();
163+
let relation_columns: Vec<_> = relation_columns
164+
.into_iter()
165+
.map(|c| format!("[{}]", c))
166+
.collect();
167+
168+
format!(
169+
"FOREIGN KEY({}) REFERENCES {}[{}]({})",
170+
columns.join(","),
171+
prefix!(schema),
172+
table,
173+
relation_columns.join(","),
174+
)
175+
}
155176
}
156177

157178
impl MsSql {

‎src/backend/mysql.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,27 @@ impl SqlGenerator for MySql {
120120
fn drop_index(name: &str) -> String {
121121
format!("DROP INDEX `{}`", name)
122122
}
123+
124+
fn add_foreign_key(
125+
columns: &[String],
126+
table: &str,
127+
relation_columns: &[String],
128+
schema: Option<&str>,
129+
) -> String {
130+
let columns: Vec<_> = columns.into_iter().map(|c| format!("`{}`", c)).collect();
131+
let relation_columns: Vec<_> = relation_columns
132+
.into_iter()
133+
.map(|c| format!("`{}`", c))
134+
.collect();
135+
136+
format!(
137+
"FOREIGN KEY({}) REFERENCES {}`{}`({})",
138+
columns.join(","),
139+
prefix!(schema),
140+
table,
141+
relation_columns.join(","),
142+
)
143+
}
123144
}
124145

125146
impl MySql {

‎src/backend/pg.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,28 @@ impl SqlGenerator for Pg {
122122
fn drop_index(name: &str) -> String {
123123
format!("DROP INDEX \"{}\"", name)
124124
}
125+
126+
fn add_foreign_key(
127+
columns: &[String],
128+
table: &str,
129+
relation_columns: &[String],
130+
schema: Option<&str>,
131+
) -> String {
132+
let columns: Vec<_> = columns.into_iter().map(|c| format!("\"{}\"", c)).collect();
133+
134+
let relation_columns: Vec<_> = relation_columns
135+
.into_iter()
136+
.map(|c| format!("\"{}\"", c))
137+
.collect();
138+
139+
format!(
140+
"FOREIGN KEY({}) REFERENCES {}\"{}\"({})",
141+
columns.join(","),
142+
prefix!(schema),
143+
table,
144+
relation_columns.join(","),
145+
)
146+
}
125147
}
126148

127149
impl Pg {

‎src/backend/sqlite3.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,27 @@ impl SqlGenerator for Sqlite {
119119
fn rename_column(_: &str, _: &str) -> String {
120120
panic!("Sqlite does not support renaming columns!")
121121
}
122+
123+
fn add_foreign_key(
124+
columns: &[String],
125+
table: &str,
126+
relation_columns: &[String],
127+
_: Option<&str>,
128+
) -> String {
129+
let columns: Vec<_> = columns.into_iter().map(|c| format!("\"{}\"", c)).collect();
130+
131+
let relation_columns: Vec<_> = relation_columns
132+
.into_iter()
133+
.map(|c| format!("\"{}\"", c))
134+
.collect();
135+
136+
format!(
137+
"FOREIGN KEY({}) REFERENCES \"{}\"({})",
138+
columns.join(","),
139+
table,
140+
relation_columns.join(","),
141+
)
142+
}
122143
}
123144

124145
impl Sqlite {

‎src/lib.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,3 +174,14 @@ pub enum IndexChange {
174174
/// Remove a multi-column index
175175
RemoveIndex(String, String),
176176
}
177+
178+
/// An enum set that represents operations done with and on foreign keys
179+
#[derive(Clone)]
180+
pub enum ForeignKeyChange {
181+
/// Add a foreign key
182+
AddForeignKey {
183+
columns: Vec<String>,
184+
table: String,
185+
relation_columns: Vec<String>,
186+
},
187+
}

‎src/migration.rs

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ impl Migration {
6060
&mut CreateTable(ref mut t, ref mut cb)
6161
| &mut CreateTableIfNotExists(ref mut t, ref mut cb) => {
6262
cb(t); // Run the user code
63-
let (cols, indices) = t.make::<T>(false, schema);
63+
let (cols, indices, foreign_keys) = t.make::<T>(false, schema);
6464

6565
let name = t.meta.name().clone();
6666
sql.push_str(&match change {
@@ -79,6 +79,20 @@ impl Migration {
7979
sql.push_str(", ");
8080
}
8181
}
82+
83+
let l = foreign_keys.len();
84+
for (i, slice) in foreign_keys.iter().enumerate() {
85+
if cols.len() > 0 && i == 0 {
86+
sql.push_str(", ")
87+
}
88+
89+
sql.push_str(slice);
90+
91+
if i < l - 1 {
92+
sql.push_str(", ")
93+
}
94+
}
95+
8296
sql.push_str(")");
8397

8498
// Add additional index columns
@@ -96,9 +110,11 @@ impl Migration {
96110
}
97111
&mut ChangeTable(ref mut t, ref mut cb) => {
98112
cb(t);
99-
let (cols, indices) = t.make::<T>(true, schema);
113+
let (cols, indices, fks) = t.make::<T>(true, schema);
114+
100115
sql.push_str(&T::alter_table(&t.meta.name(), schema));
101116
sql.push_str(" ");
117+
102118
let l = cols.len();
103119
for (i, slice) in cols.iter().enumerate() {
104120
sql.push_str(slice);
@@ -108,6 +124,20 @@ impl Migration {
108124
}
109125
}
110126

127+
let l = fks.len();
128+
for (i, slice) in fks.iter().enumerate() {
129+
if cols.len() > 0 && i == 0 {
130+
sql.push_str(", ")
131+
}
132+
133+
sql.push_str("ADD ");
134+
sql.push_str(slice);
135+
136+
if i < l - 1 {
137+
sql.push_str(", ")
138+
}
139+
}
140+
111141
// Add additional index columns
112142
if indices.len() > 0 {
113143
sql.push_str(";");

‎src/table.rs

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@
77
//! You can also change existing tables with a closure that can
88
//! then access individual columns in that table.
99
10-
use super::backend::SqlGenerator;
11-
use super::{IndexChange, TableChange};
10+
use super::{backend::SqlGenerator, ForeignKeyChange, IndexChange, TableChange};
1211
use crate::types::Type;
1312
use std::fmt::{Debug, Formatter, Result as FmtResult};
1413

@@ -24,11 +23,18 @@ impl Debug for IndexChange {
2423
}
2524
}
2625

26+
impl Debug for ForeignKeyChange {
27+
fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
28+
f.write_str("ForeignKeyChange")
29+
}
30+
}
31+
2732
#[derive(Debug, Clone)]
2833
pub struct Table {
2934
pub meta: TableMeta,
3035
columns: Vec<TableChange>,
3136
indices: Vec<IndexChange>,
37+
foreign_keys: Vec<ForeignKeyChange>,
3238
}
3339

3440
impl Table {
@@ -37,6 +43,7 @@ impl Table {
3743
meta: TableMeta::new(name.into()),
3844
columns: vec![],
3945
indices: vec![],
46+
foreign_keys: vec![],
4047
}
4148
}
4249

@@ -103,18 +110,45 @@ impl Table {
103110
));
104111
}
105112

106-
/// Generate Sql for this table, returned as two vectors
113+
pub fn add_foreign_key(
114+
&mut self,
115+
columns_on_this_side: &[&str],
116+
related_table: &str,
117+
columns_on_that_side: &[&str],
118+
) {
119+
let table = related_table.into();
120+
121+
let columns = columns_on_this_side
122+
.into_iter()
123+
.map(|c| String::from(*c))
124+
.collect();
125+
126+
let relation_columns = columns_on_that_side
127+
.into_iter()
128+
.map(|c| String::from(*c))
129+
.collect();
130+
131+
self.foreign_keys.push(ForeignKeyChange::AddForeignKey {
132+
table,
133+
columns,
134+
relation_columns,
135+
})
136+
}
137+
138+
/// Generate Sql for this table, returned as three vectors
107139
///
108140
/// The first vector (`.0`) represents all column changes done to the table,
109141
/// the second vector (`.1`) contains all index and suffix changes.
142+
/// the third vector (`.2`) contains all foreign key changes.
110143
///
111-
/// It is very well possible for either of them to be empty,
144+
/// It is very well possible for all of them to be empty,
112145
/// although both being empty *might* signify an error.
113146
pub fn make<T: SqlGenerator>(
114147
&mut self,
115148
ex: bool,
116149
schema: Option<&str>,
117-
) -> (Vec<String>, Vec<String>) {
150+
) -> (Vec<String>, Vec<String>, Vec<String>) {
151+
use ForeignKeyChange as KFC;
118152
use IndexChange as IC;
119153
use TableChange as TC;
120154

@@ -143,7 +177,24 @@ impl Table {
143177
})
144178
.collect();
145179

146-
(columns, indeces)
180+
let foreign_keys = self
181+
.foreign_keys
182+
.iter()
183+
.map(|change| match change {
184+
KFC::AddForeignKey {
185+
columns,
186+
table,
187+
relation_columns,
188+
} => T::add_foreign_key(
189+
columns.as_slice(),
190+
table,
191+
relation_columns.as_slice(),
192+
schema,
193+
),
194+
})
195+
.collect();
196+
197+
(columns, indeces, foreign_keys)
147198
}
148199
}
149200

0 commit comments

Comments
 (0)