Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy/dialects/postgresql/asyncpg.py: 46%

1# dialects/postgresql/asyncpg.py

3# file>

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

9r"""

10.. dialect:: postgresql+asyncpg

11 :name: asyncpg

12 :dbapi: asyncpg

13 :connectstring: postgresql+asyncpg://user:password@host:port/dbname[?key=value&key=value...]

14 :url: https://magicstack.github.io/asyncpg/

16The asyncpg dialect is SQLAlchemy's first Python asyncio dialect.

18Using a special asyncio mediation layer, the asyncpg dialect is usable

19as the backend for the :ref:`SQLAlchemy asyncio <asyncio_toplevel>`

20extension package.

22This dialect should normally be used only with the

23:func:`_asyncio.create_async_engine` engine creation function::

25 from sqlalchemy.ext.asyncio import create_async_engine

27 engine = create_async_engine(

28 "postgresql+asyncpg://user:pass@hostname/dbname"

29 )

31.. versionadded:: 1.4

33.. note::

35 By default asyncpg does not decode the ``json`` and ``jsonb`` types and

36 returns them as strings. SQLAlchemy sets default type decoder for ``json``

37 and ``jsonb`` types using the python builtin ``json.loads`` function.

38 The json implementation used can be changed by setting the attribute

39 ``json_deserializer`` when creating the engine with

40 :func:`create_engine` or :func:`create_async_engine`.

42.. _asyncpg_multihost:

44Multihost Connections

45--------------------------

47The asyncpg dialect features support for multiple fallback hosts in the

48same way as that of the psycopg2 and psycopg dialects. The

49syntax is the same,

50using ``host=<host>:<port>`` combinations as additional query string arguments;

51however, there is no default port, so all hosts must have a complete port number

52present, otherwise an exception is raised::

54 engine = create_async_engine(

55 "postgresql+asyncpg://user:password@/dbname?host=HostA:5432&host=HostB:5432&host=HostC:5432"

56 )

58For complete background on this syntax, see :ref:`psycopg2_multi_host`.

60.. versionadded:: 2.0.18

62.. seealso::

64 :ref:`psycopg2_multi_host`

66.. _asyncpg_prepared_statement_cache:

68Prepared Statement Cache

69--------------------------

71The asyncpg SQLAlchemy dialect makes use of ``asyncpg.connection.prepare()``

72for all statements. The prepared statement objects are cached after

73construction which appears to grant a 10% or more performance improvement for

74statement invocation. The cache is on a per-DBAPI connection basis, which

75means that the primary storage for prepared statements is within DBAPI

76connections pooled within the connection pool. The size of this cache

77defaults to 100 statements per DBAPI connection and may be adjusted using the

78``prepared_statement_cache_size`` DBAPI argument (note that while this argument

79is implemented by SQLAlchemy, it is part of the DBAPI emulation portion of the

80asyncpg dialect, therefore is handled as a DBAPI argument, not a dialect

81argument)::

84 engine = create_async_engine(

85 "postgresql+asyncpg://user:pass@hostname/dbname?prepared_statement_cache_size=500"

86 )

88To disable the prepared statement cache, use a value of zero::

90 engine = create_async_engine(

91 "postgresql+asyncpg://user:pass@hostname/dbname?prepared_statement_cache_size=0"

92 )

94.. versionadded:: 1.4.0b2 Added ``prepared_statement_cache_size`` for asyncpg.

97.. warning:: The ``asyncpg`` database driver necessarily uses caches for

98 PostgreSQL type OIDs, which become stale when custom PostgreSQL datatypes

99 such as ``ENUM`` objects are changed via DDL operations. Additionally,

100 prepared statements themselves which are optionally cached by SQLAlchemy's

101 driver as described above may also become "stale" when DDL has been emitted

102 to the PostgreSQL database which modifies the tables or other objects

103 involved in a particular prepared statement.

104

105 The SQLAlchemy asyncpg dialect will invalidate these caches within its local

106 process when statements that represent DDL are emitted on a local

107 connection, but this is only controllable within a single Python process /

108 database engine. If DDL changes are made from other database engines

109 and/or processes, a running application may encounter asyncpg exceptions

110 ``InvalidCachedStatementError`` and/or ``InternalServerError("cache lookup

111 failed for type <oid>")`` if it refers to pooled database connections which

112 operated upon the previous structures. The SQLAlchemy asyncpg dialect will

113 recover from these error cases when the driver raises these exceptions by

114 clearing its internal caches as well as those of the asyncpg driver in

115 response to them, but cannot prevent them from being raised in the first

116 place if the cached prepared statement or asyncpg type caches have gone

117 stale, nor can it retry the statement as the PostgreSQL transaction is

118 invalidated when these errors occur.

119

120.. _asyncpg_prepared_statement_name:

121

122Prepared Statement Name with PGBouncer

123--------------------------------------

124

125By default, asyncpg enumerates prepared statements in numeric order, which

126can lead to errors if a name has already been taken for another prepared

127statement. This issue can arise if your application uses database proxies

128such as PgBouncer to handle connections. One possible workaround is to

129use dynamic prepared statement names, which asyncpg now supports through

130an optional ``name`` value for the statement name. This allows you to

131generate your own unique names that won't conflict with existing ones.

132To achieve this, you can provide a function that will be called every time

133a prepared statement is prepared::

134

135 from uuid import uuid4

136

137 engine = create_async_engine(

138 "postgresql+asyncpg://user:pass@somepgbouncer/dbname",

139 poolclass=NullPool,

140 connect_args={

141 "prepared_statement_name_func": lambda: f"__asyncpg_{uuid4()}__",

142 },

143 )

144

145.. seealso::

146

147 https://github.com/MagicStack/asyncpg/issues/837

148

149 https://github.com/sqlalchemy/sqlalchemy/issues/6467

150

151.. warning:: When using PGBouncer, to prevent a buildup of useless prepared statements in

152 your application, it's important to use the :class:`.NullPool` pool

153 class, and to configure PgBouncer to use `DISCARD <https://www.postgresql.org/docs/current/sql-discard.html>`_

154 when returning connections. The DISCARD command is used to release resources held by the db connection,

155 including prepared statements. Without proper setup, prepared statements can

156 accumulate quickly and cause performance issues.

157

158Disabling the PostgreSQL JIT to improve ENUM datatype handling

159---------------------------------------------------------------

160

161Asyncpg has an `issue <https://github.com/MagicStack/asyncpg/issues/727>`_ when

162using PostgreSQL ENUM datatypes, where upon the creation of new database

163connections, an expensive query may be emitted in order to retrieve metadata

164regarding custom types which has been shown to negatively affect performance.

165To mitigate this issue, the PostgreSQL "jit" setting may be disabled from the

166client using this setting passed to :func:`_asyncio.create_async_engine`::

167

168 engine = create_async_engine(

169 "postgresql+asyncpg://user:password@localhost/tmp",

170 connect_args={"server_settings": {"jit": "off"}},

171 )

172

173.. seealso::

174

175 https://github.com/MagicStack/asyncpg/issues/727

176

177""" # noqa

178

179from __future__ import annotations

180

181from collections import deque

182import decimal

183import json as _py_json

184import re

185import time

186from types import NoneType

187from typing import Any

188from typing import Awaitable

189from typing import Callable

190from typing import NoReturn

191from typing import Optional

192from typing import Protocol

193from typing import Sequence

194from typing import Tuple

195from typing import TYPE_CHECKING

196

197from . import json

198from . import ranges

199from .array import ARRAY as PGARRAY

200from .base import _DECIMAL_TYPES

201from .base import _FLOAT_TYPES

202from .base import _INT_TYPES

203from .base import ENUM

204from .base import INTERVAL

205from .base import OID

206from .base import PGCompiler

207from .base import PGDialect

208from .base import PGExecutionContext

209from .base import PGIdentifierPreparer

210from .base import REGCLASS

211from .base import REGCONFIG

212from .bitstring import BitString

213from .types import BIT

214from .types import BYTEA

215from .types import CITEXT

216from ... import exc

217from ... import util

218from ...connectors.asyncio import AsyncAdapt_dbapi_connection

219from ...connectors.asyncio import AsyncAdapt_dbapi_cursor

220from ...connectors.asyncio import AsyncAdapt_dbapi_module

221from ...connectors.asyncio import AsyncAdapt_dbapi_ss_cursor

222from ...connectors.asyncio import AsyncAdapt_Error

223from ...connectors.asyncio import AsyncAdapt_terminate

224from ...engine import processors

225from ...sql import sqltypes

226from ...util.concurrency import await_

227

228if TYPE_CHECKING:

229 from ...engine.interfaces import _DBAPICursorDescription

230

231

232class AsyncpgARRAY(PGARRAY):

233 render_bind_cast = True

234

235

236class AsyncpgString(sqltypes.String):

237 render_bind_cast = True

238

239

240class AsyncpgREGCONFIG(REGCONFIG):

241 render_bind_cast = True

242

243

244class AsyncpgTime(sqltypes.Time):

245 render_bind_cast = True

246

247

248class AsyncpgBit(BIT):

249 render_bind_cast = True

250

251 def bind_processor(self, dialect):

252 asyncpg_BitString = dialect.dbapi.asyncpg.BitString

253

254 def to_bind(value):

255 if isinstance(value, str):

256 value = BitString(value)

257 value = asyncpg_BitString.from_int(int(value), len(value))

258 return value

259

260 return to_bind

261

262 def result_processor(self, dialect, coltype):

263 def to_result(value):

264 if value is not None:

265 value = BitString.from_int(value.to_int(), length=len(value))

266 return value

267

268 return to_result

269

270

271class AsyncpgByteA(BYTEA):

272 render_bind_cast = True

273

274

275class AsyncpgDate(sqltypes.Date):

276 render_bind_cast = True

277

278

279class AsyncpgDateTime(sqltypes.DateTime):

280 render_bind_cast = True

281

282

283class AsyncpgBoolean(sqltypes.Boolean):

284 render_bind_cast = True

285

286

287class AsyncPgInterval(INTERVAL):

288 render_bind_cast = True

289

290 @classmethod

291 def adapt_emulated_to_native(cls, interval, **kw):

292 return AsyncPgInterval(precision=interval.second_precision)

293

294

295class AsyncPgEnum(ENUM):

296 render_bind_cast = True

297

298

299class AsyncpgInteger(sqltypes.Integer):

300 render_bind_cast = True

301

302

303class AsyncpgSmallInteger(sqltypes.SmallInteger):

304 render_bind_cast = True

305

306

307class AsyncpgBigInteger(sqltypes.BigInteger):

308 render_bind_cast = True

309

310

311class AsyncpgJSONIndexType(sqltypes.JSON.JSONIndexType):

312 pass

313

314

315class AsyncpgJSONIntIndexType(sqltypes.JSON.JSONIntIndexType):

316 __visit_name__ = "json_int_index"

317

318 render_bind_cast = True

319

320

321class AsyncpgJSONStrIndexType(sqltypes.JSON.JSONStrIndexType):

322 __visit_name__ = "json_str_index"

323

324 render_bind_cast = True

325

326

327class AsyncpgJSONPathType(json.JSONPathType):

328 def bind_processor(self, dialect):

329 def process(value):

330 if isinstance(value, str):

331 # If it's already a string assume that it's in json path

332 # format. This allows using cast with json paths literals

333 return value

334 elif value:

335 tokens = [str(elem) for elem in value]

336 return tokens

337 else:

338 return []

339

340 return process

341

342

343class _AsyncpgNumericCommon(sqltypes.NumericCommon):

344 render_bind_cast = True

345

346 def bind_processor(self, dialect):

347 return None

348

349 def result_processor(self, dialect, coltype):

350 if self.asdecimal:

351 if coltype in _FLOAT_TYPES:

352 return processors.to_decimal_processor_factory(

353 decimal.Decimal, self._effective_decimal_return_scale

354 )

355 elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:

356 # pg8000 returns Decimal natively for 1700

357 return None

358 else:

359 raise exc.InvalidRequestError(

360 "Unknown PG numeric type: %d" % coltype

361 )

362 else:

363 if coltype in _FLOAT_TYPES:

364 # pg8000 returns float natively for 701

365 return None

366 elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:

367 return processors.to_float

368 else:

369 raise exc.InvalidRequestError(

370 "Unknown PG numeric type: %d" % coltype

371 )

372

373

374class AsyncpgNumeric(_AsyncpgNumericCommon, sqltypes.Numeric):

375 pass

376

377

378class AsyncpgFloat(_AsyncpgNumericCommon, sqltypes.Float):

379 pass

380

381

382class AsyncpgREGCLASS(REGCLASS):

383 render_bind_cast = True

384

385

386class AsyncpgOID(OID):

387 render_bind_cast = True

388

389

390class AsyncpgCHAR(sqltypes.CHAR):

391 render_bind_cast = True

392

393

394class _AsyncpgRange(ranges.AbstractSingleRangeImpl):

395 def bind_processor(self, dialect):

396 asyncpg_Range = dialect.dbapi.asyncpg.Range

397

398 def to_range(value):

399 if isinstance(value, ranges.Range):

400 value = asyncpg_Range(

401 value.lower,

402 value.upper,

403 lower_inc=value.bounds[0] == "[",

404 upper_inc=value.bounds[1] == "]",

405 empty=value.empty,

406 )

407 return value

408

409 return to_range

410

411 def result_processor(self, dialect, coltype):

412 def to_range(value):

413 if value is not None:

414 empty = value.isempty

415 value = ranges.Range(

416 value.lower,

417 value.upper,

418 bounds=f"{'[' if empty or value.lower_inc else '('}" # type: ignore # noqa: E501

419 f"{']' if not empty and value.upper_inc else ')'}",

420 empty=empty,

421 )

422 return value

423

424 return to_range

425

426

427class _AsyncpgMultiRange(ranges.AbstractMultiRangeImpl):

428 def bind_processor(self, dialect):

429 asyncpg_Range = dialect.dbapi.asyncpg.Range

430

431 def to_range(value):

432 if isinstance(value, (str, NoneType)):

433 return value

434

435 def to_range(value):

436 if isinstance(value, ranges.Range):

437 value = asyncpg_Range(

438 value.lower,

439 value.upper,

440 lower_inc=value.bounds[0] == "[",

441 upper_inc=value.bounds[1] == "]",

442 empty=value.empty,

443 )

444 return value

445

446 return [to_range(element) for element in value]

447

448 return to_range

449

450 def result_processor(self, dialect, coltype):

451 def to_range_array(value):

452 def to_range(rvalue):

453 if rvalue is not None:

454 empty = rvalue.isempty

455 rvalue = ranges.Range(

456 rvalue.lower,

457 rvalue.upper,

458 bounds=f"{'[' if empty or rvalue.lower_inc else '('}" # type: ignore # noqa: E501

459 f"{']' if not empty and rvalue.upper_inc else ')'}",

460 empty=empty,

461 )

462 return rvalue

463

464 if value is not None:

465 value = ranges.MultiRange(to_range(elem) for elem in value)

466

467 return value

468

469 return to_range_array

470

471

472class PGExecutionContext_asyncpg(PGExecutionContext):

473 def handle_dbapi_exception(self, e):

474 if isinstance(

475 e,

476 (

477 self.dialect.dbapi.InvalidCachedStatementError,

478 self.dialect.dbapi.InternalServerError,

479 ),

480 ):

481 self.dialect._invalidate_schema_cache()

482

483 def pre_exec(self):

484 if self.isddl:

485 self.dialect._invalidate_schema_cache()

486

487 self.cursor._invalidate_schema_cache_asof = (

488 self.dialect._invalidate_schema_cache_asof

489 )

490

491 if not self.compiled:

492 return

493

494 def create_server_side_cursor(self):

495 return self._dbapi_connection.cursor(server_side=True)

496

497

498class PGCompiler_asyncpg(PGCompiler):

499 pass

500

501

502class PGIdentifierPreparer_asyncpg(PGIdentifierPreparer):

503 pass

504

505

506class _AsyncpgTransaction(Protocol):

507 async def start(self) -> None: ...

508 async def commit(self) -> None: ...

509 async def rollback(self) -> None: ...

510

511

512class _AsyncpgConnection(Protocol):

513 async def executemany(

514 self, operation: Any, seq_of_parameters: Sequence[Tuple[Any, ...]]

515 ) -> Any: ...

516

517 async def reload_schema_state(self) -> None: ...

518

519 async def prepare(

520 self, operation: Any, *, name: Optional[str] = None

521 ) -> Any: ...

522

523 def is_closed(self) -> bool: ...

524

525 def transaction(

526 self,

527 *,

528 isolation: Optional[str] = None,

529 readonly: bool = False,

530 deferrable: bool = False,

531 ) -> _AsyncpgTransaction: ...

532

533 def fetchrow(self, operation: str) -> Any: ...

534

535 async def close(self, timeout: int = ...) -> None: ...

536

537 def terminate(self) -> None: ...

538

539

540class _AsyncpgCursor(Protocol):

541 def fetch(self, size: int) -> Any: ...

542

543

544class AsyncAdapt_asyncpg_cursor(AsyncAdapt_dbapi_cursor):

545 __slots__ = (

546 "_description",

547 "_arraysize",

548 "_rowcount",

549 "_invalidate_schema_cache_asof",

550 )

551

552 _adapt_connection: AsyncAdapt_asyncpg_connection

553 _connection: _AsyncpgConnection

554 _cursor: Optional[_AsyncpgCursor]

555 _awaitable_cursor_close: bool = False

556

557 def __init__(self, adapt_connection: AsyncAdapt_asyncpg_connection):

558 self._adapt_connection = adapt_connection

559 self._connection = adapt_connection._connection

560 self._cursor = None

561 self._rows = deque()

562 self._description = None

563 self._arraysize = 1

564 self._rowcount = -1

565 self._invalidate_schema_cache_asof = 0

566

567 def _handle_exception(self, error):

568 self._adapt_connection._handle_exception(error)

569

570 async def _prepare_and_execute(self, operation, parameters):

571 adapt_connection = self._adapt_connection

572

573 async with adapt_connection._execute_mutex:

574 if adapt_connection._transaction is None:

575 await adapt_connection._start_transaction()

576

577 if parameters is None:

578 parameters = ()

579

580 try:

581 prepared_stmt, attributes = await adapt_connection._prepare(

582 operation, self._invalidate_schema_cache_asof

583 )

584

585 if attributes:

586 self._description = [

587 (

588 attr.name,

589 attr.type.oid,

590 None,

591 None,

592 None,

593 None,

594 None,

595 )

596 for attr in attributes

597 ]

598 else:

599 self._description = None

600

601 if self.server_side:

602 self._cursor = await prepared_stmt.cursor(*parameters)

603 self._rowcount = -1

604 else:

605 self._rows = deque(await prepared_stmt.fetch(*parameters))

606 status = prepared_stmt.get_statusmsg()

607

608 reg = re.match(

609 r"(?:SELECT|UPDATE|DELETE|INSERT \d+) (\d+)",

610 status or "",

611 )

612 if reg:

613 self._rowcount = int(reg.group(1))

614 else:

615 self._rowcount = -1

616

617 except Exception as error:

618 self._handle_exception(error)

619

620 @property

621 def description(self) -> Optional[_DBAPICursorDescription]:

622 return self._description

623

624 @property

625 def rowcount(self) -> int:

626 return self._rowcount

627

628 @property

629 def arraysize(self) -> int:

630 return self._arraysize

631

632 @arraysize.setter

633 def arraysize(self, value: int) -> None:

634 self._arraysize = value

635

636 async def _executemany(self, operation, seq_of_parameters):

637 adapt_connection = self._adapt_connection

638

639 self._description = None

640 async with adapt_connection._execute_mutex:

641 await adapt_connection._check_type_cache_invalidation(

642 self._invalidate_schema_cache_asof

643 )

644

645 if adapt_connection._transaction is None:

646 await adapt_connection._start_transaction()

647

648 try:

649 return await self._connection.executemany(

650 operation, seq_of_parameters

651 )

652 except Exception as error:

653 self._handle_exception(error)

654

655 def execute(self, operation, parameters=None):

656 await_(self._prepare_and_execute(operation, parameters))

657

658 def executemany(self, operation, seq_of_parameters):

659 return await_(self._executemany(operation, seq_of_parameters))

660

661 def setinputsizes(self, *inputsizes):

662 raise NotImplementedError()

663

664

665class AsyncAdapt_asyncpg_ss_cursor(

666 AsyncAdapt_dbapi_ss_cursor, AsyncAdapt_asyncpg_cursor

667):

668 __slots__ = ("_rowbuffer",)

669

670 def __init__(self, adapt_connection):

671 super().__init__(adapt_connection)

672 self._rowbuffer = deque()

673

674 def close(self):

675 self._cursor = None

676 self._rowbuffer.clear()

677

678 def _buffer_rows(self):

679 assert self._cursor is not None

680 new_rows = await_(self._cursor.fetch(50))

681 self._rowbuffer.extend(new_rows)

682

683 def __aiter__(self):

684 return self

685

686 async def __anext__(self):

687 while True:

688 while self._rowbuffer:

689 yield self._rowbuffer.popleft()

690

691 self._buffer_rows()

692 if not self._rowbuffer:

693 break

694

695 def fetchone(self):

696 if not self._rowbuffer:

697 self._buffer_rows()

698 if not self._rowbuffer:

699 return None

700 return self._rowbuffer.popleft()

701

702 def fetchmany(self, size=None):

703 if size is None:

704 return self.fetchall()

705

706 if not self._rowbuffer:

707 self._buffer_rows()

708

709 assert self._cursor is not None

710 rb = self._rowbuffer

711 lb = len(rb)

712 if size > lb:

713 rb.extend(await_(self._cursor.fetch(size - lb)))

714

715 return [rb.popleft() for _ in range(min(size, len(rb)))]

716

717 def fetchall(self):

718 ret = list(self._rowbuffer)

719 ret.extend(await_(self._all()))

720 self._rowbuffer.clear()

721 return ret

722

723 async def _all(self):

724 rows = []

725

726 assert self._cursor is not None

727

728 # TODO: looks like we have to hand-roll some kind of batching here.

729 # hardcoding for the moment but this should be improved.

730 while True:

731 batch = await self._cursor.fetch(1000)

732 if batch:

733 rows.extend(batch)

734 continue

735 else:

736 break

737 return rows

738

739 def executemany(self, operation, seq_of_parameters):

740 raise NotImplementedError(

741 "server side cursor doesn't support executemany yet"

742 )

743

744

745class AsyncAdapt_asyncpg_connection(

746 AsyncAdapt_terminate, AsyncAdapt_dbapi_connection

747):

748 _cursor_cls = AsyncAdapt_asyncpg_cursor

749 _ss_cursor_cls = AsyncAdapt_asyncpg_ss_cursor

750

751 _connection: _AsyncpgConnection

752 _transaction: Optional[_AsyncpgTransaction]

753

754 __slots__ = (

755 "isolation_level",

756 "_isolation_setting",

757 "readonly",

758 "deferrable",

759 "_transaction",

760 "_prepared_statement_cache",

761 "_prepared_statement_name_func",

762 "_invalidate_schema_cache_asof",

763 )

764

765 def __init__(

766 self,

767 dbapi,

768 connection,

769 prepared_statement_cache_size=100,

770 prepared_statement_name_func=None,

771 ):

772 super().__init__(dbapi, connection)

773 self.isolation_level = self._isolation_setting = None

774 self.readonly = False

775 self.deferrable = False

776 self._transaction = None

777 self._invalidate_schema_cache_asof = time.time()

778

779 if prepared_statement_cache_size:

780 self._prepared_statement_cache = util.LRUCache(

781 prepared_statement_cache_size

782 )

783 else:

784 self._prepared_statement_cache = None

785

786 if prepared_statement_name_func:

787 self._prepared_statement_name_func = prepared_statement_name_func

788 else:

789 self._prepared_statement_name_func = self._default_name_func

790

791 async def _check_type_cache_invalidation(self, invalidate_timestamp):

792 if invalidate_timestamp > self._invalidate_schema_cache_asof:

793 await self._connection.reload_schema_state()

794 self._invalidate_schema_cache_asof = invalidate_timestamp

795

796 async def _prepare(self, operation, invalidate_timestamp):

797 await self._check_type_cache_invalidation(invalidate_timestamp)

798

799 cache = self._prepared_statement_cache

800 if cache is None:

801 prepared_stmt = await self._connection.prepare(

802 operation, name=self._prepared_statement_name_func()

803 )

804 attributes = prepared_stmt.get_attributes()

805 return prepared_stmt, attributes

806

807 # asyncpg uses a type cache for the "attributes" which seems to go

808 # stale independently of the PreparedStatement itself, so place that

809 # collection in the cache as well.

810 if operation in cache:

811 prepared_stmt, attributes, cached_timestamp = cache[operation]

812

813 # preparedstatements themselves also go stale for certain DDL

814 # changes such as size of a VARCHAR changing, so there is also

815 # a cross-connection invalidation timestamp

816 if cached_timestamp > invalidate_timestamp:

817 return prepared_stmt, attributes

818

819 prepared_stmt = await self._connection.prepare(

820 operation, name=self._prepared_statement_name_func()

821 )

822 attributes = prepared_stmt.get_attributes()

823 cache[operation] = (prepared_stmt, attributes, time.time())

824

825 return prepared_stmt, attributes

826

827 @classmethod

828 def _handle_exception_no_connection(

829 cls, dbapi: Any, error: Exception

830 ) -> NoReturn:

831 if not isinstance(error, AsyncAdapt_asyncpg_dbapi.Error):

832 exception_mapping = dbapi._asyncpg_error_translate

833

834 for super_ in type(error).__mro__:

835 if super_ in exception_mapping:

836 message = error.args[0]

837 translated_error = exception_mapping[super_](

838 message, error

839 )

840 raise translated_error from error

841 super()._handle_exception_no_connection(dbapi, error)

842

843 def _handle_exception(self, error: Exception) -> NoReturn:

844 if self._connection.is_closed():

845 self._transaction = None

846

847 super()._handle_exception(error)

848

849 @property

850 def autocommit(self):

851 return self.isolation_level == "autocommit"

852

853 @autocommit.setter

854 def autocommit(self, value):

855 if value:

856 self.isolation_level = "autocommit"

857 else:

858 self.isolation_level = self._isolation_setting

859

860 def ping(self):

861 try:

862 _ = await_(self._async_ping())

863 except Exception as error:

864 self._handle_exception(error)

865

866 async def _async_ping(self):

867 if self._transaction is None and self.isolation_level != "autocommit":

868 # create a transaction explicitly to support pgbouncer

869 # transaction mode. See #10226

870 tr = self._connection.transaction()

871 await tr.start()

872 try:

873 await self._connection.fetchrow(";")

874 finally:

875 await tr.rollback()

876 else:

877 await self._connection.fetchrow(";")

878

879 def set_isolation_level(self, level):

880 self.rollback()

881 self.isolation_level = self._isolation_setting = level

882

883 async def _start_transaction(self):

884 if self.isolation_level == "autocommit":

885 return

886

887 assert self._transaction is None

888 try:

889 transaction = self._connection.transaction(

890 isolation=self.isolation_level,

891 readonly=self.readonly,

892 deferrable=self.deferrable,

893 )

894 await transaction.start()

895 except Exception as error:

896 self._handle_exception(error)

897 else:

898 # only track the transaction once BEGIN has succeeded; a

899 # transaction that failed to start has nothing to roll back

900 self._transaction = transaction

901

902 async def _call_and_discard(self, fn: Callable[[], Awaitable[Any]]):

903 try:

904 await fn()

905 finally:

906 # if asyncpg fn was actually called, then whether or

907 # not it raised or succeeded, the transaction is done, discard it

908 self._transaction = None

909

910 def rollback(self):

911 if self._transaction is not None:

912 try:

913 await_(self._call_and_discard(self._transaction.rollback))

914 except Exception as error:

915 # don't dereference asyncpg transaction if we didn't

916 # actually try to call rollback() on it

917 self._handle_exception(error)

918

919 def commit(self):

920 if self._transaction is not None:

921 try:

922 await_(self._call_and_discard(self._transaction.commit))

923 except Exception as error:

924 # don't dereference asyncpg transaction if we didn't

925 # actually try to call commit() on it

926 self._handle_exception(error)

927

928 def close(self):

929 self.rollback()

930

931 await_(self._connection.close())

932

933 def _terminate_handled_exceptions(self):

934 return super()._terminate_handled_exceptions() + (

935 self.dbapi.asyncpg.PostgresError,

936 )

937

938 async def _terminate_graceful_close(self) -> None:

939 # timeout added in asyncpg 0.14.0 December 2017

940 await self._connection.close(timeout=2)

941 self._transaction = None

942

943 def _terminate_force_close(self) -> None:

944 self._connection.terminate()

945 self._transaction = None

946

947 @staticmethod

948 def _default_name_func():

949 return None

950

951

952class AsyncAdapt_asyncpg_dbapi(AsyncAdapt_dbapi_module):

953 def __init__(self, asyncpg):

954 super().__init__(asyncpg)

955 self.asyncpg = asyncpg

956 self.paramstyle = "numeric_dollar"

957

958 def connect(self, *arg, **kw):

959 creator_fn = kw.pop("async_creator_fn", self.asyncpg.connect)

960 prepared_statement_cache_size = kw.pop(

961 "prepared_statement_cache_size", 100

962 )

963 prepared_statement_name_func = kw.pop(

964 "prepared_statement_name_func", None

965 )

966

967 return await_(

968 AsyncAdapt_asyncpg_connection.create(

969 self,

970 creator_fn(*arg, **kw),

971 prepared_statement_cache_size=prepared_statement_cache_size,

972 prepared_statement_name_func=prepared_statement_name_func,

973 )

974 )

975

976 class Error(AsyncAdapt_Error):

977

978 pgcode: str | None

979

980 sqlstate: str | None

981

982 detail: str | None

983

984 def __init__(self, message, error=None):

985 super().__init__(message, error)

986 self.detail = getattr(error, "detail", None)

987 self.pgcode = self.sqlstate = getattr(error, "sqlstate", None)

988

989 class Warning(AsyncAdapt_Error): # noqa

990 pass

991

992 class InterfaceError(Error):

993 pass

994

995 class DatabaseError(Error):

996 pass

997

998 class InternalError(DatabaseError):

999 pass

1000

1001 class OperationalError(DatabaseError):

1002 pass

1003

1004 class ProgrammingError(DatabaseError):

1005 pass

1006

1007 class IntegrityError(DatabaseError):

1008 pass

1009

1010 class RestrictViolationError(IntegrityError):

1011 pass

1012

1013 class NotNullViolationError(IntegrityError):

1014 pass

1015

1016 class ForeignKeyViolationError(IntegrityError):

1017 pass

1018

1019 class UniqueViolationError(IntegrityError):

1020 pass

1021

1022 class CheckViolationError(IntegrityError):

1023 pass

1024

1025 class ExclusionViolationError(IntegrityError):

1026 pass

1027

1028 class DataError(DatabaseError):

1029 pass

1030

1031 class NotSupportedError(DatabaseError):

1032 pass

1033

1034 class InternalServerError(InternalError):

1035 pass

1036

1037 class InternalClientError(InternalError):

1038 pass

1039

1040 class InvalidCachedStatementError(NotSupportedError):

1041 def __init__(self, message, error=None):

1042 super().__init__(

1043 message + " (SQLAlchemy asyncpg dialect will now invalidate "

1044 "all prepared caches in response to this exception)",

1045 )

1046

1047 # pep-249 datatype placeholders. As of SQLAlchemy 2.0 these aren't

1048 # used, however the test suite looks for these in a few cases.

1049 STRING = util.symbol("STRING")

1050 NUMBER = util.symbol("NUMBER")

1051 DATETIME = util.symbol("DATETIME")

1052

1053 @util.memoized_property

1054 def _asyncpg_error_translate(self):

1055 import asyncpg

1056

1057 return {

1058 asyncpg.exceptions.IntegrityConstraintViolationError: self.IntegrityError, # noqa: E501

1059 asyncpg.exceptions.PostgresError: self.Error,

1060 asyncpg.exceptions.SyntaxOrAccessError: self.ProgrammingError,

1061 asyncpg.exceptions.InterfaceError: self.InterfaceError,

1062 asyncpg.exceptions.InvalidCachedStatementError: self.InvalidCachedStatementError, # noqa: E501

1063 asyncpg.exceptions.InternalServerError: self.InternalServerError,

1064 asyncpg.exceptions.RestrictViolationError: self.RestrictViolationError, # noqa: E501

1065 asyncpg.exceptions.NotNullViolationError: self.NotNullViolationError, # noqa: E501

1066 asyncpg.exceptions.ForeignKeyViolationError: self.ForeignKeyViolationError, # noqa: E501

1067 asyncpg.exceptions.UniqueViolationError: self.UniqueViolationError,

1068 asyncpg.exceptions.CheckViolationError: self.CheckViolationError,

1069 asyncpg.exceptions.ExclusionViolationError: self.ExclusionViolationError, # noqa: E501

1070 asyncpg.exceptions.InternalClientError: self.InternalClientError,

1071 }

1072

1073 def Binary(self, value):

1074 return value

1075

1076

1077class PGDialect_asyncpg(PGDialect):

1078 driver = "asyncpg"

1079 supports_statement_cache = True

1080

1081 supports_server_side_cursors = True

1082

1083 render_bind_cast = True

1084 has_terminate = True

1085

1086 default_paramstyle = "numeric_dollar"

1087 supports_sane_multi_rowcount = False

1088 execution_ctx_cls = PGExecutionContext_asyncpg

1089 statement_compiler = PGCompiler_asyncpg

1090 preparer = PGIdentifierPreparer_asyncpg

1091

1092 supports_native_json_serialization = False

1093 supports_native_json_deserialization = True

1094 dialect_injects_custom_json_deserializer = True

1095

1096 colspecs = util.update_copy(

1097 PGDialect.colspecs,

1098 {

1099 sqltypes.String: AsyncpgString,

1100 sqltypes.ARRAY: AsyncpgARRAY,

1101 BIT: AsyncpgBit,

1102 CITEXT: CITEXT,

1103 REGCONFIG: AsyncpgREGCONFIG,

1104 sqltypes.Time: AsyncpgTime,

1105 sqltypes.Date: AsyncpgDate,

1106 sqltypes.DateTime: AsyncpgDateTime,

1107 sqltypes.Interval: AsyncPgInterval,

1108 INTERVAL: AsyncPgInterval,

1109 sqltypes.Boolean: AsyncpgBoolean,

1110 sqltypes.Integer: AsyncpgInteger,

1111 sqltypes.SmallInteger: AsyncpgSmallInteger,

1112 sqltypes.BigInteger: AsyncpgBigInteger,

1113 sqltypes.Numeric: AsyncpgNumeric,

1114 sqltypes.Float: AsyncpgFloat,

1115 sqltypes.LargeBinary: AsyncpgByteA,

1116 sqltypes.JSON.JSONPathType: AsyncpgJSONPathType,

1117 sqltypes.JSON.JSONIndexType: AsyncpgJSONIndexType,

1118 sqltypes.JSON.JSONIntIndexType: AsyncpgJSONIntIndexType,

1119 sqltypes.JSON.JSONStrIndexType: AsyncpgJSONStrIndexType,

1120 sqltypes.Enum: AsyncPgEnum,

1121 OID: AsyncpgOID,

1122 REGCLASS: AsyncpgREGCLASS,

1123 sqltypes.CHAR: AsyncpgCHAR,

1124 ranges.AbstractSingleRange: _AsyncpgRange,

1125 ranges.AbstractMultiRange: _AsyncpgMultiRange,

1126 },

1127 )

1128 is_async = True

1129 _invalidate_schema_cache_asof = 0

1130

1131 def _invalidate_schema_cache(self):

1132 self._invalidate_schema_cache_asof = time.time()

1133

1134 def retrieve_dbapi_version(self, dbapi):

1135 # dbapi is the AsyncAdapt_asyncpg_dbapi wrapper; the version is on

1136 # the asyncpg module itself, which is ``.driver``

1137 driver = getattr(dbapi, "driver", None)

1138 return util.parse_version_string(getattr(driver, "__version__", None))

1139

1140 @classmethod

1141 def import_dbapi(cls):

1142 return AsyncAdapt_asyncpg_dbapi(__import__("asyncpg"))

1143

1144 @util.memoized_property

1145 def _isolation_lookup(self):

1146 return {

1147 "AUTOCOMMIT": "autocommit",

1148 "READ COMMITTED": "read_committed",

1149 "REPEATABLE READ": "repeatable_read",

1150 "SERIALIZABLE": "serializable",

1151 }

1152

1153 def get_isolation_level_values(self, dbapi_connection):

1154 return list(self._isolation_lookup)

1155

1156 def set_isolation_level(self, dbapi_connection, level):

1157 dbapi_connection.set_isolation_level(self._isolation_lookup[level])

1158

1159 def detect_autocommit_setting(self, dbapi_conn) -> bool:

1160 return bool(dbapi_conn.autocommit)

1161

1162 def set_readonly(self, connection, value):

1163 connection.readonly = value

1164

1165 def get_readonly(self, connection):

1166 return connection.readonly

1167

1168 def set_deferrable(self, connection, value):

1169 connection.deferrable = value

1170

1171 def get_deferrable(self, connection):

1172 return connection.deferrable

1173

1174 def do_terminate(self, dbapi_connection) -> None:

1175 dbapi_connection.terminate()

1176

1177 def create_connect_args(self, url):

1178 opts = url.translate_connect_args(username="user")

1179 multihosts, multiports = self._split_multihost_from_url(url)

1180

1181 opts.update(url.query)

1182

1183 if multihosts:

1184 assert multiports

1185 if len(multihosts) == 1:

1186 opts["host"] = multihosts[0]

1187 if multiports[0] is not None:

1188 opts["port"] = multiports[0]

1189 elif not all(multihosts):

1190 raise exc.ArgumentError(

1191 "All hosts are required to be present"

1192 " for asyncpg multiple host URL"

1193 )

1194 elif not all(multiports):

1195 raise exc.ArgumentError(

1196 "All ports are required to be present"

1197 " for asyncpg multiple host URL"

1198 )

1199 else:

1200 opts["host"] = list(multihosts)

1201 opts["port"] = list(multiports)

1202 else:

1203 util.coerce_kw_type(opts, "port", int)

1204 util.coerce_kw_type(opts, "prepared_statement_cache_size", int)

1205 return ([], opts)

1206

1207 def do_ping(self, dbapi_connection):

1208 dbapi_connection.ping()

1209 return True

1210

1211 def is_disconnect(self, e, connection, cursor):

1212 if connection:

1213 return connection._connection.is_closed()

1214 else:

1215 return isinstance(

1216 e, self.dbapi.InterfaceError

1217 ) and "connection is closed" in str(e)

1218

1219 async def setup_asyncpg_json_codec(self, conn):

1220 """set up JSON codec for asyncpg.

1221

1222 This occurs for all new connections and

1223 can be overridden by third party dialects.

1224

1225 .. versionadded:: 1.4.27

1226

1227 """

1228

1229 asyncpg_connection = conn._connection

1230 deserializer = self._json_deserializer or _py_json.loads

1231

1232 def _json_decoder(bin_value):

1233 return deserializer(bin_value.decode())

1234

1235 await asyncpg_connection.set_type_codec(

1236 "json",

1237 encoder=str.encode,

1238 decoder=_json_decoder,

1239 schema="pg_catalog",

1240 format="binary",

1241 )

1242

1243 async def setup_asyncpg_jsonb_codec(self, conn):

1244 """set up JSONB codec for asyncpg.

1245

1246 This occurs for all new connections and

1247 can be overridden by third party dialects.

1248

1249 .. versionadded:: 1.4.27

1250

1251 """

1252

1253 asyncpg_connection = conn._connection

1254

1255 def _jsonb_encoder(str_value):

1256 # \x01 is the prefix for jsonb used by PostgreSQL.

1257 # asyncpg requires it when format='binary'

1258 return b"\x01" + str_value.encode()

1259

1260 deserializer = self._json_deserializer or _py_json.loads

1261

1262 def _jsonb_decoder(bin_value):

1263 # the byte is the \x01 prefix for jsonb used by PostgreSQL.

1264 # asyncpg returns it when format='binary'

1265 return deserializer(bin_value[1:].decode())

1266

1267 await asyncpg_connection.set_type_codec(

1268 "jsonb",

1269 encoder=_jsonb_encoder,

1270 decoder=_jsonb_decoder,

1271 schema="pg_catalog",

1272 format="binary",

1273 )

1274

1275 async def _disable_asyncpg_inet_codecs(self, conn):

1276 asyncpg_connection = conn._connection

1277

1278 await asyncpg_connection.set_type_codec(

1279 "inet",

1280 encoder=lambda s: s,

1281 decoder=lambda s: s,

1282 schema="pg_catalog",

1283 format="text",

1284 )

1285

1286 await asyncpg_connection.set_type_codec(

1287 "cidr",

1288 encoder=lambda s: s,

1289 decoder=lambda s: s,

1290 schema="pg_catalog",

1291 format="text",

1292 )

1293

1294 def on_connect(self):

1295 """on_connect for asyncpg

1296

1297 A major component of this for asyncpg is to set up type decoders at the

1298 asyncpg level.

1299

1300 See https://github.com/MagicStack/asyncpg/issues/623 for

1301 notes on JSON/JSONB implementation.

1302

1303 """

1304

1305 super_connect = super().on_connect()

1306

1307 def connect(conn):

1308 await_(self.setup_asyncpg_json_codec(conn))

1309 await_(self.setup_asyncpg_jsonb_codec(conn))

1310

1311 if self._native_inet_types is False:

1312 await_(self._disable_asyncpg_inet_codecs(conn))

1313 if super_connect is not None:

1314 super_connect(conn)

1315

1316 return connect

1317

1318 def get_driver_connection(self, connection):

1319 return connection._connection

1320

1321

1322dialect = PGDialect_asyncpg