Skip to content

Commit 090be97

Browse files
committed
Fix sa_column/sa_type lost when using Annotated with validators
When using Annotated[type, Field(sa_column=...), Validator(...)], Pydantic V2 creates a new pydantic.fields.FieldInfo that doesn't preserve SQLModel-specific attributes like sa_column and sa_type. This caused timezone-aware datetime columns defined with DateTime(timezone=True) to lose their timezone setting. The fix extracts the SQLModel FieldInfo from the original Annotated type's metadata, preserving sa_column and sa_type attributes even when Pydantic V2 merges the field info. Fixes timezone-aware datetime columns losing timezone=True when using Annotated types with Pydantic validators.
1 parent 8c2e4c4 commit 090be97

File tree

2 files changed

+150
-5
lines changed

2 files changed

+150
-5
lines changed

sqlmodel/main.py

Lines changed: 56 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@
5454
from sqlalchemy.orm.instrumentation import is_instrumented
5555
from sqlalchemy.sql.schema import MetaData
5656
from sqlalchemy.sql.sqltypes import LargeBinary, Time, Uuid
57-
from typing_extensions import Literal, TypeAlias, deprecated, get_origin
57+
from typing_extensions import Annotated, Literal, TypeAlias, deprecated, get_args, get_origin
5858

5959
from ._compat import ( # type: ignore[attr-defined]
6060
IS_PYDANTIC_V2,
@@ -562,7 +562,8 @@ def get_config(name: str) -> Any:
562562
# If it was passed by kwargs, ensure it's also set in config
563563
set_config_value(model=new_cls, parameter="table", value=config_table)
564564
for k, v in get_model_fields(new_cls).items():
565-
col = get_column_from_field(v)
565+
original_annotation = new_cls.__annotations__.get(k)
566+
col = get_column_from_field(v, original_annotation)
566567
setattr(new_cls, k, col)
567568
# Set a config flag to tell FastAPI that this should be read with a field
568569
# in orm_mode instead of preemptively converting it to a dict.
@@ -646,12 +647,44 @@ def __init__(
646647
ModelMetaclass.__init__(cls, classname, bases, dict_, **kw)
647648

648649

649-
def get_sqlalchemy_type(field: Any) -> Any:
650+
def _get_sqlmodel_field_info_from_annotation(annotation: Any) -> Optional["FieldInfo"]:
651+
"""Extract SQLModel FieldInfo from an Annotated type's metadata.
652+
653+
When using Annotated[type, Field(...), Validator(...)], Pydantic V2 may create
654+
a new pydantic.fields.FieldInfo that doesn't preserve SQLModel-specific attributes
655+
like sa_column and sa_type. This function looks through the Annotated metadata
656+
to find the original SQLModel FieldInfo.
657+
"""
658+
if get_origin(annotation) is not Annotated:
659+
return None
660+
for arg in get_args(annotation)[1:]: # Skip the first arg (the actual type)
661+
if isinstance(arg, FieldInfo):
662+
return arg
663+
return None
664+
665+
666+
def get_sqlalchemy_type(field: Any, original_annotation: Any = None) -> Any:
650667
if IS_PYDANTIC_V2:
651668
field_info = field
652669
else:
653670
field_info = field.field_info
654671
sa_type = getattr(field_info, "sa_type", Undefined) # noqa: B009
672+
# If sa_type not found on field_info, check if it's in the Annotated metadata
673+
# This handles the case where Pydantic V2 creates a new FieldInfo losing SQLModel attrs
674+
if sa_type is Undefined and IS_PYDANTIC_V2:
675+
# First try field_info.annotation (may be unpacked by Pydantic)
676+
annotation = getattr(field_info, "annotation", None)
677+
if annotation is not None:
678+
sqlmodel_field_info = _get_sqlmodel_field_info_from_annotation(annotation)
679+
if sqlmodel_field_info is not None:
680+
sa_type = getattr(sqlmodel_field_info, "sa_type", Undefined)
681+
# If still not found, try the original annotation from the class
682+
if sa_type is Undefined and original_annotation is not None:
683+
sqlmodel_field_info = _get_sqlmodel_field_info_from_annotation(
684+
original_annotation
685+
)
686+
if sqlmodel_field_info is not None:
687+
sa_type = getattr(sqlmodel_field_info, "sa_type", Undefined)
655688
if sa_type is not Undefined:
656689
return sa_type
657690

@@ -703,15 +736,33 @@ def get_sqlalchemy_type(field: Any) -> Any:
703736
raise ValueError(f"{type_} has no matching SQLAlchemy type")
704737

705738

706-
def get_column_from_field(field: Any) -> Column: # type: ignore
739+
def get_column_from_field(
740+
field: Any, original_annotation: Any = None
741+
) -> Column: # type: ignore
707742
if IS_PYDANTIC_V2:
708743
field_info = field
709744
else:
710745
field_info = field.field_info
711746
sa_column = getattr(field_info, "sa_column", Undefined)
747+
# If sa_column not found on field_info, check if it's in the Annotated metadata
748+
# This handles the case where Pydantic V2 creates a new FieldInfo losing SQLModel attrs
749+
if sa_column is Undefined and IS_PYDANTIC_V2:
750+
# First try field_info.annotation (may be unpacked by Pydantic)
751+
annotation = getattr(field_info, "annotation", None)
752+
if annotation is not None:
753+
sqlmodel_field_info = _get_sqlmodel_field_info_from_annotation(annotation)
754+
if sqlmodel_field_info is not None:
755+
sa_column = getattr(sqlmodel_field_info, "sa_column", Undefined)
756+
# If still not found, try the original annotation from the class
757+
if sa_column is Undefined and original_annotation is not None:
758+
sqlmodel_field_info = _get_sqlmodel_field_info_from_annotation(
759+
original_annotation
760+
)
761+
if sqlmodel_field_info is not None:
762+
sa_column = getattr(sqlmodel_field_info, "sa_column", Undefined)
712763
if isinstance(sa_column, Column):
713764
return sa_column
714-
sa_type = get_sqlalchemy_type(field)
765+
sa_type = get_sqlalchemy_type(field, original_annotation)
715766
primary_key = getattr(field_info, "primary_key", Undefined)
716767
if primary_key is Undefined:
717768
primary_key = False

tests/test_annotated_sa_column.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""Tests for Annotated fields with sa_column and Pydantic validators.
2+
3+
When using Annotated[type, Field(sa_column=...), Validator(...)], Pydantic V2 may
4+
create a new FieldInfo that doesn't preserve SQLModel-specific attributes like
5+
sa_column. These tests ensure the sa_column is properly extracted from the
6+
Annotated metadata.
7+
"""
8+
9+
from datetime import datetime
10+
from typing import Annotated, Optional
11+
12+
from pydantic import AfterValidator, BeforeValidator
13+
from sqlalchemy import Column, DateTime, String
14+
from sqlmodel import Field, SQLModel
15+
16+
17+
def test_annotated_sa_column_with_validators() -> None:
18+
"""Test that sa_column is preserved when using Annotated with validators."""
19+
20+
def before_validate(v: datetime) -> datetime:
21+
return v
22+
23+
def after_validate(v: datetime) -> datetime:
24+
return v
25+
26+
class Position(SQLModel, table=True):
27+
id: Optional[int] = Field(default=None, primary_key=True)
28+
timestamp: Annotated[
29+
datetime,
30+
Field(
31+
sa_column=Column(
32+
DateTime(timezone=True), nullable=False, index=True
33+
)
34+
),
35+
BeforeValidator(before_validate),
36+
AfterValidator(after_validate),
37+
]
38+
39+
# Verify the column type has timezone=True
40+
assert Position.__table__.c.timestamp.type.timezone is True
41+
assert Position.__table__.c.timestamp.nullable is False
42+
assert Position.__table__.c.timestamp.index is True
43+
44+
45+
def test_annotated_sa_column_with_single_validator() -> None:
46+
"""Test sa_column with just one validator."""
47+
48+
def validate_name(v: str) -> str:
49+
return v.strip()
50+
51+
class Item(SQLModel, table=True):
52+
id: Optional[int] = Field(default=None, primary_key=True)
53+
name: Annotated[
54+
str,
55+
Field(sa_column=Column(String(100), nullable=False, unique=True)),
56+
AfterValidator(validate_name),
57+
]
58+
59+
assert isinstance(Item.__table__.c.name.type, String)
60+
assert Item.__table__.c.name.type.length == 100
61+
assert Item.__table__.c.name.nullable is False
62+
assert Item.__table__.c.name.unique is True
63+
64+
65+
def test_annotated_sa_column_without_validators() -> None:
66+
"""Test that sa_column still works with Annotated but no validators."""
67+
68+
class Record(SQLModel, table=True):
69+
id: Optional[int] = Field(default=None, primary_key=True)
70+
created_at: Annotated[
71+
datetime,
72+
Field(sa_column=Column(DateTime(timezone=True), nullable=False)),
73+
]
74+
75+
assert Record.__table__.c.created_at.type.timezone is True
76+
assert Record.__table__.c.created_at.nullable is False
77+
78+
79+
def test_annotated_sa_type_with_validators() -> None:
80+
"""Test that sa_type is preserved when using Annotated with validators."""
81+
82+
def validate_timestamp(v: datetime) -> datetime:
83+
return v
84+
85+
class Event(SQLModel, table=True):
86+
id: Optional[int] = Field(default=None, primary_key=True)
87+
occurred_at: Annotated[
88+
datetime,
89+
Field(sa_type=DateTime(timezone=True)),
90+
AfterValidator(validate_timestamp),
91+
]
92+
93+
# Verify the column type has timezone=True
94+
assert Event.__table__.c.occurred_at.type.timezone is True

0 commit comments

Comments
 (0)