Fixed check constraint handling with alternate schemas (fixes agronholm#3)

agronholm · agronholm · commit 6ff3604a40ec · 2013-09-01T17:12:22.000+03:00
diff --git a/sqlacodegen/codegen.py b/sqlacodegen/codegen.py
@@ -20,6 +20,7 @@
 
 
 _re_boolean_check_constraint = re.compile(r"(?:(?:.*?)\.)?(.*?) IN \(0, 1\)")
+_re_column_name = re.compile(r'(?:(["`])(?:.*)\1\.)?(["`]?)(.*)\2')
 _re_enum_check_constraint = re.compile(r"(?:(?:.*?)\.)?(.*?) IN \((.+)\)")
 _re_enum_item = re.compile(r"'(.*?)(?<!\\)'")
 
@@ -113,7 +114,7 @@ def _render_column(column, show_name):
     is_unique = any(isinstance(c, UniqueConstraint) and set(c.columns) == set([column])
                     for c in column.table.constraints)
     has_index = any(set(i.columns) == set([column]) for i in column.table.indexes)
-    
+
     # Render the column type if there are no foreign keys on it or any of them points back to itself
     render_coltype = not dedicated_fks or any(fk.column is column for fk in dedicated_fks)
 
@@ -483,15 +484,15 @@ def __init__(self, metadata, noindexes=False, noconstraints=False, nojoined=Fals
                         # Turn any integer-like column with a CheckConstraint like "column IN (0, 1)" into a Boolean
                         match = _re_boolean_check_constraint.match(sqltext)
                         if match:
-                            colname = match.group(1).strip('`"')
+                            colname = _re_column_name.match(match.group(1)).group(3)
                             table.constraints.remove(constraint)
                             table.c[colname].type = Boolean()
                             continue
 
                         # Turn any string-type column with a CheckConstraint like "column IN (...)" into an Enum
                         match = _re_enum_check_constraint.match(sqltext)
                         if match:
-                            colname = match.group(1).strip('`"')
+                            colname = _re_column_name.match(match.group(1)).group(3)
                             items = match.group(2)
                             if isinstance(table.c[colname].type, String):
                                 table.constraints.remove(constraint)
diff --git a/test/test_codegen.py b/test/test_codegen.py
@@ -820,13 +820,13 @@ class SimpleItem(Base):
 """)
 
     def test_table_args_kwargs(self):
-        Table(
+        simple_items = Table(
             'simple_items', self.metadata,
             Column('id', INTEGER, primary_key=True),
             Column('name', VARCHAR),
-            Index('testidx', 'id', 'name'),
             schema='testschema'
         )
+        simple_items.indexes.add(Index('testidx', simple_items.c.id, simple_items.c.name))
 
         eq_(self.generate_code(), """\
 # coding: utf-8
@@ -889,4 +889,44 @@ def test_foreign_key_options(self):
     'simple_items', metadata,
     Column('name', String, ForeignKey('simple_items.name', deferrable=True, initially='DEFERRED'))
 )
-""")
+""")
+
+    def test_foreign_key_schema(self):
+        Table(
+            'simple_items', self.metadata,
+            Column('id', INTEGER, primary_key=True),
+            Column('other_item_id', INTEGER),
+            ForeignKeyConstraint(['other_item_id'], ['otherschema.other_items.id'])
+        )
+        Table(
+            'other_items', self.metadata,
+            Column('id', INTEGER, primary_key=True),
+            schema='otherschema'
+        )
+
+        eq_(self.generate_code(), """\
+# coding: utf-8
+from sqlalchemy import Column, ForeignKey, Integer
+from sqlalchemy.orm import relationship
+from sqlalchemy.ext.declarative import declarative_base
+
+
+Base = declarative_base()
+metadata = Base.metadata
+
+
+class SimpleItem(Base):
+    __tablename__ = 'simple_items'
+
+    id = Column(Integer, primary_key=True)
+    other_item_id = Column(ForeignKey('otherschema.other_items.id'))
+
+    other_item = relationship('OtherItem')
+
+
+class OtherItem(Base):
+    __tablename__ = 'other_items'
+    __table_args__ = {'schema': 'otherschema'}
+
+    id = Column(Integer, primary_key=True)
+""")
