codekingpro/portable-devtools
115k
1# dialects/postgresql/psycopg2.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS 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# mypy: ignore-errors
8
9r"""
10.. dialect:: postgresql+psycopg2
11 :name: psycopg2
12 :dbapi: psycopg2
13 :connectstring: postgresql+psycopg2://user:password@host:port/dbname[?key=value&key=value...]
14 :url: https://pypi.org/project/psycopg2/
15
16.. _psycopg2_toplevel:
17
18psycopg2 Connect Arguments
19--------------------------
20
21Keyword arguments that are specific to the SQLAlchemy psycopg2 dialect
22may be passed to :func:`_sa.create_engine()`, and include the following:
23
24
25* ``isolation_level``: This option, available for all PostgreSQL dialects,
26 includes the ``AUTOCOMMIT`` isolation level when using the psycopg2
27 dialect. This option sets the **default** isolation level for the
28 connection that is set immediately upon connection to the database before
29 the connection is pooled. This option is generally superseded by the more
30 modern :paramref:`_engine.Connection.execution_options.isolation_level`
31 execution option, detailed at :ref:`dbapi_autocommit`.
32
33 .. seealso::
34
35 :ref:`psycopg2_isolation_level`
36
37 :ref:`dbapi_autocommit`
38
39
40* ``client_encoding``: sets the client encoding in a libpq-agnostic way,
41 using psycopg2's ``set_client_encoding()`` method.
42
43 .. seealso::
44
45 :ref:`psycopg2_unicode`
46
47
48* ``executemany_mode``, ``executemany_batch_page_size``,
49 ``executemany_values_page_size``: Allows use of psycopg2
50 extensions for optimizing "executemany"-style queries. See the referenced
51 section below for details.
52
53 .. seealso::
54
55 :ref:`psycopg2_executemany_mode`
56
57.. tip::
58
59 The above keyword arguments are **dialect** keyword arguments, meaning
60 that they are passed as explicit keyword arguments to :func:`_sa.create_engine()`::
61
62 engine = create_engine(
63 "postgresql+psycopg2://scott:tiger@localhost/test",
64 isolation_level="SERIALIZABLE",
65 )
66
67 These should not be confused with **DBAPI** connect arguments, which
68 are passed as part of the :paramref:`_sa.create_engine.connect_args`
69 dictionary and/or are passed in the URL query string, as detailed in
70 the section :ref:`custom_dbapi_args`.
71
72.. _psycopg2_ssl:
73
74SSL Connections
75---------------
76
77The psycopg2 module has a connection argument named ``sslmode`` for
78controlling its behavior regarding secure (SSL) connections. The default is
79``sslmode=prefer``; it will attempt an SSL connection and if that fails it
80will fall back to an unencrypted connection. ``sslmode=require`` may be used
81to ensure that only secure connections are established. Consult the
82psycopg2 / libpq documentation for further options that are available.
83
84Note that ``sslmode`` is specific to psycopg2 so it is included in the
85connection URI::
86
87 engine = sa.create_engine(
88 "postgresql+psycopg2://scott:tiger@192.168.0.199:5432/test?sslmode=require"
89 )
90
91Unix Domain Connections
92------------------------
93
94psycopg2 supports connecting via Unix domain connections. When the ``host``
95portion of the URL is omitted, SQLAlchemy passes ``None`` to psycopg2,
96which specifies Unix-domain communication rather than TCP/IP communication::
97
98 create_engine("postgresql+psycopg2://user:password@/dbname")
99
100By default, the socket file used is to connect to a Unix-domain socket
101in ``/tmp``, or whatever socket directory was specified when PostgreSQL
102was built. This value can be overridden by passing a pathname to psycopg2,
103using ``host`` as an additional keyword argument::
104
105 create_engine(
106 "postgresql+psycopg2://user:password@/dbname?host=/var/lib/postgresql"
107 )
108
109.. warning:: The format accepted here allows for a hostname in the main URL
110 in addition to the "host" query string argument. **When using this URL
111 format, the initial host is silently ignored**. That is, this URL::
112
113 engine = create_engine(
114 "postgresql+psycopg2://user:password@myhost1/dbname?host=myhost2"
115 )
116
117 Above, the hostname ``myhost1`` is **silently ignored and discarded.** The
118 host which is connected is the ``myhost2`` host.
119
120 This is to maintain some degree of compatibility with PostgreSQL's own URL
121 format which has been tested to behave the same way and for which tools like
122 PifPaf hardcode two hostnames.
123
124.. seealso::
125
126 `PQconnectdbParams \
127 <https://www.postgresql.org/docs/current/static/libpq-connect.html#LIBPQ-PQCONNECTDBPARAMS>`_
128
129.. _psycopg2_multi_host:
130
131Specifying multiple fallback hosts
132-----------------------------------
133
134psycopg2 supports multiple connection points in the connection string.
135When the ``host`` parameter is used multiple times in the query section of
136the URL, SQLAlchemy will create a single string of the host and port
137information provided to make the connections. Tokens may consist of
138``host::port`` or just ``host``; in the latter case, the default port
139is selected by libpq. In the example below, three host connections
140are specified, for ``HostA::PortA``, ``HostB`` connecting to the default port,
141and ``HostC::PortC``::
142
143 create_engine(
144 "postgresql+psycopg2://user:password@/dbname?host=HostA:PortA&host=HostB&host=HostC:PortC"
145 )
146
147As an alternative, libpq query string format also may be used; this specifies
148``host`` and ``port`` as single query string arguments with comma-separated
149lists - the default port can be chosen by indicating an empty value
150in the comma separated list::
151
152 create_engine(
153 "postgresql+psycopg2://user:password@/dbname?host=HostA,HostB,HostC&port=PortA,,PortC"
154 )
155
156With either URL style, connections to each host is attempted based on a
157configurable strategy, which may be configured using the libpq
158``target_session_attrs`` parameter. Per libpq this defaults to ``any``
159which indicates a connection to each host is then attempted until a connection is successful.
160Other strategies include ``primary``, ``prefer-standby``, etc. The complete
161list is documented by PostgreSQL at
162`libpq connection strings <https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING>`_.
163
164For example, to indicate two hosts using the ``primary`` strategy::
165
166 create_engine(
167 "postgresql+psycopg2://user:password@/dbname?host=HostA:PortA&host=HostB&host=HostC:PortC&target_session_attrs=primary"
168 )
169
170.. versionchanged:: 1.4.40 Port specification in psycopg2 multiple host format
171 is repaired, previously ports were not correctly interpreted in this context.
172 libpq comma-separated format is also now supported.
173
174.. versionadded:: 1.3.20 Support for multiple hosts in PostgreSQL connection
175 string.
176
177.. seealso::
178
179 `libpq connection strings <https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING>`_ - please refer
180 to this section in the libpq documentation for complete background on multiple host support.
181
182
183Empty DSN Connections / Environment Variable Connections
184---------------------------------------------------------
185
186The psycopg2 DBAPI can connect to PostgreSQL by passing an empty DSN to the
187libpq client library, which by default indicates to connect to a localhost
188PostgreSQL database that is open for "trust" connections. This behavior can be
189further tailored using a particular set of environment variables which are
190prefixed with ``PG_...``, which are consumed by ``libpq`` to take the place of
191any or all elements of the connection string.
192
193For this form, the URL can be passed without any elements other than the
194initial scheme::
195
196 engine = create_engine("postgresql+psycopg2://")
197
198In the above form, a blank "dsn" string is passed to the ``psycopg2.connect()``
199function which in turn represents an empty DSN passed to libpq.
200
201.. versionadded:: 1.3.2 support for parameter-less connections with psycopg2.
202
203.. seealso::
204
205 `Environment Variables\
206 <https://www.postgresql.org/docs/current/libpq-envars.html>`_ -
207 PostgreSQL documentation on how to use ``PG_...``
208 environment variables for connections.
209
210.. _psycopg2_execution_options:
211
212Per-Statement/Connection Execution Options
213-------------------------------------------
214
215The following DBAPI-specific options are respected when used with
216:meth:`_engine.Connection.execution_options`,
217:meth:`.Executable.execution_options`,
218:meth:`_query.Query.execution_options`,
219in addition to those not specific to DBAPIs:
220
221* ``isolation_level`` - Set the transaction isolation level for the lifespan
222 of a :class:`_engine.Connection` (can only be set on a connection,
223 not a statement
224 or query). See :ref:`psycopg2_isolation_level`.
225
226* ``stream_results`` - Enable or disable usage of psycopg2 server side
227 cursors - this feature makes use of "named" cursors in combination with
228 special result handling methods so that result rows are not fully buffered.
229 Defaults to False, meaning cursors are buffered by default.
230
231* ``max_row_buffer`` - when using ``stream_results``, an integer value that
232 specifies the maximum number of rows to buffer at a time. This is
233 interpreted by the :class:`.BufferedRowCursorResult`, and if omitted the
234 buffer will grow to ultimately store 1000 rows at a time.
235
236 .. versionchanged:: 1.4 The ``max_row_buffer`` size can now be greater than
237 1000, and the buffer will grow to that size.
238
239.. _psycopg2_batch_mode:
240
241.. _psycopg2_executemany_mode:
242
243Psycopg2 Fast Execution Helpers
244-------------------------------
245
246Modern versions of psycopg2 include a feature known as
247`Fast Execution Helpers \
248<https://www.psycopg.org/docs/extras.html#fast-execution-helpers>`_, which
249have been shown in benchmarking to improve psycopg2's executemany()
250performance, primarily with INSERT statements, by at least
251an order of magnitude.
252
253SQLAlchemy implements a native form of the "insert many values"
254handler that will rewrite a single-row INSERT statement to accommodate for
255many values at once within an extended VALUES clause; this handler is
256equivalent to psycopg2's ``execute_values()`` handler; an overview of this
257feature and its configuration are at :ref:`engine_insertmanyvalues`.
258
259.. versionadded:: 2.0 Replaced psycopg2's ``execute_values()`` fast execution
260 helper with a native SQLAlchemy mechanism known as
261 :ref:`insertmanyvalues <engine_insertmanyvalues>`.
262
263The psycopg2 dialect retains the ability to use the psycopg2-specific
264``execute_batch()`` feature, although it is not expected that this is a widely
265used feature. The use of this extension may be enabled using the
266``executemany_mode`` flag which may be passed to :func:`_sa.create_engine`::
267
268 engine = create_engine(
269 "postgresql+psycopg2://scott:tiger@host/dbname",
270 executemany_mode="values_plus_batch",
271 )
272
273Possible options for ``executemany_mode`` include:
274
275* ``values_only`` - this is the default value. SQLAlchemy's native
276 :ref:`insertmanyvalues <engine_insertmanyvalues>` handler is used for qualifying
277 INSERT statements, assuming
278 :paramref:`_sa.create_engine.use_insertmanyvalues` is left at
279 its default value of ``True``. This handler rewrites simple
280 INSERT statements to include multiple VALUES clauses so that many
281 parameter sets can be inserted with one statement.
282
283* ``'values_plus_batch'``- SQLAlchemy's native
284 :ref:`insertmanyvalues <engine_insertmanyvalues>` handler is used for qualifying
285 INSERT statements, assuming
286 :paramref:`_sa.create_engine.use_insertmanyvalues` is left at its default
287 value of ``True``. Then, psycopg2's ``execute_batch()`` handler is used for
288 qualifying UPDATE and DELETE statements when executed with multiple parameter
289 sets. When using this mode, the :attr:`_engine.CursorResult.rowcount`
290 attribute will not contain a value for executemany-style executions against
291 UPDATE and DELETE statements.
292
293.. versionchanged:: 2.0 Removed the ``'batch'`` and ``'None'`` options
294 from psycopg2 ``executemany_mode``. Control over batching for INSERT
295 statements is now configured via the
296 :paramref:`_sa.create_engine.use_insertmanyvalues` engine-level parameter.
297
298The term "qualifying statements" refers to the statement being executed
299being a Core :func:`_expression.insert`, :func:`_expression.update`
300or :func:`_expression.delete` construct, and **not** a plain textual SQL
301string or one constructed using :func:`_expression.text`. It also may **not** be
302a special "extension" statement such as an "ON CONFLICT" "upsert" statement.
303When using the ORM, all insert/update/delete statements used by the ORM flush process
304are qualifying.
305
306The "page size" for the psycopg2 "batch" strategy can be affected
307by using the ``executemany_batch_page_size`` parameter, which defaults to
308100.
309
310For the "insertmanyvalues" feature, the page size can be controlled using the
311:paramref:`_sa.create_engine.insertmanyvalues_page_size` parameter,
312which defaults to 1000. An example of modifying both parameters
313is below::
314
315 engine = create_engine(
316 "postgresql+psycopg2://scott:tiger@host/dbname",
317 executemany_mode="values_plus_batch",
318 insertmanyvalues_page_size=5000,
319 executemany_batch_page_size=500,
320 )
321
322.. seealso::
323
324 :ref:`engine_insertmanyvalues` - background on "insertmanyvalues"
325
326 :ref:`tutorial_multiple_parameters` - General information on using the
327 :class:`_engine.Connection`
328 object to execute statements in such a way as to make
329 use of the DBAPI ``.executemany()`` method.
330
331
332.. _psycopg2_unicode:
333
334Unicode with Psycopg2
335----------------------
336
337The psycopg2 DBAPI driver supports Unicode data transparently.
338
339The client character encoding can be controlled for the psycopg2 dialect
340in the following ways:
341
342* For PostgreSQL 9.1 and above, the ``client_encoding`` parameter may be
343 passed in the database URL; this parameter is consumed by the underlying
344 ``libpq`` PostgreSQL client library::
345
346 engine = create_engine(
347 "postgresql+psycopg2://user:pass@host/dbname?client_encoding=utf8"
348 )
349
350 Alternatively, the above ``client_encoding`` value may be passed using
351 :paramref:`_sa.create_engine.connect_args` for programmatic establishment with
352 ``libpq``::
353
354 engine = create_engine(
355 "postgresql+psycopg2://user:pass@host/dbname",
356 connect_args={"client_encoding": "utf8"},
357 )
358
359* For all PostgreSQL versions, psycopg2 supports a client-side encoding
360 value that will be passed to database connections when they are first
361 established. The SQLAlchemy psycopg2 dialect supports this using the
362 ``client_encoding`` parameter passed to :func:`_sa.create_engine`::
363
364 engine = create_engine(
365 "postgresql+psycopg2://user:pass@host/dbname", client_encoding="utf8"
366 )
367
368 .. tip:: The above ``client_encoding`` parameter admittedly is very similar
369 in appearance to usage of the parameter within the
370 :paramref:`_sa.create_engine.connect_args` dictionary; the difference
371 above is that the parameter is consumed by psycopg2 and is
372 passed to the database connection using ``SET client_encoding TO
373 'utf8'``; in the previously mentioned style, the parameter is instead
374 passed through psycopg2 and consumed by the ``libpq`` library.
375
376* A common way to set up client encoding with PostgreSQL databases is to
377 ensure it is configured within the server-side postgresql.conf file;
378 this is the recommended way to set encoding for a server that is
379 consistently of one encoding in all databases::
380
381 # postgresql.conf file
382
383 # client_encoding = sql_ascii # actually, defaults to database
384 # encoding
385 client_encoding = utf8
386
387Transactions
388------------
389
390The psycopg2 dialect fully supports SAVEPOINT and two-phase commit operations.
391
392.. _psycopg2_isolation_level:
393
394Psycopg2 Transaction Isolation Level
395-------------------------------------
396
397As discussed in :ref:`postgresql_isolation_level`,
398all PostgreSQL dialects support setting of transaction isolation level
399both via the ``isolation_level`` parameter passed to :func:`_sa.create_engine`
400,
401as well as the ``isolation_level`` argument used by
402:meth:`_engine.Connection.execution_options`. When using the psycopg2 dialect
403, these
404options make use of psycopg2's ``set_isolation_level()`` connection method,
405rather than emitting a PostgreSQL directive; this is because psycopg2's
406API-level setting is always emitted at the start of each transaction in any
407case.
408
409The psycopg2 dialect supports these constants for isolation level:
410
411* ``READ COMMITTED``
412* ``READ UNCOMMITTED``
413* ``REPEATABLE READ``
414* ``SERIALIZABLE``
415* ``AUTOCOMMIT``
416
417.. seealso::
418
419 :ref:`postgresql_isolation_level`
420
421 :ref:`pg8000_isolation_level`
422
423
424NOTICE logging
425---------------
426
427The psycopg2 dialect will log PostgreSQL NOTICE messages
428via the ``sqlalchemy.dialects.postgresql`` logger. When this logger
429is set to the ``logging.INFO`` level, notice messages will be logged::
430
431 import logging
432
433 logging.getLogger("sqlalchemy.dialects.postgresql").setLevel(logging.INFO)
434
435Above, it is assumed that logging is configured externally. If this is not
436the case, configuration such as ``logging.basicConfig()`` must be utilized::
437
438 import logging
439
440 logging.basicConfig() # log messages to stdout
441 logging.getLogger("sqlalchemy.dialects.postgresql").setLevel(logging.INFO)
442
443.. seealso::
444
445 `Logging HOWTO <https://docs.python.org/3/howto/logging.html>`_ - on the python.org website
446
447.. _psycopg2_hstore:
448
449HSTORE type
450------------
451
452The ``psycopg2`` DBAPI includes an extension to natively handle marshalling of
453the HSTORE type. The SQLAlchemy psycopg2 dialect will enable this extension
454by default when psycopg2 version 2.4 or greater is used, and
455it is detected that the target database has the HSTORE type set up for use.
456In other words, when the dialect makes the first
457connection, a sequence like the following is performed:
458
4591. Request the available HSTORE oids using
460 ``psycopg2.extras.HstoreAdapter.get_oids()``.
461 If this function returns a list of HSTORE identifiers, we then determine
462 that the ``HSTORE`` extension is present.
463 This function is **skipped** if the version of psycopg2 installed is
464 less than version 2.4.
465
4662. If the ``use_native_hstore`` flag is at its default of ``True``, and
467 we've detected that ``HSTORE`` oids are available, the
468 ``psycopg2.extensions.register_hstore()`` extension is invoked for all
469 connections.
470
471The ``register_hstore()`` extension has the effect of **all Python
472dictionaries being accepted as parameters regardless of the type of target
473column in SQL**. The dictionaries are converted by this extension into a
474textual HSTORE expression. If this behavior is not desired, disable the
475use of the hstore extension by setting ``use_native_hstore`` to ``False`` as
476follows::
477
478 engine = create_engine(
479 "postgresql+psycopg2://scott:tiger@localhost/test",
480 use_native_hstore=False,
481 )
482
483The ``HSTORE`` type is **still supported** when the
484``psycopg2.extensions.register_hstore()`` extension is not used. It merely
485means that the coercion between Python dictionaries and the HSTORE
486string format, on both the parameter side and the result side, will take
487place within SQLAlchemy's own marshalling logic, and not that of ``psycopg2``
488which may be more performant.
489
490""" # noqa
491from __future__ import annotations
492
493import collections.abc as collections_abc
494import logging
495import re
496from typing import cast
497
498from . import ranges
499from ._psycopg_common import _PGDialect_common_psycopg
500from ._psycopg_common import _PGExecutionContext_common_psycopg
501from .base import PGIdentifierPreparer
502from .json import JSON
503from .json import JSONB
504from ... import types as sqltypes
505from ... import util
506from ...util import FastIntFlag
507from ...util import parse_user_argument_for_enum
508
509logger = logging.getLogger("sqlalchemy.dialects.postgresql")
510
511
512class _PGJSON(JSON):
513 def result_processor(self, dialect, coltype):
514 return None
515
516
517class _PGJSONB(JSONB):
518 def result_processor(self, dialect, coltype):
519 return None
520
521
522class _Psycopg2Range(ranges.AbstractSingleRangeImpl):
523 _psycopg2_range_cls = "none"
524
525 def bind_processor(self, dialect):
526 psycopg2_Range = getattr(
527 cast(PGDialect_psycopg2, dialect)._psycopg2_extras,
528 self._psycopg2_range_cls,
529 )
530
531 def to_range(value):
532 if isinstance(value, ranges.Range):
533 value = psycopg2_Range(
534 value.lower, value.upper, value.bounds, value.empty
535 )
536 return value
537
538 return to_range
539
540 def result_processor(self, dialect, coltype):
541 def to_range(value):
542 if value is not None:
543 value = ranges.Range(
544 value._lower,
545 value._upper,
546 bounds=value._bounds if value._bounds else "[)",
547 empty=not value._bounds,
548 )
549 return value
550
551 return to_range
552
553
554class _Psycopg2NumericRange(_Psycopg2Range):
555 _psycopg2_range_cls = "NumericRange"
556
557
558class _Psycopg2DateRange(_Psycopg2Range):
559 _psycopg2_range_cls = "DateRange"
560
561
562class _Psycopg2DateTimeRange(_Psycopg2Range):
563 _psycopg2_range_cls = "DateTimeRange"
564
565
566class _Psycopg2DateTimeTZRange(_Psycopg2Range):
567 _psycopg2_range_cls = "DateTimeTZRange"
568
569
570class PGExecutionContext_psycopg2(_PGExecutionContext_common_psycopg):
571 _psycopg2_fetched_rows = None
572
573 def post_exec(self):
574 self._log_notices(self.cursor)
575
576 def _log_notices(self, cursor):
577 # check also that notices is an iterable, after it's already
578 # established that we will be iterating through it. This is to get
579 # around test suites such as SQLAlchemy's using a Mock object for
580 # cursor
581 if not cursor.connection.notices or not isinstance(
582 cursor.connection.notices, collections_abc.Iterable
583 ):
584 return
585
586 for notice in cursor.connection.notices:
587 # NOTICE messages have a
588 # newline character at the end
589 logger.info(notice.rstrip())
590
591 cursor.connection.notices[:] = []
592
593
594class PGIdentifierPreparer_psycopg2(PGIdentifierPreparer):
595 pass
596
597
598class ExecutemanyMode(FastIntFlag):
599 EXECUTEMANY_VALUES = 0
600 EXECUTEMANY_VALUES_PLUS_BATCH = 1
601
602
603(
604 EXECUTEMANY_VALUES,
605 EXECUTEMANY_VALUES_PLUS_BATCH,
606) = ExecutemanyMode.__members__.values()
607
608
609class PGDialect_psycopg2(_PGDialect_common_psycopg):
610 driver = "psycopg2"
611
612 supports_statement_cache = True
613 supports_server_side_cursors = True
614
615 default_paramstyle = "pyformat"
616 # set to true based on psycopg2 version
617 supports_sane_multi_rowcount = False
618 execution_ctx_cls = PGExecutionContext_psycopg2
619 preparer = PGIdentifierPreparer_psycopg2
620 psycopg2_version = (0, 0)
621 use_insertmanyvalues_wo_returning = True
622
623 returns_native_bytes = False
624
625 _has_native_hstore = True
626
627 colspecs = util.update_copy(
628 _PGDialect_common_psycopg.colspecs,
629 {
630 JSON: _PGJSON,
631 sqltypes.JSON: _PGJSON,
632 JSONB: _PGJSONB,
633 ranges.INT4RANGE: _Psycopg2NumericRange,
634 ranges.INT8RANGE: _Psycopg2NumericRange,
635 ranges.NUMRANGE: _Psycopg2NumericRange,
636 ranges.DATERANGE: _Psycopg2DateRange,
637 ranges.TSRANGE: _Psycopg2DateTimeRange,
638 ranges.TSTZRANGE: _Psycopg2DateTimeTZRange,
639 },
640 )
641
642 def __init__(
643 self,
644 executemany_mode="values_only",
645 executemany_batch_page_size=100,
646 **kwargs,
647 ):
648 _PGDialect_common_psycopg.__init__(self, **kwargs)
649
650 if self._native_inet_types:
651 raise NotImplementedError(
652 "The psycopg2 dialect does not implement "
653 "ipaddress type handling; native_inet_types cannot be set "
654 "to ``True`` when using this dialect."
655 )
656
657 # Parse executemany_mode argument, allowing it to be only one of the
658 # symbol names
659 self.executemany_mode = parse_user_argument_for_enum(
660 executemany_mode,
661 {
662 EXECUTEMANY_VALUES: ["values_only"],
663 EXECUTEMANY_VALUES_PLUS_BATCH: ["values_plus_batch"],
664 },
665 "executemany_mode",
666 )
667
668 self.executemany_batch_page_size = executemany_batch_page_size
669
670 if self.dbapi and hasattr(self.dbapi, "__version__"):
671 m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", self.dbapi.__version__)
672 if m:
673 self.psycopg2_version = tuple(
674 int(x) for x in m.group(1, 2, 3) if x is not None
675 )
676
677 if self.psycopg2_version < (2, 7):
678 raise ImportError(
679 "psycopg2 version 2.7 or higher is required."
680 )
681
682 def initialize(self, connection):
683 super().initialize(connection)
684 self._has_native_hstore = (
685 self.use_native_hstore
686 and self._hstore_oids(connection.connection.dbapi_connection)
687 is not None
688 )
689
690 self.supports_sane_multi_rowcount = (
691 self.executemany_mode is not EXECUTEMANY_VALUES_PLUS_BATCH
692 )
693
694 @classmethod
695 def import_dbapi(cls):
696 import psycopg2
697
698 return psycopg2
699
700 @util.memoized_property
701 def _psycopg2_extensions(cls):
702 from psycopg2 import extensions
703
704 return extensions
705
706 @util.memoized_property
707 def _psycopg2_extras(cls):
708 from psycopg2 import extras
709
710 return extras
711
712 @util.memoized_property
713 def _isolation_lookup(self):
714 extensions = self._psycopg2_extensions
715 return {
716 "AUTOCOMMIT": extensions.ISOLATION_LEVEL_AUTOCOMMIT,
717 "READ COMMITTED": extensions.ISOLATION_LEVEL_READ_COMMITTED,
718 "READ UNCOMMITTED": extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
719 "REPEATABLE READ": extensions.ISOLATION_LEVEL_REPEATABLE_READ,
720 "SERIALIZABLE": extensions.ISOLATION_LEVEL_SERIALIZABLE,
721 }
722
723 def set_isolation_level(self, dbapi_connection, level):
724 dbapi_connection.set_isolation_level(self._isolation_lookup[level])
725
726 def set_readonly(self, connection, value):
727 connection.readonly = value
728
729 def get_readonly(self, connection):
730 return connection.readonly
731
732 def set_deferrable(self, connection, value):
733 connection.deferrable = value
734
735 def get_deferrable(self, connection):
736 return connection.deferrable
737
738 def on_connect(self):
739 extras = self._psycopg2_extras
740
741 fns = []
742 if self.client_encoding is not None:
743
744 def on_connect(dbapi_conn):
745 dbapi_conn.set_client_encoding(self.client_encoding)
746
747 fns.append(on_connect)
748
749 if self.dbapi:
750
751 def on_connect(dbapi_conn):
752 extras.register_uuid(None, dbapi_conn)
753
754 fns.append(on_connect)
755
756 if self.dbapi and self.use_native_hstore:
757
758 def on_connect(dbapi_conn):
759 hstore_oids = self._hstore_oids(dbapi_conn)
760 if hstore_oids is not None:
761 oid, array_oid = hstore_oids
762 kw = {"oid": oid}
763 kw["array_oid"] = array_oid
764 extras.register_hstore(dbapi_conn, **kw)
765
766 fns.append(on_connect)
767
768 if self.dbapi and self._json_deserializer:
769
770 def on_connect(dbapi_conn):
771 extras.register_default_json(
772 dbapi_conn, loads=self._json_deserializer
773 )
774 extras.register_default_jsonb(
775 dbapi_conn, loads=self._json_deserializer
776 )
777
778 fns.append(on_connect)
779
780 if fns:
781
782 def on_connect(dbapi_conn):
783 for fn in fns:
784 fn(dbapi_conn)
785
786 return on_connect
787 else:
788 return None
789
790 def do_executemany(self, cursor, statement, parameters, context=None):
791 if self.executemany_mode is EXECUTEMANY_VALUES_PLUS_BATCH:
792 if self.executemany_batch_page_size:
793 kwargs = {"page_size": self.executemany_batch_page_size}
794 else:
795 kwargs = {}
796 self._psycopg2_extras.execute_batch(
797 cursor, statement, parameters, **kwargs
798 )
799 else:
800 cursor.executemany(statement, parameters)
801
802 def do_begin_twophase(self, connection, xid):
803 connection.connection.tpc_begin(xid)
804
805 def do_prepare_twophase(self, connection, xid):
806 connection.connection.tpc_prepare()
807
808 def _do_twophase(self, dbapi_conn, operation, xid, recover=False):
809 if recover:
810 if dbapi_conn.status != self._psycopg2_extensions.STATUS_READY:
811 dbapi_conn.rollback()
812 operation(xid)
813 else:
814 operation()
815
816 def do_rollback_twophase(
817 self, connection, xid, is_prepared=True, recover=False
818 ):
819 dbapi_conn = connection.connection.dbapi_connection
820 self._do_twophase(
821 dbapi_conn, dbapi_conn.tpc_rollback, xid, recover=recover
822 )
823
824 def do_commit_twophase(
825 self, connection, xid, is_prepared=True, recover=False
826 ):
827 dbapi_conn = connection.connection.dbapi_connection
828 self._do_twophase(
829 dbapi_conn, dbapi_conn.tpc_commit, xid, recover=recover
830 )
831
832 @util.memoized_instancemethod
833 def _hstore_oids(self, dbapi_connection):
834 extras = self._psycopg2_extras
835 oids = extras.HstoreAdapter.get_oids(dbapi_connection)
836 if oids is not None and oids[0]:
837 return oids[0:2]
838 else:
839 return None
840
841 def is_disconnect(self, e, connection, cursor):
842 if isinstance(e, self.dbapi.Error):
843 # check the "closed" flag. this might not be
844 # present on old psycopg2 versions. Also,
845 # this flag doesn't actually help in a lot of disconnect
846 # situations, so don't rely on it.
847 if getattr(connection, "closed", False):
848 return True
849
850 # checks based on strings. in the case that .closed
851 # didn't cut it, fall back onto these.
852 str_e = str(e).partition("\n")[0]
853 for msg in self._is_disconnect_messages:
854 idx = str_e.find(msg)
855 if idx >= 0 and '"' not in str_e[:idx]:
856 return True
857 return False
858
859 @util.memoized_property
860 def _is_disconnect_messages(self):
861 return (
862 # these error messages from libpq: interfaces/libpq/fe-misc.c
863 # and interfaces/libpq/fe-secure.c.
864 "terminating connection",
865 "closed the connection",
866 "connection not open",
867 "could not receive data from server",
868 "could not send data to server",
869 # psycopg2 client errors, psycopg2/connection.h,
870 # psycopg2/cursor.h
871 "connection already closed",
872 "cursor already closed",
873 # not sure where this path is originally from, it may
874 # be obsolete. It really says "losed", not "closed".
875 "losed the connection unexpectedly",
876 # these can occur in newer SSL
877 "connection has been closed unexpectedly",
878 "SSL error: decryption failed or bad record mac",
879 "SSL SYSCALL error: Bad file descriptor",
880 "SSL SYSCALL error: EOF detected",
881 "SSL SYSCALL error: Operation timed out",
882 "SSL SYSCALL error: Bad address",
883 # This can occur in OpenSSL 1 when an unexpected EOF occurs.
884 # https://www.openssl.org/docs/man1.1.1/man3/SSL_get_error.html#BUGS
885 # It may also occur in newer OpenSSL for a non-recoverable I/O
886 # error as a result of a system call that does not set 'errno'
887 # in libc.
888 "SSL SYSCALL error: Success",
889 )
890
891
892dialect = PGDialect_psycopg2
893