Skip to content

Commit

Permalink
feat: implement SQLA factory (#369)
Browse files Browse the repository at this point in the history
* feat: add sqlalchemy factory (initial)

* feat: Add Sqlalchemy factory (#342)

* feat: Add SQLAlchemyFactory table column parsing

* docs: Add SQLAlchemy reference

* docs: amend link

* doc: amend link

* Revert type updates

* feat: add SQLAlchemy relationship configuration

* feat: improve mapping for SQLAlchemy ARRAY

* feat: SQLAlchemy test hints for 3.8

* feat: SQLAlchemy test hints for 3.8

* feat: SQLAlchemy test hints for 3.8

* feat: SQLAlchemy factory hints for 3.8

* feat: add extra attributes for SQLAlchemyFactory

* feat: resolve SQLAlchemyFactory type issues

* feat: fix rebase errors

* feat: Add SQLA persistence handles, update class attributes

* docs: Add SQLAFactory docs

* test: add missing __init__ files, fix imports, fix rebase errors

* docs: fix SQLA class ref

---------

Co-authored-by: Na'aman Hirschfeld <nhirschfeld@gmail.com>
  • Loading branch information
adhtruong and Goldziher committed Sep 18, 2023
1 parent 8e41372 commit c76ffc9
Show file tree
Hide file tree
Showing 16 changed files with 806 additions and 5 deletions.
2 changes: 2 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ repos:
pydantic>=2,
pytest,
sphinx,
sqlalchemy>=2,
]
- repo: https://github.com/RobertCraigie/pyright-python
rev: v1.1.324
Expand All @@ -68,6 +69,7 @@ repos:
pydantic>=2,
pytest,
sphinx,
sqlalchemy>=2,
]
- repo: https://github.com/sphinx-contrib/sphinx-lint
rev: "v0.6.8"
Expand Down
2 changes: 2 additions & 0 deletions docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,8 @@
(PY_CLASS, "BeanieDocumentFactory"),
(PY_CLASS, "OdmanticModelFactory"),
(PY_CLASS, "ModelField"),
(PY_CLASS, "Session"),
(PY_CLASS, "AsyncSession"),
]
nitpick_ignore_regex = [
(PY_RE, r"typing_extensions.*"),
Expand Down
Empty file.
Empty file.
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
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]):
__model__ = Author


def test_sqla_factory() -> None:
author = AuthorFactory.build()
assert isinstance(author, Author)
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
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]):
__model__ = Author


class AuthorFactoryWithRelationship(SQLAlchemyFactory[Author]):
__model__ = 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)
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
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]):
__model__ = 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
1 change: 1 addition & 0 deletions docs/reference/factories/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,4 @@ factories
odmantic_odm_factory
beanie_odm_factory
attrs_factory
sqlalchemy_factory
4 changes: 4 additions & 0 deletions docs/reference/factories/sqlalchemy_factory.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
sqlalchemy_factory
==================

.. automodule:: polyfactory.factories.sqlalchemy_factory
2 changes: 1 addition & 1 deletion docs/usage/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ Usage Guide
:titlesonly:
:maxdepth: 1

library_factories
library_factories/index
declaring_factories
configuration
fields
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ These include:
:class:`AttrsFactory <polyfactory.factories.attrs_factory.AttrsFactory>`
a base factory for `attrs <https://www.attrs.org/en/stable/index.html>`_ models.

:class:`SQLAlchemyFactory <polyfactory.factories.sqlalchemy_factory.SQLAlchemyFactory>`
a base factory for `SQLAlchemy <https://www.sqlalchemy.org/>`_ models.

.. note::
All factories exported from ``polyfactory.factories`` do not require any additional dependencies. The other factories,
such as :class:`ModelFactory <polyfactory.factories.pydantic_factory.ModelFactory>`, require an additional but optional
Expand All @@ -34,3 +37,9 @@ These include:

.. note::
We will be adding additional factories to this package, so make sure to checkout the above list from time to time.


.. toctree::
:maxdepth: 1

sqlalchemy_factory
39 changes: 39 additions & 0 deletions docs/usage/library_factories/sqlalchemy_factory.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
SQLAlchemyFactory
===================

Basic usage is like other factories

.. literalinclude:: /examples/library_factories/sqlalchemy_factory/test_example_1.py
:caption: Declaring a factory for a SQLAlchemy model
:language: python

Configurations
------------------------------

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

.. literalinclude:: /examples/library_factories/sqlalchemy_factory/test_example_2.py
:caption: Setting relationships
:language: python

.. 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.

.. literalinclude:: /examples/library_factories/sqlalchemy_factory/test_example_3.py
:caption: Using persistence
:language: python

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``.


API reference
------------------------------
Full API docs are available :class:`here <polyfactory.factories.sqlalchemy_factory.SQLAlchemyFactory>`.
Loading

0 comments on commit c76ffc9

Please sign in to comment.