SQLAlchemyDTO: column/relationship type inference. (#1879)

If type annotations aren't available for a given column/relationship we make an effort
to infer a type annotation from the relevant SQLAlchemy object.

Closes #1853

Author: peterschutt

litestar/contrib/sqlalchemy/dto.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import singledispatchmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar
 
 from sqlalchemy import Column, inspect, orm, sql
 from sqlalchemy.ext.associationproxy import AssociationProxy, AssociationProxyExtensionType
@@ -13,6 +13,7 @@
     Mapped,
     NotExtension,
     QueryableAttribute,
+    RelationshipDirection,
     RelationshipProperty,
 )
 
@@ -22,6 +23,7 @@
 from litestar.dto.factory.utils import get_model_type_hints
 from litestar.exceptions import ImproperlyConfiguredException
 from litestar.types.empty import Empty
+from litestar.typing import ParsedType
 from litestar.utils.helpers import get_fully_qualified_class_name
 from litestar.utils.signature import ParsedSignature
 
@@ -30,8 +32,6 @@
 
     from typing_extensions import TypeAlias
 
-    from litestar.typing import ParsedType
-
 __all__ = ("SQLAlchemyDTO",)
 
 T = TypeVar("T", bound="DeclarativeBase | Collection[DeclarativeBase]")
@@ -89,10 +89,8 @@ def _(
                 (parsed_type,) = parsed_type.inner_types
             else:
                 raise NotImplementedError(f"Expected 'Mapped' origin, got: '{parsed_type.origin}'")
-        except KeyError as e:
-            raise ImproperlyConfiguredException(
-                f"No type information found for '{orm_descriptor}'. Has a type annotation been added to the column?"
-            ) from e
+        except KeyError:
+            parsed_type = parse_type_from_element(elem)
 
         return [
             FieldDefinition(
@@ -221,6 +219,59 @@ def default_factory(d: Any = sqla_default) -> Any:
         else:
             raise ValueError("Unexpected default type")
     else:
-        if getattr(elem, "nullable", False):
+        if (
+            isinstance(elem, RelationshipProperty)
+            and detect_nullable_relationship(elem)
+            or getattr(elem, "nullable", False)
+        ):
             default = None
+
     return default, default_factory
+
+
+def parse_type_from_element(elem: ElementType) -> ParsedType:
+    """Parses a type from a SQLAlchemy element.
+
+    Args:
+        elem: The SQLAlchemy element to parse.
+
+    Returns:
+        ParsedType: The parsed type.
+
+    Raises:
+        ImproperlyConfiguredException: If the type cannot be parsed.
+    """
+
+    if isinstance(elem, Column):
+        if elem.nullable:
+            return ParsedType(Optional[elem.type.python_type])
+        return ParsedType(elem.type.python_type)
+
+    if isinstance(elem, RelationshipProperty):
+        if elem.direction in (RelationshipDirection.ONETOMANY, RelationshipDirection.MANYTOMANY):
+            collection_type = ParsedType(elem.collection_class or list)
+            return ParsedType(collection_type.safe_generic_origin[elem.mapper.class_])
+
+        if detect_nullable_relationship(elem):
+            return ParsedType(Optional[elem.mapper.class_])
+
+        return ParsedType(elem.mapper.class_)
+
+    raise ImproperlyConfiguredException(
+        f"Unable to parse type from element '{elem}'. Consider adding a type hint.",
+    )
+
+
+def detect_nullable_relationship(elem: RelationshipProperty) -> bool:
+    """Detects if a relationship is nullable.
+
+    This attempts to decide if we should allow a ``None`` default value for a relationship by looking at the
+    foreign key fields. If all foreign key fields are nullable, then we allow a ``None`` default value.
+
+    Args:
+        elem: The relationship to check.
+
+    Returns:
+        bool: ``True`` if the relationship is nullable, ``False`` otherwise.
+    """
+    return elem.direction == RelationshipDirection.MANYTOONE and all(c.nullable for c in elem.local_columns)

litestar/dto/factory/_backends/utils.py

@@ -279,7 +279,7 @@ def transfer_type_data(
     if isinstance(transfer_type, CollectionType):
         if transfer_type.has_nested:
             return transfer_nested_collection_type_data(
-                transfer_type.parsed_type.origin, transfer_type, dto_for, source_value
+                transfer_type.parsed_type.instantiable_origin, transfer_type, dto_for, source_value
             )
         return transfer_type.parsed_type.instantiable_origin(source_value)
     return source_value

tests/unit/test_contrib/test_sqlalchemy/test_dto.py

@@ -7,11 +7,11 @@
 
 import pytest
 import sqlalchemy
-from sqlalchemy import func
-from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, declared_attr, mapped_column
+from sqlalchemy import ForeignKey, func
+from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, declared_attr, mapped_column, relationship
 from typing_extensions import Annotated
 
-from litestar.contrib.sqlalchemy.dto import SQLAlchemyDTO
+from litestar.contrib.sqlalchemy.dto import SQLAlchemyDTO, parse_type_from_element
 from litestar.dto.factory import DTOConfig, DTOField, Mark
 from litestar.dto.factory.field import DTO_FIELD_META_KEY
 from litestar.dto.interface import ConnectionContext, HandlerContext
@@ -553,10 +553,76 @@ class A(Base):
     assert vars(model)["c"] == [1, 2, 3]
 
 
-async def test_no_type_hints(base: type[DeclarativeBase], connection_context: ConnectionContext) -> None:
+async def test_no_type_hint_column(base: type[DeclarativeBase], connection_context: ConnectionContext) -> None:
     class Model(base):
-        field = mapped_column(sqlalchemy.String)
+        nullable_field = mapped_column(sqlalchemy.String)
+        not_nullable_field = mapped_column(sqlalchemy.String, nullable=False, default="")
 
     dto_type = SQLAlchemyDTO[Annotated[Model, DTOConfig()]]
-    with pytest.raises(ImproperlyConfiguredException, match="No type information found for 'Model.field'"):
-        await get_model_from_dto(dto_type, Model, connection_context, b"")
+    model = await get_model_from_dto(dto_type, Model, connection_context, b"{}")
+    assert model.nullable_field is None
+    assert model.not_nullable_field == ""
+
+
+async def test_no_type_hint_scalar_relationship_with_nullable_fk(
+    base: type[DeclarativeBase], connection_context: ConnectionContext
+) -> None:
+    class Child(base):
+        ...
+
+    class Model(base):
+        child_id = mapped_column(ForeignKey("child.id"))
+        child = relationship(Child)
+
+    dto_type = SQLAlchemyDTO[Annotated[Model, DTOConfig(exclude={"child_id"})]]
+    model = await get_model_from_dto(dto_type, Model, connection_context, b"{}")
+    assert model.child is None
+
+
+async def test_no_type_hint_scalar_relationship_with_not_nullable_fk(
+    base: type[DeclarativeBase], connection_context: ConnectionContext
+) -> None:
+    class Child(base):
+        ...
+
+    class Model(base):
+        child_id = mapped_column(ForeignKey("child.id"), nullable=False)
+        child = relationship(Child)
+
+    dto_type = SQLAlchemyDTO[Annotated[Model, DTOConfig(exclude={"child_id"})]]
+    model = await get_model_from_dto(dto_type, Model, connection_context, b'{"child": {}}')
+    assert isinstance(model.child, Child)
+
+
+async def test_no_type_hint_collection_relationship(
+    base: type[DeclarativeBase], connection_context: ConnectionContext
+) -> None:
+    class Child(base):
+        model_id = mapped_column(ForeignKey("model.id"))
+
+    class Model(base):
+        children = relationship(Child)
+
+    dto_type = SQLAlchemyDTO[Annotated[Model, DTOConfig()]]
+    model = await get_model_from_dto(dto_type, Model, connection_context, b'{"children": []}')
+    assert model.children == []
+
+
+async def test_no_type_hint_collection_relationship_alt_collection_class(
+    base: type[DeclarativeBase], connection_context: ConnectionContext
+) -> None:
+    class Child(base):
+        model_id = mapped_column(ForeignKey("model.id"))
+
+    class Model(base):
+        children = relationship(Child, collection_class=set)
+
+    dto_type = SQLAlchemyDTO[Annotated[Model, DTOConfig()]]
+    model = await get_model_from_dto(dto_type, Model, connection_context, b'{"children": []}')
+    assert model.children == set()
+
+
+def test_parse_type_from_element_failure() -> None:
+    with pytest.raises(ImproperlyConfiguredException) as exc:
+        parse_type_from_element(1)
+    assert str(exc.value) == "500: Unable to parse type from element '1'. Consider adding a type hint."