You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1050 lines
33KB

  1. # postgresql/asyncpg.py
  2. # Copyright (C) 2005-2021 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: http://www.opensource.org/licenses/mit-license.php
  7. r"""
  8. .. dialect:: postgresql+asyncpg
  9. :name: asyncpg
  10. :dbapi: asyncpg
  11. :connectstring: postgresql+asyncpg://user:password@host:port/dbname[?key=value&key=value...]
  12. :url: https://magicstack.github.io/asyncpg/
  13. The asyncpg dialect is SQLAlchemy's first Python asyncio dialect.
  14. Using a special asyncio mediation layer, the asyncpg 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("postgresql+asyncpg://user:pass@hostname/dbname")
  21. The dialect can also be run as a "synchronous" dialect within the
  22. :func:`_sa.create_engine` function, which will pass "await" calls into
  23. an ad-hoc event loop. This mode of operation is of **limited use**
  24. and is for special testing scenarios only. The mode can be enabled by
  25. adding the SQLAlchemy-specific flag ``async_fallback`` to the URL
  26. in conjunction with :func:`_sa.create_engine`::
  27. # for testing purposes only; do not use in production!
  28. engine = create_engine("postgresql+asyncpg://user:pass@hostname/dbname?async_fallback=true")
  29. .. versionadded:: 1.4
  30. .. note::
  31. By default asyncpg does not decode the ``json`` and ``jsonb`` types and
  32. returns them as strings. SQLAlchemy sets default type decoder for ``json``
  33. and ``jsonb`` types using the python builtin ``json.loads`` function.
  34. The json implementation used can be changed by setting the attribute
  35. ``json_deserializer`` when creating the engine with
  36. :func:`create_engine` or :func:`create_async_engine`.
  37. .. _asyncpg_prepared_statement_cache:
  38. Prepared Statement Cache
  39. --------------------------
  40. The asyncpg SQLAlchemy dialect makes use of ``asyncpg.connection.prepare()``
  41. for all statements. The prepared statement objects are cached after
  42. construction which appears to grant a 10% or more performance improvement for
  43. statement invocation. The cache is on a per-DBAPI connection basis, which
  44. means that the primary storage for prepared statements is within DBAPI
  45. connections pooled within the connection pool. The size of this cache
  46. defaults to 100 statements per DBAPI connection and may be adjusted using the
  47. ``prepared_statement_cache_size`` DBAPI argument (note that while this argument
  48. is implemented by SQLAlchemy, it is part of the DBAPI emulation portion of the
  49. asyncpg dialect, therefore is handled as a DBAPI argument, not a dialect
  50. argument)::
  51. engine = create_async_engine("postgresql+asyncpg://user:pass@hostname/dbname?prepared_statement_cache_size=500")
  52. To disable the prepared statement cache, use a value of zero::
  53. engine = create_async_engine("postgresql+asyncpg://user:pass@hostname/dbname?prepared_statement_cache_size=0")
  54. .. versionadded:: 1.4.0b2 Added ``prepared_statement_cache_size`` for asyncpg.
  55. .. warning:: The ``asyncpg`` database driver necessarily uses caches for
  56. PostgreSQL type OIDs, which become stale when custom PostgreSQL datatypes
  57. such as ``ENUM`` objects are changed via DDL operations. Additionally,
  58. prepared statements themselves which are optionally cached by SQLAlchemy's
  59. driver as described above may also become "stale" when DDL has been emitted
  60. to the PostgreSQL database which modifies the tables or other objects
  61. involved in a particular prepared statement.
  62. The SQLAlchemy asyncpg dialect will invalidate these caches within its local
  63. process when statements that represent DDL are emitted on a local
  64. connection, but this is only controllable within a single Python process /
  65. database engine. If DDL changes are made from other database engines
  66. and/or processes, a running application may encounter asyncpg exceptions
  67. ``InvalidCachedStatementError`` and/or ``InternalServerError("cache lookup
  68. failed for type <oid>")`` if it refers to pooled database connections which
  69. operated upon the previous structures. The SQLAlchemy asyncpg dialect will
  70. recover from these error cases when the driver raises these exceptions by
  71. clearing its internal caches as well as those of the asyncpg driver in
  72. response to them, but cannot prevent them from being raised in the first
  73. place if the cached prepared statement or asyncpg type caches have gone
  74. stale, nor can it retry the statement as the PostgreSQL transaction is
  75. invalidated when these errors occur.
  76. """ # noqa
  77. import collections
  78. import decimal
  79. import json as _py_json
  80. import re
  81. import time
  82. from . import json
  83. from .base import _DECIMAL_TYPES
  84. from .base import _FLOAT_TYPES
  85. from .base import _INT_TYPES
  86. from .base import ENUM
  87. from .base import INTERVAL
  88. from .base import OID
  89. from .base import PGCompiler
  90. from .base import PGDialect
  91. from .base import PGExecutionContext
  92. from .base import PGIdentifierPreparer
  93. from .base import REGCLASS
  94. from .base import UUID
  95. from ... import exc
  96. from ... import pool
  97. from ... import processors
  98. from ... import util
  99. from ...sql import sqltypes
  100. from ...util.concurrency import asyncio
  101. from ...util.concurrency import await_fallback
  102. from ...util.concurrency import await_only
  103. try:
  104. from uuid import UUID as _python_UUID # noqa
  105. except ImportError:
  106. _python_UUID = None
  107. class AsyncpgTime(sqltypes.Time):
  108. def get_dbapi_type(self, dbapi):
  109. return dbapi.TIME
  110. class AsyncpgDate(sqltypes.Date):
  111. def get_dbapi_type(self, dbapi):
  112. return dbapi.DATE
  113. class AsyncpgDateTime(sqltypes.DateTime):
  114. def get_dbapi_type(self, dbapi):
  115. if self.timezone:
  116. return dbapi.TIMESTAMP_W_TZ
  117. else:
  118. return dbapi.TIMESTAMP
  119. class AsyncpgBoolean(sqltypes.Boolean):
  120. def get_dbapi_type(self, dbapi):
  121. return dbapi.BOOLEAN
  122. class AsyncPgInterval(INTERVAL):
  123. def get_dbapi_type(self, dbapi):
  124. return dbapi.INTERVAL
  125. @classmethod
  126. def adapt_emulated_to_native(cls, interval, **kw):
  127. return AsyncPgInterval(precision=interval.second_precision)
  128. class AsyncPgEnum(ENUM):
  129. def get_dbapi_type(self, dbapi):
  130. return dbapi.ENUM
  131. class AsyncpgInteger(sqltypes.Integer):
  132. def get_dbapi_type(self, dbapi):
  133. return dbapi.INTEGER
  134. class AsyncpgBigInteger(sqltypes.BigInteger):
  135. def get_dbapi_type(self, dbapi):
  136. return dbapi.BIGINTEGER
  137. class AsyncpgJSON(json.JSON):
  138. def get_dbapi_type(self, dbapi):
  139. return dbapi.JSON
  140. def result_processor(self, dialect, coltype):
  141. return None
  142. class AsyncpgJSONB(json.JSONB):
  143. def get_dbapi_type(self, dbapi):
  144. return dbapi.JSONB
  145. def result_processor(self, dialect, coltype):
  146. return None
  147. class AsyncpgJSONIndexType(sqltypes.JSON.JSONIndexType):
  148. def get_dbapi_type(self, dbapi):
  149. raise NotImplementedError("should not be here")
  150. class AsyncpgJSONIntIndexType(sqltypes.JSON.JSONIntIndexType):
  151. def get_dbapi_type(self, dbapi):
  152. return dbapi.INTEGER
  153. class AsyncpgJSONStrIndexType(sqltypes.JSON.JSONStrIndexType):
  154. def get_dbapi_type(self, dbapi):
  155. return dbapi.STRING
  156. class AsyncpgJSONPathType(json.JSONPathType):
  157. def bind_processor(self, dialect):
  158. def process(value):
  159. assert isinstance(value, util.collections_abc.Sequence)
  160. tokens = [util.text_type(elem) for elem in value]
  161. return tokens
  162. return process
  163. class AsyncpgUUID(UUID):
  164. def get_dbapi_type(self, dbapi):
  165. return dbapi.UUID
  166. def bind_processor(self, dialect):
  167. if not self.as_uuid and dialect.use_native_uuid:
  168. def process(value):
  169. if value is not None:
  170. value = _python_UUID(value)
  171. return value
  172. return process
  173. def result_processor(self, dialect, coltype):
  174. if not self.as_uuid and dialect.use_native_uuid:
  175. def process(value):
  176. if value is not None:
  177. value = str(value)
  178. return value
  179. return process
  180. class AsyncpgNumeric(sqltypes.Numeric):
  181. def bind_processor(self, dialect):
  182. return None
  183. def result_processor(self, dialect, coltype):
  184. if self.asdecimal:
  185. if coltype in _FLOAT_TYPES:
  186. return processors.to_decimal_processor_factory(
  187. decimal.Decimal, self._effective_decimal_return_scale
  188. )
  189. elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
  190. # pg8000 returns Decimal natively for 1700
  191. return None
  192. else:
  193. raise exc.InvalidRequestError(
  194. "Unknown PG numeric type: %d" % coltype
  195. )
  196. else:
  197. if coltype in _FLOAT_TYPES:
  198. # pg8000 returns float natively for 701
  199. return None
  200. elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
  201. return processors.to_float
  202. else:
  203. raise exc.InvalidRequestError(
  204. "Unknown PG numeric type: %d" % coltype
  205. )
  206. class AsyncpgREGCLASS(REGCLASS):
  207. def get_dbapi_type(self, dbapi):
  208. return dbapi.STRING
  209. class AsyncpgOID(OID):
  210. def get_dbapi_type(self, dbapi):
  211. return dbapi.INTEGER
  212. class PGExecutionContext_asyncpg(PGExecutionContext):
  213. def handle_dbapi_exception(self, e):
  214. if isinstance(
  215. e,
  216. (
  217. self.dialect.dbapi.InvalidCachedStatementError,
  218. self.dialect.dbapi.InternalServerError,
  219. ),
  220. ):
  221. self.dialect._invalidate_schema_cache()
  222. def pre_exec(self):
  223. if self.isddl:
  224. self.dialect._invalidate_schema_cache()
  225. self.cursor._invalidate_schema_cache_asof = (
  226. self.dialect._invalidate_schema_cache_asof
  227. )
  228. if not self.compiled:
  229. return
  230. # we have to exclude ENUM because "enum" not really a "type"
  231. # we can cast to, it has to be the name of the type itself.
  232. # for now we just omit it from casting
  233. self.exclude_set_input_sizes = {AsyncAdapt_asyncpg_dbapi.ENUM}
  234. def create_server_side_cursor(self):
  235. return self._dbapi_connection.cursor(server_side=True)
  236. class PGCompiler_asyncpg(PGCompiler):
  237. pass
  238. class PGIdentifierPreparer_asyncpg(PGIdentifierPreparer):
  239. pass
  240. class AsyncAdapt_asyncpg_cursor:
  241. __slots__ = (
  242. "_adapt_connection",
  243. "_connection",
  244. "_rows",
  245. "description",
  246. "arraysize",
  247. "rowcount",
  248. "_inputsizes",
  249. "_cursor",
  250. "_invalidate_schema_cache_asof",
  251. )
  252. server_side = False
  253. def __init__(self, adapt_connection):
  254. self._adapt_connection = adapt_connection
  255. self._connection = adapt_connection._connection
  256. self._rows = []
  257. self._cursor = None
  258. self.description = None
  259. self.arraysize = 1
  260. self.rowcount = -1
  261. self._inputsizes = None
  262. self._invalidate_schema_cache_asof = 0
  263. def close(self):
  264. self._rows[:] = []
  265. def _handle_exception(self, error):
  266. self._adapt_connection._handle_exception(error)
  267. def _parameter_placeholders(self, params):
  268. if not self._inputsizes:
  269. return tuple("$%d" % idx for idx, _ in enumerate(params, 1))
  270. else:
  271. return tuple(
  272. "$%d::%s" % (idx, typ) if typ else "$%d" % idx
  273. for idx, typ in enumerate(
  274. (_pg_types.get(typ) for typ in self._inputsizes), 1
  275. )
  276. )
  277. async def _prepare_and_execute(self, operation, parameters):
  278. adapt_connection = self._adapt_connection
  279. async with adapt_connection._execute_mutex:
  280. if not adapt_connection._started:
  281. await adapt_connection._start_transaction()
  282. if parameters is not None:
  283. operation = operation % self._parameter_placeholders(
  284. parameters
  285. )
  286. else:
  287. parameters = ()
  288. try:
  289. prepared_stmt, attributes = await adapt_connection._prepare(
  290. operation, self._invalidate_schema_cache_asof
  291. )
  292. if attributes:
  293. self.description = [
  294. (
  295. attr.name,
  296. attr.type.oid,
  297. None,
  298. None,
  299. None,
  300. None,
  301. None,
  302. )
  303. for attr in attributes
  304. ]
  305. else:
  306. self.description = None
  307. if self.server_side:
  308. self._cursor = await prepared_stmt.cursor(*parameters)
  309. self.rowcount = -1
  310. else:
  311. self._rows = await prepared_stmt.fetch(*parameters)
  312. status = prepared_stmt.get_statusmsg()
  313. reg = re.match(
  314. r"(?:UPDATE|DELETE|INSERT \d+) (\d+)", status
  315. )
  316. if reg:
  317. self.rowcount = int(reg.group(1))
  318. else:
  319. self.rowcount = -1
  320. except Exception as error:
  321. self._handle_exception(error)
  322. async def _executemany(self, operation, seq_of_parameters):
  323. adapt_connection = self._adapt_connection
  324. async with adapt_connection._execute_mutex:
  325. await adapt_connection._check_type_cache_invalidation(
  326. self._invalidate_schema_cache_asof
  327. )
  328. if not adapt_connection._started:
  329. await adapt_connection._start_transaction()
  330. operation = operation % self._parameter_placeholders(
  331. seq_of_parameters[0]
  332. )
  333. try:
  334. return await self._connection.executemany(
  335. operation, seq_of_parameters
  336. )
  337. except Exception as error:
  338. self._handle_exception(error)
  339. def execute(self, operation, parameters=None):
  340. self._adapt_connection.await_(
  341. self._prepare_and_execute(operation, parameters)
  342. )
  343. def executemany(self, operation, seq_of_parameters):
  344. return self._adapt_connection.await_(
  345. self._executemany(operation, seq_of_parameters)
  346. )
  347. def setinputsizes(self, *inputsizes):
  348. self._inputsizes = inputsizes
  349. def __iter__(self):
  350. while self._rows:
  351. yield self._rows.pop(0)
  352. def fetchone(self):
  353. if self._rows:
  354. return self._rows.pop(0)
  355. else:
  356. return None
  357. def fetchmany(self, size=None):
  358. if size is None:
  359. size = self.arraysize
  360. retval = self._rows[0:size]
  361. self._rows[:] = self._rows[size:]
  362. return retval
  363. def fetchall(self):
  364. retval = self._rows[:]
  365. self._rows[:] = []
  366. return retval
  367. class AsyncAdapt_asyncpg_ss_cursor(AsyncAdapt_asyncpg_cursor):
  368. server_side = True
  369. __slots__ = ("_rowbuffer",)
  370. def __init__(self, adapt_connection):
  371. super(AsyncAdapt_asyncpg_ss_cursor, self).__init__(adapt_connection)
  372. self._rowbuffer = None
  373. def close(self):
  374. self._cursor = None
  375. self._rowbuffer = None
  376. def _buffer_rows(self):
  377. new_rows = self._adapt_connection.await_(self._cursor.fetch(50))
  378. self._rowbuffer = collections.deque(new_rows)
  379. def __aiter__(self):
  380. return self
  381. async def __anext__(self):
  382. if not self._rowbuffer:
  383. self._buffer_rows()
  384. while True:
  385. while self._rowbuffer:
  386. yield self._rowbuffer.popleft()
  387. self._buffer_rows()
  388. if not self._rowbuffer:
  389. break
  390. def fetchone(self):
  391. if not self._rowbuffer:
  392. self._buffer_rows()
  393. if not self._rowbuffer:
  394. return None
  395. return self._rowbuffer.popleft()
  396. def fetchmany(self, size=None):
  397. if size is None:
  398. return self.fetchall()
  399. if not self._rowbuffer:
  400. self._buffer_rows()
  401. buf = list(self._rowbuffer)
  402. lb = len(buf)
  403. if size > lb:
  404. buf.extend(
  405. self._adapt_connection.await_(self._cursor.fetch(size - lb))
  406. )
  407. result = buf[0:size]
  408. self._rowbuffer = collections.deque(buf[size:])
  409. return result
  410. def fetchall(self):
  411. ret = list(self._rowbuffer) + list(
  412. self._adapt_connection.await_(self._all())
  413. )
  414. self._rowbuffer.clear()
  415. return ret
  416. async def _all(self):
  417. rows = []
  418. # TODO: looks like we have to hand-roll some kind of batching here.
  419. # hardcoding for the moment but this should be improved.
  420. while True:
  421. batch = await self._cursor.fetch(1000)
  422. if batch:
  423. rows.extend(batch)
  424. continue
  425. else:
  426. break
  427. return rows
  428. def executemany(self, operation, seq_of_parameters):
  429. raise NotImplementedError(
  430. "server side cursor doesn't support executemany yet"
  431. )
  432. class AsyncAdapt_asyncpg_connection:
  433. __slots__ = (
  434. "dbapi",
  435. "_connection",
  436. "isolation_level",
  437. "_isolation_setting",
  438. "readonly",
  439. "deferrable",
  440. "_transaction",
  441. "_started",
  442. "_prepared_statement_cache",
  443. "_invalidate_schema_cache_asof",
  444. "_execute_mutex",
  445. )
  446. await_ = staticmethod(await_only)
  447. def __init__(self, dbapi, connection, prepared_statement_cache_size=100):
  448. self.dbapi = dbapi
  449. self._connection = connection
  450. self.isolation_level = self._isolation_setting = "read_committed"
  451. self.readonly = False
  452. self.deferrable = False
  453. self._transaction = None
  454. self._started = False
  455. self._invalidate_schema_cache_asof = time.time()
  456. self._execute_mutex = asyncio.Lock()
  457. if prepared_statement_cache_size:
  458. self._prepared_statement_cache = util.LRUCache(
  459. prepared_statement_cache_size
  460. )
  461. else:
  462. self._prepared_statement_cache = None
  463. async def _check_type_cache_invalidation(self, invalidate_timestamp):
  464. if invalidate_timestamp > self._invalidate_schema_cache_asof:
  465. await self._connection.reload_schema_state()
  466. self._invalidate_schema_cache_asof = invalidate_timestamp
  467. async def _prepare(self, operation, invalidate_timestamp):
  468. await self._check_type_cache_invalidation(invalidate_timestamp)
  469. cache = self._prepared_statement_cache
  470. if cache is None:
  471. prepared_stmt = await self._connection.prepare(operation)
  472. attributes = prepared_stmt.get_attributes()
  473. return prepared_stmt, attributes
  474. # asyncpg uses a type cache for the "attributes" which seems to go
  475. # stale independently of the PreparedStatement itself, so place that
  476. # collection in the cache as well.
  477. if operation in cache:
  478. prepared_stmt, attributes, cached_timestamp = cache[operation]
  479. # preparedstatements themselves also go stale for certain DDL
  480. # changes such as size of a VARCHAR changing, so there is also
  481. # a cross-connection invalidation timestamp
  482. if cached_timestamp > invalidate_timestamp:
  483. return prepared_stmt, attributes
  484. prepared_stmt = await self._connection.prepare(operation)
  485. attributes = prepared_stmt.get_attributes()
  486. cache[operation] = (prepared_stmt, attributes, time.time())
  487. return prepared_stmt, attributes
  488. def _handle_exception(self, error):
  489. if self._connection.is_closed():
  490. self._transaction = None
  491. self._started = False
  492. if not isinstance(error, AsyncAdapt_asyncpg_dbapi.Error):
  493. exception_mapping = self.dbapi._asyncpg_error_translate
  494. for super_ in type(error).__mro__:
  495. if super_ in exception_mapping:
  496. translated_error = exception_mapping[super_](
  497. "%s: %s" % (type(error), error)
  498. )
  499. translated_error.pgcode = (
  500. translated_error.sqlstate
  501. ) = getattr(error, "sqlstate", None)
  502. raise translated_error from error
  503. else:
  504. raise error
  505. else:
  506. raise error
  507. @property
  508. def autocommit(self):
  509. return self.isolation_level == "autocommit"
  510. @autocommit.setter
  511. def autocommit(self, value):
  512. if value:
  513. self.isolation_level = "autocommit"
  514. else:
  515. self.isolation_level = self._isolation_setting
  516. def set_isolation_level(self, level):
  517. if self._started:
  518. self.rollback()
  519. self.isolation_level = self._isolation_setting = level
  520. async def _start_transaction(self):
  521. if self.isolation_level == "autocommit":
  522. return
  523. try:
  524. self._transaction = self._connection.transaction(
  525. isolation=self.isolation_level,
  526. readonly=self.readonly,
  527. deferrable=self.deferrable,
  528. )
  529. await self._transaction.start()
  530. except Exception as error:
  531. self._handle_exception(error)
  532. else:
  533. self._started = True
  534. def cursor(self, server_side=False):
  535. if server_side:
  536. return AsyncAdapt_asyncpg_ss_cursor(self)
  537. else:
  538. return AsyncAdapt_asyncpg_cursor(self)
  539. def rollback(self):
  540. if self._started:
  541. try:
  542. self.await_(self._transaction.rollback())
  543. except Exception as error:
  544. self._handle_exception(error)
  545. finally:
  546. self._transaction = None
  547. self._started = False
  548. def commit(self):
  549. if self._started:
  550. try:
  551. self.await_(self._transaction.commit())
  552. except Exception as error:
  553. self._handle_exception(error)
  554. finally:
  555. self._transaction = None
  556. self._started = False
  557. def close(self):
  558. self.rollback()
  559. self.await_(self._connection.close())
  560. class AsyncAdaptFallback_asyncpg_connection(AsyncAdapt_asyncpg_connection):
  561. __slots__ = ()
  562. await_ = staticmethod(await_fallback)
  563. class AsyncAdapt_asyncpg_dbapi:
  564. def __init__(self, asyncpg):
  565. self.asyncpg = asyncpg
  566. self.paramstyle = "format"
  567. def connect(self, *arg, **kw):
  568. async_fallback = kw.pop("async_fallback", False)
  569. prepared_statement_cache_size = kw.pop(
  570. "prepared_statement_cache_size", 100
  571. )
  572. if util.asbool(async_fallback):
  573. return AsyncAdaptFallback_asyncpg_connection(
  574. self,
  575. await_fallback(self.asyncpg.connect(*arg, **kw)),
  576. prepared_statement_cache_size=prepared_statement_cache_size,
  577. )
  578. else:
  579. return AsyncAdapt_asyncpg_connection(
  580. self,
  581. await_only(self.asyncpg.connect(*arg, **kw)),
  582. prepared_statement_cache_size=prepared_statement_cache_size,
  583. )
  584. class Error(Exception):
  585. pass
  586. class Warning(Exception): # noqa
  587. pass
  588. class InterfaceError(Error):
  589. pass
  590. class DatabaseError(Error):
  591. pass
  592. class InternalError(DatabaseError):
  593. pass
  594. class OperationalError(DatabaseError):
  595. pass
  596. class ProgrammingError(DatabaseError):
  597. pass
  598. class IntegrityError(DatabaseError):
  599. pass
  600. class DataError(DatabaseError):
  601. pass
  602. class NotSupportedError(DatabaseError):
  603. pass
  604. class InternalServerError(InternalError):
  605. pass
  606. class InvalidCachedStatementError(NotSupportedError):
  607. def __init__(self, message):
  608. super(
  609. AsyncAdapt_asyncpg_dbapi.InvalidCachedStatementError, self
  610. ).__init__(
  611. message + " (SQLAlchemy asyncpg dialect will now invalidate "
  612. "all prepared caches in response to this exception)",
  613. )
  614. @util.memoized_property
  615. def _asyncpg_error_translate(self):
  616. import asyncpg
  617. return {
  618. asyncpg.exceptions.IntegrityConstraintViolationError: self.IntegrityError, # noqa: E501
  619. asyncpg.exceptions.PostgresError: self.Error,
  620. asyncpg.exceptions.SyntaxOrAccessError: self.ProgrammingError,
  621. asyncpg.exceptions.InterfaceError: self.InterfaceError,
  622. asyncpg.exceptions.InvalidCachedStatementError: self.InvalidCachedStatementError, # noqa: E501
  623. asyncpg.exceptions.InternalServerError: self.InternalServerError,
  624. }
  625. def Binary(self, value):
  626. return value
  627. STRING = util.symbol("STRING")
  628. TIMESTAMP = util.symbol("TIMESTAMP")
  629. TIMESTAMP_W_TZ = util.symbol("TIMESTAMP_W_TZ")
  630. TIME = util.symbol("TIME")
  631. DATE = util.symbol("DATE")
  632. INTERVAL = util.symbol("INTERVAL")
  633. NUMBER = util.symbol("NUMBER")
  634. FLOAT = util.symbol("FLOAT")
  635. BOOLEAN = util.symbol("BOOLEAN")
  636. INTEGER = util.symbol("INTEGER")
  637. BIGINTEGER = util.symbol("BIGINTEGER")
  638. BYTES = util.symbol("BYTES")
  639. DECIMAL = util.symbol("DECIMAL")
  640. JSON = util.symbol("JSON")
  641. JSONB = util.symbol("JSONB")
  642. ENUM = util.symbol("ENUM")
  643. UUID = util.symbol("UUID")
  644. BYTEA = util.symbol("BYTEA")
  645. DATETIME = TIMESTAMP
  646. BINARY = BYTEA
  647. _pg_types = {
  648. AsyncAdapt_asyncpg_dbapi.STRING: "varchar",
  649. AsyncAdapt_asyncpg_dbapi.TIMESTAMP: "timestamp",
  650. AsyncAdapt_asyncpg_dbapi.TIMESTAMP_W_TZ: "timestamp with time zone",
  651. AsyncAdapt_asyncpg_dbapi.DATE: "date",
  652. AsyncAdapt_asyncpg_dbapi.TIME: "time",
  653. AsyncAdapt_asyncpg_dbapi.INTERVAL: "interval",
  654. AsyncAdapt_asyncpg_dbapi.NUMBER: "numeric",
  655. AsyncAdapt_asyncpg_dbapi.FLOAT: "float",
  656. AsyncAdapt_asyncpg_dbapi.BOOLEAN: "bool",
  657. AsyncAdapt_asyncpg_dbapi.INTEGER: "integer",
  658. AsyncAdapt_asyncpg_dbapi.BIGINTEGER: "bigint",
  659. AsyncAdapt_asyncpg_dbapi.BYTES: "bytes",
  660. AsyncAdapt_asyncpg_dbapi.DECIMAL: "decimal",
  661. AsyncAdapt_asyncpg_dbapi.JSON: "json",
  662. AsyncAdapt_asyncpg_dbapi.JSONB: "jsonb",
  663. AsyncAdapt_asyncpg_dbapi.ENUM: "enum",
  664. AsyncAdapt_asyncpg_dbapi.UUID: "uuid",
  665. AsyncAdapt_asyncpg_dbapi.BYTEA: "bytea",
  666. }
  667. class PGDialect_asyncpg(PGDialect):
  668. driver = "asyncpg"
  669. supports_statement_cache = True
  670. supports_unicode_statements = True
  671. supports_server_side_cursors = True
  672. supports_unicode_binds = True
  673. default_paramstyle = "format"
  674. supports_sane_multi_rowcount = False
  675. execution_ctx_cls = PGExecutionContext_asyncpg
  676. statement_compiler = PGCompiler_asyncpg
  677. preparer = PGIdentifierPreparer_asyncpg
  678. use_setinputsizes = True
  679. use_native_uuid = True
  680. colspecs = util.update_copy(
  681. PGDialect.colspecs,
  682. {
  683. sqltypes.Time: AsyncpgTime,
  684. sqltypes.Date: AsyncpgDate,
  685. sqltypes.DateTime: AsyncpgDateTime,
  686. sqltypes.Interval: AsyncPgInterval,
  687. INTERVAL: AsyncPgInterval,
  688. UUID: AsyncpgUUID,
  689. sqltypes.Boolean: AsyncpgBoolean,
  690. sqltypes.Integer: AsyncpgInteger,
  691. sqltypes.BigInteger: AsyncpgBigInteger,
  692. sqltypes.Numeric: AsyncpgNumeric,
  693. sqltypes.JSON: AsyncpgJSON,
  694. json.JSONB: AsyncpgJSONB,
  695. sqltypes.JSON.JSONPathType: AsyncpgJSONPathType,
  696. sqltypes.JSON.JSONIndexType: AsyncpgJSONIndexType,
  697. sqltypes.JSON.JSONIntIndexType: AsyncpgJSONIntIndexType,
  698. sqltypes.JSON.JSONStrIndexType: AsyncpgJSONStrIndexType,
  699. sqltypes.Enum: AsyncPgEnum,
  700. OID: AsyncpgOID,
  701. REGCLASS: AsyncpgREGCLASS,
  702. },
  703. )
  704. is_async = True
  705. _invalidate_schema_cache_asof = 0
  706. def _invalidate_schema_cache(self):
  707. self._invalidate_schema_cache_asof = time.time()
  708. @util.memoized_property
  709. def _dbapi_version(self):
  710. if self.dbapi and hasattr(self.dbapi, "__version__"):
  711. return tuple(
  712. [
  713. int(x)
  714. for x in re.findall(
  715. r"(\d+)(?:[-\.]?|$)", self.dbapi.__version__
  716. )
  717. ]
  718. )
  719. else:
  720. return (99, 99, 99)
  721. @classmethod
  722. def dbapi(cls):
  723. return AsyncAdapt_asyncpg_dbapi(__import__("asyncpg"))
  724. @util.memoized_property
  725. def _isolation_lookup(self):
  726. return {
  727. "AUTOCOMMIT": "autocommit",
  728. "READ COMMITTED": "read_committed",
  729. "REPEATABLE READ": "repeatable_read",
  730. "SERIALIZABLE": "serializable",
  731. }
  732. def set_isolation_level(self, connection, level):
  733. try:
  734. level = self._isolation_lookup[level.replace("_", " ")]
  735. except KeyError as err:
  736. util.raise_(
  737. exc.ArgumentError(
  738. "Invalid value '%s' for isolation_level. "
  739. "Valid isolation levels for %s are %s"
  740. % (level, self.name, ", ".join(self._isolation_lookup))
  741. ),
  742. replace_context=err,
  743. )
  744. connection.set_isolation_level(level)
  745. def set_readonly(self, connection, value):
  746. connection.readonly = value
  747. def get_readonly(self, connection):
  748. return connection.readonly
  749. def set_deferrable(self, connection, value):
  750. connection.deferrable = value
  751. def get_deferrable(self, connection):
  752. return connection.deferrable
  753. def create_connect_args(self, url):
  754. opts = url.translate_connect_args(username="user")
  755. opts.update(url.query)
  756. util.coerce_kw_type(opts, "prepared_statement_cache_size", int)
  757. util.coerce_kw_type(opts, "port", int)
  758. return ([], opts)
  759. @classmethod
  760. def get_pool_class(cls, url):
  761. async_fallback = url.query.get("async_fallback", False)
  762. if util.asbool(async_fallback):
  763. return pool.FallbackAsyncAdaptedQueuePool
  764. else:
  765. return pool.AsyncAdaptedQueuePool
  766. def is_disconnect(self, e, connection, cursor):
  767. if connection:
  768. return connection._connection.is_closed()
  769. else:
  770. return isinstance(
  771. e, self.dbapi.InterfaceError
  772. ) and "connection is closed" in str(e)
  773. def do_set_input_sizes(self, cursor, list_of_tuples, context):
  774. if self.positional:
  775. cursor.setinputsizes(
  776. *[dbtype for key, dbtype, sqltype in list_of_tuples]
  777. )
  778. else:
  779. cursor.setinputsizes(
  780. **{
  781. key: dbtype
  782. for key, dbtype, sqltype in list_of_tuples
  783. if dbtype
  784. }
  785. )
  786. def on_connect(self):
  787. super_connect = super(PGDialect_asyncpg, self).on_connect()
  788. def _jsonb_encoder(str_value):
  789. # \x01 is the prefix for jsonb used by PostgreSQL.
  790. # asyncpg requires it when format='binary'
  791. return b"\x01" + str_value.encode()
  792. deserializer = self._json_deserializer or _py_json.loads
  793. def _json_decoder(bin_value):
  794. return deserializer(bin_value.decode())
  795. def _jsonb_decoder(bin_value):
  796. # the byte is the \x01 prefix for jsonb used by PostgreSQL.
  797. # asyncpg returns it when format='binary'
  798. return deserializer(bin_value[1:].decode())
  799. async def _setup_type_codecs(conn):
  800. """set up type decoders at the asyncpg level.
  801. these are set_type_codec() calls to normalize
  802. There was a tentative decoder for the "char" datatype here
  803. to have it return strings however this type is actually a binary
  804. type that other drivers are likely mis-interpreting.
  805. See https://github.com/MagicStack/asyncpg/issues/623 for reference
  806. on why it's set up this way.
  807. """
  808. await conn._connection.set_type_codec(
  809. "json",
  810. encoder=str.encode,
  811. decoder=_json_decoder,
  812. schema="pg_catalog",
  813. format="binary",
  814. )
  815. await conn._connection.set_type_codec(
  816. "jsonb",
  817. encoder=_jsonb_encoder,
  818. decoder=_jsonb_decoder,
  819. schema="pg_catalog",
  820. format="binary",
  821. )
  822. def connect(conn):
  823. conn.await_(_setup_type_codecs(conn))
  824. if super_connect is not None:
  825. super_connect(conn)
  826. return connect
  827. dialect = PGDialect_asyncpg