Add database constraints

butane_core/src/codegen/migration.rs

@@ -1,5 +1,8 @@
-use super::*;
-use crate::migrations::adb::{create_many_table, AColumn, ATable};
+use super::{
+    dbobj, fields, get_default, get_deferred_sql_type, get_many_sql_type, is_auto, is_foreign_key,
+    is_many_to_many, is_option, is_row_field, is_unique, pk_field,
+};
+use crate::migrations::adb::{create_many_table, AColumn, ARef, ATable, DeferredSqlType, TypeKey};
 use crate::migrations::{MigrationMut, MigrationsMut};
 use crate::Result;
 use syn::{Field, ItemStruct};
@@ -43,15 +46,20 @@ fn create_atables(ast_struct: &ItemStruct, config: &dbobj::Config) -> Vec<ATable
             .expect("db object fields must be named")
             .to_string();
         if is_row_field(f) {
-            let col = AColumn::new(
+            let deferred_type = get_deferred_sql_type(&f.ty);
+            let mut col = AColumn::new(
                 name,
-                get_deferred_sql_type(&f.ty),
+                deferred_type.clone(),
                 is_nullable(f),
                 f == &pk,
                 is_auto(f),
                 is_unique(f),
                 get_default(f).expect("Malformed default attribute"),
+                None,
             );
+            if is_foreign_key(f) {
+                col.add_reference(&ARef::Deferred(deferred_type))
+            }
             table.add_column(col);
         } else if is_many_to_many(f) {
             result.push(many_table(&table.name, f, &pk));
@@ -69,9 +77,20 @@ fn many_table(main_table_name: &str, many_field: &Field, pk_field: &Field) -> AT
         .to_string();
     let many_field_type = get_many_sql_type(many_field)
         .unwrap_or_else(|| panic!("Mis-identified Many field {field_name}"));
+    let pk_field_name = pk_field
+        .ident
+        .as_ref()
+        .expect("fields must be named")
+        .to_string();
     let pk_field_type = get_deferred_sql_type(&pk_field.ty);
 
-    create_many_table(main_table_name, &field_name, many_field_type, pk_field_type)
+    create_many_table(
+        main_table_name,
+        &field_name,
+        many_field_type,
+        &pk_field_name,
+        pk_field_type,
+    )
 }
 
 fn is_nullable(field: &Field) -> bool {

butane_core/src/codegen/mod.rs

@@ -289,6 +289,19 @@ fn get_option_sql_type(ty: &syn::Type) -> Option<DeferredSqlType> {
     })
 }
 
+fn get_foreign_key_sql_type(field: &Field) -> Option<DeferredSqlType> {
+    if let Some(inner_type_path) = get_type_argument(&field.ty, &OPTION_TYNAMES) {
+        let inner_ty: syn::Type = syn::TypePath {
+            qself: None,
+            path: inner_type_path.clone(),
+        }
+        .into();
+
+        return get_foreign_sql_type(&inner_ty, &FKEY_TYNAMES);
+    }
+    get_foreign_sql_type(&field.ty, &FKEY_TYNAMES)
+}
+
 fn get_many_sql_type(field: &Field) -> Option<DeferredSqlType> {
     get_foreign_sql_type(&field.ty, &MANY_TYNAMES)
 }
@@ -309,6 +322,10 @@ fn is_many_to_many(field: &Field) -> bool {
     get_many_sql_type(field).is_some()
 }
 
+fn is_foreign_key(field: &Field) -> bool {
+    get_foreign_key_sql_type(field).is_some()
+}
+
 fn is_option(field: &Field) -> bool {
     get_type_argument(&field.ty, &OPTION_TYNAMES).is_some()
 }
@@ -598,6 +615,7 @@ impl From<TokenStream2> for CompilerErrorMsg {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use syn::parse::Parser;
 
     #[test]
     fn test_get_type_argument_option() {
@@ -647,4 +665,37 @@ mod tests {
         let rv = get_type_argument(&typ, &MANY_TYNAMES);
         assert!(rv.is_none());
     }
+
+    #[test]
+    fn test_is_foreign_key() {
+        let tokens = quote::quote! {
+           foos: butane::ForeignKey<Foo>
+        };
+        let field = syn::Field::parse_named.parse2(tokens).unwrap();
+        assert!(is_foreign_key(&field));
+
+        let tokens = quote::quote! {
+           foos: Option<butane::ForeignKey<Foo>>
+        };
+        let field = syn::Field::parse_named.parse2(tokens).unwrap();
+        assert!(is_foreign_key(&field));
+
+        let tokens = quote::quote! {
+            foos: i8
+        };
+        let field = syn::Field::parse_named.parse2(tokens).unwrap();
+        assert!(!is_foreign_key(&field));
+
+        let tokens = quote::quote! {
+            foos: Option<i8>
+        };
+        let field = syn::Field::parse_named.parse2(tokens).unwrap();
+        assert!(!is_foreign_key(&field));
+
+        let tokens = quote::quote! {
+           foos: Option<SomethingElse>
+        };
+        let field = syn::Field::parse_named.parse2(tokens).unwrap();
+        assert!(!is_foreign_key(&field));
+    }
 }

butane_core/src/db/pg.rs

@@ -3,7 +3,7 @@ use super::connmethods::VecRows;
 use super::helper;
 use super::*;
 use crate::custom::{SqlTypeCustom, SqlValRefCustom};
-use crate::migrations::adb::{AColumn, ATable, Operation, TypeIdentifier, ADB};
+use crate::migrations::adb::{AColumn, ARef, ATable, Operation, TypeIdentifier, ADB};
 use crate::{debug, query};
 use crate::{Result, SqlType, SqlVal, SqlValRef};
 use bytes::BufMut;
@@ -546,6 +546,7 @@ where
 fn sql_for_op(current: &mut ADB, op: &Operation) -> Result<String> {
     match op {
         Operation::AddTable(table) => Ok(create_table(table, false)?),
+        Operation::AddTableConstraints(table) => Ok(create_table_constraints(table)),
         Operation::AddTableIfNotExists(table) => Ok(create_table(table, true)?),
         Operation::RemoveTable(name) => Ok(drop_table(name)),
         Operation::AddColumn(tbl, col) => add_column(tbl, col),
@@ -570,6 +571,16 @@ fn create_table(table: &ATable, allow_exists: bool) -> Result<String> {
     ))
 }
 
+fn create_table_constraints(table: &ATable) -> String {
+    table
+        .columns
+        .iter()
+        .filter(|column| column.reference().is_some())
+        .map(|column| define_constraint(table, column))
+        .collect::<Vec<String>>()
+        .join("\n")
+}
+
 fn define_column(col: &AColumn) -> Result<String> {
     let mut constraints: Vec<String> = Vec::new();
     if !col.nullable() {
@@ -589,6 +600,25 @@ fn define_column(col: &AColumn) -> Result<String> {
     ))
 }
 
+fn define_constraint(table: &ATable, column: &AColumn) -> String {
+    let reference = column
+        .reference()
+        .as_ref()
+        .expect("must have a references value");
+    match reference {
+        ARef::Literal(literal) => {
+            format!(
+                "ALTER TABLE {} ADD FOREIGN KEY ({}) REFERENCES {}({});",
+                helper::quote_reserved_word(&table.name),
+                helper::quote_reserved_word(column.name()),
+                helper::quote_reserved_word(literal.table_name()),
+                helper::quote_reserved_word(literal.column_name()),
+            )
+        }
+        _ => panic!(),
+    }
+}
+
 fn col_sqltype(col: &AColumn) -> Result<Cow<str>> {
     match col.typeid()? {
         TypeIdentifier::Name(name) => Ok(Cow::Owned(name)),

butane_core/src/db/sqlite.rs

@@ -6,7 +6,7 @@ use super::helper;
 use super::*;
 use crate::db::connmethods::BackendRows;
 use crate::debug;
-use crate::migrations::adb::{AColumn, ATable, Operation, TypeIdentifier, ADB};
+use crate::migrations::adb::{AColumn, ARef, ATable, Operation, TypeIdentifier, ADB};
 use crate::query;
 use crate::query::Order;
 use crate::{Result, SqlType, SqlVal, SqlValRef};
@@ -493,6 +493,7 @@ fn sql_valref_from_rusqlite<'a>(
 fn sql_for_op(current: &mut ADB, op: &Operation) -> Result<String> {
     match op {
         Operation::AddTable(table) => Ok(create_table(table, false)),
+        Operation::AddTableConstraints(_table) => Ok("".to_owned()),
         Operation::AddTableIfNotExists(table) => Ok(create_table(table, true)),
         Operation::RemoveTable(name) => Ok(drop_table(name)),
         Operation::AddColumn(tbl, col) => add_column(tbl, col),
@@ -509,7 +510,24 @@ fn create_table(table: &ATable, allow_exists: bool) -> String {
         .collect::<Vec<String>>()
         .join(",\n");
     let modifier = if allow_exists { "IF NOT EXISTS " } else { "" };
-    format!("CREATE TABLE {}{} (\n{}\n);", modifier, table.name, coldefs)
+    let mut constraints = create_table_constraints(table);
+    if !constraints.is_empty() {
+        constraints = ",\n".to_owned() + &constraints;
+    }
+    format!(
+        "CREATE TABLE {}{} (\n{}{}\n);",
+        modifier, table.name, coldefs, constraints
+    )
+}
+
+fn create_table_constraints(table: &ATable) -> String {
+    table
+        .columns
+        .iter()
+        .filter(|column| column.reference().is_some())
+        .map(define_constraint)
+        .collect::<Vec<String>>()
+        .join("\n")
 }
 
 fn define_column(col: &AColumn) -> String {
@@ -536,6 +554,24 @@ fn define_column(col: &AColumn) -> String {
     )
 }
 
+fn define_constraint(column: &AColumn) -> String {
+    let reference = column
+        .reference()
+        .as_ref()
+        .expect("must have a references value");
+    match reference {
+        ARef::Literal(literal) => {
+            format!(
+                "FOREIGN KEY ({}) REFERENCES {}({})",
+                helper::quote_reserved_word(column.name()),
+                helper::quote_reserved_word(literal.table_name()),
+                helper::quote_reserved_word(literal.column_name()),
+            )
+        }
+        _ => panic!(),
+    }
+}
+
 fn col_sqltype(col: &AColumn) -> Cow<str> {
     match col.typeid() {
         Ok(TypeIdentifier::Ty(ty)) => Cow::Borrowed(sqltype(&ty)),