Repository: FactoryBoy/factory_boy
Pull request #658 (2019-10-31 22:54:19)

Add SQLAlchemy get_or_create support
This is a naive implementation purely translated from django  `get_or_create` functionality as requested in [#99](https://github.com/FactoryBoy/factory_boy/issues/99).

federicobond: Hi! Thank you for your contribution. Code looks good but we would need some tests before we can merge this.
Gonzasestopal: Sure! Tests have been added.

I found this issue regarding unicode fields while executing DB operations on python2.7 , but  `StringModel` solved it for now. Maybe consider migrating from unicode fields to string fields, or actually mocking with unicode data.
Gonzasestopal: Any update on this?
sutera: Any plans to merge this? Seems like a widely requested feature in the issues tab.
br-dr3: Any update on this?
francoisfreitag: Thank you! :tada: 

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 3a22d92..9bccbb5 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -11,6 +11,7 @@ ChangeLog
     - Report misconfiguration when a :py:class:`~factory.Factory` is used as the :py:attr:`~factory.Factory.model` for another :py:class:`~factory.Factory`.
     - Allow configuring the color palette of :py:class:`~factory.django.ImageField`.
     - :py:meth:`get_random_state()` now represents the state of Faker and ``factory_boy`` fuzzy attributes.
+    - Add SQLAlchemy ``get_or_create`` support
 
 *Bugfix:*
 
diff --git a/docs/orms.rst b/docs/orms.rst
index 0883c6c..e67eb13 100644
--- a/docs/orms.rst
+++ b/docs/orms.rst
@@ -325,6 +325,43 @@ To work, this class needs an `SQLAlchemy`_ session object affected to the :attr:
 
         The default value is ``None``.
 
+    .. attribute:: sqlalchemy_get_or_create
+
+        .. versionadded:: 3.0.0
+
+        Fields whose name are passed in this list will be used to perform a
+        :meth:`Model.query.one_or_none() <sqlalchemy.orm.query.Query.one_or_none>`
+        or the usual :meth:`Session.add() <sqlalchemy.orm.session.Session.add>`:
+
+        .. code-block:: python
+
+            class UserFactory(factory.alchemy.SQLAlchemyModelFactory):
+                class Meta:
+                    model = User
+                    sqlalchemy_session = session
+                    sqlalchemy_get_or_create = ('username',)
+
+                username = 'john'
+
+        .. code-block:: pycon
+
+            >>> User.query.all()
+            []
+            >>> UserFactory()                   # Creates a new user
+            <User: john>
+            >>> User.query.all()
+            [<User: john>]
+
+            >>> UserFactory()                   # Fetches the existing user
+            <User: john>
+            >>> User.query.all()                # No new user!
+            [<User: john>]
+
+            >>> UserFactory(username='jack')    # Creates another user
+            <User: jack>
+            >>> User.query.all()
+            [<User: john>, <User: jack>]
+
 
 A (very) simple example:
 
diff --git a/factory/alchemy.py b/factory/alchemy.py
index 862f89e..beccdcc 100644
--- a/factory/alchemy.py
+++ b/factory/alchemy.py
@@ -1,6 +1,9 @@
 # Copyright: See the LICENSE file.
 
-from . import base
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm.exc import NoResultFound
+
+from . import base, errors
 
 SESSION_PERSISTENCE_COMMIT = 'commit'
 SESSION_PERSISTENCE_FLUSH = 'flush'
@@ -21,6 +24,7 @@ class SQLAlchemyOptions(base.FactoryOptions):
 
     def _build_default_options(self):
         return super()._build_default_options() + [
+            base.OptionDefault('sqlalchemy_get_or_create', (), inherit=True),
             base.OptionDefault('sqlalchemy_session', None, inherit=True),
             base.OptionDefault(
                 'sqlalchemy_session_persistence',
@@ -39,15 +43,66 @@ class SQLAlchemyModelFactory(base.Factory):
     class Meta:
         abstract = True
 
+    @classmethod
+    def _generate(cls, strategy, params):
+        # Original params are used in _get_or_create if it cannot build an
+        # object initially due to an IntegrityError being raised
+        cls._original_params = params
+        return super(SQLAlchemyModelFactory, cls)._generate(strategy, params)
+
+    @classmethod
+    def _get_or_create(cls, model_class, session, *args, **kwargs):
+        key_fields = {}
+        for field in cls._meta.sqlalchemy_get_or_create:
+            if field not in kwargs:
+                raise errors.FactoryError(
+                    "sqlalchemy_get_or_create - "
+                    "Unable to find initialization value for '%s' in factory %s" %
+                    (field, cls.__name__))
+            key_fields[field] = kwargs.pop(field)
+
+        obj = session.query(model_class).filter_by(
+            *args, **key_fields).one_or_none()
+
+        if not obj:
+            try:
+                obj = cls._save(model_class, session, *args, **key_fields, **kwargs)
+            except IntegrityError as e:
+                session.rollback()
+                get_or_create_params = {
+                    lookup: value
+                    for lookup, value in cls._original_params.items()
+                    if lookup in cls._meta.sqlalchemy_get_or_create
+                }
+                if get_or_create_params:
+                    try:
+                        obj = session.query(model_class).filter_by(
+                            **get_or_create_params).one()
+                    except NoResultFound:
+                        # Original params are not a valid lookup and triggered a create(),
+                        # that resulted in an IntegrityError.
+                        raise e
+                else:
+                    raise e
+
+        return obj
+
     @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
-        session_persistence = cls._meta.sqlalchemy_session_persistence
 
-        obj = model_class(*args, **kwargs)
         if session is None:
             raise RuntimeError("No session provided.")
+        if cls._meta.sqlalchemy_get_or_create:
+            return cls._get_or_create(model_class, session, *args, **kwargs)
+        return cls._save(model_class, session, *args, **kwargs)
+
+    @classmethod
+    def _save(cls, model_class, session, *args, **kwargs):
+        session_persistence = cls._meta.sqlalchemy_session_persistence
+
+        obj = model_class(*args, **kwargs)
         session.add(obj)
         if session_persistence == SESSION_PERSISTENCE_FLUSH:
             session.flush()


Test patch:
diff --git a/tests/alchemyapp/models.py b/tests/alchemyapp/models.py
index e31e903..013a56a 100644
--- a/tests/alchemyapp/models.py
+++ b/tests/alchemyapp/models.py
@@ -20,6 +20,23 @@ class StandardModel(Base):
     foo = Column(Unicode(20))
 
 
+class MultiFieldModel(Base):
+    __tablename__ = 'MultiFieldModelTable'
+
+    id = Column(Integer(), primary_key=True)
+    foo = Column(Unicode(20))
+    slug = Column(Unicode(20), unique=True)
+
+
+class MultifieldUniqueModel(Base):
+    __tablename__ = 'MultiFieldUniqueModelTable'
+
+    id = Column(Integer(), primary_key=True)
+    slug = Column(Unicode(20), unique=True)
+    text = Column(Unicode(20), unique=True)
+    title = Column(Unicode(20), unique=True)
+
+
 class NonIntegerPk(Base):
     __tablename__ = 'NonIntegerPk'
 
diff --git a/tests/test_alchemy.py b/tests/test_alchemy.py
index 1f05c31..2205ca1 100644
--- a/tests/test_alchemy.py
+++ b/tests/test_alchemy.py
@@ -5,6 +5,8 @@
 import unittest
 from unittest import mock
 
+import sqlalchemy
+
 import factory
 from factory.alchemy import SQLAlchemyModelFactory
 
@@ -36,8 +38,41 @@ class NoSessionFactory(SQLAlchemyModelFactory):
     id = factory.Sequence(lambda n: n)
 
 
-class SQLAlchemyPkSequenceTestCase(unittest.TestCase):
+class MultifieldModelFactory(SQLAlchemyModelFactory):
+    class Meta:
+        model = models.MultiFieldModel
+        sqlalchemy_get_or_create = ('slug',)
+        sqlalchemy_session = models.session
+        sqlalchemy_session_persistence = 'commit'
+
+    id = factory.Sequence(lambda n: n)
+    foo = factory.Sequence(lambda n: 'foo%d' % n)
 
+
+class WithGetOrCreateFieldFactory(SQLAlchemyModelFactory):
+    class Meta:
+        model = models.StandardModel
+        sqlalchemy_get_or_create = ('foo',)
+        sqlalchemy_session = models.session
+        sqlalchemy_session_persistence = 'commit'
+
+    id = factory.Sequence(lambda n: n)
+    foo = factory.Sequence(lambda n: 'foo%d' % n)
+
+
+class WithMultipleGetOrCreateFieldsFactory(SQLAlchemyModelFactory):
+    class Meta:
+        model = models.MultifieldUniqueModel
+        sqlalchemy_get_or_create = ("slug", "text",)
+        sqlalchemy_session = models.session
+        sqlalchemy_session_persistence = 'commit'
+
+    id = factory.Sequence(lambda n: n)
+    slug = factory.Sequence(lambda n: "slug%s" % n)
+    text = factory.Sequence(lambda n: "text%s" % n)
+
+
+class SQLAlchemyPkSequenceTestCase(unittest.TestCase):
     def setUp(self):
         super().setUp()
         StandardFactory.reset_sequence(1)
@@ -74,6 +109,61 @@ class SQLAlchemyPkSequenceTestCase(unittest.TestCase):
         self.assertEqual(0, std2.id)
 
 
+class SQLAlchemyGetOrCreateTests(unittest.TestCase):
+    def setUp(self):
+        models.session.rollback()
+
+    def test_simple_call(self):
+        obj1 = WithGetOrCreateFieldFactory(foo='foo1')
+        obj2 = WithGetOrCreateFieldFactory(foo='foo1')
+        self.assertEqual(obj1, obj2)
+
+    def test_missing_arg(self):
+        with self.assertRaises(factory.FactoryError):
+            MultifieldModelFactory()
+
+    def test_raises_exception_when_existing_objs(self):
+        StandardFactory.create_batch(2, foo='foo')
+        with self.assertRaises(sqlalchemy.orm.exc.MultipleResultsFound):
+            WithGetOrCreateFieldFactory(foo='foo')
+
+    def test_multicall(self):
+        objs = MultifieldModelFactory.create_batch(
+            6,
+            slug=factory.Iterator(['main', 'alt']),
+        )
+        self.assertEqual(6, len(objs))
+        self.assertEqual(2, len(set(objs)))
+        self.assertEqual(
+            list(
+                obj.slug for obj in models.session.query(
+                    models.MultiFieldModel.slug
+                )
+            ),
+            ["alt", "main"],
+        )
+
+
+class MultipleGetOrCreateFieldsTest(unittest.TestCase):
+    def setUp(self):
+        models.session.rollback()
+
+    def test_one_defined(self):
+        obj1 = WithMultipleGetOrCreateFieldsFactory()
+        obj2 = WithMultipleGetOrCreateFieldsFactory(slug=obj1.slug)
+        self.assertEqual(obj1, obj2)
+
+    def test_both_defined(self):
+        obj1 = WithMultipleGetOrCreateFieldsFactory()
+        with self.assertRaises(sqlalchemy.exc.IntegrityError):
+            WithMultipleGetOrCreateFieldsFactory(slug=obj1.slug, text="alt")
+
+    def test_unique_field_not_in_get_or_create(self):
+        WithMultipleGetOrCreateFieldsFactory(title='Title')
+        with self.assertRaises(sqlalchemy.exc.IntegrityError):
+            WithMultipleGetOrCreateFieldsFactory(title='Title')
+
+
 class SQLAlchemySessionPersistenceTestCase(unittest.TestCase):
     def setUp(self):
         super().setUp()