main.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005
  1. import ipaddress
  2. import uuid
  3. import weakref
  4. from datetime import date, datetime, time, timedelta
  5. from decimal import Decimal
  6. from enum import Enum
  7. from pathlib import Path
  8. from typing import (
  9. TYPE_CHECKING,
  10. AbstractSet,
  11. Any,
  12. Callable,
  13. ClassVar,
  14. Dict,
  15. List,
  16. Mapping,
  17. Optional,
  18. Sequence,
  19. Set,
  20. Tuple,
  21. Type,
  22. TypeVar,
  23. Union,
  24. cast,
  25. overload,
  26. )
  27. from pydantic import BaseModel, EmailStr
  28. from pydantic.fields import FieldInfo as PydanticFieldInfo
  29. from sqlalchemy import (
  30. Boolean,
  31. Column,
  32. Date,
  33. DateTime,
  34. Float,
  35. ForeignKey,
  36. Integer,
  37. Interval,
  38. Numeric,
  39. inspect,
  40. )
  41. from sqlalchemy import Enum as sa_Enum
  42. from sqlalchemy.orm import (
  43. Mapped,
  44. RelationshipProperty,
  45. declared_attr,
  46. registry,
  47. relationship,
  48. )
  49. from sqlalchemy.orm.attributes import set_attribute
  50. from sqlalchemy.orm.decl_api import DeclarativeMeta
  51. from sqlalchemy.orm.instrumentation import is_instrumented
  52. from sqlalchemy.sql.schema import MetaData
  53. from sqlalchemy.sql.sqltypes import LargeBinary, Time, Uuid
  54. from typing_extensions import Literal, deprecated, get_origin
  55. from ._compat import ( # type: ignore[attr-defined]
  56. IS_PYDANTIC_V2,
  57. PYDANTIC_VERSION,
  58. BaseConfig,
  59. ModelField,
  60. ModelMetaclass,
  61. Representation,
  62. SQLModelConfig,
  63. Undefined,
  64. UndefinedType,
  65. _calculate_keys,
  66. finish_init,
  67. get_annotations,
  68. get_config_value,
  69. get_field_metadata,
  70. get_model_fields,
  71. get_relationship_to,
  72. get_sa_type_from_field,
  73. init_pydantic_private_attrs,
  74. is_field_noneable,
  75. is_table_model_class,
  76. post_init_field_info,
  77. set_config_value,
  78. sqlmodel_init,
  79. sqlmodel_validate,
  80. )
  81. from .sql.sqltypes import AutoString
  82. if TYPE_CHECKING:
  83. from pydantic._internal._model_construction import ModelMetaclass as ModelMetaclass
  84. from pydantic._internal._repr import Representation as Representation
  85. from pydantic_core import PydanticUndefined as Undefined
  86. from pydantic_core import PydanticUndefinedType as UndefinedType
  87. _T = TypeVar("_T")
  88. NoArgAnyCallable = Callable[[], Any]
  89. IncEx = Union[Set[int], Set[str], Dict[int, Any], Dict[str, Any], None]
  90. OnDeleteType = Literal["CASCADE", "SET NULL", "RESTRICT"]
  91. def __dataclass_transform__(
  92. *,
  93. eq_default: bool = True,
  94. order_default: bool = False,
  95. kw_only_default: bool = False,
  96. field_descriptors: Tuple[Union[type, Callable[..., Any]], ...] = (()),
  97. ) -> Callable[[_T], _T]:
  98. return lambda a: a
  99. class FieldInfo(PydanticFieldInfo):
  100. def __init__(self, default: Any = Undefined, **kwargs: Any) -> None:
  101. primary_key = kwargs.pop("primary_key", False)
  102. nullable = kwargs.pop("nullable", Undefined)
  103. foreign_key = kwargs.pop("foreign_key", Undefined)
  104. ondelete = kwargs.pop("ondelete", Undefined)
  105. unique = kwargs.pop("unique", False)
  106. index = kwargs.pop("index", Undefined)
  107. sa_type = kwargs.pop("sa_type", Undefined)
  108. sa_column = kwargs.pop("sa_column", Undefined)
  109. sa_column_args = kwargs.pop("sa_column_args", Undefined)
  110. sa_column_kwargs = kwargs.pop("sa_column_kwargs", Undefined)
  111. if sa_column is not Undefined:
  112. if sa_column_args is not Undefined:
  113. raise RuntimeError(
  114. "Passing sa_column_args is not supported when "
  115. "also passing a sa_column"
  116. )
  117. if sa_column_kwargs is not Undefined:
  118. raise RuntimeError(
  119. "Passing sa_column_kwargs is not supported when "
  120. "also passing a sa_column"
  121. )
  122. if primary_key is not Undefined:
  123. raise RuntimeError(
  124. "Passing primary_key is not supported when "
  125. "also passing a sa_column"
  126. )
  127. if nullable is not Undefined:
  128. raise RuntimeError(
  129. "Passing nullable is not supported when also passing a sa_column"
  130. )
  131. if foreign_key is not Undefined:
  132. raise RuntimeError(
  133. "Passing foreign_key is not supported when "
  134. "also passing a sa_column"
  135. )
  136. if ondelete is not Undefined:
  137. raise RuntimeError(
  138. "Passing ondelete is not supported when also passing a sa_column"
  139. )
  140. if unique is not Undefined:
  141. raise RuntimeError(
  142. "Passing unique is not supported when also passing a sa_column"
  143. )
  144. if index is not Undefined:
  145. raise RuntimeError(
  146. "Passing index is not supported when also passing a sa_column"
  147. )
  148. if sa_type is not Undefined:
  149. raise RuntimeError(
  150. "Passing sa_type is not supported when also passing a sa_column"
  151. )
  152. if ondelete is not Undefined:
  153. if foreign_key is Undefined:
  154. raise RuntimeError("ondelete can only be used with foreign_key")
  155. super().__init__(default=default, **kwargs)
  156. self.primary_key = primary_key
  157. self.nullable = nullable
  158. self.foreign_key = foreign_key
  159. self.ondelete = ondelete
  160. self.unique = unique
  161. self.index = index
  162. self.sa_type = sa_type
  163. self.sa_column = sa_column
  164. self.sa_column_args = sa_column_args
  165. self.sa_column_kwargs = sa_column_kwargs
  166. class RelationshipInfo(Representation):
  167. def __init__(
  168. self,
  169. *,
  170. back_populates: Optional[str] = None,
  171. cascade_delete: Optional[bool] = False,
  172. passive_deletes: Optional[Union[bool, Literal["all"]]] = False,
  173. link_model: Optional[Any] = None,
  174. sa_relationship: Optional[RelationshipProperty] = None, # type: ignore
  175. sa_relationship_args: Optional[Sequence[Any]] = None,
  176. sa_relationship_kwargs: Optional[Mapping[str, Any]] = None,
  177. ) -> None:
  178. if sa_relationship is not None:
  179. if sa_relationship_args is not None:
  180. raise RuntimeError(
  181. "Passing sa_relationship_args is not supported when "
  182. "also passing a sa_relationship"
  183. )
  184. if sa_relationship_kwargs is not None:
  185. raise RuntimeError(
  186. "Passing sa_relationship_kwargs is not supported when "
  187. "also passing a sa_relationship"
  188. )
  189. self.back_populates = back_populates
  190. self.cascade_delete = cascade_delete
  191. self.passive_deletes = passive_deletes
  192. self.link_model = link_model
  193. self.sa_relationship = sa_relationship
  194. self.sa_relationship_args = sa_relationship_args
  195. self.sa_relationship_kwargs = sa_relationship_kwargs
  196. # include sa_type, sa_column_args, sa_column_kwargs
  197. @overload
  198. def Field(
  199. default: Any = Undefined,
  200. *,
  201. default_factory: Optional[NoArgAnyCallable] = None,
  202. alias: Optional[str] = None,
  203. title: Optional[str] = None,
  204. description: Optional[str] = None,
  205. exclude: Union[
  206. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  207. ] = None,
  208. include: Union[
  209. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  210. ] = None,
  211. const: Optional[bool] = None,
  212. gt: Optional[float] = None,
  213. ge: Optional[float] = None,
  214. lt: Optional[float] = None,
  215. le: Optional[float] = None,
  216. multiple_of: Optional[float] = None,
  217. max_digits: Optional[int] = None,
  218. decimal_places: Optional[int] = None,
  219. min_items: Optional[int] = None,
  220. max_items: Optional[int] = None,
  221. unique_items: Optional[bool] = None,
  222. min_length: Optional[int] = None,
  223. max_length: Optional[int] = None,
  224. allow_mutation: bool = True,
  225. regex: Optional[str] = None,
  226. discriminator: Optional[str] = None,
  227. repr: bool = True,
  228. primary_key: Union[bool, UndefinedType] = Undefined,
  229. foreign_key: Any = Undefined,
  230. unique: Union[bool, UndefinedType] = Undefined,
  231. nullable: Union[bool, UndefinedType] = Undefined,
  232. index: Union[bool, UndefinedType] = Undefined,
  233. sa_type: Union[Type[Any], UndefinedType] = Undefined,
  234. sa_column_args: Union[Sequence[Any], UndefinedType] = Undefined,
  235. sa_column_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
  236. schema_extra: Optional[Dict[str, Any]] = None,
  237. ) -> Any: ...
  238. # When foreign_key is str, include ondelete
  239. # include sa_type, sa_column_args, sa_column_kwargs
  240. @overload
  241. def Field(
  242. default: Any = Undefined,
  243. *,
  244. default_factory: Optional[NoArgAnyCallable] = None,
  245. alias: Optional[str] = None,
  246. title: Optional[str] = None,
  247. description: Optional[str] = None,
  248. exclude: Union[
  249. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  250. ] = None,
  251. include: Union[
  252. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  253. ] = None,
  254. const: Optional[bool] = None,
  255. gt: Optional[float] = None,
  256. ge: Optional[float] = None,
  257. lt: Optional[float] = None,
  258. le: Optional[float] = None,
  259. multiple_of: Optional[float] = None,
  260. max_digits: Optional[int] = None,
  261. decimal_places: Optional[int] = None,
  262. min_items: Optional[int] = None,
  263. max_items: Optional[int] = None,
  264. unique_items: Optional[bool] = None,
  265. min_length: Optional[int] = None,
  266. max_length: Optional[int] = None,
  267. allow_mutation: bool = True,
  268. regex: Optional[str] = None,
  269. discriminator: Optional[str] = None,
  270. repr: bool = True,
  271. primary_key: Union[bool, UndefinedType] = Undefined,
  272. foreign_key: str,
  273. ondelete: Union[OnDeleteType, UndefinedType] = Undefined,
  274. unique: Union[bool, UndefinedType] = Undefined,
  275. nullable: Union[bool, UndefinedType] = Undefined,
  276. index: Union[bool, UndefinedType] = Undefined,
  277. sa_type: Union[Type[Any], UndefinedType] = Undefined,
  278. sa_column_args: Union[Sequence[Any], UndefinedType] = Undefined,
  279. sa_column_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
  280. schema_extra: Optional[Dict[str, Any]] = None,
  281. ) -> Any: ...
  282. # Include sa_column, don't include
  283. # primary_key
  284. # foreign_key
  285. # ondelete
  286. # unique
  287. # nullable
  288. # index
  289. # sa_type
  290. # sa_column_args
  291. # sa_column_kwargs
  292. @overload
  293. def Field(
  294. default: Any = Undefined,
  295. *,
  296. default_factory: Optional[NoArgAnyCallable] = None,
  297. alias: Optional[str] = None,
  298. title: Optional[str] = None,
  299. description: Optional[str] = None,
  300. exclude: Union[
  301. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  302. ] = None,
  303. include: Union[
  304. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  305. ] = None,
  306. const: Optional[bool] = None,
  307. gt: Optional[float] = None,
  308. ge: Optional[float] = None,
  309. lt: Optional[float] = None,
  310. le: Optional[float] = None,
  311. multiple_of: Optional[float] = None,
  312. max_digits: Optional[int] = None,
  313. decimal_places: Optional[int] = None,
  314. min_items: Optional[int] = None,
  315. max_items: Optional[int] = None,
  316. unique_items: Optional[bool] = None,
  317. min_length: Optional[int] = None,
  318. max_length: Optional[int] = None,
  319. allow_mutation: bool = True,
  320. regex: Optional[str] = None,
  321. discriminator: Optional[str] = None,
  322. repr: bool = True,
  323. sa_column: Union[Column, UndefinedType] = Undefined, # type: ignore
  324. schema_extra: Optional[Dict[str, Any]] = None,
  325. ) -> Any: ...
  326. def Field(
  327. default: Any = Undefined,
  328. *,
  329. default_factory: Optional[NoArgAnyCallable] = None,
  330. alias: Optional[str] = None,
  331. title: Optional[str] = None,
  332. description: Optional[str] = None,
  333. exclude: Union[
  334. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  335. ] = None,
  336. include: Union[
  337. AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
  338. ] = None,
  339. const: Optional[bool] = None,
  340. gt: Optional[float] = None,
  341. ge: Optional[float] = None,
  342. lt: Optional[float] = None,
  343. le: Optional[float] = None,
  344. multiple_of: Optional[float] = None,
  345. max_digits: Optional[int] = None,
  346. decimal_places: Optional[int] = None,
  347. min_items: Optional[int] = None,
  348. max_items: Optional[int] = None,
  349. unique_items: Optional[bool] = None,
  350. min_length: Optional[int] = None,
  351. max_length: Optional[int] = None,
  352. allow_mutation: bool = True,
  353. regex: Optional[str] = None,
  354. discriminator: Optional[str] = None,
  355. repr: bool = True,
  356. primary_key: Union[bool, UndefinedType] = Undefined,
  357. foreign_key: Any = Undefined,
  358. ondelete: Union[OnDeleteType, UndefinedType] = Undefined,
  359. unique: Union[bool, UndefinedType] = Undefined,
  360. nullable: Union[bool, UndefinedType] = Undefined,
  361. index: Union[bool, UndefinedType] = Undefined,
  362. sa_type: Union[Type[Any], UndefinedType] = Undefined,
  363. sa_column: Union[Column, UndefinedType] = Undefined, # type: ignore
  364. sa_column_args: Union[Sequence[Any], UndefinedType] = Undefined,
  365. sa_column_kwargs: Union[Mapping[str, Any], UndefinedType] = Undefined,
  366. schema_extra: Optional[Dict[str, Any]] = None,
  367. ) -> Any:
  368. current_schema_extra = schema_extra or {}
  369. field_info = FieldInfo(
  370. default,
  371. default_factory=default_factory,
  372. alias=alias,
  373. title=title,
  374. description=description,
  375. exclude=exclude,
  376. include=include,
  377. const=const,
  378. gt=gt,
  379. ge=ge,
  380. lt=lt,
  381. le=le,
  382. multiple_of=multiple_of,
  383. max_digits=max_digits,
  384. decimal_places=decimal_places,
  385. min_items=min_items,
  386. max_items=max_items,
  387. unique_items=unique_items,
  388. min_length=min_length,
  389. max_length=max_length,
  390. allow_mutation=allow_mutation,
  391. regex=regex,
  392. discriminator=discriminator,
  393. repr=repr,
  394. primary_key=primary_key,
  395. foreign_key=foreign_key,
  396. ondelete=ondelete,
  397. unique=unique,
  398. nullable=nullable,
  399. index=index,
  400. sa_type=sa_type,
  401. sa_column=sa_column,
  402. sa_column_args=sa_column_args,
  403. sa_column_kwargs=sa_column_kwargs,
  404. **current_schema_extra,
  405. )
  406. post_init_field_info(field_info)
  407. return field_info
  408. @overload
  409. def Relationship(
  410. *,
  411. back_populates: Optional[str] = None,
  412. cascade_delete: Optional[bool] = False,
  413. passive_deletes: Optional[Union[bool, Literal["all"]]] = False,
  414. link_model: Optional[Any] = None,
  415. sa_relationship_args: Optional[Sequence[Any]] = None,
  416. sa_relationship_kwargs: Optional[Mapping[str, Any]] = None,
  417. ) -> Any: ...
  418. @overload
  419. def Relationship(
  420. *,
  421. back_populates: Optional[str] = None,
  422. cascade_delete: Optional[bool] = False,
  423. passive_deletes: Optional[Union[bool, Literal["all"]]] = False,
  424. link_model: Optional[Any] = None,
  425. sa_relationship: Optional[RelationshipProperty[Any]] = None,
  426. ) -> Any: ...
  427. def Relationship(
  428. *,
  429. back_populates: Optional[str] = None,
  430. cascade_delete: Optional[bool] = False,
  431. passive_deletes: Optional[Union[bool, Literal["all"]]] = False,
  432. link_model: Optional[Any] = None,
  433. sa_relationship: Optional[RelationshipProperty[Any]] = None,
  434. sa_relationship_args: Optional[Sequence[Any]] = None,
  435. sa_relationship_kwargs: Optional[Mapping[str, Any]] = None,
  436. ) -> Any:
  437. relationship_info = RelationshipInfo(
  438. back_populates=back_populates,
  439. cascade_delete=cascade_delete,
  440. passive_deletes=passive_deletes,
  441. link_model=link_model,
  442. sa_relationship=sa_relationship,
  443. sa_relationship_args=sa_relationship_args,
  444. sa_relationship_kwargs=sa_relationship_kwargs,
  445. )
  446. return relationship_info
  447. @__dataclass_transform__(kw_only_default=True, field_descriptors=(Field, FieldInfo))
  448. class SQLModelMetaclass(ModelMetaclass, DeclarativeMeta):
  449. __sqlmodel_relationships__: Dict[str, RelationshipInfo]
  450. model_config: SQLModelConfig
  451. model_fields: Dict[str, FieldInfo]
  452. __config__: Type[SQLModelConfig]
  453. __fields__: Dict[str, ModelField] # type: ignore[assignment]
  454. # Replicate SQLAlchemy
  455. def __setattr__(cls, name: str, value: Any) -> None:
  456. if is_table_model_class(cls):
  457. DeclarativeMeta.__setattr__(cls, name, value)
  458. else:
  459. super().__setattr__(name, value)
  460. def __delattr__(cls, name: str) -> None:
  461. if is_table_model_class(cls):
  462. DeclarativeMeta.__delattr__(cls, name)
  463. else:
  464. super().__delattr__(name)
  465. # From Pydantic
  466. def __new__(
  467. cls,
  468. name: str,
  469. bases: Tuple[Type[Any], ...],
  470. class_dict: Dict[str, Any],
  471. **kwargs: Any,
  472. ) -> Any:
  473. relationships: Dict[str, RelationshipInfo] = {}
  474. dict_for_pydantic = {}
  475. original_annotations = get_annotations(class_dict)
  476. pydantic_annotations = {}
  477. relationship_annotations = {}
  478. for k, v in class_dict.items():
  479. if isinstance(v, RelationshipInfo):
  480. relationships[k] = v
  481. else:
  482. dict_for_pydantic[k] = v
  483. for k, v in original_annotations.items():
  484. if k in relationships:
  485. relationship_annotations[k] = v
  486. else:
  487. pydantic_annotations[k] = v
  488. dict_used = {
  489. **dict_for_pydantic,
  490. "__weakref__": None,
  491. "__sqlmodel_relationships__": relationships,
  492. "__annotations__": pydantic_annotations,
  493. }
  494. # Duplicate logic from Pydantic to filter config kwargs because if they are
  495. # passed directly including the registry Pydantic will pass them over to the
  496. # superclass causing an error
  497. allowed_config_kwargs: Set[str] = {
  498. key
  499. for key in dir(BaseConfig)
  500. if not (
  501. key.startswith("__") and key.endswith("__")
  502. ) # skip dunder methods and attributes
  503. }
  504. config_kwargs = {
  505. key: kwargs[key] for key in kwargs.keys() & allowed_config_kwargs
  506. }
  507. new_cls = super().__new__(cls, name, bases, dict_used, **config_kwargs)
  508. new_cls.__annotations__ = {
  509. **relationship_annotations,
  510. **pydantic_annotations,
  511. **new_cls.__annotations__,
  512. }
  513. def get_config(name: str) -> Any:
  514. config_class_value = get_config_value(
  515. model=new_cls, parameter=name, default=Undefined
  516. )
  517. if config_class_value is not Undefined:
  518. return config_class_value
  519. kwarg_value = kwargs.get(name, Undefined)
  520. if kwarg_value is not Undefined:
  521. return kwarg_value
  522. return Undefined
  523. config_table = get_config("table")
  524. if config_table is True:
  525. # If it was passed by kwargs, ensure it's also set in config
  526. set_config_value(model=new_cls, parameter="table", value=config_table)
  527. for k, v in get_model_fields(new_cls).items():
  528. col = get_column_from_field(v)
  529. setattr(new_cls, k, col)
  530. # Set a config flag to tell FastAPI that this should be read with a field
  531. # in orm_mode instead of preemptively converting it to a dict.
  532. # This could be done by reading new_cls.model_config['table'] in FastAPI, but
  533. # that's very specific about SQLModel, so let's have another config that
  534. # other future tools based on Pydantic can use.
  535. set_config_value(
  536. model=new_cls, parameter="read_from_attributes", value=True
  537. )
  538. # For compatibility with older versions
  539. # TODO: remove this in the future
  540. set_config_value(model=new_cls, parameter="read_with_orm_mode", value=True)
  541. config_registry = get_config("registry")
  542. if config_registry is not Undefined:
  543. config_registry = cast(registry, config_registry)
  544. # If it was passed by kwargs, ensure it's also set in config
  545. set_config_value(model=new_cls, parameter="registry", value=config_table)
  546. setattr(new_cls, "_sa_registry", config_registry) # noqa: B010
  547. setattr(new_cls, "metadata", config_registry.metadata) # noqa: B010
  548. setattr(new_cls, "__abstract__", True) # noqa: B010
  549. return new_cls
  550. # Override SQLAlchemy, allow both SQLAlchemy and plain Pydantic models
  551. def __init__(
  552. cls, classname: str, bases: Tuple[type, ...], dict_: Dict[str, Any], **kw: Any
  553. ) -> None:
  554. # Only one of the base classes (or the current one) should be a table model
  555. # this allows FastAPI cloning a SQLModel for the response_model without
  556. # trying to create a new SQLAlchemy, for a new table, with the same name, that
  557. # triggers an error
  558. base_is_table = any(is_table_model_class(base) for base in bases)
  559. if is_table_model_class(cls) and not base_is_table:
  560. for rel_name, rel_info in cls.__sqlmodel_relationships__.items():
  561. if rel_info.sa_relationship:
  562. # There's a SQLAlchemy relationship declared, that takes precedence
  563. # over anything else, use that and continue with the next attribute
  564. setattr(cls, rel_name, rel_info.sa_relationship) # Fix #315
  565. continue
  566. raw_ann = cls.__annotations__[rel_name]
  567. origin = get_origin(raw_ann)
  568. if origin is Mapped:
  569. ann = raw_ann.__args__[0]
  570. else:
  571. ann = raw_ann
  572. # Plain forward references, for models not yet defined, are not
  573. # handled well by SQLAlchemy without Mapped, so, wrap the
  574. # annotations in Mapped here
  575. cls.__annotations__[rel_name] = Mapped[ann] # type: ignore[valid-type]
  576. relationship_to = get_relationship_to(
  577. name=rel_name, rel_info=rel_info, annotation=ann
  578. )
  579. rel_kwargs: Dict[str, Any] = {}
  580. if rel_info.back_populates:
  581. rel_kwargs["back_populates"] = rel_info.back_populates
  582. if rel_info.cascade_delete:
  583. rel_kwargs["cascade"] = "all, delete-orphan"
  584. if rel_info.passive_deletes:
  585. rel_kwargs["passive_deletes"] = rel_info.passive_deletes
  586. if rel_info.link_model:
  587. ins = inspect(rel_info.link_model)
  588. local_table = getattr(ins, "local_table") # noqa: B009
  589. if local_table is None:
  590. raise RuntimeError(
  591. "Couldn't find the secondary table for "
  592. f"model {rel_info.link_model}"
  593. )
  594. rel_kwargs["secondary"] = local_table
  595. rel_args: List[Any] = []
  596. if rel_info.sa_relationship_args:
  597. rel_args.extend(rel_info.sa_relationship_args)
  598. if rel_info.sa_relationship_kwargs:
  599. rel_kwargs.update(rel_info.sa_relationship_kwargs)
  600. rel_value = relationship(relationship_to, *rel_args, **rel_kwargs)
  601. setattr(cls, rel_name, rel_value) # Fix #315
  602. # SQLAlchemy no longer uses dict_
  603. # Ref: https://github.com/sqlalchemy/sqlalchemy/commit/428ea01f00a9cc7f85e435018565eb6da7af1b77
  604. # Tag: 1.4.36
  605. DeclarativeMeta.__init__(cls, classname, bases, dict_, **kw)
  606. else:
  607. ModelMetaclass.__init__(cls, classname, bases, dict_, **kw)
  608. def get_sqlalchemy_type(field: Any) -> Any:
  609. if IS_PYDANTIC_V2:
  610. field_info = field
  611. else:
  612. field_info = field.field_info
  613. sa_type = getattr(field_info, "sa_type", Undefined) # noqa: B009
  614. if sa_type is not Undefined:
  615. return sa_type
  616. type_ = get_sa_type_from_field(field)
  617. metadata = get_field_metadata(field)
  618. # Check enums first as an enum can also be a str, needed by Pydantic/FastAPI
  619. if issubclass(type_, Enum):
  620. return sa_Enum(type_)
  621. if issubclass(
  622. type_,
  623. (
  624. str,
  625. ipaddress.IPv4Address,
  626. ipaddress.IPv4Network,
  627. ipaddress.IPv6Address,
  628. ipaddress.IPv6Network,
  629. Path,
  630. EmailStr,
  631. ),
  632. ):
  633. max_length = getattr(metadata, "max_length", None)
  634. if max_length:
  635. return AutoString(length=max_length)
  636. return AutoString
  637. if issubclass(type_, float):
  638. return Float
  639. if issubclass(type_, bool):
  640. return Boolean
  641. if issubclass(type_, int):
  642. return Integer
  643. if issubclass(type_, datetime):
  644. return DateTime
  645. if issubclass(type_, date):
  646. return Date
  647. if issubclass(type_, timedelta):
  648. return Interval
  649. if issubclass(type_, time):
  650. return Time
  651. if issubclass(type_, bytes):
  652. return LargeBinary
  653. if issubclass(type_, Decimal):
  654. return Numeric(
  655. precision=getattr(metadata, "max_digits", None),
  656. scale=getattr(metadata, "decimal_places", None),
  657. )
  658. if issubclass(type_, uuid.UUID):
  659. return Uuid
  660. raise ValueError(f"{type_} has no matching SQLAlchemy type")
  661. def get_column_from_field(field: Any) -> Column: # type: ignore
  662. if IS_PYDANTIC_V2:
  663. field_info = field
  664. else:
  665. field_info = field.field_info
  666. sa_column = getattr(field_info, "sa_column", Undefined)
  667. if isinstance(sa_column, Column):
  668. return sa_column
  669. sa_type = get_sqlalchemy_type(field)
  670. primary_key = getattr(field_info, "primary_key", Undefined)
  671. if primary_key is Undefined:
  672. primary_key = False
  673. index = getattr(field_info, "index", Undefined)
  674. if index is Undefined:
  675. index = False
  676. nullable = not primary_key and is_field_noneable(field)
  677. # Override derived nullability if the nullable property is set explicitly
  678. # on the field
  679. field_nullable = getattr(field_info, "nullable", Undefined) # noqa: B009
  680. if field_nullable is not Undefined:
  681. assert not isinstance(field_nullable, UndefinedType)
  682. nullable = field_nullable
  683. args = []
  684. foreign_key = getattr(field_info, "foreign_key", Undefined)
  685. if foreign_key is Undefined:
  686. foreign_key = None
  687. unique = getattr(field_info, "unique", Undefined)
  688. if unique is Undefined:
  689. unique = False
  690. if foreign_key:
  691. if field_info.ondelete == "SET NULL" and not nullable:
  692. raise RuntimeError('ondelete="SET NULL" requires nullable=True')
  693. assert isinstance(foreign_key, str)
  694. ondelete = getattr(field_info, "ondelete", Undefined)
  695. if ondelete is Undefined:
  696. ondelete = None
  697. assert isinstance(ondelete, (str, type(None))) # for typing
  698. args.append(ForeignKey(foreign_key, ondelete=ondelete))
  699. kwargs = {
  700. "primary_key": primary_key,
  701. "nullable": nullable,
  702. "index": index,
  703. "unique": unique,
  704. }
  705. sa_default = Undefined
  706. if field_info.default_factory:
  707. sa_default = field_info.default_factory
  708. elif field_info.default is not Undefined:
  709. sa_default = field_info.default
  710. if sa_default is not Undefined:
  711. kwargs["default"] = sa_default
  712. sa_column_args = getattr(field_info, "sa_column_args", Undefined)
  713. if sa_column_args is not Undefined:
  714. args.extend(list(cast(Sequence[Any], sa_column_args)))
  715. sa_column_kwargs = getattr(field_info, "sa_column_kwargs", Undefined)
  716. if sa_column_kwargs is not Undefined:
  717. kwargs.update(cast(Dict[Any, Any], sa_column_kwargs))
  718. return Column(sa_type, *args, **kwargs) # type: ignore
  719. class_registry = weakref.WeakValueDictionary() # type: ignore
  720. default_registry = registry()
  721. _TSQLModel = TypeVar("_TSQLModel", bound="SQLModel")
  722. class SQLModel(BaseModel, metaclass=SQLModelMetaclass, registry=default_registry):
  723. # SQLAlchemy needs to set weakref(s), Pydantic will set the other slots values
  724. __slots__ = ("__weakref__",)
  725. __tablename__: ClassVar[Union[str, Callable[..., str]]]
  726. __sqlmodel_relationships__: ClassVar[Dict[str, RelationshipProperty[Any]]]
  727. __name__: ClassVar[str]
  728. metadata: ClassVar[MetaData]
  729. __allow_unmapped__ = True # https://docs.sqlalchemy.org/en/20/changelog/migration_20.html#migration-20-step-six
  730. if IS_PYDANTIC_V2:
  731. model_config = SQLModelConfig(from_attributes=True)
  732. else:
  733. class Config:
  734. orm_mode = True
  735. def __new__(cls, *args: Any, **kwargs: Any) -> Any:
  736. new_object = super().__new__(cls)
  737. # SQLAlchemy doesn't call __init__ on the base class when querying from DB
  738. # Ref: https://docs.sqlalchemy.org/en/14/orm/constructors.html
  739. # Set __fields_set__ here, that would have been set when calling __init__
  740. # in the Pydantic model so that when SQLAlchemy sets attributes that are
  741. # added (e.g. when querying from DB) to the __fields_set__, this already exists
  742. init_pydantic_private_attrs(new_object)
  743. return new_object
  744. def __init__(__pydantic_self__, **data: Any) -> None:
  745. # Uses something other than `self` the first arg to allow "self" as a
  746. # settable attribute
  747. # SQLAlchemy does very dark black magic and modifies the __init__ method in
  748. # sqlalchemy.orm.instrumentation._generate_init()
  749. # so, to make SQLAlchemy work, it's needed to explicitly call __init__ to
  750. # trigger all the SQLAlchemy logic, it doesn't work using cls.__new__, setting
  751. # attributes obj.__dict__, etc. The __init__ method has to be called. But
  752. # there are cases where calling all the default logic is not ideal, e.g.
  753. # when calling Model.model_validate(), as the validation is done outside
  754. # of instance creation.
  755. # At the same time, __init__ is what users would normally call, by creating
  756. # a new instance, which should have validation and all the default logic.
  757. # So, to be able to set up the internal SQLAlchemy logic alone without
  758. # executing the rest, and support things like Model.model_validate(), we
  759. # use a contextvar to know if we should execute everything.
  760. if finish_init.get():
  761. sqlmodel_init(self=__pydantic_self__, data=data)
  762. def __setattr__(self, name: str, value: Any) -> None:
  763. if name in {"_sa_instance_state"}:
  764. self.__dict__[name] = value
  765. return
  766. else:
  767. # Set in SQLAlchemy, before Pydantic to trigger events and updates
  768. if is_table_model_class(self.__class__) and is_instrumented(self, name): # type: ignore[no-untyped-call]
  769. set_attribute(self, name, value)
  770. # Set in Pydantic model to trigger possible validation changes, only for
  771. # non relationship values
  772. if name not in self.__sqlmodel_relationships__:
  773. super().__setattr__(name, value)
  774. def __repr_args__(self) -> Sequence[Tuple[Optional[str], Any]]:
  775. # Don't show SQLAlchemy private attributes
  776. return [
  777. (k, v)
  778. for k, v in super().__repr_args__()
  779. if not (isinstance(k, str) and k.startswith("_sa_"))
  780. ]
  781. @declared_attr # type: ignore
  782. def __tablename__(cls) -> str:
  783. return cls.__name__.lower()
  784. @classmethod
  785. def model_validate(
  786. cls: Type[_TSQLModel],
  787. obj: Any,
  788. *,
  789. strict: Union[bool, None] = None,
  790. from_attributes: Union[bool, None] = None,
  791. context: Union[Dict[str, Any], None] = None,
  792. update: Union[Dict[str, Any], None] = None,
  793. ) -> _TSQLModel:
  794. return sqlmodel_validate(
  795. cls=cls,
  796. obj=obj,
  797. strict=strict,
  798. from_attributes=from_attributes,
  799. context=context,
  800. update=update,
  801. )
  802. def model_dump(
  803. self,
  804. *,
  805. mode: Union[Literal["json", "python"], str] = "python",
  806. include: IncEx = None,
  807. exclude: IncEx = None,
  808. context: Union[Dict[str, Any], None] = None,
  809. by_alias: bool = False,
  810. exclude_unset: bool = False,
  811. exclude_defaults: bool = False,
  812. exclude_none: bool = False,
  813. round_trip: bool = False,
  814. warnings: Union[bool, Literal["none", "warn", "error"]] = True,
  815. serialize_as_any: bool = False,
  816. ) -> Dict[str, Any]:
  817. if PYDANTIC_VERSION >= "2.7.0":
  818. extra_kwargs: Dict[str, Any] = {
  819. "context": context,
  820. "serialize_as_any": serialize_as_any,
  821. }
  822. else:
  823. extra_kwargs = {}
  824. if IS_PYDANTIC_V2:
  825. return super().model_dump(
  826. mode=mode,
  827. include=include,
  828. exclude=exclude,
  829. by_alias=by_alias,
  830. exclude_unset=exclude_unset,
  831. exclude_defaults=exclude_defaults,
  832. exclude_none=exclude_none,
  833. round_trip=round_trip,
  834. warnings=warnings,
  835. **extra_kwargs,
  836. )
  837. else:
  838. return super().dict(
  839. include=include,
  840. exclude=exclude,
  841. by_alias=by_alias,
  842. exclude_unset=exclude_unset,
  843. exclude_defaults=exclude_defaults,
  844. exclude_none=exclude_none,
  845. )
  846. @deprecated(
  847. """
  848. 🚨 `obj.dict()` was deprecated in SQLModel 0.0.14, you should
  849. instead use `obj.model_dump()`.
  850. """
  851. )
  852. def dict(
  853. self,
  854. *,
  855. include: IncEx = None,
  856. exclude: IncEx = None,
  857. by_alias: bool = False,
  858. exclude_unset: bool = False,
  859. exclude_defaults: bool = False,
  860. exclude_none: bool = False,
  861. ) -> Dict[str, Any]:
  862. return self.model_dump(
  863. include=include,
  864. exclude=exclude,
  865. by_alias=by_alias,
  866. exclude_unset=exclude_unset,
  867. exclude_defaults=exclude_defaults,
  868. exclude_none=exclude_none,
  869. )
  870. @classmethod
  871. @deprecated(
  872. """
  873. 🚨 `obj.from_orm(data)` was deprecated in SQLModel 0.0.14, you should
  874. instead use `obj.model_validate(data)`.
  875. """
  876. )
  877. def from_orm(
  878. cls: Type[_TSQLModel], obj: Any, update: Optional[Dict[str, Any]] = None
  879. ) -> _TSQLModel:
  880. return cls.model_validate(obj, update=update)
  881. @classmethod
  882. @deprecated(
  883. """
  884. 🚨 `obj.parse_obj(data)` was deprecated in SQLModel 0.0.14, you should
  885. instead use `obj.model_validate(data)`.
  886. """
  887. )
  888. def parse_obj(
  889. cls: Type[_TSQLModel], obj: Any, update: Optional[Dict[str, Any]] = None
  890. ) -> _TSQLModel:
  891. if not IS_PYDANTIC_V2:
  892. obj = cls._enforce_dict_if_root(obj) # type: ignore[attr-defined] # noqa
  893. return cls.model_validate(obj, update=update)
  894. # From Pydantic, override to only show keys from fields, omit SQLAlchemy attributes
  895. @deprecated(
  896. """
  897. 🚨 You should not access `obj._calculate_keys()` directly.
  898. It is only useful for Pydantic v1.X, you should probably upgrade to
  899. Pydantic v2.X.
  900. """,
  901. category=None,
  902. )
  903. def _calculate_keys(
  904. self,
  905. include: Optional[Mapping[Union[int, str], Any]],
  906. exclude: Optional[Mapping[Union[int, str], Any]],
  907. exclude_unset: bool,
  908. update: Optional[Dict[str, Any]] = None,
  909. ) -> Optional[AbstractSet[str]]:
  910. return _calculate_keys(
  911. self,
  912. include=include,
  913. exclude=exclude,
  914. exclude_unset=exclude_unset,
  915. update=update,
  916. )
  917. def sqlmodel_update(
  918. self: _TSQLModel,
  919. obj: Union[Dict[str, Any], BaseModel],
  920. *,
  921. update: Union[Dict[str, Any], None] = None,
  922. ) -> _TSQLModel:
  923. use_update = (update or {}).copy()
  924. if isinstance(obj, dict):
  925. for key, value in {**obj, **use_update}.items():
  926. if key in get_model_fields(self):
  927. setattr(self, key, value)
  928. elif isinstance(obj, BaseModel):
  929. for key in get_model_fields(obj):
  930. if key in use_update:
  931. value = use_update.pop(key)
  932. else:
  933. value = getattr(obj, key)
  934. setattr(self, key, value)
  935. for remaining_key in use_update:
  936. if remaining_key in get_model_fields(self):
  937. value = use_update.pop(remaining_key)
  938. setattr(self, remaining_key, value)
  939. else:
  940. raise ValueError(
  941. "Can't use sqlmodel_update() with something that "
  942. f"is not a dict or SQLModel or Pydantic model: {obj}"
  943. )
  944. return self