asyncmy.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. # dialects/mysql/asyncmy.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+asyncmy
  9. :name: asyncmy
  10. :dbapi: asyncmy
  11. :connectstring: mysql+asyncmy://user:password@host:port/dbname[?key=value&key=value...]
  12. :url: https://github.com/long2ice/asyncmy
  13. Using a special asyncio mediation layer, the asyncmy dialect is usable
  14. as the backend for the :ref:`SQLAlchemy asyncio <asyncio_toplevel>`
  15. extension package.
  16. This dialect should normally be used only with the
  17. :func:`_asyncio.create_async_engine` engine creation function::
  18. from sqlalchemy.ext.asyncio import create_async_engine
  19. engine = create_async_engine(
  20. "mysql+asyncmy://user:pass@hostname/dbname?charset=utf8mb4"
  21. )
  22. """ # noqa
  23. from __future__ import annotations
  24. from types import ModuleType
  25. from typing import Any
  26. from typing import NoReturn
  27. from typing import Optional
  28. from typing import TYPE_CHECKING
  29. from typing import Union
  30. from .pymysql import MySQLDialect_pymysql
  31. from ... import pool
  32. from ... import util
  33. from ...connectors.asyncio import AsyncAdapt_dbapi_connection
  34. from ...connectors.asyncio import AsyncAdapt_dbapi_cursor
  35. from ...connectors.asyncio import AsyncAdapt_dbapi_module
  36. from ...connectors.asyncio import AsyncAdapt_dbapi_ss_cursor
  37. from ...connectors.asyncio import AsyncAdapt_terminate
  38. from ...util.concurrency import await_fallback
  39. from ...util.concurrency import await_only
  40. if TYPE_CHECKING:
  41. from ...connectors.asyncio import AsyncIODBAPIConnection
  42. from ...connectors.asyncio import AsyncIODBAPICursor
  43. from ...engine.interfaces import ConnectArgsType
  44. from ...engine.interfaces import DBAPIConnection
  45. from ...engine.interfaces import DBAPICursor
  46. from ...engine.interfaces import DBAPIModule
  47. from ...engine.interfaces import PoolProxiedConnection
  48. from ...engine.url import URL
  49. class AsyncAdapt_asyncmy_cursor(AsyncAdapt_dbapi_cursor):
  50. __slots__ = ()
  51. class AsyncAdapt_asyncmy_ss_cursor(
  52. AsyncAdapt_dbapi_ss_cursor, AsyncAdapt_asyncmy_cursor
  53. ):
  54. __slots__ = ()
  55. def _make_new_cursor(
  56. self, connection: AsyncIODBAPIConnection
  57. ) -> AsyncIODBAPICursor:
  58. return connection.cursor(
  59. self._adapt_connection.dbapi.asyncmy.cursors.SSCursor
  60. )
  61. class AsyncAdapt_asyncmy_connection(
  62. AsyncAdapt_terminate, AsyncAdapt_dbapi_connection
  63. ):
  64. __slots__ = ()
  65. _cursor_cls = AsyncAdapt_asyncmy_cursor
  66. _ss_cursor_cls = AsyncAdapt_asyncmy_ss_cursor
  67. def _handle_exception(self, error: Exception) -> NoReturn:
  68. if isinstance(error, AttributeError):
  69. raise self.dbapi.InternalError(
  70. "network operation failed due to asyncmy attribute error"
  71. )
  72. raise error
  73. def ping(self, reconnect: bool) -> None:
  74. assert not reconnect
  75. return self.await_(self._do_ping())
  76. async def _do_ping(self) -> None:
  77. try:
  78. async with self._execute_mutex:
  79. await self._connection.ping(False)
  80. except Exception as error:
  81. self._handle_exception(error)
  82. def character_set_name(self) -> Optional[str]:
  83. return self._connection.character_set_name() # type: ignore[no-any-return] # noqa: E501
  84. def autocommit(self, value: Any) -> None:
  85. self.await_(self._connection.autocommit(value))
  86. def get_autocommit(self) -> bool:
  87. return self._connection.get_autocommit() # type: ignore
  88. def close(self) -> None:
  89. self.await_(self._connection.ensure_closed())
  90. async def _terminate_graceful_close(self) -> None:
  91. await self._connection.ensure_closed()
  92. def _terminate_force_close(self) -> None:
  93. # it's not awaitable.
  94. self._connection.close()
  95. class AsyncAdaptFallback_asyncmy_connection(AsyncAdapt_asyncmy_connection):
  96. __slots__ = ()
  97. await_ = staticmethod(await_fallback)
  98. class AsyncAdapt_asyncmy_dbapi(AsyncAdapt_dbapi_module):
  99. def __init__(self, asyncmy: ModuleType):
  100. self.asyncmy = asyncmy
  101. self.paramstyle = "format"
  102. self._init_dbapi_attributes()
  103. def _init_dbapi_attributes(self) -> None:
  104. for name in (
  105. "Warning",
  106. "Error",
  107. "InterfaceError",
  108. "DataError",
  109. "DatabaseError",
  110. "OperationalError",
  111. "InterfaceError",
  112. "IntegrityError",
  113. "ProgrammingError",
  114. "InternalError",
  115. "NotSupportedError",
  116. ):
  117. setattr(self, name, getattr(self.asyncmy.errors, name))
  118. STRING = util.symbol("STRING")
  119. NUMBER = util.symbol("NUMBER")
  120. BINARY = util.symbol("BINARY")
  121. DATETIME = util.symbol("DATETIME")
  122. TIMESTAMP = util.symbol("TIMESTAMP")
  123. Binary = staticmethod(bytes)
  124. def connect(self, *arg: Any, **kw: Any) -> AsyncAdapt_asyncmy_connection:
  125. async_fallback = kw.pop("async_fallback", False)
  126. creator_fn = kw.pop("async_creator_fn", self.asyncmy.connect)
  127. if util.asbool(async_fallback):
  128. return AsyncAdaptFallback_asyncmy_connection(
  129. self,
  130. await_fallback(creator_fn(*arg, **kw)),
  131. )
  132. else:
  133. return AsyncAdapt_asyncmy_connection(
  134. self,
  135. await_only(creator_fn(*arg, **kw)),
  136. )
  137. class MySQLDialect_asyncmy(MySQLDialect_pymysql):
  138. driver = "asyncmy"
  139. supports_statement_cache = True
  140. supports_server_side_cursors = True
  141. _sscursor = AsyncAdapt_asyncmy_ss_cursor
  142. is_async = True
  143. has_terminate = True
  144. @classmethod
  145. def import_dbapi(cls) -> DBAPIModule:
  146. return AsyncAdapt_asyncmy_dbapi(__import__("asyncmy"))
  147. @classmethod
  148. def get_pool_class(cls, url: URL) -> type:
  149. async_fallback = url.query.get("async_fallback", False)
  150. if util.asbool(async_fallback):
  151. return pool.FallbackAsyncAdaptedQueuePool
  152. else:
  153. return pool.AsyncAdaptedQueuePool
  154. def do_terminate(self, dbapi_connection: DBAPIConnection) -> None:
  155. dbapi_connection.terminate()
  156. def create_connect_args(self, url: URL) -> ConnectArgsType: # type: ignore[override] # noqa: E501
  157. return super().create_connect_args(
  158. url, _translate_args=dict(username="user", database="db")
  159. )
  160. def is_disconnect(
  161. self,
  162. e: DBAPIModule.Error,
  163. connection: Optional[Union[PoolProxiedConnection, DBAPIConnection]],
  164. cursor: Optional[DBAPICursor],
  165. ) -> bool:
  166. if super().is_disconnect(e, connection, cursor):
  167. return True
  168. else:
  169. str_e = str(e).lower()
  170. return (
  171. "not connected" in str_e or "network operation failed" in str_e
  172. )
  173. def _found_rows_client_flag(self) -> int:
  174. from asyncmy.constants import CLIENT # type: ignore
  175. return CLIENT.FOUND_ROWS # type: ignore[no-any-return]
  176. def get_driver_connection(
  177. self, connection: DBAPIConnection
  178. ) -> AsyncIODBAPIConnection:
  179. return connection._connection # type: ignore[no-any-return]
  180. dialect = MySQLDialect_asyncmy