litestar-org/polyfactory

SQLAlchemyFactory with foreign keys

Closed this issue ยท 3 comments

Hi, I'm facing some trouble while creating factories that involve ForeignKeys in SQLAlchemy...

class EmailCategory(SoftDeleteMixin, BaseModal):

    __tablename__ = "email_categories"

    id = Column(Integer, primary_key=True)
    name = Column(String(255), unique=True, nullable=False)
    email_templates = relationship("EmailTemplate", back_populates="category")

class EmailTemplate(SoftDeleteMixin, BaseModal):

    __tablename__ = "email_templates"

    id = Column(Integer, primary_key=True)
    name = Column(String(255), unique=True, nullable=False)
    subject = Column(String(255), nullable=False)
    body = Column(Text, nullable=False)
    category_id = Column(
        Integer, ForeignKey("email_categories.id", ondelete="CASCADE"), nullable=False
    )
    category = relationship("EmailCategory", back_populates="email_templates")


class AsyncPersistenceHandler(AsyncPersistenceProtocol[T]):
    async def save(self, data: T) -> T:
        async with get_db_context() as session:
            session.add(data)
            await session.commit()

    async def save_many(self, data: List[T]) -> List[T]:
        async with get_db_context() as session:
            session.add_all(data)
            await session.commit()


class BaseFactory(SQLAlchemyFactory):
    """Base Faactory"""

    __is_base_factory__ = True
    __async_persistence__ = AsyncPersistenceHandler

    id = Ignore()
    deleted_at = None
    created_at = Use(lambda: datetime.now())


class EmailCategoryFactory(BaseFactory):
    """Factory for generating EmailCategory dummy data."""

    __model__ = EmailCategory
    name = Use(lambda: EmailCategoryFactory.__faker__.user_name())


class EmailTemplateFactory(BaseFactory):
    """Factory for generating EmailTemplate dummy data."""

    __model__ = EmailTemplate
    category = EmailCategoryFactory


async def _main():
    template = await EmailTemplateFactory.create_async()
    print(template)


if __name__ == "__main__":
    asyncio.run(_main())

When I try to create an EmailCategory model instance, I get an IntegrityError mentioning that the random integer assigned to foreign key field is not present in the table "email_categories" (which is expected). Right now I'm passing the foreign_key values manually, like this

    category = await EmailCategoryFactory.create_async()
    template = await EmailTemplateFactory.create_async(category_id=category.id)
    print(template)

Is there any other way to implement the same?


Note

While we are open for sponsoring on GitHub Sponsors and
OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh dashboard

  • If you would like to see an issue prioritized, make a pledge towards it!
  • We receive the pledge once the issue is completed & verified
  • This, along with engagement in the community, helps us know which features are a priority to our users.
Fund with Polar

Hi @vickypalani, could you try adding adding __set_relationships__? Docs here. By default these, relationships are not set as these causes some issues with recursive relationships originally.

guacs commented

@vickypalani I think the solution mentioned above should work. If not, feel free to reopen the issue.

Hi @adhtruong thanks for the response, @guacs sorry for not updating the status of the issue. I was able to resolve this after going through the documentation last week. This is how my factory looks right now,

class EmailTemplateFactory(BaseFactory):
    """Factory for generating EmailTemplate dummy data."""

    __model__ = EmailTemplate
    __set_relationships__ = True
    category = EmailCategoryFactory
    created_by = UserFactory