Use AwareDateTime when generating field from DateTime(timezone=True) (#407)

src/marshmallow_sqlalchemy/convert.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime as dt
 import functools
 import inspect
 import uuid
@@ -93,6 +94,16 @@ def _enum_field_factory(
     )
 
 
+def _datetime_field_factory(
+    converter: ModelConverter, data_type: sa.DateTime
+) -> type[fields.Field]:
+    if getattr(data_type, "timezone", False) or isinstance(
+        data_type, mssql.DATETIMEOFFSET
+    ):
+        return fields.AwareDateTime
+    return converter.type_mapping.get(dt.datetime, fields.DateTime)
+
+
 class ModelConverter:
     """Converts a SQLAlchemy model into a dictionary of corresponding
     marshmallow `Fields <marshmallow.fields.Field>`.
@@ -105,6 +116,7 @@ class ModelConverter:
         sa.JSON: fields.Raw,
         sa.ARRAY: _list_field_factory,
         sa.PickleType: fields.Raw,
+        sa.DateTime: _datetime_field_factory,
         postgresql.BIT: fields.Integer,
         postgresql.OID: fields.Integer,
         postgresql.UUID: fields.UUID,
@@ -123,7 +135,7 @@ class ModelConverter:
         mysql.SET: fields.List,
         mysql.ENUM: fields.Field,
         mysql.INTEGER: fields.Integer,
-        mysql.DATETIME: fields.DateTime,
+        mysql.DATETIME: _datetime_field_factory,
         mssql.BIT: fields.Integer,
         mssql.UNIQUEIDENTIFIER: fields.UUID,
     }

tests/test_conversion.py

@@ -7,7 +7,7 @@
 import sqlalchemy as sa
 from marshmallow import Schema, fields, validate
 from sqlalchemy import Integer, String
-from sqlalchemy.dialects import mysql, postgresql
+from sqlalchemy.dialects import mssql, mysql, postgresql
 from sqlalchemy.orm import Mapped, Session, column_property
 
 from marshmallow_sqlalchemy import (
@@ -205,6 +205,26 @@ def test_convert_types(self, converter, sa_type, field_type):
         field = converter.property2field(prop)
         assert type(field) is field_type
 
+    @pytest.mark.parametrize(
+        ("sa_type_instance", "field_type"),
+        (
+            (sa.DateTime(timezone=True), fields.AwareDateTime),
+            (sa.DateTime(timezone=False), fields.DateTime),
+            (sa.DateTime(), fields.DateTime),
+            (postgresql.TIMESTAMP(timezone=True), fields.AwareDateTime),
+            (postgresql.TIMESTAMP(timezone=False), fields.DateTime),
+            (mysql.DATETIME(timezone=True), fields.AwareDateTime),
+            (mysql.DATETIME(timezone=False), fields.DateTime),
+            (mssql.DATETIMEOFFSET(), fields.AwareDateTime),
+            (sa.TIMESTAMP(timezone=True), fields.AwareDateTime),
+            (sa.TIMESTAMP(timezone=False), fields.DateTime),
+        ),
+    )
+    def test_convert_datetime_types(self, converter, sa_type_instance, field_type):
+        prop = make_property(sa_type_instance)
+        field = converter.property2field(prop)
+        assert type(field) is field_type
+
     def test_convert_Numeric(self, converter):
         prop = make_property(sa.Numeric(scale=2))
         field = converter.property2field(prop)

tests/test_sqlalchemy_schema.py

@@ -750,3 +750,20 @@ class Meta:
     assert schema.dump(teacher) == {
         "data": {"foo": "bar"},
     }
+
+
+def test_auto_field_aware_datetime(Base):
+    class Item(Base):
+        __tablename__ = "item_aware_datetime"
+        id = sa.Column(sa.Integer, primary_key=True)
+        created_at = sa.Column(sa.DateTime(timezone=True), nullable=False)
+        updated_at = sa.Column(sa.DateTime(timezone=False), nullable=True)
+
+    class ItemSchema(SQLAlchemyAutoSchema):
+        class Meta:
+            model = Item
+
+    schema = ItemSchema()
+    assert isinstance(schema.fields["created_at"], fields.AwareDateTime)
+    assert isinstance(schema.fields["updated_at"], fields.DateTime)
+    assert not isinstance(schema.fields["updated_at"], fields.AwareDateTime)