overhaul to schema, addition of ForeignKeyConstraint/

PrimaryKeyConstraint objects (also UniqueConstraint not
completed yet).  table creation and reflection modified
to be more oriented towards these new table-level objects.
reflection for sqlite/postgres/mysql supports composite
foreign keys; oracle/mssql/firebird not converted yet.

7 files changed, 59 insertions, 93 deletions

diff --git a/lib/sqlalchemy/databases/firebird.py b/lib/sqlalchemy/databases/firebird.py
index 0039333d5..085d8cf44 100644
--- a/lib/sqlalchemy/databases/firebird.py
+++ b/lib/sqlalchemy/databases/firebird.py
@@ -293,7 +293,7 @@ class FBCompiler(ansisql.ANSICompiler):
         return ""
 
 class FBSchemaGenerator(ansisql.ANSISchemaGenerator):
-    def get_column_specification(self, column, override_pk=False, **kwargs):
+    def get_column_specification(self, column, **kwargs):
         colspec = column.name 
         colspec += " " + column.type.engine_impl(self.engine).get_col_spec()
         default = self.get_column_default_string(column)
@@ -302,10 +302,6 @@ class FBSchemaGenerator(ansisql.ANSISchemaGenerator):
 
         if not column.nullable:
             colspec += " NOT NULL"
-        if column.primary_key and not override_pk:
-            colspec += " PRIMARY KEY"
-        if column.foreign_key:
-            colspec += " REFERENCES %s(%s)" % (column.foreign_key.column.table.name, column.foreign_key.column.name) 
         return colspec
 
     def visit_sequence(self, sequence):
diff --git a/lib/sqlalchemy/databases/information_schema.py b/lib/sqlalchemy/databases/information_schema.py
index 08236f799..296db2de5 100644
--- a/lib/sqlalchemy/databases/information_schema.py
+++ b/lib/sqlalchemy/databases/information_schema.py
@@ -54,6 +54,7 @@ pg_key_constraints = schema.Table("key_column_usage", ischema,
     Column("table_name", String),
     Column("column_name", String),
     Column("constraint_name", String),
+    Column("ordinal_position", Integer),
     schema="information_schema")
 
 #mysql_key_constraints = schema.Table("key_column_usage", ischema,
@@ -100,13 +101,9 @@ class ISchema(object):
         return self.cache[name]
 
 
-def reflecttable(connection, table, ischema_names, use_mysql=False):
+def reflecttable(connection, table, ischema_names):
     
-    if use_mysql:
-        # no idea which INFORMATION_SCHEMA spec is correct, mysql or postgres
-        key_constraints = mysql_key_constraints
-    else:
-        key_constraints = pg_key_constraints
+    key_constraints = pg_key_constraints
         
     if table.schema is not None:
         current_schema = table.schema
@@ -152,39 +149,50 @@ def reflecttable(connection, table, ischema_names, use_mysql=False):
     if not found_table:
         raise exceptions.NoSuchTableError(table.name)
 
-    s = select([constraints.c.constraint_name, constraints.c.constraint_type, constraints.c.table_name, key_constraints], use_labels=True, from_obj=[constraints.join(column_constraints, column_constraints.c.constraint_name==constraints.c.constraint_name).join(key_constraints, key_constraints.c.constraint_name==column_constraints.c.constraint_name)])
-    if not use_mysql:
-        s.append_column(column_constraints)
-        s.append_whereclause(constraints.c.table_name==table.name)
-        s.append_whereclause(constraints.c.table_schema==current_schema)
-        colmap = [constraints.c.constraint_type, key_constraints.c.column_name, column_constraints.c.table_schema, column_constraints.c.table_name, column_constraints.c.column_name]
-    else:
-        # this doesnt seem to pick up any foreign keys with mysql
-        s.append_whereclause(key_constraints.c.table_name==constraints.c.table_name)
-        s.append_whereclause(key_constraints.c.table_schema==constraints.c.table_schema)
-        s.append_whereclause(constraints.c.table_name==table.name)
-        s.append_whereclause(constraints.c.table_schema==current_schema)
-        colmap = [constraints.c.constraint_type, key_constraints.c.column_name, key_constraints.c.referenced_table_schema, key_constraints.c.referenced_table_name, key_constraints.c.referenced_column_name]
+    # we are relying on the natural ordering of the constraint_column_usage table to return the referenced columns
+    # in an order that corresponds to the ordinal_position in the key_constraints table, otherwise composite foreign keys
+    # wont reflect properly.  dont see a way around this based on whats available from information_schema
+    s = select([constraints.c.constraint_name, constraints.c.constraint_type, constraints.c.table_name, key_constraints], use_labels=True, from_obj=[constraints.join(column_constraints, column_constraints.c.constraint_name==constraints.c.constraint_name).join(key_constraints, key_constraints.c.constraint_name==column_constraints.c.constraint_name)], order_by=[key_constraints.c.ordinal_position])
+    s.append_column(column_constraints)
+    s.append_whereclause(constraints.c.table_name==table.name)
+    s.append_whereclause(constraints.c.table_schema==current_schema)
+    colmap = [constraints.c.constraint_type, key_constraints.c.column_name, column_constraints.c.table_schema, column_constraints.c.table_name, column_constraints.c.column_name, constraints.c.constraint_name, key_constraints.c.ordinal_position]
     c = connection.execute(s)
 
+    fks = {}
     while True:
         row = c.fetchone()
         if row is None:
             break
-#        continue
-        (type, constrained_column, referred_schema, referred_table, referred_column) = (
+        (type, constrained_column, referred_schema, referred_table, referred_column, constraint_name, ordinal_position) = (
             row[colmap[0]],
             row[colmap[1]],
             row[colmap[2]],
             row[colmap[3]],
-            row[colmap[4]]
+            row[colmap[4]],
+            row[colmap[5]],
+            row[colmap[6]]
         )
         #print "type %s on column %s to remote %s.%s.%s" % (type, constrained_column, referred_schema, referred_table, referred_column) 
         if type=='PRIMARY KEY':
             table.c[constrained_column]._set_primary_key()
         elif type=='FOREIGN KEY':
+            try:
+                fk = fks[constraint_name]
+            except KeyError:
+                fk = ([],[])
+                fks[constraint_name] = fk
             if current_schema == referred_schema:
                 referred_schema = table.schema
-            remotetable = Table(referred_table, table.metadata, autoload=True, autoload_with=connection, schema=referred_schema)
-            table.c[constrained_column].append_item(schema.ForeignKey(remotetable.c[referred_column]))
+            if referred_schema is not None:
+                refspec = ".".join([referred_schema, referred_table, referred_column])
+            else:
+                refspec = ".".join([referred_table, referred_column])
+            if constrained_column not in fk[0]:
+                fk[0].append(constrained_column)
+            if refspec not in fk[1]:
+                fk[1].append(refspec)
+    
+    for name, value in fks.iteritems():
+        table.append_item(ForeignKeyConstraint(value[0], value[1], name=name))    
             
diff --git a/lib/sqlalchemy/databases/mssql.py b/lib/sqlalchemy/databases/mssql.py
index c297195ca..9d51d535d 100644
--- a/lib/sqlalchemy/databases/mssql.py
+++ b/lib/sqlalchemy/databases/mssql.py
@@ -511,7 +511,7 @@ class MSSQLCompiler(ansisql.ANSICompiler):
 
         
 class MSSQLSchemaGenerator(ansisql.ANSISchemaGenerator):
-    def get_column_specification(self, column, override_pk=False, first_pk=False):
+    def get_column_specification(self, column, **kwargs):
         colspec = column.name + " " + column.type.engine_impl(self.engine).get_col_spec()
 
         # install a IDENTITY Sequence if we have an implicit IDENTITY column
@@ -528,12 +528,6 @@ class MSSQLSchemaGenerator(ansisql.ANSISchemaGenerator):
             default = self.get_column_default_string(column)
             if default is not None:
                 colspec += " DEFAULT " + default
-
-        if column.primary_key:
-            if not override_pk:
-                colspec += " PRIMARY KEY"
-        if column.foreign_key:
-            colspec += " REFERENCES %s(%s)" % (column.foreign_key.column.table.fullname, column.foreign_key.column.name)
         
         return colspec
 
diff --git a/lib/sqlalchemy/databases/mysql.py b/lib/sqlalchemy/databases/mysql.py
index 997010f1c..1d587ff7c 100644
--- a/lib/sqlalchemy/databases/mysql.py
+++ b/lib/sqlalchemy/databases/mysql.py
@@ -309,8 +309,6 @@ class MySQLDialect(ansisql.ANSIDialect):
                 break
             #print "row! " + repr(row)
             if not found_table:
-                tabletype, foreignkeyD = self.moretableinfo(connection, table=table)
-                table.kwargs['mysql_engine'] = tabletype
                 found_table = True
 
             (name, type, nullable, primary_key, default) = (row[0], row[1], row[2] == 'YES', row[3] == 'PRI', row[4])
@@ -338,16 +336,15 @@ class MySQLDialect(ansisql.ANSIDialect):
                     argslist = re.findall(r'(\d+)', args)
                     coltype = coltype(*[int(a) for a in argslist], **kw)
             
-            arglist = []
-            fkey = foreignkeyD.get(name)
-            if fkey is not None:
-                arglist.append(schema.ForeignKey(fkey))
-    
-            table.append_item(schema.Column(name, coltype, *arglist,
+            table.append_item(schema.Column(name, coltype, 
                                             **dict(primary_key=primary_key,
                                                    nullable=nullable,
                                                    default=default
                                                    )))
+
+        tabletype = self.moretableinfo(connection, table=table)
+        table.kwargs['mysql_engine'] = tabletype
+
         if not found_table:
             raise exceptions.NoSuchTableError(table.name)
     
@@ -368,15 +365,15 @@ class MySQLDialect(ansisql.ANSIDialect):
             match = re.search(r'\b(?:TYPE|ENGINE)=(?P<ttype>.+)\b', desc[lastparen.start():], re.I)
             if match:
                 tabletype = match.group('ttype')
-        foreignkeyD = {}
-        fkpat = (r'FOREIGN KEY\s*\(`?(?P<name>.+?)`?\)'
-                 r'\s*REFERENCES\s*`?(?P<reftable>.+?)`?'
-                 r'\s*\(`?(?P<refcol>.+?)`?\)'
-                )
+
+        fkpat = r'CONSTRAINT `(?P<name>.+?)` FOREIGN KEY \((?P<columns>.+?)\) REFERENCES `(?P<reftable>.+?)` \((?P<refcols>.+?)\)'
         for match in re.finditer(fkpat, desc):
-            foreignkeyD[match.group('name')] = match.group('reftable') + '.' + match.group('refcol')
+            columns = re.findall(r'`(.+?)`', match.group('columns'))
+            refcols = [match.group('reftable') + "." + x for x in re.findall(r'`(.+?)`', match.group('refcols'))]
+            constraint = schema.ForeignKeyConstraint(columns, refcols, name=match.group('name'))
+            table.append_item(constraint)
 
-        return (tabletype, foreignkeyD)
+        return tabletype
         
 
 class MySQLCompiler(ansisql.ANSICompiler):
@@ -411,12 +408,8 @@ class MySQLSchemaGenerator(ansisql.ANSISchemaGenerator):
         if not column.nullable:
             colspec += " NOT NULL"
         if column.primary_key:
-            if not override_pk:
-                colspec += " PRIMARY KEY"
             if not column.foreign_key and first_pk and isinstance(column.type, sqltypes.Integer):
                 colspec += " AUTO_INCREMENT"
-        if column.foreign_key:
-            colspec += ", FOREIGN KEY (%s) REFERENCES %s(%s)" % (column.name, column.foreign_key.column.table.name, column.foreign_key.column.name) 
         return colspec
 
     def post_create_table(self, table):
diff --git a/lib/sqlalchemy/databases/oracle.py b/lib/sqlalchemy/databases/oracle.py
index bf6c1fd8d..d184291fd 100644
--- a/lib/sqlalchemy/databases/oracle.py
+++ b/lib/sqlalchemy/databases/oracle.py
@@ -320,7 +320,7 @@ class OracleCompiler(ansisql.ANSICompiler):
         return ""
 
 class OracleSchemaGenerator(ansisql.ANSISchemaGenerator):
-    def get_column_specification(self, column, override_pk=False, **kwargs):
+    def get_column_specification(self, column, **kwargs):
         colspec = column.name
         colspec += " " + column.type.engine_impl(self.engine).get_col_spec()
         default = self.get_column_default_string(column)
@@ -329,10 +329,6 @@ class OracleSchemaGenerator(ansisql.ANSISchemaGenerator):
 
         if not column.nullable:
             colspec += " NOT NULL"
-        if column.primary_key and not override_pk:
-            colspec += " PRIMARY KEY"
-        if column.foreign_key:
-            colspec += " REFERENCES %s(%s)" % (column.foreign_key.column.table.name, column.foreign_key.column.name) 
         return colspec
 
     def visit_sequence(self, sequence):
diff --git a/lib/sqlalchemy/databases/postgres.py b/lib/sqlalchemy/databases/postgres.py
index de21bd570..decccba58 100644
--- a/lib/sqlalchemy/databases/postgres.py
+++ b/lib/sqlalchemy/databases/postgres.py
@@ -329,7 +329,7 @@ class PGCompiler(ansisql.ANSICompiler):
         
 class PGSchemaGenerator(ansisql.ANSISchemaGenerator):
         
-    def get_column_specification(self, column, override_pk=False, **kwargs):
+    def get_column_specification(self, column, **kwargs):
         colspec = column.name
         if column.primary_key and not column.foreign_key and isinstance(column.type, sqltypes.Integer) and (column.default is None or (isinstance(column.default, schema.Sequence) and column.default.optional)):
             colspec += " SERIAL"
@@ -341,10 +341,6 @@ class PGSchemaGenerator(ansisql.ANSISchemaGenerator):
 
         if not column.nullable:
             colspec += " NOT NULL"
-        if column.primary_key and not override_pk:
-            colspec += " PRIMARY KEY"
-        if column.foreign_key:
-            colspec += " REFERENCES %s(%s)" % (column.foreign_key.column.table.fullname, column.foreign_key.column.name) 
         return colspec
 
     def visit_sequence(self, sequence):
diff --git a/lib/sqlalchemy/databases/sqlite.py b/lib/sqlalchemy/databases/sqlite.py
index c07952ff2..c703cd81e 100644
--- a/lib/sqlalchemy/databases/sqlite.py
+++ b/lib/sqlalchemy/databases/sqlite.py
@@ -257,7 +257,7 @@ class SQLiteCompiler(ansisql.ANSICompiler):
             return ansisql.ANSICompiler.binary_operator_string(self, binary)
 
 class SQLiteSchemaGenerator(ansisql.ANSISchemaGenerator):
-    def get_column_specification(self, column, override_pk=False, **kwargs):
+    def get_column_specification(self, column, **kwargs):
         colspec = column.name + " " + column.type.engine_impl(self.engine).get_col_spec()
         default = self.get_column_default_string(column)
         if default is not None:
@@ -265,34 +265,17 @@ class SQLiteSchemaGenerator(ansisql.ANSISchemaGenerator):
 
         if not column.nullable:
             colspec += " NOT NULL"
-        if column.primary_key and not override_pk:
-            colspec += " PRIMARY KEY"
-        if column.foreign_key:
-            colspec += " REFERENCES %s(%s)" % (column.foreign_key.column.table.name, column.foreign_key.column.name) 
         return colspec
-    def visit_table(self, table):
-        """sqlite is going to create multi-primary keys with just a UNIQUE index."""
-        self.append("\nCREATE TABLE " + table.fullname + "(")
-
-        separator = "\n"
-
-        have_pk = False
-        use_pks = len(table.primary_key) == 1
-        for column in table.columns:
-            self.append(separator)
-            separator = ", \n"
-            self.append("\t" + self.get_column_specification(column, override_pk=not use_pks))
-                
-        if len(table.primary_key) > 1:
-            self.append(", \n")
-            # put all PRIMARY KEYS in a UNIQUE index
-            self.append("\tUNIQUE (%s)" % string.join([c.name for c in table.primary_key],', '))
-
-        self.append("\n)\n\n")
-        self.execute()        
-        if hasattr(table, 'indexes'):
-            for index in table.indexes:
-                self.visit_index(index)
 
+    # this doesnt seem to be needed, although i suspect older versions of sqlite might still
+    # not directly support composite primary keys
+    #def visit_primary_key_constraint(self, constraint):
+    #    if len(constraint) > 1:
+    #        self.append(", \n")
+    #        # put all PRIMARY KEYS in a UNIQUE index
+    #        self.append("\tUNIQUE (%s)" % string.join([c.name for c in constraint],', '))
+    #    else:
+    #        super(SQLiteSchemaGenerator, self).visit_primary_key_constraint(constraint)
+            
 dialect = SQLiteDialect
 poolclass = pool.SingletonThreadPool