psycopg2.py 31 KB

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