Added optional onupdate and ondelete params to Operations.create_check_constraint

2 files changed, 34 insertions, 4 deletions

diff --git a/alembic/operations.py b/alembic/operations.py
index 1bbe5a1..22f260c 100644
--- a/alembic/operations.py
+++ b/alembic/operations.py
@@ -51,7 +51,8 @@ class Operations(object):
         alembic.op._remove_proxy()
 
     def _foreign_key_constraint(self, name, source, referent, 
-                                    local_cols, remote_cols):
+                                    local_cols, remote_cols,
+                                    onupdate=None, ondelete=None):
         m = schema.MetaData()
         t1 = schema.Table(source, m, 
                 *[schema.Column(n, NULLTYPE) for n in local_cols])
@@ -61,7 +62,9 @@ class Operations(object):
         f = schema.ForeignKeyConstraint(local_cols, 
                                             ["%s.%s" % (referent, n) 
                                             for n in remote_cols],
-                                            name=name
+                                            name=name,
+                                            onupdate=onupdate,
+                                            ondelete=ondelete
                                             )
         t1.append_constraint(f)
 
@@ -315,7 +318,7 @@ class Operations(object):
 
 
     def create_foreign_key(self, name, source, referent, local_cols, 
-                                    remote_cols):
+                           remote_cols, onupdate=None, ondelete=None):
         """Issue a "create foreign key" instruction using the 
         current migration context.
 
@@ -349,12 +352,19 @@ class Operations(object):
          source table.
         :param remote_cols: a list of string column names in the
          remote table.
+        :param onupdate: Optional string. If set, emit ON UPDATE <value> when
+         issuing DDL for this constraint. Typical values include CASCADE,
+         DELETE and RESTRICT.
+        :param ondelete: Optional string. If set, emit ON DELETE <value> when
+         issuing DDL for this constraint. Typical values include CASCADE,
+         DELETE and RESTRICT.
 
         """
 
         self.impl.add_constraint(
                     self._foreign_key_constraint(name, source, referent, 
-                            local_cols, remote_cols)
+                            local_cols, remote_cols,
+                            onupdate=onupdate, ondelete=ondelete)
                 )
 
     def create_unique_constraint(self, name, source, local_cols, **kw):
diff --git a/tests/test_op.py b/tests/test_op.py
index 8af4711..58f19cb 100644
--- a/tests/test_op.py
+++ b/tests/test_op.py
@@ -154,6 +154,26 @@ def test_add_foreign_key():
             "REFERENCES t2 (bat, hoho)"
     )
 
+def test_add_foreign_key_onupdate():
+    context = op_fixture()
+    op.create_foreign_key('fk_test', 't1', 't2',
+                    ['foo', 'bar'], ['bat', 'hoho'],
+                    onupdate='CASCADE')
+    context.assert_(
+        "ALTER TABLE t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) "
+            "REFERENCES t2 (bat, hoho) ON UPDATE CASCADE"
+    )
+
+def test_add_foreign_key_ondelete():
+    context = op_fixture()
+    op.create_foreign_key('fk_test', 't1', 't2',
+                    ['foo', 'bar'], ['bat', 'hoho'],
+                    ondelete='CASCADE')
+    context.assert_(
+        "ALTER TABLE t1 ADD CONSTRAINT fk_test FOREIGN KEY(foo, bar) "
+            "REFERENCES t2 (bat, hoho) ON DELETE CASCADE"
+    )
+
 def test_add_check_constraint():
     context = op_fixture()
     op.create_check_constraint(