asyncio.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. # connectors/asyncio.py
  2. # Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. # mypy: ignore-errors
  8. """generic asyncio-adapted versions of DBAPI connection and cursor"""
  9. from __future__ import annotations
  10. import collections
  11. import itertools
  12. from ..engine import AdaptedConnection
  13. from ..util.concurrency import asyncio
  14. from ..util.concurrency import await_fallback
  15. from ..util.concurrency import await_only
  16. class AsyncAdapt_dbapi_cursor:
  17. server_side = False
  18. __slots__ = (
  19. "_adapt_connection",
  20. "_connection",
  21. "await_",
  22. "_cursor",
  23. "_rows",
  24. )
  25. def __init__(self, adapt_connection):
  26. self._adapt_connection = adapt_connection
  27. self._connection = adapt_connection._connection
  28. self.await_ = adapt_connection.await_
  29. cursor = self._connection.cursor()
  30. self._cursor = self._aenter_cursor(cursor)
  31. self._rows = collections.deque()
  32. def _aenter_cursor(self, cursor):
  33. return self.await_(cursor.__aenter__())
  34. @property
  35. def description(self):
  36. return self._cursor.description
  37. @property
  38. def rowcount(self):
  39. return self._cursor.rowcount
  40. @property
  41. def arraysize(self):
  42. return self._cursor.arraysize
  43. @arraysize.setter
  44. def arraysize(self, value):
  45. self._cursor.arraysize = value
  46. @property
  47. def lastrowid(self):
  48. return self._cursor.lastrowid
  49. def close(self):
  50. # note we aren't actually closing the cursor here,
  51. # we are just letting GC do it. see notes in aiomysql dialect
  52. self._rows.clear()
  53. def execute(self, operation, parameters=None):
  54. return self.await_(self._execute_async(operation, parameters))
  55. def executemany(self, operation, seq_of_parameters):
  56. return self.await_(
  57. self._executemany_async(operation, seq_of_parameters)
  58. )
  59. async def _execute_async(self, operation, parameters):
  60. async with self._adapt_connection._execute_mutex:
  61. result = await self._cursor.execute(operation, parameters or ())
  62. if self._cursor.description and not self.server_side:
  63. self._rows = collections.deque(await self._cursor.fetchall())
  64. return result
  65. async def _executemany_async(self, operation, seq_of_parameters):
  66. async with self._adapt_connection._execute_mutex:
  67. return await self._cursor.executemany(operation, seq_of_parameters)
  68. def nextset(self):
  69. self.await_(self._cursor.nextset())
  70. if self._cursor.description and not self.server_side:
  71. self._rows = collections.deque(
  72. self.await_(self._cursor.fetchall())
  73. )
  74. def setinputsizes(self, *inputsizes):
  75. # NOTE: this is overrridden in aioodbc due to
  76. # see https://github.com/aio-libs/aioodbc/issues/451
  77. # right now
  78. return self.await_(self._cursor.setinputsizes(*inputsizes))
  79. def __iter__(self):
  80. while self._rows:
  81. yield self._rows.popleft()
  82. def fetchone(self):
  83. if self._rows:
  84. return self._rows.popleft()
  85. else:
  86. return None
  87. def fetchmany(self, size=None):
  88. if size is None:
  89. size = self.arraysize
  90. rr = iter(self._rows)
  91. retval = list(itertools.islice(rr, 0, size))
  92. self._rows = collections.deque(rr)
  93. return retval
  94. def fetchall(self):
  95. retval = list(self._rows)
  96. self._rows.clear()
  97. return retval
  98. class AsyncAdapt_dbapi_ss_cursor(AsyncAdapt_dbapi_cursor):
  99. __slots__ = ()
  100. server_side = True
  101. def __init__(self, adapt_connection):
  102. self._adapt_connection = adapt_connection
  103. self._connection = adapt_connection._connection
  104. self.await_ = adapt_connection.await_
  105. cursor = self._connection.cursor()
  106. self._cursor = self.await_(cursor.__aenter__())
  107. def close(self):
  108. if self._cursor is not None:
  109. self.await_(self._cursor.close())
  110. self._cursor = None
  111. def fetchone(self):
  112. return self.await_(self._cursor.fetchone())
  113. def fetchmany(self, size=None):
  114. return self.await_(self._cursor.fetchmany(size=size))
  115. def fetchall(self):
  116. return self.await_(self._cursor.fetchall())
  117. class AsyncAdapt_dbapi_connection(AdaptedConnection):
  118. _cursor_cls = AsyncAdapt_dbapi_cursor
  119. _ss_cursor_cls = AsyncAdapt_dbapi_ss_cursor
  120. await_ = staticmethod(await_only)
  121. __slots__ = ("dbapi", "_execute_mutex")
  122. def __init__(self, dbapi, connection):
  123. self.dbapi = dbapi
  124. self._connection = connection
  125. self._execute_mutex = asyncio.Lock()
  126. def ping(self, reconnect):
  127. return self.await_(self._connection.ping(reconnect))
  128. def add_output_converter(self, *arg, **kw):
  129. self._connection.add_output_converter(*arg, **kw)
  130. def character_set_name(self):
  131. return self._connection.character_set_name()
  132. @property
  133. def autocommit(self):
  134. return self._connection.autocommit
  135. @autocommit.setter
  136. def autocommit(self, value):
  137. # https://github.com/aio-libs/aioodbc/issues/448
  138. # self._connection.autocommit = value
  139. self._connection._conn.autocommit = value
  140. def cursor(self, server_side=False):
  141. if server_side:
  142. return self._ss_cursor_cls(self)
  143. else:
  144. return self._cursor_cls(self)
  145. def rollback(self):
  146. self.await_(self._connection.rollback())
  147. def commit(self):
  148. self.await_(self._connection.commit())
  149. def close(self):
  150. self.await_(self._connection.close())
  151. class AsyncAdaptFallback_dbapi_connection(AsyncAdapt_dbapi_connection):
  152. __slots__ = ()
  153. await_ = staticmethod(await_fallback)