aiomysql.py 9.7 KB

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