SQLAlchemyFactory

Basic usage is like other factories

Declaring a factory for a SQLAlchemy model

from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory


class Base(DeclarativeBase): ...


class Author(Base):
    __tablename__ = "authors"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]


class AuthorFactory(SQLAlchemyFactory[Author]): ...


def test_sqla_factory() -> None:
    author = AuthorFactory.build()
    assert isinstance(author, Author)

Note

The examples here require SQLAlchemy 2 to be installed. The factory itself supports both 1.4 and 2.

Configuration

SQLAlchemyFactory allows to override some configuration attributes so that a described factory can use a behavior from SQLAlchemy ORM such as relationship() or Association Proxy.

Relationship

By default, __set_relationships__ is set to False. If it is True, all fields with the SQLAlchemy relationship() will be included in the result created by build method.

Setting relationships

from typing import List

from sqlalchemy import ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory


class Base(DeclarativeBase): ...


class Author(Base):
    __tablename__ = "authors"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    books: Mapped[List["Book"]] = relationship("Book", uselist=True)


class Book(Base):
    __tablename__ = "books"

    id: Mapped[int] = mapped_column(primary_key=True)
    author_id: Mapped[int] = mapped_column(ForeignKey(Author.id))


class AuthorFactory(SQLAlchemyFactory[Author]): ...


class AuthorFactoryWithRelationship(SQLAlchemyFactory[Author]):
    __set_relationships__ = True


def test_sqla_factory() -> None:
    author = AuthorFactory.build()
    assert author.books == []


def test_sqla_factory_with_relationship() -> None:
    author = AuthorFactoryWithRelationship.build()
    assert isinstance(author, Author)
    assert isinstance(author.books[0], Book)

Note

If __set_relationships__ = True, ForeignKey fields associated with relationship() will be automatically generated by build method because __set_foreign_keys__ is set to True by default. But their values will be overwritten by using create_sync/ create_async methods, so SQLAlchemy ORM creates them.

Association Proxy

By default, __set_association_proxy__ is set to False. If it is True, all SQLAlchemy fields mapped to ORM Association Proxy class will be included in the result created by build method.

Setting association_proxy

from __future__ import annotations

from sqlalchemy import ForeignKey
from sqlalchemy.ext.associationproxy import AssociationProxy, association_proxy
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory


class Base(DeclarativeBase): ...


class User(Base):
    __tablename__ = "users"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    user_keyword_associations: Mapped[list["UserKeywordAssociation"]] = relationship(
        back_populates="user",
    )
    keywords: AssociationProxy[list["Keyword"]] = association_proxy(
        "user_keyword_associations",
        "keyword",
        creator=lambda keyword_obj: UserKeywordAssociation(keyword=keyword_obj),
    )


class UserKeywordAssociation(Base):
    __tablename__ = "user_keyword"
    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"), primary_key=True)
    keyword_id: Mapped[int] = mapped_column(ForeignKey("keywords.id"), primary_key=True)

    user: Mapped[User] = relationship(back_populates="user_keyword_associations")
    keyword: Mapped["Keyword"] = relationship()


class Keyword(Base):
    __tablename__ = "keywords"
    id: Mapped[int] = mapped_column(primary_key=True)
    keyword: Mapped[str]


class UserFactory(SQLAlchemyFactory[User]): ...


class UserFactoryWithAssociation(SQLAlchemyFactory[User]):
    __set_association_proxy__ = True


def test_sqla_factory() -> None:
    user = UserFactory.build()
    assert not user.user_keyword_associations
    assert not user.keywords


def test_sqla_factory_with_association() -> None:
    user = UserFactoryWithAssociation.build()
    assert isinstance(user.user_keyword_associations[0], UserKeywordAssociation)
    assert isinstance(user.keywords[0], Keyword)

Note

If __set_relationships__ = True, the Polyfactory will create both fields from a particular SQLAlchemy model (association_proxy and its relationship), but eventually a relationship field will be overwritten by using create_sync/ create_async methods via SQLAlchemy ORM with a proper instance from an Association Proxy relation.

Persistence

A handler is provided to allow persistence. This can be used by setting __session__ attribute on a factory.

Using persistence

from typing import List

from sqlalchemy import ForeignKey, create_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory


class Base(DeclarativeBase): ...


class Author(Base):
    __tablename__ = "authors"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    books: Mapped[List["Book"]] = relationship("Book", uselist=True)


class Book(Base):
    __tablename__ = "books"

    id: Mapped[int] = mapped_column(primary_key=True)
    author_id: Mapped[int] = mapped_column(ForeignKey(Author.id))


class AuthorFactory(SQLAlchemyFactory[Author]):
    __set_relationships__ = True


def test_sqla_factory_persistence() -> None:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine)
    session = Session(engine)

    AuthorFactory.__session__ = session  # Or using a callable that returns a session

    author = AuthorFactory.create_sync()
    assert author.id is not None
    assert author.id == author.books[0].author_id

By default, this will add generated models to the session and then commit. This can be customised further by setting __sync_persistence__.

Similarly for __async_session__ and create_async.

Adding global overrides

By combining the above and using other settings, a global base factory can be set up for other factories.

Using persistence

from typing import List

from sqlalchemy import ForeignKey, create_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory, T


class Base(DeclarativeBase): ...


class Author(Base):
    __tablename__ = "authors"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str]

    books: Mapped[List["Book"]] = relationship(
        "Book",
        uselist=True,
        back_populates="author",
    )


class Book(Base):
    __tablename__ = "books"

    id: Mapped[int] = mapped_column(primary_key=True)
    author_id: Mapped[int] = mapped_column(ForeignKey(Author.id), nullable=False)
    author: Mapped[Author] = relationship(
        "Author",
        uselist=False,
        back_populates="books",
    )


class BaseFactory(SQLAlchemyFactory[T]):
    __is_base_factory__ = True
    __set_relationships__ = True
    __randomize_collection_length__ = True
    __min_collection_length__ = 3


def test_custom_sqla_factory() -> None:
    engine = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(engine)
    session = Session(engine)

    BaseFactory.__session__ = session  # Or using a callable that returns a session

    author = BaseFactory.create_factory(Author).create_sync()
    assert author.id is not None
    assert author.id == author.books[0].author_id

    book = BaseFactory.create_factory(Book).create_sync()
    assert book.id is not None
    assert book.author.books == [book]

Add column type mapping

Columns types and normally automatically mapped to Python type. This can be overridden for cases where Python type is not available or need to provide extra information to correctly generate.

Overriding SQLA column type mapping

from __future__ import annotations

from collections.abc import Callable
from decimal import Decimal
from typing import Any

from sqlalchemy import types
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from polyfactory.factories.sqlalchemy_factory import SQLAlchemyFactory


class Location(types.TypeEngine):
    cache_ok = True


class Base(DeclarativeBase): ...


class Place(Base):
    __tablename__ = "location"

    id: Mapped[int] = mapped_column(primary_key=True)
    location: Mapped[tuple[Decimal, Decimal]] = mapped_column(Location)


class PlaceFactory(SQLAlchemyFactory[Place]):
    @classmethod
    def get_sqlalchemy_types(cls) -> dict[Any, Callable[[], Any]]:
        mapping = super().get_sqlalchemy_types()
        mapping[Location] = cls.__faker__.latlng
        return mapping


def test_custom_sqla_factory() -> None:
    result = PlaceFactory.build()
    assert isinstance(result.location, tuple)

API reference

Full API docs are available here.