asyncmy.py 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330
  1. # dialects/mysql/asyncmy.py
  2. # Copyright (C) 2005-2024 the SQLAlchemy authors and contributors <see AUTHORS
  3. # 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. r"""
  9. .. dialect:: mysql+asyncmy
  10. :name: asyncmy
  11. :dbapi: asyncmy
  12. :connectstring: mysql+asyncmy://user:password@host:port/dbname[?key=value&key=value...]
  13. :url: https://github.com/long2ice/asyncmy
  14. Using a special asyncio mediation layer, the asyncmy dialect is usable
  15. as the backend for the :ref:`SQLAlchemy asyncio <asyncio_toplevel>`
  16. extension package.
  17. This dialect should normally be used only with the
  18. :func:`_asyncio.create_async_engine` engine creation function::
  19. from sqlalchemy.ext.asyncio import create_async_engine
  20. engine = create_async_engine("mysql+asyncmy://user:pass@hostname/dbname?charset=utf8mb4")
  21. """ # noqa
  22. from contextlib import asynccontextmanager
  23. from .pymysql import MySQLDialect_pymysql
  24. from ... import pool
  25. from ... import util
  26. from ...engine import AdaptedConnection
  27. from ...util.concurrency import asyncio
  28. from ...util.concurrency import await_fallback
  29. from ...util.concurrency import await_only
  30. class AsyncAdapt_asyncmy_cursor:
  31. # TODO: base on connectors/asyncio.py
  32. # see #10415
  33. server_side = False
  34. __slots__ = (
  35. "_adapt_connection",
  36. "_connection",
  37. "await_",
  38. "_cursor",
  39. "_rows",
  40. )
  41. def __init__(self, adapt_connection):
  42. self._adapt_connection = adapt_connection
  43. self._connection = adapt_connection._connection
  44. self.await_ = adapt_connection.await_
  45. cursor = self._connection.cursor()
  46. self._cursor = self.await_(cursor.__aenter__())
  47. self._rows = []
  48. @property
  49. def description(self):
  50. return self._cursor.description
  51. @property
  52. def rowcount(self):
  53. return self._cursor.rowcount
  54. @property
  55. def arraysize(self):
  56. return self._cursor.arraysize
  57. @arraysize.setter
  58. def arraysize(self, value):
  59. self._cursor.arraysize = value
  60. @property
  61. def lastrowid(self):
  62. return self._cursor.lastrowid
  63. def close(self):
  64. # note we aren't actually closing the cursor here,
  65. # we are just letting GC do it. to allow this to be async
  66. # we would need the Result to change how it does "Safe close cursor".
  67. # MySQL "cursors" don't actually have state to be "closed" besides
  68. # exhausting rows, which we already have done for sync cursor.
  69. # another option would be to emulate aiosqlite dialect and assign
  70. # cursor only if we are doing server side cursor operation.
  71. self._rows[:] = []
  72. def execute(self, operation, parameters=None):
  73. return self.await_(self._execute_async(operation, parameters))
  74. def executemany(self, operation, seq_of_parameters):
  75. return self.await_(
  76. self._executemany_async(operation, seq_of_parameters)
  77. )
  78. async def _execute_async(self, operation, parameters):
  79. async with self._adapt_connection._mutex_and_adapt_errors():
  80. if parameters is None:
  81. result = await self._cursor.execute(operation)
  82. else:
  83. result = await self._cursor.execute(operation, parameters)
  84. if not self.server_side:
  85. # asyncmy has a "fake" async result, so we have to pull it out
  86. # of that here since our default result is not async.
  87. # we could just as easily grab "_rows" here and be done with it
  88. # but this is safer.
  89. self._rows = list(await self._cursor.fetchall())
  90. return result
  91. async def _executemany_async(self, operation, seq_of_parameters):
  92. async with self._adapt_connection._mutex_and_adapt_errors():
  93. return await self._cursor.executemany(operation, seq_of_parameters)
  94. def setinputsizes(self, *inputsizes):
  95. pass
  96. def __iter__(self):
  97. while self._rows:
  98. yield self._rows.pop(0)
  99. def fetchone(self):
  100. if self._rows:
  101. return self._rows.pop(0)
  102. else:
  103. return None
  104. def fetchmany(self, size=None):
  105. if size is None:
  106. size = self.arraysize
  107. retval = self._rows[0:size]
  108. self._rows[:] = self._rows[size:]
  109. return retval
  110. def fetchall(self):
  111. retval = self._rows[:]
  112. self._rows[:] = []
  113. return retval
  114. class AsyncAdapt_asyncmy_ss_cursor(AsyncAdapt_asyncmy_cursor):
  115. # TODO: base on connectors/asyncio.py
  116. # see #10415
  117. __slots__ = ()
  118. server_side = True
  119. def __init__(self, adapt_connection):
  120. self._adapt_connection = adapt_connection
  121. self._connection = adapt_connection._connection
  122. self.await_ = adapt_connection.await_
  123. cursor = self._connection.cursor(
  124. adapt_connection.dbapi.asyncmy.cursors.SSCursor
  125. )
  126. self._cursor = self.await_(cursor.__aenter__())
  127. def close(self):
  128. if self._cursor is not None:
  129. self.await_(self._cursor.close())
  130. self._cursor = None
  131. def fetchone(self):
  132. return self.await_(self._cursor.fetchone())
  133. def fetchmany(self, size=None):
  134. return self.await_(self._cursor.fetchmany(size=size))
  135. def fetchall(self):
  136. return self.await_(self._cursor.fetchall())
  137. class AsyncAdapt_asyncmy_connection(AdaptedConnection):
  138. # TODO: base on connectors/asyncio.py
  139. # see #10415
  140. await_ = staticmethod(await_only)
  141. __slots__ = ("dbapi", "_execute_mutex")
  142. def __init__(self, dbapi, connection):
  143. self.dbapi = dbapi
  144. self._connection = connection
  145. self._execute_mutex = asyncio.Lock()
  146. @asynccontextmanager
  147. async def _mutex_and_adapt_errors(self):
  148. async with self._execute_mutex:
  149. try:
  150. yield
  151. except AttributeError:
  152. raise self.dbapi.InternalError(
  153. "network operation failed due to asyncmy attribute error"
  154. )
  155. def ping(self, reconnect):
  156. assert not reconnect
  157. return self.await_(self._do_ping())
  158. async def _do_ping(self):
  159. async with self._mutex_and_adapt_errors():
  160. return await self._connection.ping(False)
  161. def character_set_name(self):
  162. return self._connection.character_set_name()
  163. def autocommit(self, value):
  164. self.await_(self._connection.autocommit(value))
  165. def cursor(self, server_side=False):
  166. if server_side:
  167. return AsyncAdapt_asyncmy_ss_cursor(self)
  168. else:
  169. return AsyncAdapt_asyncmy_cursor(self)
  170. def rollback(self):
  171. self.await_(self._connection.rollback())
  172. def commit(self):
  173. self.await_(self._connection.commit())
  174. def close(self):
  175. # it's not awaitable.
  176. self._connection.close()
  177. class AsyncAdaptFallback_asyncmy_connection(AsyncAdapt_asyncmy_connection):
  178. __slots__ = ()
  179. await_ = staticmethod(await_fallback)
  180. def _Binary(x):
  181. """Return x as a binary type."""
  182. return bytes(x)
  183. class AsyncAdapt_asyncmy_dbapi:
  184. def __init__(self, asyncmy):
  185. self.asyncmy = asyncmy
  186. self.paramstyle = "format"
  187. self._init_dbapi_attributes()
  188. def _init_dbapi_attributes(self):
  189. for name in (
  190. "Warning",
  191. "Error",
  192. "InterfaceError",
  193. "DataError",
  194. "DatabaseError",
  195. "OperationalError",
  196. "InterfaceError",
  197. "IntegrityError",
  198. "ProgrammingError",
  199. "InternalError",
  200. "NotSupportedError",
  201. ):
  202. setattr(self, name, getattr(self.asyncmy.errors, name))
  203. STRING = util.symbol("STRING")
  204. NUMBER = util.symbol("NUMBER")
  205. BINARY = util.symbol("BINARY")
  206. DATETIME = util.symbol("DATETIME")
  207. TIMESTAMP = util.symbol("TIMESTAMP")
  208. Binary = staticmethod(_Binary)
  209. def connect(self, *arg, **kw):
  210. async_fallback = kw.pop("async_fallback", False)
  211. creator_fn = kw.pop("async_creator_fn", self.asyncmy.connect)
  212. if util.asbool(async_fallback):
  213. return AsyncAdaptFallback_asyncmy_connection(
  214. self,
  215. await_fallback(creator_fn(*arg, **kw)),
  216. )
  217. else:
  218. return AsyncAdapt_asyncmy_connection(
  219. self,
  220. await_only(creator_fn(*arg, **kw)),
  221. )
  222. class MySQLDialect_asyncmy(MySQLDialect_pymysql):
  223. driver = "asyncmy"
  224. supports_statement_cache = True
  225. supports_server_side_cursors = True
  226. _sscursor = AsyncAdapt_asyncmy_ss_cursor
  227. is_async = True
  228. @classmethod
  229. def import_dbapi(cls):
  230. return AsyncAdapt_asyncmy_dbapi(__import__("asyncmy"))
  231. @classmethod
  232. def get_pool_class(cls, url):
  233. async_fallback = url.query.get("async_fallback", False)
  234. if util.asbool(async_fallback):
  235. return pool.FallbackAsyncAdaptedQueuePool
  236. else:
  237. return pool.AsyncAdaptedQueuePool
  238. def create_connect_args(self, url):
  239. return super().create_connect_args(
  240. url, _translate_args=dict(username="user", database="db")
  241. )
  242. def is_disconnect(self, e, connection, cursor):
  243. if super().is_disconnect(e, connection, cursor):
  244. return True
  245. else:
  246. str_e = str(e).lower()
  247. return (
  248. "not connected" in str_e or "network operation failed" in str_e
  249. )
  250. def _found_rows_client_flag(self):
  251. from asyncmy.constants import CLIENT
  252. return CLIENT.FOUND_ROWS
  253. def get_driver_connection(self, connection):
  254. return connection._connection
  255. dialect = MySQLDialect_asyncmy