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

ForeignKey

By default, __set_foreign_keys__ is set to False. If it is True, all fields with the SQLAlchemy ForeignKey will be included in the resulting mock dictionary created by build method.

Setting foreign_keys

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)
    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 BookFactory(SQLAlchemyFactory[Book]): ...


class BookFactoryWithForeignKey(SQLAlchemyFactory[Book]):
    __set_foreign_keys__ = True


def test_sqla_factory() -> None:
    book = BookFactory.build()
    assert not book.author_id


def test_sqla_factory_with_foreign_keys() -> None:
    book = BookFactoryWithForeignKey.build()
    assert book.author_id
    assert isinstance(book.author_id, int)

Relationship

By default, relationships will not be set. This can be overridden via __set_relationships__.

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_without_relationship() -> None:
    author = AuthorFactory.build()
    assert author.books == []


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

Note

In general, foreign keys are not automatically generated by .build. This can be resolved by setting the fields yourself and/or using create_sync/ create_async so models can be added to a SQLA session so these are set.

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]

API reference

Full API docs are available here.