Small next sequence refactoring

2 files changed, 8 insertions, 6 deletions

diff --git a/factory/__init__.py b/factory/__init__.py
index 058fd76..f90f40a 100644
--- a/factory/__init__.py
+++ b/factory/__init__.py
@@ -23,6 +23,7 @@
 __version__ = '2.1.0-dev'
 __author__ = 'Raphaël Barrois <raphael.barrois+fboy@polytechnique.org>'
 
+
 from .base import (
     Factory,
     BaseDictFactory,
@@ -39,7 +40,6 @@ from .base import (
 
 from .mogo import MogoFactory
 from .django import DjangoModelFactory
-from .alchemy import SQLAlchemyModelFactory
 
 from .declarations import (
     LazyAttribute,
diff --git a/factory/alchemy.py b/factory/alchemy.py
index 2bfaf81..ca7aefa 100644
--- a/factory/alchemy.py
+++ b/factory/alchemy.py
@@ -19,6 +19,8 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 from __future__ import unicode_literals
+from sqlalchemy.sql.functions import max
+
 from . import base
 
 
@@ -30,12 +32,12 @@ class SQLAlchemyModelFactory(base.Factory):
     @classmethod
     def _setup_next_sequence(cls, *args, **kwargs):
         """Compute the next available PK, based on the 'pk' database field."""
-        from sqlalchemy.sql.functions import max
         session = cls.FACTORY_SESSION
-        pk = cls.FACTORY_FOR.__table__.primary_key.columns.values()[0].key
-        max_pk = session.query(max(getattr(cls.FACTORY_FOR, pk))).one()
-        if isinstance(max_pk[0], int):
-            return max_pk[0] + 1 if max_pk[0] else 1
+        model = cls.FACTORY_FOR
+        pk = getattr(model, model.__mapper__.primary_key[0].name)
+        max_pk = session.query(max(pk)).one()[0]
+        if isinstance(max_pk, int):
+            return max_pk + 1 if max_pk else 1
         else:
             return 1