Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 47 additions & 4 deletions polyfactory/factories/sqlalchemy_factory.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import enum
import inspect
from collections.abc import Collection, Mapping
from dataclasses import is_dataclass
from datetime import date, datetime
Expand All @@ -17,21 +18,24 @@
cast,
)

from sqlalchemy.sql.schema import ColumnDefault
from sqlalchemy.util.langhelpers import duck_type_collection

from polyfactory.exceptions import ConfigurationException, MissingDependencyException, ParameterException
from polyfactory.factories.base import BaseFactory
from polyfactory.factories.base import BuildContext as BaseBuildContext
from polyfactory.field_meta import Constraints, FieldMeta
from polyfactory.field_meta import Constraints, FieldMeta, Null
from polyfactory.persistence import AsyncPersistenceProtocol, SyncPersistenceProtocol
from polyfactory.utils.types import Frozendict

try:
from sqlalchemy import ARRAY, Column, Numeric, String, inspect, types
from sqlalchemy import ARRAY, Column, Numeric, String, types
from sqlalchemy import inspect as sqla_inspect
from sqlalchemy.dialects import mssql, mysql, postgresql, sqlite
from sqlalchemy.exc import NoInspectionAvailable
from sqlalchemy.ext.associationproxy import AssociationProxy
from sqlalchemy.orm import InstanceState, Mapper, RelationshipProperty
from sqlalchemy.sql.elements import TextClause
except ImportError as e:
msg = "sqlalchemy is not installed"
raise MissingDependencyException(msg) from e
Expand Down Expand Up @@ -178,6 +182,21 @@ async def create_async(cls, **kwargs: Any) -> T:
kwargs["_build_context"] = build_context
return await super().create_async(**kwargs)

@classmethod
def build(cls, *_: Any, **kwargs: Any) -> T:
"""Build an instance of the factory's __model__ and apply field defaults."""
data = super().process_kwargs(**kwargs)

if cls.__use_defaults__:
for field_meta in cls.get_model_fields():
if field_meta.name in data:
continue

if cls.should_use_default_value(field_meta):
data[field_meta.name] = field_meta.default

return cls.__model__(**data)

@classmethod
def get_sqlalchemy_types(cls) -> dict[Any, Callable[[], Any]]:
"""Get mapping of types where column type should be used directly.
Expand Down Expand Up @@ -227,7 +246,7 @@ def get_provider_map(cls) -> dict[Any, Callable[[], Any]]:
@classmethod
def is_supported_type(cls, value: Any) -> TypeGuard[type[T]]:
try:
inspected = inspect(value)
inspected = sqla_inspect(value)
except NoInspectionAvailable:
return False
return isinstance(inspected, (Mapper, InstanceState))
Expand Down Expand Up @@ -365,15 +384,39 @@ def _get_association_proxy_type(cls, table: Mapper, proxy: AssociationProxy) ->
class_ = target_attr.entity.class_
return class_ if not target_collection.uselist else list[class_] # type: ignore[valid-type]

@classmethod
def _get_column_default(cls, col: Column) -> Any:
if (default := col.default) is None:
return Null

if not isinstance(default, ColumnDefault):
return Null

if default.is_scalar:
return default.arg

if default.is_clause_element and isinstance(default.arg, TextClause):
return default.arg.text

if default.is_callable:
sig = inspect.signature(default.arg)
# Skip context-sensitive default functions as there is no
# obvious and correct way to construct the SQLAlchemy context manually.
if not sig.parameters:
return default.arg(None) # type: ignore[arg-type]

return Null

@classmethod
def get_model_fields(cls) -> list[FieldMeta]:
fields_meta: list[FieldMeta] = []

table: Mapper = inspect(cls.__model__) # type: ignore[assignment]
table: Mapper = sqla_inspect(cls.__model__) # type: ignore[assignment]
fields_meta.extend(
FieldMeta.from_type(
annotation=cls.get_type_from_column(column),
name=name,
default=cls._get_column_default(column),
)
for name, column in table.columns.items()
if cls.should_column_be_set(column)
Expand Down
20 changes: 20 additions & 0 deletions tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import text

try:
from sqlalchemy.orm.collections import (
Expand Down Expand Up @@ -202,6 +203,25 @@ class ModelFactory(SQLAlchemyFactory[Model]):
assert result.optional_field is None


def test_default_fields() -> None:
class Model(Base):
__tablename__ = "model_with_default_fields"

id = Column(Integer(), primary_key=True)
default_scalar = Column(Integer(), default=10)
default_callable = Column(String(), default=lambda: "callable")
default_sql = Column(String(), default=text("0"))

class ModelFactory(SQLAlchemyFactory[Model]):
__model__ = Model
__use_defaults__ = True

result = ModelFactory.build()
assert result.default_scalar == 10
assert result.default_callable == "callable"
assert result.default_sql == "0"


@pytest.mark.parametrize("set_primary_key", (True, False))
def test_set_primary_key(set_primary_key: bool) -> None:
class AuthorFactory(SQLAlchemyFactory[Author]):
Expand Down
Loading