session.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. from typing import (
  2. Any,
  3. Dict,
  4. Mapping,
  5. Optional,
  6. Sequence,
  7. TypeVar,
  8. Union,
  9. overload,
  10. )
  11. from sqlalchemy import util
  12. from sqlalchemy.engine.interfaces import _CoreAnyExecuteParams
  13. from sqlalchemy.engine.result import Result, ScalarResult, TupleResult
  14. from sqlalchemy.orm import Query as _Query
  15. from sqlalchemy.orm import Session as _Session
  16. from sqlalchemy.orm._typing import OrmExecuteOptionsParameter
  17. from sqlalchemy.sql._typing import _ColumnsClauseArgument
  18. from sqlalchemy.sql.base import Executable as _Executable
  19. from sqlmodel.sql.base import Executable
  20. from sqlmodel.sql.expression import Select, SelectOfScalar
  21. from typing_extensions import deprecated
  22. _TSelectParam = TypeVar("_TSelectParam", bound=Any)
  23. class Session(_Session):
  24. @overload
  25. def exec(
  26. self,
  27. statement: Select[_TSelectParam],
  28. *,
  29. params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
  30. execution_options: Mapping[str, Any] = util.EMPTY_DICT,
  31. bind_arguments: Optional[Dict[str, Any]] = None,
  32. _parent_execute_state: Optional[Any] = None,
  33. _add_event: Optional[Any] = None,
  34. ) -> TupleResult[_TSelectParam]: ...
  35. @overload
  36. def exec(
  37. self,
  38. statement: SelectOfScalar[_TSelectParam],
  39. *,
  40. params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
  41. execution_options: Mapping[str, Any] = util.EMPTY_DICT,
  42. bind_arguments: Optional[Dict[str, Any]] = None,
  43. _parent_execute_state: Optional[Any] = None,
  44. _add_event: Optional[Any] = None,
  45. ) -> ScalarResult[_TSelectParam]: ...
  46. def exec(
  47. self,
  48. statement: Union[
  49. Select[_TSelectParam],
  50. SelectOfScalar[_TSelectParam],
  51. Executable[_TSelectParam],
  52. ],
  53. *,
  54. params: Optional[Union[Mapping[str, Any], Sequence[Mapping[str, Any]]]] = None,
  55. execution_options: Mapping[str, Any] = util.EMPTY_DICT,
  56. bind_arguments: Optional[Dict[str, Any]] = None,
  57. _parent_execute_state: Optional[Any] = None,
  58. _add_event: Optional[Any] = None,
  59. ) -> Union[TupleResult[_TSelectParam], ScalarResult[_TSelectParam]]:
  60. results = super().execute(
  61. statement,
  62. params=params,
  63. execution_options=execution_options,
  64. bind_arguments=bind_arguments,
  65. _parent_execute_state=_parent_execute_state,
  66. _add_event=_add_event,
  67. )
  68. if isinstance(statement, SelectOfScalar):
  69. return results.scalars()
  70. return results # type: ignore
  71. @deprecated(
  72. """
  73. 🚨 You probably want to use `session.exec()` instead of `session.execute()`.
  74. This is the original SQLAlchemy `session.execute()` method that returns objects
  75. of type `Row`, and that you have to call `scalars()` to get the model objects.
  76. For example:
  77. ```Python
  78. heroes = session.execute(select(Hero)).scalars().all()
  79. ```
  80. instead you could use `exec()`:
  81. ```Python
  82. heroes = session.exec(select(Hero)).all()
  83. ```
  84. """,
  85. category=None,
  86. )
  87. def execute( # type: ignore
  88. self,
  89. statement: _Executable,
  90. params: Optional[_CoreAnyExecuteParams] = None,
  91. *,
  92. execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
  93. bind_arguments: Optional[Dict[str, Any]] = None,
  94. _parent_execute_state: Optional[Any] = None,
  95. _add_event: Optional[Any] = None,
  96. ) -> Result[Any]:
  97. """
  98. 🚨 You probably want to use `session.exec()` instead of `session.execute()`.
  99. This is the original SQLAlchemy `session.execute()` method that returns objects
  100. of type `Row`, and that you have to call `scalars()` to get the model objects.
  101. For example:
  102. ```Python
  103. heroes = session.execute(select(Hero)).scalars().all()
  104. ```
  105. instead you could use `exec()`:
  106. ```Python
  107. heroes = session.exec(select(Hero)).all()
  108. ```
  109. """
  110. return super().execute(
  111. statement,
  112. params=params,
  113. execution_options=execution_options,
  114. bind_arguments=bind_arguments,
  115. _parent_execute_state=_parent_execute_state,
  116. _add_event=_add_event,
  117. )
  118. @deprecated(
  119. """
  120. 🚨 You probably want to use `session.exec()` instead of `session.query()`.
  121. `session.exec()` is SQLModel's own short version with increased type
  122. annotations.
  123. Or otherwise you might want to use `session.execute()` instead of
  124. `session.query()`.
  125. """
  126. )
  127. def query( # type: ignore
  128. self, *entities: _ColumnsClauseArgument[Any], **kwargs: Any
  129. ) -> _Query[Any]:
  130. """
  131. 🚨 You probably want to use `session.exec()` instead of `session.query()`.
  132. `session.exec()` is SQLModel's own short version with increased type
  133. annotations.
  134. Or otherwise you might want to use `session.execute()` instead of
  135. `session.query()`.
  136. """
  137. return super().query(*entities, **kwargs)