Render the ondelete, onupdate and match options in foreign key constraints

agronholm · agronholm · commit 29561b961aaf · 2014-01-03T04:44:30.000+02:00
diff --git a/setup.cfg b/setup.cfg
@@ -4,6 +4,7 @@ exclude = .tox
 
 [pytest]
 pep8maxlinelength = 120
+addopts = --tb=short
 
 [wheel]
 universal = 1
diff --git a/sqlacodegen/__init__.py b/sqlacodegen/__init__.py
@@ -1 +1 @@
-version = '1.1.4'
+version = '1.1.4.post1'
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -9,6 +9,7 @@
 
 from sqlalchemy import (Enum, ForeignKeyConstraint, PrimaryKeyConstraint, CheckConstraint, UniqueConstraint, Table,
                         Column)
+from sqlalchemy.schema import ForeignKey
 from sqlalchemy.util import OrderedDict
 from sqlalchemy.types import Boolean, String
 import sqlalchemy
@@ -109,15 +110,6 @@ def _render_column_type(coltype):
     return text
 
 
-def _render_fk(fk):
-    text = "ForeignKey('%s.%s'" % (fk.column.table.fullname, fk.column.name)
-    if fk.deferrable:
-        text += ', deferrable=True'
-    if fk.initially:
-        text += ', initially=%r' % fk.initially
-    return text + ')'
-
-
 def _render_column(column, show_name):
     kwarg = []
     is_sole_pk = column.primary_key and len(column.table.primary_key) == 1
@@ -149,17 +141,29 @@ def _render_column(column, show_name):
     return 'Column({0})'.format(', '.join(
         ([repr(column.name)] if show_name else []) +
         ([_render_column_type(column.type)] if render_coltype else []) +
-        [_render_fk(x) for x in dedicated_fks] +
+        [_render_constraint(x) for x in dedicated_fks] +
         [repr(x) for x in column.constraints] +
         ['{0}={1}'.format(k, repr(getattr(column, k))) for k in kwarg]))
 
 
 def _render_constraint(constraint):
-    if isinstance(constraint, ForeignKeyConstraint):
+    def render_fk_options(*opts):
+        opts = [repr(opt) for opt in opts]
+        for attr in 'ondelete', 'onupdate', 'deferrable', 'initially', 'match':
+            value = getattr(constraint, attr)
+            if value:
+                opts.append('{0}={1!r}'.format(attr, value))
+
+        return ', '.join(opts)
+
+    if isinstance(constraint, ForeignKey):
+        remote_column = '{0}.{1}'.format(constraint.column.table.fullname, constraint.column.name)
+        return 'ForeignKey({0})'.format(render_fk_options(remote_column))
+    elif isinstance(constraint, ForeignKeyConstraint):
         local_columns = constraint.columns
         remote_columns = ['{0}.{1}'.format(fk.column.table.fullname, fk.column.name)
                           for fk in constraint.elements]
-        return 'ForeignKeyConstraint({0!r}, {1!r})'.format(local_columns, remote_columns)
+        return 'ForeignKeyConstraint({0})'.format(render_fk_options(local_columns, remote_columns))
     elif isinstance(constraint, CheckConstraint):
         return 'CheckConstraint({0!r})'.format(_get_compiled_expression(constraint.sqltext))
     elif isinstance(constraint, UniqueConstraint):
@@ -465,7 +469,7 @@ def __init__(self, metadata, noindexes=False, noconstraints=False, nojoined=Fals
         for table in metadata.tables.values():
             # Link tables have exactly two foreign key constraints and all columns are involved in them
             fk_constraints = [constr for constr in table.constraints if isinstance(constr, ForeignKeyConstraint)]
-            if (len(fk_constraints) == 2 and all(col.foreign_keys for col in table.columns)):
+            if len(fk_constraints) == 2 and all(col.foreign_keys for col in table.columns):
                 association_tables.add(table.name)
                 tablename = sorted(fk_constraints, key=_get_constraint_sort_key)[0].elements[0].column.table.name
                 links[tablename].append(table)
diff --git a/test/test_codegen.py b/test/test_codegen.py
@@ -414,7 +414,8 @@ def test_onetomany_composite(self):
             Column('id', INTEGER, primary_key=True),
             Column('container_id1', INTEGER),
             Column('container_id2', INTEGER),
-            ForeignKeyConstraint(['container_id1', 'container_id2'], ['simple_containers.id1', 'simple_containers.id2'])
+            ForeignKeyConstraint(['container_id1', 'container_id2'], ['simple_containers.id1', 'simple_containers.id2'],
+                                 ondelete='CASCADE', onupdate='CASCADE', match='PARTIAL')
         )
         Table(
             'simple_containers', self.metadata,
@@ -443,7 +444,8 @@ class SimpleContainer(Base):
 class SimpleItem(Base):
     __tablename__ = 'simple_items'
     __table_args__ = (
-        ForeignKeyConstraint(['container_id1', 'container_id2'], ['simple_containers.id1', 'simple_containers.id2']),
+        ForeignKeyConstraint(['container_id1', 'container_id2'], ['simple_containers.id1', 'simple_containers.id2'], \
+ondelete='CASCADE', onupdate='CASCADE', match='PARTIAL'),
     )
 
     id = Column(Integer, primary_key=True)
@@ -899,7 +901,8 @@ def test_schema_boolean(self):
     def test_foreign_key_options(self):
         Table(
             'simple_items', self.metadata,
-            Column('name', VARCHAR, ForeignKey('simple_items.name', deferrable=True, initially='DEFERRED'))
+            Column('name', VARCHAR, ForeignKey('simple_items.name', ondelete='CASCADE', onupdate='CASCADE',
+                                               deferrable=True, initially='DEFERRED', match='PARTIAL'))
         )
 
         assert self.generate_code() == """\
@@ -912,7 +915,8 @@ def test_foreign_key_options(self):
 
 t_simple_items = Table(
     'simple_items', metadata,
-    Column('name', String, ForeignKey('simple_items.name', deferrable=True, initially='DEFERRED'))
+    Column('name', String, ForeignKey('simple_items.name', ondelete='CASCADE', onupdate='CASCADE', \
+deferrable=True, initially='DEFERRED', match='PARTIAL'))
 )
 """
 

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-version = '1.1.4'`
	`1`	`+version = '1.1.4.post1'`