1 files changed, 3 insertions, 18 deletions
diff --git a/factory/alchemy.py b/factory/alchemy.py
index 3c91411..a9aab23 100644
--- a/factory/alchemy.py
+++ b/factory/alchemy.py
@@ -19,7 +19,6 @@
 # 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
 
@@ -28,6 +27,7 @@ class SQLAlchemyOptions(base.FactoryOptions):
     def _build_default_options(self):
         return super(SQLAlchemyOptions, self)._build_default_options() + [
             base.OptionDefault('sqlalchemy_session', None, inherit=True),
+            base.OptionDefault('force_flush', False, inherit=True),
         ]
 
 
@@ -38,27 +38,12 @@ class SQLAlchemyModelFactory(base.Factory):
     class Meta:
         abstract = True
 
-    _OLDSTYLE_ATTRIBUTES = base.Factory._OLDSTYLE_ATTRIBUTES.copy()
-    _OLDSTYLE_ATTRIBUTES.update({
-        'FACTORY_SESSION': 'sqlalchemy_session',
-    })
-
-    @classmethod
-    def _setup_next_sequence(cls, *args, **kwargs):
-        """Compute the next available PK, based on the 'pk' database field."""
-        session = cls._meta.sqlalchemy_session
-        model = cls._meta.model
-        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
-
     @classmethod
     def _create(cls, model_class, *args, **kwargs):
         """Create an instance of the model, and save it to the database."""
         session = cls._meta.sqlalchemy_session
         obj = model_class(*args, **kwargs)
         session.add(obj)
+        if cls._meta.force_flush:
+            session.flush()
         return obj