Implemented CHECK constraint reflection for SQLite and PostgreSQL

Co-Authored-By: Mike Bayer <mike_mp@zzzcomputing.com>
Change-Id: Ie6cf2d2958d1c567324db9e08fef2d3186e97350
Pull-request: https://bitbucket.org/zzzeek/sqlalchemy/pull-requests/80

4 files changed, 117 insertions, 0 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 9d019b56e..fe3d29450 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -2497,6 +2497,31 @@ class PGDialect(default.DefaultDialect):
             for name, uc in uniques.items()
         ]
 
+    @reflection.cache
+    def get_check_constraints(
+            self, connection, table_name, schema=None, **kw):
+        table_oid = self.get_table_oid(connection, table_name, schema,
+                                       info_cache=kw.get('info_cache'))
+
+        CHECK_SQL = """
+            SELECT
+                cons.conname as name,
+                cons.consrc as src
+            FROM
+                pg_catalog.pg_constraint cons
+            WHERE
+                cons.conrelid = :table_oid AND
+                cons.contype = 'c'
+        """
+
+        c = connection.execute(sql.text(CHECK_SQL), table_oid=table_oid)
+
+        return [
+            {'name': name,
+             'sqltext': src[1:-1]}
+            for name, src in c.fetchall()
+            ]
+
     def _load_enums(self, connection, schema=None):
         schema = schema or self.default_schema_name
         if not self.supports_native_enum:
diff --git a/lib/sqlalchemy/dialects/sqlite/base.py b/lib/sqlalchemy/dialects/sqlite/base.py
index ddd869448..2fe10556b 100644
--- a/lib/sqlalchemy/dialects/sqlite/base.py
+++ b/lib/sqlalchemy/dialects/sqlite/base.py
@@ -1474,6 +1474,32 @@ class SQLiteDialect(default.DefaultDialect):
         return unique_constraints
 
     @reflection.cache
+    def get_check_constraints(self, connection, table_name,
+                              schema=None, **kw):
+        table_data = self._get_table_sql(
+            connection, table_name, schema=schema, **kw)
+        if not table_data:
+            return []
+
+        CHECK_PATTERN = (
+            '(?:CONSTRAINT (\w+) +)?'
+            'CHECK *\( *(.+) *\),? *'
+        )
+        check_constraints = []
+        # NOTE: we aren't using re.S here because we actually are
+        # taking advantage of each CHECK constraint being all on one
+        # line in the table definition in order to delineate.  This
+        # necessarily makes assumptions as to how the CREATE TABLE
+        # was emitted.
+        for match in re.finditer(CHECK_PATTERN, table_data, re.I):
+            check_constraints.append({
+                'sqltext': match.group(2),
+                'name': match.group(1)
+            })
+
+        return check_constraints
+
+    @reflection.cache
     def get_indexes(self, connection, table_name, schema=None, **kw):
         pragma_indexes = self._get_table_pragma(
             connection, "index_list", table_name, schema=schema)
diff --git a/lib/sqlalchemy/engine/interfaces.py b/lib/sqlalchemy/engine/interfaces.py
index 26731f9a5..13e8bf1f4 100644
--- a/lib/sqlalchemy/engine/interfaces.py
+++ b/lib/sqlalchemy/engine/interfaces.py
@@ -399,6 +399,29 @@ class Dialect(object):
 
         raise NotImplementedError()
 
+    def get_check_constraints(
+            self, connection, table_name, schema=None, **kw):
+        """Return information about check constraints in `table_name`.
+
+        Given a string `table_name` and an optional string `schema`, return
+        check constraint information as a list of dicts with these keys:
+
+        name
+          the check constraint's name
+
+        sqltext
+          the check constraint's SQL expression
+
+        \**kw
+          other options passed to the dialect's get_check_constraints()
+          method.
+
+        .. versionadded:: 1.1.0
+
+        """
+
+        raise NotImplementedError()
+
     def normalize_name(self, name):
         """convert the given name to lowercase if it is detected as
         case insensitive.
diff --git a/lib/sqlalchemy/engine/reflection.py b/lib/sqlalchemy/engine/reflection.py
index 2d524978d..14e98cecf 100644
--- a/lib/sqlalchemy/engine/reflection.py
+++ b/lib/sqlalchemy/engine/reflection.py
@@ -506,6 +506,32 @@ class Inspector(object):
         return self.dialect.get_unique_constraints(
             self.bind, table_name, schema, info_cache=self.info_cache, **kw)
 
+    def get_check_constraints(self, table_name, schema=None, **kw):
+        """Return information about check constraints in `table_name`.
+
+        Given a string `table_name` and an optional string `schema`, return
+        check constraint information as a list of dicts with these keys:
+
+        name
+          the check constraint's name
+
+        sqltext
+          the check constraint's SQL expression
+
+        :param table_name: string name of the table.  For special quoting,
+         use :class:`.quoted_name`.
+
+        :param schema: string schema name; if omitted, uses the default schema
+         of the database connection.  For special quoting,
+         use :class:`.quoted_name`.
+
+        .. versionadded:: 1.1.0
+
+        """
+
+        return self.dialect.get_check_constraints(
+            self.bind, table_name, schema, info_cache=self.info_cache, **kw)
+
     def reflecttable(self, table, include_columns, exclude_columns=()):
         """Given a Table object, load its internal constructs based on
         introspection.
@@ -586,6 +612,10 @@ class Inspector(object):
             table_name, schema, table, cols_by_orig_name,
             include_columns, exclude_columns, reflection_options)
 
+        self._reflect_check_constraints(
+            table_name, schema, table, cols_by_orig_name,
+            include_columns, exclude_columns, reflection_options)
+
     def _reflect_column(
         self, table, col_d, include_columns,
             exclude_columns, cols_by_orig_name):
@@ -788,3 +818,16 @@ class Inspector(object):
                     constrained_cols.append(constrained_col)
             table.append_constraint(
                 sa_schema.UniqueConstraint(*constrained_cols, name=conname))
+
+    def _reflect_check_constraints(
+            self, table_name, schema, table, cols_by_orig_name,
+            include_columns, exclude_columns, reflection_options):
+        try:
+            constraints = self.get_check_constraints(table_name, schema)
+        except NotImplementedError:
+            # optional dialect feature
+            return
+
+        for const_d in constraints:
+            table.append_constraint(
+                sa_schema.CheckConstraint(**const_d))