aiomysql.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250
  1. # dialects/mysql/aiomysql.py
  2. # Copyright (C) 2005-2026 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. The aiomysql dialect is SQLAlchemy's second Python asyncio dialect.
  14. Using a special asyncio mediation layer, the aiomysql 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(
  21. "mysql+aiomysql://user:pass@hostname/dbname?charset=utf8mb4"
  22. )
  23. """ # noqa
  24. from __future__ import annotations
  25. from types import ModuleType
  26. from typing import Any
  27. from typing import Dict
  28. from typing import Optional
  29. from typing import Tuple
  30. from typing import TYPE_CHECKING
  31. from typing import Union
  32. from .pymysql import MySQLDialect_pymysql
  33. from ... import pool
  34. from ... import util
  35. from ...connectors.asyncio import AsyncAdapt_dbapi_connection
  36. from ...connectors.asyncio import AsyncAdapt_dbapi_cursor
  37. from ...connectors.asyncio import AsyncAdapt_dbapi_module
  38. from ...connectors.asyncio import AsyncAdapt_dbapi_ss_cursor
  39. from ...connectors.asyncio import AsyncAdapt_terminate
  40. from ...util.concurrency import await_fallback
  41. from ...util.concurrency import await_only
  42. if TYPE_CHECKING:
  43. from ...connectors.asyncio import AsyncIODBAPIConnection
  44. from ...connectors.asyncio import AsyncIODBAPICursor
  45. from ...engine.interfaces import ConnectArgsType
  46. from ...engine.interfaces import DBAPIConnection
  47. from ...engine.interfaces import DBAPICursor
  48. from ...engine.interfaces import DBAPIModule
  49. from ...engine.interfaces import PoolProxiedConnection
  50. from ...engine.url import URL
  51. class AsyncAdapt_aiomysql_cursor(AsyncAdapt_dbapi_cursor):
  52. __slots__ = ()
  53. def _make_new_cursor(
  54. self, connection: AsyncIODBAPIConnection
  55. ) -> AsyncIODBAPICursor:
  56. return connection.cursor(self._adapt_connection.dbapi.Cursor)
  57. class AsyncAdapt_aiomysql_ss_cursor(
  58. AsyncAdapt_dbapi_ss_cursor, AsyncAdapt_aiomysql_cursor
  59. ):
  60. __slots__ = ()
  61. def _make_new_cursor(
  62. self, connection: AsyncIODBAPIConnection
  63. ) -> AsyncIODBAPICursor:
  64. return connection.cursor(
  65. self._adapt_connection.dbapi.aiomysql.cursors.SSCursor
  66. )
  67. class AsyncAdapt_aiomysql_connection(
  68. AsyncAdapt_terminate, AsyncAdapt_dbapi_connection
  69. ):
  70. __slots__ = ()
  71. _cursor_cls = AsyncAdapt_aiomysql_cursor
  72. _ss_cursor_cls = AsyncAdapt_aiomysql_ss_cursor
  73. def ping(self, reconnect: bool) -> None:
  74. assert not reconnect
  75. self.await_(self._connection.ping(reconnect))
  76. def character_set_name(self) -> Optional[str]:
  77. return self._connection.character_set_name() # type: ignore[no-any-return] # noqa: E501
  78. def autocommit(self, value: Any) -> None:
  79. self.await_(self._connection.autocommit(value))
  80. def get_autocommit(self) -> bool:
  81. return self._connection.get_autocommit() # type: ignore
  82. def close(self) -> None:
  83. self.await_(self._connection.ensure_closed())
  84. async def _terminate_graceful_close(self) -> None:
  85. await self._connection.ensure_closed()
  86. def _terminate_force_close(self) -> None:
  87. # it's not awaitable.
  88. self._connection.close()
  89. class AsyncAdaptFallback_aiomysql_connection(AsyncAdapt_aiomysql_connection):
  90. __slots__ = ()
  91. await_ = staticmethod(await_fallback)
  92. class AsyncAdapt_aiomysql_dbapi(AsyncAdapt_dbapi_module):
  93. def __init__(self, aiomysql: ModuleType, pymysql: ModuleType):
  94. self.aiomysql = aiomysql
  95. self.pymysql = pymysql
  96. self.paramstyle = "format"
  97. self._init_dbapi_attributes()
  98. self.Cursor, self.SSCursor = self._init_cursors_subclasses()
  99. def _init_dbapi_attributes(self) -> None:
  100. for name in (
  101. "Warning",
  102. "Error",
  103. "InterfaceError",
  104. "DataError",
  105. "DatabaseError",
  106. "OperationalError",
  107. "InterfaceError",
  108. "IntegrityError",
  109. "ProgrammingError",
  110. "InternalError",
  111. "NotSupportedError",
  112. ):
  113. setattr(self, name, getattr(self.aiomysql, name))
  114. for name in (
  115. "NUMBER",
  116. "STRING",
  117. "DATETIME",
  118. "BINARY",
  119. "TIMESTAMP",
  120. "Binary",
  121. ):
  122. setattr(self, name, getattr(self.pymysql, name))
  123. def connect(self, *arg: Any, **kw: Any) -> AsyncAdapt_aiomysql_connection:
  124. async_fallback = kw.pop("async_fallback", False)
  125. creator_fn = kw.pop("async_creator_fn", self.aiomysql.connect)
  126. if util.asbool(async_fallback):
  127. return AsyncAdaptFallback_aiomysql_connection(
  128. self,
  129. await_fallback(creator_fn(*arg, **kw)),
  130. )
  131. else:
  132. return AsyncAdapt_aiomysql_connection(
  133. self,
  134. await_only(creator_fn(*arg, **kw)),
  135. )
  136. def _init_cursors_subclasses(
  137. self,
  138. ) -> Tuple[AsyncIODBAPICursor, AsyncIODBAPICursor]:
  139. # suppress unconditional warning emitted by aiomysql
  140. class Cursor(self.aiomysql.Cursor): # type: ignore[misc, name-defined]
  141. async def _show_warnings(
  142. self, conn: AsyncIODBAPIConnection
  143. ) -> None:
  144. pass
  145. class SSCursor(self.aiomysql.SSCursor): # type: ignore[misc, name-defined] # noqa: E501
  146. async def _show_warnings(
  147. self, conn: AsyncIODBAPIConnection
  148. ) -> None:
  149. pass
  150. return Cursor, SSCursor # type: ignore[return-value]
  151. class MySQLDialect_aiomysql(MySQLDialect_pymysql):
  152. driver = "aiomysql"
  153. supports_statement_cache = True
  154. supports_server_side_cursors = True
  155. _sscursor = AsyncAdapt_aiomysql_ss_cursor
  156. is_async = True
  157. has_terminate = True
  158. @classmethod
  159. def import_dbapi(cls) -> AsyncAdapt_aiomysql_dbapi:
  160. return AsyncAdapt_aiomysql_dbapi(
  161. __import__("aiomysql"), __import__("pymysql")
  162. )
  163. @classmethod
  164. def get_pool_class(cls, url: URL) -> type:
  165. async_fallback = url.query.get("async_fallback", False)
  166. if util.asbool(async_fallback):
  167. return pool.FallbackAsyncAdaptedQueuePool
  168. else:
  169. return pool.AsyncAdaptedQueuePool
  170. def do_terminate(self, dbapi_connection: DBAPIConnection) -> None:
  171. dbapi_connection.terminate()
  172. def create_connect_args(
  173. self, url: URL, _translate_args: Optional[Dict[str, Any]] = None
  174. ) -> ConnectArgsType:
  175. return super().create_connect_args(
  176. url, _translate_args=dict(username="user", database="db")
  177. )
  178. def is_disconnect(
  179. self,
  180. e: DBAPIModule.Error,
  181. connection: Optional[Union[PoolProxiedConnection, DBAPIConnection]],
  182. cursor: Optional[DBAPICursor],
  183. ) -> bool:
  184. if super().is_disconnect(e, connection, cursor):
  185. return True
  186. else:
  187. str_e = str(e).lower()
  188. return "not connected" in str_e
  189. def _found_rows_client_flag(self) -> int:
  190. from pymysql.constants import CLIENT # type: ignore
  191. return CLIENT.FOUND_ROWS # type: ignore[no-any-return]
  192. def get_driver_connection(
  193. self, connection: DBAPIConnection
  194. ) -> AsyncIODBAPIConnection:
  195. return connection._connection # type: ignore[no-any-return]
  196. dialect = MySQLDialect_aiomysql