In work projects in the past, my team would try to avoid dealing with the complexity of SQLAlchemy database sessions by making one global session that every module referenced. It made things easy and straightforward and it played nice with FactoryBoy - whose factories seem to work quite well under those conditions.
Going into building lobit.io, I tried to do things the same way, but as the code base and the number of unit tests grew, I kept facing what seemed to be an unquashable avalanche of OperationalErrors griping about "Too many connections".
I implemented an overhaul of the codebase, wrapping every piece of ORM code in a separate session and creating a DbUtils class to handle sessions consistently without too much boilerplate:
from contextlib import contextmanager from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from sqlalchemy.orm.session import Session from config import config class DbUtils: Session = sessionmaker(bind=create_engine(config.db_connection_url)) @classmethod @contextmanager def session_scope(cls) -> Session: """Provide a transactional scope around a series of operations.""" session = cls.Session() try: yield session session.commit() except: session.rollback() raise finally: session.close()
Fixing the test factories
When it came to unit tests, I was at a loss for how to best handle all of the factories I was using. They were all tied into the global session. Furthermore, factories took in the session at the time of class creation, not instantiation, e.g.
import factory from models.post import Post from tests.factories.post_thread_factory import PostThreadFactory class PostFactory(factory.alchemy.SQLAlchemyModelFactory): class Meta: model = Post sqlalchemy_session_persistence = "commit" sqlalchemy_session = session text = factory.Faker("text", max_nb_chars=100) external_id = factory.Faker("text", max_nb_chars=50) post_thread = factory.SubFactory(PostThreadFactory)
Wrapping factory class creation in a function did the trick:
import factory from models.post import Post from tests.factories.post_thread_factory import PostThreadFactory def PostFactory(session): class _PostFactory(factory.alchemy.SQLAlchemyModelFactory): class Meta: model = Post sqlalchemy_session_persistence = "commit" sqlalchemy_session = session text = factory.Faker("text", max_nb_chars=100) external_id = factory.Faker("text", max_nb_chars=50) post_thread = factory.SubFactory(PostThreadFactory(session)) return _PostFactory
While simple wrapping worked, it made my tests look God-awful. Manually managing sessions had already added boilerplate overhead to every unit test that hit the database.
Before refactoring, I could create a post using
PostFactory.create(). Now it was
PostFactory(session).create(). Most of the unit tests hit multiple factories several times each, so it was also less efficient to keep creating the same classes over and over.
I took to creating a single instance of each required factory at the top of each session:
def test_post_stuff(self): with DbUtils.session_scope() as session: post_factory = PostFactory(session) post_thread_factory = PostThreadFactory(session) ...
Cleaning it up
A given test could have nearly a dozen factories defined at the top - which looked terrible even if it didn't trigger a "too-many-locals" warning from Pylint (which it almost always did).
So I made one more change and hired a manager to keep all the factories running smoothly.
from tests.factories.post_factory import PostFactory from tests.factories.post_thread_factory import PostThreadFactory ... class FactoryManager: known_factories = [ PostFactory, PostThreadFactory, ... ] def __init__(self, session): self.session = session for factory_func in self.known_factories: setattr(self, factory_func.__name__, factory_func(self.session))
The FactoryManager didn't require a lot of code, but let me leave my tests largely unchanged. Instead of initializing several objects at the top of each test, I could initialize one.
def test_post_stuff(self): with DbUtils.session_scope() as session: fm = FactoryManager(session)
Now I could prefix all factory calls with
fm. and get something almost as clean as before. At least,
fm.PostFactory.create() looks better to me than
PostFactory(session).create() or littering each test with a bunch of locals.
Prefer to catch my posts elsewhere?
Originally published on my blog: http://typenil.com/sqlalchemy-session-wrapping-factories/