PostgreSQL foreign key inspection includes options

2 files changed, 105 insertions, 5 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index ecbbf1674..71d9dc2f6 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1883,6 +1883,15 @@ class PGDialect(default.DefaultDialect):
                 n.oid = c.relnamespace
           ORDER BY 1
         """
+        # http://www.postgresql.org/docs/9.0/static/sql-createtable.html
+        FK_REGEX = re.compile(
+            r'FOREIGN KEY \((.*?)\) REFERENCES (?:(.*?)\.)?(.*?)\((.*?)\)'
+            r'[\s]?(MATCH (FULL|PARTIAL|SIMPLE)+)?'
+            r'[\s]?(ON UPDATE (CASCADE|RESTRICT|NO ACTION|SET NULL|SET DEFAULT)+)?'
+            r'[\s]?(ON DELETE (CASCADE|RESTRICT|NO ACTION|SET NULL|SET DEFAULT)+)?'
+            r'[\s]?(DEFERRABLE|NOT DEFERRABLE)?'
+            r'[\s]?(INITIALLY (DEFERRED|IMMEDIATE)+)?'
+        )
 
         t = sql.text(FK_SQL, typemap={
                                 'conname': sqltypes.Unicode,
@@ -1890,15 +1899,18 @@ class PGDialect(default.DefaultDialect):
         c = connection.execute(t, table=table_oid)
         fkeys = []
         for conname, condef, conschema in c.fetchall():
-            m = re.search('FOREIGN KEY \((.*?)\) REFERENCES '
-                            '(?:(.*?)\.)?(.*?)\((.*?)\)', condef).groups()
+            m = re.search(FK_REGEX, condef).groups()
             constrained_columns, referred_schema, \
-                    referred_table, referred_columns = m
+                    referred_table, referred_columns, \
+                    _, match, _, onupdate, _, ondelete, \
+                    deferrable, _, initially = m
+            if deferrable is not None:
+                deferrable = True if deferrable == 'DEFERRABLE' else False
             constrained_columns = [preparer._unquote_identifier(x)
                         for x in re.split(r'\s*,\s*', constrained_columns)]
 
             if referred_schema:
-                referred_schema =\
+                referred_schema = \
                                 preparer._unquote_identifier(referred_schema)
             elif schema is not None and schema == conschema:
                 # no schema was returned by pg_get_constraintdef().  This
@@ -1916,7 +1928,14 @@ class PGDialect(default.DefaultDialect):
                 'constrained_columns': constrained_columns,
                 'referred_schema': referred_schema,
                 'referred_table': referred_table,
-                'referred_columns': referred_columns
+                'referred_columns': referred_columns,
+                'options': {
+                    'onupdate': onupdate,
+                    'ondelete': ondelete,
+                    'deferrable': deferrable,
+                    'initially': initially,
+                    'match': match
+                }
             }
             fkeys.append(fkey_d)
         return fkeys
diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index 3d48230f3..6dfb84f40 100644
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -17,6 +17,7 @@ from sqlalchemy.dialects.postgresql import base as postgresql
 import logging
 import logging.handlers
 from sqlalchemy.testing.mock import Mock
+from sqlalchemy.engine.reflection import Inspector
 
 class MiscTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
 
@@ -218,3 +219,83 @@ class MiscTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
                 ddl_compiler.get_column_specification(t.c.c),
                 "c %s NOT NULL" % expected
             )
+
+class InspectionTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
+
+    __only_on__ = 'postgresql'
+
+    def test_get_foreign_keys(self):
+        """
+        PGDialect.get_foreign_keys()
+        """
+        metadata = MetaData(testing.db)
+        person = Table('person', metadata,
+            Column('id', String(length=32), nullable=False, primary_key=True),
+            Column('company_id', ForeignKey('company.id',
+                name='person_company_id_fkey',
+                match='FULL', onupdate='RESTRICT', ondelete='RESTRICT',
+                deferrable=True, initially='DEFERRED'
+                )
+            )
+        )
+        company = Table('company', metadata,
+            Column('id', String(length=32), nullable=False, primary_key=True),
+            Column('name', String(length=255)),
+            Column('industry_id', ForeignKey('industry.id',
+                name='company_industry_id_fkey',
+                onupdate='CASCADE', ondelete='CASCADE',
+                deferrable=False, # PG default
+                initially='IMMEDIATE' # PG default
+                )
+            )
+        )
+        industry = Table('industry', metadata,
+            Column('id', Integer(), nullable=False, primary_key=True),
+            Column('name', String(length=255))
+        )
+        fk_ref = {
+            'person_company_id_fkey': {
+                'name': 'person_company_id_fkey',
+                'constrained_columns': ['company_id'],
+                'referred_columns': ['id'],
+                'referred_table': 'company',
+                'referred_schema': None,
+                'options': {
+                    'onupdate': 'RESTRICT',
+                    'deferrable': True,
+                    'ondelete': 'RESTRICT',
+                    'initially': 'DEFERRED',
+                    'match': 'FULL'
+                    }
+                },
+            'company_industry_id_fkey': {
+                'name': 'company_industry_id_fkey',
+                'constrained_columns': ['industry_id'],
+                'referred_columns': ['id'],
+                'referred_table': 'industry',
+                'referred_schema': None,
+                'options': {
+                    'onupdate': 'CASCADE',
+                    'deferrable': None,
+                    'ondelete': 'CASCADE',
+                    'initially': None,
+                    'match': None
+                }
+            }
+        }
+        try:
+            connection = testing.db.connect()
+            industry.create()
+            company.create()
+            person.create()
+            inspector = Inspector.from_engine(connection)
+            fks = inspector.get_foreign_keys('person') + \
+                inspector.get_foreign_keys('company')
+            for fk in fks:
+                ref = fk_ref[fk['name']]
+                for key, val in fk.items():
+                    eq_(val, ref[key])
+        finally:
+            person.drop()
+            company.drop()
+            industry.drop()