Add support for constraints, refs #161

Add support for foreign key constraints in has_index and has_unique_index. Remove has_indexed_foreign_key function (superceded now by more versatile has_index function).

CHANGES.rst

@@ -4,6 +4,15 @@ Changelog
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
 
+0.31.0 (2015-09-17)
+^^^^^^^^^^^^^^^^^^^
+
+- Made has_index allow fk constraint as parameter
+- Made has_unique_index allow fk constraint as parameter
+- Made the extra packages in setup.py to be returned in deterministic order
+- Removed is_indexed_foreign_key (superceded by more versatile has_index)
+
+
 0.30.17 (2015-08-16)
 ^^^^^^^^^^^^^^^^^^^^
 

sqlalchemy_utils/__init__.py

@@ -22,6 +22,7 @@ from .functions import (  # noqa
     get_column_key,
     get_columns,
     get_declarative_base,
+    get_fk_constraint_for_columns,
     get_hybrid_properties,
     get_mapper,
     get_primary_keys,
@@ -92,4 +93,4 @@ from .types import (  # noqa
     WeekDaysType
 )
 
-__version__ = '0.30.17'
+__version__ = '0.31.0'

sqlalchemy_utils/functions/__init__.py

@@ -11,9 +11,9 @@ from .database import (  # noqa
 )
 from .foreign_keys import (  # noqa
     dependent_objects,
+    get_fk_constraint_for_columns,
     get_referencing_foreign_keys,
     group_foreign_keys,
-    is_indexed_foreign_key,
     merge_references,
     non_indexed_foreign_keys
 )

sqlalchemy_utils/functions/database.py

@@ -9,6 +9,7 @@ from sqlalchemy.exc import OperationalError, ProgrammingError
 
 from sqlalchemy_utils.expressions import explain_analyze
 
+from ..utils import starts_with
 from .orm import quote
 
 
@@ -187,16 +188,23 @@ def json_sql(value, scalars_to_json=True):
     return value
 
 
-def has_index(column):
+def has_index(column_or_constraint):
     """
-    Return whether or not given column has an index. A column has an index if
+    Return whether or not given column or the columns of given foreign key
-    it has a single column index or it is the first column in compound column
+    constraint have an index. A column has an index if it has a single column
-    index.
+    index or it is the first column in compound column index.
 
-    :param column: SQLAlchemy Column object
+    A foreign key constraint has an index if the constraint columns are the
+    first columns in compound column index.
+
+    :param column_or_constraint:
+        SQLAlchemy Column object or SA ForeignKeyConstraint object
 
     .. versionadded: 0.26.2
 
+    .. versionchanged: 0.30.18
+        Added support for foreign key constaints.
+
     ::
 
         from sqlalchemy_utils import has_index
@@ -240,34 +248,82 @@ def has_index(column):
 
         has_index(table.c.locale)   # False
         has_index(table.c.id)       # True
+
+
+    This function supports foreign key constraints as well:
+
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index',
+                    author_first_name,
+                    author_last_name
+                )
+            )
+
+        table = Article.__table__
+        constraint = list(table.foreign_keys)[0].constraint
+
+        has_index(constraint)  # True
     """
-    table = column.table
+    table = column_or_constraint.table
     if not isinstance(table, sa.Table):
         raise TypeError(
             'Only columns belonging to Table objects are supported. Given '
             'column belongs to %r.' % table
         )
     primary_keys = table.primary_key.columns.values()
+    if isinstance(column_or_constraint, sa.ForeignKeyConstraint):
+        columns = list(column_or_constraint.columns.values())
+    else:
+        columns = [column_or_constraint]
+
     return (
-        (primary_keys and column is primary_keys[0])
+        (primary_keys and starts_with(primary_keys, columns))
         or
         any(
-            index.columns.values()[0] is column
+            starts_with(index.columns.values(), columns)
             for index in table.indexes
         )
     )
 
 
-def has_unique_index(column):
+def has_unique_index(column_or_constraint):
     """
-    Return whether or not given column has a unique index. A column has a
+    Return whether or not given column or given foreign key constraint has a
-    unique index if it has a single column primary key index or it has a
+    unique index.
-    single column UniqueConstraint.
+
+    A column has a unique index if it has a single column primary key index or
+    it has a single column UniqueConstraint.
+
+    A foreign key constraint has a unique index if the columns of the
+    constraint are the same as the columns of table primary key or the coluns
+    of any unique index or any unique constraint of the given table.
 
     :param column: SQLAlchemy Column object
 
     .. versionadded: 0.27.1
 
+    .. versionchanged: 0.30.18
+        Added support for foreign key constaints.
+
+        Fixed support for unique indexes (previously only worked for unique
+        constraints)
+
     ::
 
         from sqlalchemy_utils import has_unique_index
@@ -289,31 +345,69 @@ def has_unique_index(column):
         has_unique_index(table.c.id)           # True
 
 
+    This function supports foreign key constraints as well:
+
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index',
+                    author_first_name,
+                    author_last_name,
+                    unique=True
+                )
+            )
+
+        table = Article.__table__
+        constraint = list(table.foreign_keys)[0].constraint
+
+        has_unique_index(constraint)  # True
+
+
     :raises TypeError: if given column does not belong to a Table object
     """
-    table = column.table
+    table = column_or_constraint.table
     if not isinstance(table, sa.Table):
         raise TypeError(
             'Only columns belonging to Table objects are supported. Given '
             'column belongs to %r.' % table
         )
-    pks = table.primary_key.columns
+    primary_keys = list(table.primary_key.columns.values())
+    if isinstance(column_or_constraint, sa.ForeignKeyConstraint):
+        columns = list(column_or_constraint.columns.values())
+    else:
+        columns = [column_or_constraint]
+
     return (
-        (column is pks.values()[0] and len(pks) == 1)
+        (columns == primary_keys)
         or
         any(
-            match_columns(constraint.columns.values()[0], column) and
+            columns == list(constraint.columns.values())
-            len(constraint.columns) == 1
+            for constraint in table.constraints
-            for constraint in column.table.constraints
             if isinstance(constraint, sa.sql.schema.UniqueConstraint)
         )
+        or
+        any(
+            columns == list(index.columns.values())
+            for index in table.indexes
+            if index.unique
+        )
     )
 
 
-def match_columns(column, column2):
-    return column.table is column2.table and column.name == column2.name
-
-
 def is_auto_assigned_date_column(column):
     """
     Returns whether or not given SQLAlchemy Column object's is auto assigned

sqlalchemy_utils/functions/foreign_keys.py

@@ -8,6 +8,7 @@ from sqlalchemy.orm import object_session
 from sqlalchemy.schema import ForeignKeyConstraint, MetaData, Table
 
 from ..query_chain import QueryChain
+from .database import has_index
 from .orm import get_column_key, get_mapper, get_tables
 
 
@@ -344,22 +345,13 @@ def non_indexed_foreign_keys(metadata, engine=None):
             if not isinstance(constraint, ForeignKeyConstraint):
                 continue
 
-            if not is_indexed_foreign_key(constraint):
+            if not has_index(constraint):
                 constraints[table.name].append(constraint)
 
     return dict(constraints)
 
 
-def is_indexed_foreign_key(constraint):
+def get_fk_constraint_for_columns(table, *columns):
-    """
+    for constraint in table.constraints:
-    Whether or not given foreign key constraint's columns have been indexed.
+        if list(constraint.columns.values()) == list(columns):
-
+            return constraint
-    :param constraint: ForeignKeyConstraint object to check the indexes
-    """
-    return any(
-        set(constraint.columns.keys())
-        ==
-        set(column.name for column in index.columns)
-        for index
-        in constraint.table.indexes
-    )

sqlalchemy_utils/utils.py

@@ -20,3 +20,10 @@ def is_sequence(value):
     return (
         isinstance(value, Iterable) and not isinstance(value, six.string_types)
     )
+
+
+def starts_with(iterable, prefix):
+    """
+    Returns whether or not given iterable starts with given prefix.
+    """
+    return list(iterable)[0:len(prefix)] == list(prefix)

tests/functions/test_has_index.py

@@ -2,7 +2,7 @@ import sqlalchemy as sa
 from pytest import raises
 from sqlalchemy.ext.declarative import declarative_base
 
-from sqlalchemy_utils import has_index
+from sqlalchemy_utils import get_fk_constraint_for_columns, has_index
 
 
 class TestHasIndex(object):
@@ -47,3 +47,98 @@ class TestHasIndex(object):
             sa.Column('name', sa.String)
         )
         assert not has_index(article.c.name)
+
+
+class TestHasIndexWithFKConstraint(object):
+    def test_composite_fk_without_index(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert not has_index(constraint)
+
+    def test_composite_fk_with_index(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index', author_first_name, author_last_name
+                )
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert has_index(constraint)
+
+    def test_composite_fk_with_partial_index_match(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index',
+                    author_first_name,
+                    author_last_name,
+                    id
+                )
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert has_index(constraint)

tests/functions/test_has_unique_index.py

@@ -2,7 +2,7 @@ import sqlalchemy as sa
 from pytest import raises
 from sqlalchemy.ext.declarative import declarative_base
 
-from sqlalchemy_utils import has_unique_index
+from sqlalchemy_utils import get_fk_constraint_for_columns, has_unique_index
 
 
 class TestHasUniqueIndex(object):
@@ -50,3 +50,102 @@ class TestHasUniqueIndex(object):
     def test_compound_column_unique_index(self):
         assert not has_unique_index(self.article_translations.c.is_published)
         assert not has_unique_index(self.article_translations.c.is_archived)
+
+
+class TestHasUniqueIndexWithFKConstraint(object):
+    def test_composite_fk_without_index(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert not has_unique_index(constraint)
+
+    def test_composite_fk_with_index(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index',
+                    author_first_name,
+                    author_last_name,
+                    unique=True
+                )
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert has_unique_index(constraint)
+
+    def test_composite_fk_with_partial_index_match(self):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            first_name = sa.Column(sa.Unicode(255), primary_key=True)
+            last_name = sa.Column(sa.Unicode(255), primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_first_name = sa.Column(sa.Unicode(255))
+            author_last_name = sa.Column(sa.Unicode(255))
+            __table_args__ = (
+                sa.ForeignKeyConstraint(
+                    [author_first_name, author_last_name],
+                    [User.first_name, User.last_name]
+                ),
+                sa.Index(
+                    'my_index',
+                    author_first_name,
+                    author_last_name,
+                    id,
+                    unique=True
+                )
+            )
+
+        table = Article.__table__
+        constraint = get_fk_constraint_for_columns(
+            table,
+            table.c.author_first_name,
+            table.c.author_last_name
+        )
+        assert not has_unique_index(constraint)