litestar-org / polyfactory

Simple and powerful factories for mock data generation
https://polyfactory.litestar.dev/
MIT License
1.06k stars 83 forks source link

Enhancement: Add SQLModelFactory #160

Closed mrharpo closed 1 year ago

mrharpo commented 1 year ago

Description

Adapting the pets example from the documentation to SQLModel, does not generate sub-factories as expected when using Relationship, despite explicitly calling Use. Foreign keys are auto-generated correctly.

Passing in a pre-built sub-model to the parent factory works as expected.

URL to code causing the issue

No response

MCVE

Working

from sqlmodel import SQLModel, Field
from polyfactory import Use
from polyfactory.factories.pydantic_factory import ModelFactory
from typing import List, Optional

class Pet(SQLModel, table=True):
    id: Optional[int] = Field(primary_key=True, default=None)
    name: str = Field(index=True)
    sound: str = Field(index=True)

class Person(SQLModel):
    id: Optional[int] = Field(primary_key=True, default=None)
    name: str = Field(index=True)
    pets: List[Pet] = Field(index=True)

class PetFactory(ModelFactory[Pet]):
    __model__ = Pet
    __allow_none_optionals___ = False

    name = Use(ModelFactory.__random__.choice, ["Ralph", "Roxy"])

class PersonFactory(ModelFactory[Person]):
    __model__ = Person
    __allow_none_optionals___ = False

    pets = Use(PetFactory.batch, size=2)

if __name__ == '__main__':
    f = PersonFactory()
    print(f.build())

# => id=None name='amCvFfBbBknGOJkgMyjz' pets=[Pet(id=7021, name='Roxy', sound='QDzzuwnCzZlMTlibiesY'), Pet(id=1446, name='Ralph', sound='hmArwGgEkGqjJpuVOOBi')]

Not working

from sqlmodel import SQLModel, Field, Relationship
from polyfactory import Use
from polyfactory.factories.pydantic_factory import ModelFactory
from typing import List, Optional

class Pet(SQLModel, table=True):
    id: Optional[int] = Field(primary_key=True, default=None)
    name: str = Field(index=True)
    sound: str = Field(index=True)
    person_id: int = Field(foreign_key='person.id', default=None)
    person: 'Person' = Relationship(back_populates='pets')

class Person(SQLModel, table=True):
    id: Optional[int] = Field(primary_key=True, default=None)
    name: str = Field(index=True)
    pets: List[Pet] = Relationship(back_populates='person')

class PetFactory(ModelFactory[Pet]):
    __model__ = Pet
    __allow_none_optionals___ = False

    name = Use(ModelFactory.__random__.choice, ["Ralph", "Roxy"])

class PersonFactory(ModelFactory[Person]):
    __model__ = Person
    __allow_none_optionals___ = False

    pets = Use(PetFactory().batch, size=2)

if __name__ == '__main__':
    f = PersonFactory()
    print(f.build())
    # => id=None name='RbVNzbPaSyukYfQfOCZM'
    p = PetFactory()
    print('manual', f.build(pets=p.batch(size=2)))
    # => manual id=912 name='nxthozUfRKEbzRbHOezx' pets=[Pet(id=None, name='Ralph', sound='VURZSWbctAUgQOrCNsHt', person_id=2655, person=Person(id=912, name='nxthozUfRKEbzRbHOezx', pets=[...])), Pet(id=None, name='Ralph', sound='iBvnxluOSPHyYHPMPTwj', person_id=4501, person=Person(id=912, name='nxthozUfRKEbzRbHOezx', pets=[...]))]

Steps to reproduce

1. Create SQLModel models
2. See sub-factories working
3. Attach relationships
4. Sub-factories don't generate automatically

Screenshots

No response

Logs

No response

Starlite Version

Platform

Fund with Polar

Goldziher commented 1 year ago

This is not a bug - you are expecting SQLModel relations to be created, but the pydantic factory doesnt know what these are - its not a pydantic entity at all, but an SQLModel specific thing. There needs to be a special SQLModel factory for this to work, which currently there isnt. It can be a feature request of course.

guacs commented 1 year ago

Closed by #369.