session.py 4.8 KB

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