diff options
| -rw-r--r-- | factory/__init__.py | 2 | ||||
| -rw-r--r-- | factory/alchemy.py | 12 |
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 |