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

1061 lines
38KB

  1. # postgresql/psycopg2.py
  2. # Copyright (C) 2005-2021 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: http://www.opensource.org/licenses/mit-license.php
  7. r"""
  8. .. dialect:: postgresql+psycopg2
  9. :name: psycopg2
  10. :dbapi: psycopg2
  11. :connectstring: postgresql+psycopg2://user:password@host:port/dbname[?key=value&key=value...]
  12. :url: http://pypi.python.org/pypi/psycopg2/
  13. psycopg2 Connect Arguments
  14. --------------------------
  15. Keyword arguments that are specific to the SQLAlchemy psycopg2 dialect
  16. may be passed to :func:`_sa.create_engine()`, and include the following:
  17. * ``isolation_level``: This option, available for all PostgreSQL dialects,
  18. includes the ``AUTOCOMMIT`` isolation level when using the psycopg2
  19. dialect. This option sets the **default** isolation level for the
  20. connection that is set immediately upon connection to the database before
  21. the connection is pooled. This option is generally superseded by the more
  22. modern :paramref:`_engine.Connection.execution_options.isolation_level`
  23. execution option, detailed at :ref:`dbapi_autocommit`.
  24. .. seealso::
  25. :ref:`psycopg2_isolation_level`
  26. :ref:`dbapi_autocommit`
  27. * ``client_encoding``: sets the client encoding in a libpq-agnostic way,
  28. using psycopg2's ``set_client_encoding()`` method.
  29. .. seealso::
  30. :ref:`psycopg2_unicode`
  31. * ``use_native_unicode``: Under Python 2 only, this can be set to False to
  32. disable the use of psycopg2's native Unicode support.
  33. .. seealso::
  34. :ref:`psycopg2_disable_native_unicode`
  35. * ``executemany_mode``, ``executemany_batch_page_size``,
  36. ``executemany_values_page_size``: Allows use of psycopg2
  37. extensions for optimizing "executemany"-stye queries. See the referenced
  38. section below for details.
  39. .. seealso::
  40. :ref:`psycopg2_executemany_mode`
  41. .. tip::
  42. The above keyword arguments are **dialect** keyword arguments, meaning
  43. that they are passed as explicit keyword arguments to :func:`_sa.create_engine()`::
  44. engine = create_engine(
  45. "postgresql+psycopg2://scott:tiger@localhost/test",
  46. isolation_level="SERIALIZABLE",
  47. )
  48. These should not be confused with **DBAPI** connect arguments, which
  49. are passed as part of the :paramref:`_sa.create_engine.connect_args`
  50. dictionary and/or are passed in the URL query string, as detailed in
  51. the section :ref:`custom_dbapi_args`.
  52. .. _psycopg2_ssl:
  53. SSL Connections
  54. ---------------
  55. The psycopg2 module has a connection argument named ``sslmode`` for
  56. controlling its behavior regarding secure (SSL) connections. The default is
  57. ``sslmode=prefer``; it will attempt an SSL connection and if that fails it
  58. will fall back to an unencrypted connection. ``sslmode=require`` may be used
  59. to ensure that only secure connections are established. Consult the
  60. psycopg2 / libpq documentation for further options that are available.
  61. Note that ``sslmode`` is specific to psycopg2 so it is included in the
  62. connection URI::
  63. engine = sa.create_engine(
  64. "postgresql+psycopg2://scott:tiger@192.168.0.199:5432/test?sslmode=require"
  65. )
  66. Unix Domain Connections
  67. ------------------------
  68. psycopg2 supports connecting via Unix domain connections. When the ``host``
  69. portion of the URL is omitted, SQLAlchemy passes ``None`` to psycopg2,
  70. which specifies Unix-domain communication rather than TCP/IP communication::
  71. create_engine("postgresql+psycopg2://user:password@/dbname")
  72. By default, the socket file used is to connect to a Unix-domain socket
  73. in ``/tmp``, or whatever socket directory was specified when PostgreSQL
  74. was built. This value can be overridden by passing a pathname to psycopg2,
  75. using ``host`` as an additional keyword argument::
  76. create_engine("postgresql+psycopg2://user:password@/dbname?host=/var/lib/postgresql")
  77. .. seealso::
  78. `PQconnectdbParams \
  79. <http://www.postgresql.org/docs/9.1/static/libpq-connect.html#LIBPQ-PQCONNECTDBPARAMS>`_
  80. .. _psycopg2_multi_host:
  81. Specifying multiple fallback hosts
  82. -----------------------------------
  83. psycopg2 supports multiple connection points in the connection string.
  84. When the ``host`` parameter is used multiple times in the query section of
  85. the URL, SQLAlchemy will create a single string of the host and port
  86. information provided to make the connections::
  87. create_engine(
  88. "postgresql+psycopg2://user:password@/dbname?host=HostA:port1&host=HostB&host=HostC"
  89. )
  90. A connection to each host is then attempted until either a connection is successful
  91. or all connections are unsuccessful in which case an error is raised.
  92. .. versionadded:: 1.3.20 Support for multiple hosts in PostgreSQL connection
  93. string.
  94. .. seealso::
  95. `PQConnString \
  96. <https://www.postgresql.org/docs/10/libpq-connect.html#LIBPQ-CONNSTRING>`_
  97. Empty DSN Connections / Environment Variable Connections
  98. ---------------------------------------------------------
  99. The psycopg2 DBAPI can connect to PostgreSQL by passing an empty DSN to the
  100. libpq client library, which by default indicates to connect to a localhost
  101. PostgreSQL database that is open for "trust" connections. This behavior can be
  102. further tailored using a particular set of environment variables which are
  103. prefixed with ``PG_...``, which are consumed by ``libpq`` to take the place of
  104. any or all elements of the connection string.
  105. For this form, the URL can be passed without any elements other than the
  106. initial scheme::
  107. engine = create_engine('postgresql+psycopg2://')
  108. In the above form, a blank "dsn" string is passed to the ``psycopg2.connect()``
  109. function which in turn represents an empty DSN passed to libpq.
  110. .. versionadded:: 1.3.2 support for parameter-less connections with psycopg2.
  111. .. seealso::
  112. `Environment Variables\
  113. <https://www.postgresql.org/docs/current/libpq-envars.html>`_ -
  114. PostgreSQL documentation on how to use ``PG_...``
  115. environment variables for connections.
  116. .. _psycopg2_execution_options:
  117. Per-Statement/Connection Execution Options
  118. -------------------------------------------
  119. The following DBAPI-specific options are respected when used with
  120. :meth:`_engine.Connection.execution_options`,
  121. :meth:`.Executable.execution_options`,
  122. :meth:`_query.Query.execution_options`,
  123. in addition to those not specific to DBAPIs:
  124. * ``isolation_level`` - Set the transaction isolation level for the lifespan
  125. of a :class:`_engine.Connection` (can only be set on a connection,
  126. not a statement
  127. or query). See :ref:`psycopg2_isolation_level`.
  128. * ``stream_results`` - Enable or disable usage of psycopg2 server side
  129. cursors - this feature makes use of "named" cursors in combination with
  130. special result handling methods so that result rows are not fully buffered.
  131. Defaults to False, meaning cursors are buffered by default.
  132. * ``max_row_buffer`` - when using ``stream_results``, an integer value that
  133. specifies the maximum number of rows to buffer at a time. This is
  134. interpreted by the :class:`.BufferedRowCursorResult`, and if omitted the
  135. buffer will grow to ultimately store 1000 rows at a time.
  136. .. versionchanged:: 1.4 The ``max_row_buffer`` size can now be greater than
  137. 1000, and the buffer will grow to that size.
  138. .. _psycopg2_batch_mode:
  139. .. _psycopg2_executemany_mode:
  140. Psycopg2 Fast Execution Helpers
  141. -------------------------------
  142. Modern versions of psycopg2 include a feature known as
  143. `Fast Execution Helpers \
  144. <http://initd.org/psycopg/docs/extras.html#fast-execution-helpers>`_, which
  145. have been shown in benchmarking to improve psycopg2's executemany()
  146. performance, primarily with INSERT statements, by multiple orders of magnitude.
  147. SQLAlchemy internally makes use of these extensions for ``executemany()`` style
  148. calls, which correspond to lists of parameters being passed to
  149. :meth:`_engine.Connection.execute` as detailed in :ref:`multiple parameter
  150. sets <execute_multiple>`. The ORM also uses this mode internally whenever
  151. possible.
  152. The two available extensions on the psycopg2 side are the ``execute_values()``
  153. and ``execute_batch()`` functions. The psycopg2 dialect defaults to using the
  154. ``execute_values()`` extension for all qualifying INSERT statements.
  155. .. versionchanged:: 1.4 The psycopg2 dialect now defaults to a new mode
  156. ``"values_only"`` for ``executemany_mode``, which allows an order of
  157. magnitude performance improvement for INSERT statements, but does not
  158. include "batch" mode for UPDATE and DELETE statements which removes the
  159. ability of ``cursor.rowcount`` to function correctly.
  160. The use of these extensions is controlled by the ``executemany_mode`` flag
  161. which may be passed to :func:`_sa.create_engine`::
  162. engine = create_engine(
  163. "postgresql+psycopg2://scott:tiger@host/dbname",
  164. executemany_mode='values_plus_batch')
  165. Possible options for ``executemany_mode`` include:
  166. * ``values_only`` - this is the default value. the psycopg2 execute_values()
  167. extension is used for qualifying INSERT statements, which rewrites the INSERT
  168. to include multiple VALUES clauses so that many parameter sets can be
  169. inserted with one statement.
  170. .. versionadded:: 1.4 Added ``"values_only"`` setting for ``executemany_mode``
  171. which is also now the default.
  172. * ``None`` - No psycopg2 extensions are not used, and the usual
  173. ``cursor.executemany()`` method is used when invoking statements with
  174. multiple parameter sets.
  175. * ``'batch'`` - Uses ``psycopg2.extras.execute_batch`` for all qualifying
  176. INSERT, UPDATE and DELETE statements, so that multiple copies
  177. of a SQL query, each one corresponding to a parameter set passed to
  178. ``executemany()``, are joined into a single SQL string separated by a
  179. semicolon. When using this mode, the :attr:`_engine.CursorResult.rowcount`
  180. attribute will not contain a value for executemany-style executions.
  181. * ``'values_plus_batch'``- ``execute_values`` is used for qualifying INSERT
  182. statements, ``execute_batch`` is used for UPDATE and DELETE.
  183. When using this mode, the :attr:`_engine.CursorResult.rowcount`
  184. attribute will not contain a value for executemany-style executions against
  185. UPDATE and DELETE statements.
  186. By "qualifying statements", we mean that the statement being executed
  187. must be a Core :func:`_expression.insert`, :func:`_expression.update`
  188. or :func:`_expression.delete` construct, and not a plain textual SQL
  189. string or one constructed using :func:`_expression.text`. When using the
  190. ORM, all insert/update/delete statements used by the ORM flush process
  191. are qualifying.
  192. The "page size" for the "values" and "batch" strategies can be affected
  193. by using the ``executemany_batch_page_size`` and
  194. ``executemany_values_page_size`` engine parameters. These
  195. control how many parameter sets
  196. should be represented in each execution. The "values" page size defaults
  197. to 1000, which is different that psycopg2's default. The "batch" page
  198. size defaults to 100. These can be affected by passing new values to
  199. :func:`_engine.create_engine`::
  200. engine = create_engine(
  201. "postgresql+psycopg2://scott:tiger@host/dbname",
  202. executemany_mode='values',
  203. executemany_values_page_size=10000, executemany_batch_page_size=500)
  204. .. versionchanged:: 1.4
  205. The default for ``executemany_values_page_size`` is now 1000, up from
  206. 100.
  207. .. seealso::
  208. :ref:`execute_multiple` - General information on using the
  209. :class:`_engine.Connection`
  210. object to execute statements in such a way as to make
  211. use of the DBAPI ``.executemany()`` method.
  212. .. _psycopg2_unicode:
  213. Unicode with Psycopg2
  214. ----------------------
  215. The psycopg2 DBAPI driver supports Unicode data transparently. Under Python 2
  216. only, the SQLAlchemy psycopg2 dialect will enable the
  217. ``psycopg2.extensions.UNICODE`` extension by default to ensure Unicode is
  218. handled properly; under Python 3, this is psycopg2's default behavior.
  219. The client character encoding can be controlled for the psycopg2 dialect
  220. in the following ways:
  221. * For PostgreSQL 9.1 and above, the ``client_encoding`` parameter may be
  222. passed in the database URL; this parameter is consumed by the underlying
  223. ``libpq`` PostgreSQL client library::
  224. engine = create_engine("postgresql+psycopg2://user:pass@host/dbname?client_encoding=utf8")
  225. Alternatively, the above ``client_encoding`` value may be passed using
  226. :paramref:`_sa.create_engine.connect_args` for programmatic establishment with
  227. ``libpq``::
  228. engine = create_engine(
  229. "postgresql+psycopg2://user:pass@host/dbname",
  230. connect_args={'client_encoding': 'utf8'}
  231. )
  232. * For all PostgreSQL versions, psycopg2 supports a client-side encoding
  233. value that will be passed to database connections when they are first
  234. established. The SQLAlchemy psycopg2 dialect supports this using the
  235. ``client_encoding`` parameter passed to :func:`_sa.create_engine`::
  236. engine = create_engine(
  237. "postgresql+psycopg2://user:pass@host/dbname",
  238. client_encoding="utf8"
  239. )
  240. .. tip:: The above ``client_encoding`` parameter admittedly is very similar
  241. in appearance to usage of the parameter within the
  242. :paramref:`_sa.create_engine.connect_args` dictionary; the difference
  243. above is that the parameter is consumed by psycopg2 and is
  244. passed to the database connection using ``SET client_encoding TO
  245. 'utf8'``; in the previously mentioned style, the parameter is instead
  246. passed through psycopg2 and consumed by the ``libpq`` library.
  247. * A common way to set up client encoding with PostgreSQL databases is to
  248. ensure it is configured within the server-side postgresql.conf file;
  249. this is the recommended way to set encoding for a server that is
  250. consistently of one encoding in all databases::
  251. # postgresql.conf file
  252. # client_encoding = sql_ascii # actually, defaults to database
  253. # encoding
  254. client_encoding = utf8
  255. .. _psycopg2_disable_native_unicode:
  256. Disabling Native Unicode
  257. ^^^^^^^^^^^^^^^^^^^^^^^^
  258. Under Python 2 only, SQLAlchemy can also be instructed to skip the usage of the
  259. psycopg2 ``UNICODE`` extension and to instead utilize its own unicode
  260. encode/decode services, which are normally reserved only for those DBAPIs that
  261. don't fully support unicode directly. Passing ``use_native_unicode=False`` to
  262. :func:`_sa.create_engine` will disable usage of ``psycopg2.extensions.
  263. UNICODE``. SQLAlchemy will instead encode data itself into Python bytestrings
  264. on the way in and coerce from bytes on the way back, using the value of the
  265. :func:`_sa.create_engine` ``encoding`` parameter, which defaults to ``utf-8``.
  266. SQLAlchemy's own unicode encode/decode functionality is steadily becoming
  267. obsolete as most DBAPIs now support unicode fully.
  268. Transactions
  269. ------------
  270. The psycopg2 dialect fully supports SAVEPOINT and two-phase commit operations.
  271. .. _psycopg2_isolation_level:
  272. Psycopg2 Transaction Isolation Level
  273. -------------------------------------
  274. As discussed in :ref:`postgresql_isolation_level`,
  275. all PostgreSQL dialects support setting of transaction isolation level
  276. both via the ``isolation_level`` parameter passed to :func:`_sa.create_engine`
  277. ,
  278. as well as the ``isolation_level`` argument used by
  279. :meth:`_engine.Connection.execution_options`. When using the psycopg2 dialect
  280. , these
  281. options make use of psycopg2's ``set_isolation_level()`` connection method,
  282. rather than emitting a PostgreSQL directive; this is because psycopg2's
  283. API-level setting is always emitted at the start of each transaction in any
  284. case.
  285. The psycopg2 dialect supports these constants for isolation level:
  286. * ``READ COMMITTED``
  287. * ``READ UNCOMMITTED``
  288. * ``REPEATABLE READ``
  289. * ``SERIALIZABLE``
  290. * ``AUTOCOMMIT``
  291. .. seealso::
  292. :ref:`postgresql_isolation_level`
  293. :ref:`pg8000_isolation_level`
  294. NOTICE logging
  295. ---------------
  296. The psycopg2 dialect will log PostgreSQL NOTICE messages
  297. via the ``sqlalchemy.dialects.postgresql`` logger. When this logger
  298. is set to the ``logging.INFO`` level, notice messages will be logged::
  299. import logging
  300. logging.getLogger('sqlalchemy.dialects.postgresql').setLevel(logging.INFO)
  301. Above, it is assumed that logging is configured externally. If this is not
  302. the case, configuration such as ``logging.basicConfig()`` must be utilized::
  303. import logging
  304. logging.basicConfig() # log messages to stdout
  305. logging.getLogger('sqlalchemy.dialects.postgresql').setLevel(logging.INFO)
  306. .. seealso::
  307. `Logging HOWTO <https://docs.python.org/3/howto/logging.html>`_ - on the python.org website
  308. .. _psycopg2_hstore:
  309. HSTORE type
  310. ------------
  311. The ``psycopg2`` DBAPI includes an extension to natively handle marshalling of
  312. the HSTORE type. The SQLAlchemy psycopg2 dialect will enable this extension
  313. by default when psycopg2 version 2.4 or greater is used, and
  314. it is detected that the target database has the HSTORE type set up for use.
  315. In other words, when the dialect makes the first
  316. connection, a sequence like the following is performed:
  317. 1. Request the available HSTORE oids using
  318. ``psycopg2.extras.HstoreAdapter.get_oids()``.
  319. If this function returns a list of HSTORE identifiers, we then determine
  320. that the ``HSTORE`` extension is present.
  321. This function is **skipped** if the version of psycopg2 installed is
  322. less than version 2.4.
  323. 2. If the ``use_native_hstore`` flag is at its default of ``True``, and
  324. we've detected that ``HSTORE`` oids are available, the
  325. ``psycopg2.extensions.register_hstore()`` extension is invoked for all
  326. connections.
  327. The ``register_hstore()`` extension has the effect of **all Python
  328. dictionaries being accepted as parameters regardless of the type of target
  329. column in SQL**. The dictionaries are converted by this extension into a
  330. textual HSTORE expression. If this behavior is not desired, disable the
  331. use of the hstore extension by setting ``use_native_hstore`` to ``False`` as
  332. follows::
  333. engine = create_engine("postgresql+psycopg2://scott:tiger@localhost/test",
  334. use_native_hstore=False)
  335. The ``HSTORE`` type is **still supported** when the
  336. ``psycopg2.extensions.register_hstore()`` extension is not used. It merely
  337. means that the coercion between Python dictionaries and the HSTORE
  338. string format, on both the parameter side and the result side, will take
  339. place within SQLAlchemy's own marshalling logic, and not that of ``psycopg2``
  340. which may be more performant.
  341. """ # noqa
  342. from __future__ import absolute_import
  343. import decimal
  344. import logging
  345. import re
  346. from uuid import UUID as _python_UUID
  347. from .base import _DECIMAL_TYPES
  348. from .base import _FLOAT_TYPES
  349. from .base import _INT_TYPES
  350. from .base import ENUM
  351. from .base import PGCompiler
  352. from .base import PGDialect
  353. from .base import PGExecutionContext
  354. from .base import PGIdentifierPreparer
  355. from .base import UUID
  356. from .hstore import HSTORE
  357. from .json import JSON
  358. from .json import JSONB
  359. from ... import exc
  360. from ... import processors
  361. from ... import types as sqltypes
  362. from ... import util
  363. from ...engine import cursor as _cursor
  364. from ...sql import elements
  365. from ...util import collections_abc
  366. logger = logging.getLogger("sqlalchemy.dialects.postgresql")
  367. class _PGNumeric(sqltypes.Numeric):
  368. def bind_processor(self, dialect):
  369. return None
  370. def result_processor(self, dialect, coltype):
  371. if self.asdecimal:
  372. if coltype in _FLOAT_TYPES:
  373. return processors.to_decimal_processor_factory(
  374. decimal.Decimal, self._effective_decimal_return_scale
  375. )
  376. elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
  377. # pg8000 returns Decimal natively for 1700
  378. return None
  379. else:
  380. raise exc.InvalidRequestError(
  381. "Unknown PG numeric type: %d" % coltype
  382. )
  383. else:
  384. if coltype in _FLOAT_TYPES:
  385. # pg8000 returns float natively for 701
  386. return None
  387. elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
  388. return processors.to_float
  389. else:
  390. raise exc.InvalidRequestError(
  391. "Unknown PG numeric type: %d" % coltype
  392. )
  393. class _PGEnum(ENUM):
  394. def result_processor(self, dialect, coltype):
  395. if util.py2k and self._expect_unicode is True:
  396. # for py2k, if the enum type needs unicode data (which is set up as
  397. # part of the Enum() constructor based on values passed as py2k
  398. # unicode objects) we have to use our own converters since
  399. # psycopg2's don't work, a rare exception to the "modern DBAPIs
  400. # support unicode everywhere" theme of deprecating
  401. # convert_unicode=True. Use the special "force_nocheck" directive
  402. # which forces unicode conversion to happen on the Python side
  403. # without an isinstance() check. in py3k psycopg2 does the right
  404. # thing automatically.
  405. self._expect_unicode = "force_nocheck"
  406. return super(_PGEnum, self).result_processor(dialect, coltype)
  407. class _PGHStore(HSTORE):
  408. def bind_processor(self, dialect):
  409. if dialect._has_native_hstore:
  410. return None
  411. else:
  412. return super(_PGHStore, self).bind_processor(dialect)
  413. def result_processor(self, dialect, coltype):
  414. if dialect._has_native_hstore:
  415. return None
  416. else:
  417. return super(_PGHStore, self).result_processor(dialect, coltype)
  418. class _PGJSON(JSON):
  419. def result_processor(self, dialect, coltype):
  420. return None
  421. class _PGJSONB(JSONB):
  422. def result_processor(self, dialect, coltype):
  423. return None
  424. class _PGUUID(UUID):
  425. def bind_processor(self, dialect):
  426. if not self.as_uuid and dialect.use_native_uuid:
  427. def process(value):
  428. if value is not None:
  429. value = _python_UUID(value)
  430. return value
  431. return process
  432. def result_processor(self, dialect, coltype):
  433. if not self.as_uuid and dialect.use_native_uuid:
  434. def process(value):
  435. if value is not None:
  436. value = str(value)
  437. return value
  438. return process
  439. _server_side_id = util.counter()
  440. class PGExecutionContext_psycopg2(PGExecutionContext):
  441. _psycopg2_fetched_rows = None
  442. def create_server_side_cursor(self):
  443. # use server-side cursors:
  444. # http://lists.initd.org/pipermail/psycopg/2007-January/005251.html
  445. ident = "c_%s_%s" % (hex(id(self))[2:], hex(_server_side_id())[2:])
  446. return self._dbapi_connection.cursor(ident)
  447. def post_exec(self):
  448. if (
  449. self._psycopg2_fetched_rows
  450. and self.compiled
  451. and self.compiled.returning
  452. ):
  453. # psycopg2 execute_values will provide for a real cursor where
  454. # cursor.description works correctly. however, it executes the
  455. # INSERT statement multiple times for multiple pages of rows, so
  456. # while this cursor also supports calling .fetchall() directly, in
  457. # order to get the list of all rows inserted across multiple pages,
  458. # we have to retrieve the aggregated list from the execute_values()
  459. # function directly.
  460. strat_cls = _cursor.FullyBufferedCursorFetchStrategy
  461. self.cursor_fetch_strategy = strat_cls(
  462. self.cursor, initial_buffer=self._psycopg2_fetched_rows
  463. )
  464. self._log_notices(self.cursor)
  465. def _log_notices(self, cursor):
  466. # check also that notices is an iterable, after it's already
  467. # established that we will be iterating through it. This is to get
  468. # around test suites such as SQLAlchemy's using a Mock object for
  469. # cursor
  470. if not cursor.connection.notices or not isinstance(
  471. cursor.connection.notices, collections_abc.Iterable
  472. ):
  473. return
  474. for notice in cursor.connection.notices:
  475. # NOTICE messages have a
  476. # newline character at the end
  477. logger.info(notice.rstrip())
  478. cursor.connection.notices[:] = []
  479. class PGCompiler_psycopg2(PGCompiler):
  480. def visit_bindparam(self, bindparam, skip_bind_expression=False, **kw):
  481. text = super(PGCompiler_psycopg2, self).visit_bindparam(
  482. bindparam, skip_bind_expression=skip_bind_expression, **kw
  483. )
  484. # note that if the type has a bind_expression(), we will get a
  485. # double compile here
  486. if not skip_bind_expression and (
  487. bindparam.type._is_array or bindparam.type._is_type_decorator
  488. ):
  489. typ = bindparam.type._unwrapped_dialect_impl(self.dialect)
  490. if typ._is_array:
  491. text += "::%s" % (
  492. elements.TypeClause(typ)._compiler_dispatch(
  493. self, skip_bind_expression=skip_bind_expression, **kw
  494. ),
  495. )
  496. return text
  497. class PGIdentifierPreparer_psycopg2(PGIdentifierPreparer):
  498. pass
  499. EXECUTEMANY_PLAIN = util.symbol("executemany_plain", canonical=0)
  500. EXECUTEMANY_BATCH = util.symbol("executemany_batch", canonical=1)
  501. EXECUTEMANY_VALUES = util.symbol("executemany_values", canonical=2)
  502. EXECUTEMANY_VALUES_PLUS_BATCH = util.symbol(
  503. "executemany_values_plus_batch",
  504. canonical=EXECUTEMANY_BATCH | EXECUTEMANY_VALUES,
  505. )
  506. class PGDialect_psycopg2(PGDialect):
  507. driver = "psycopg2"
  508. supports_statement_cache = True
  509. if util.py2k:
  510. # turn off supports_unicode_statements for Python 2. psycopg2 supports
  511. # unicode statements in Py2K. But! it does not support unicode *bound
  512. # parameter names* because it uses the Python "%" operator to
  513. # interpolate these into the string, and this fails. So for Py2K, we
  514. # have to use full-on encoding for statements and parameters before
  515. # passing to cursor.execute().
  516. supports_unicode_statements = False
  517. supports_server_side_cursors = True
  518. default_paramstyle = "pyformat"
  519. # set to true based on psycopg2 version
  520. supports_sane_multi_rowcount = False
  521. execution_ctx_cls = PGExecutionContext_psycopg2
  522. statement_compiler = PGCompiler_psycopg2
  523. preparer = PGIdentifierPreparer_psycopg2
  524. psycopg2_version = (0, 0)
  525. _has_native_hstore = True
  526. engine_config_types = PGDialect.engine_config_types.union(
  527. {"use_native_unicode": util.asbool}
  528. )
  529. colspecs = util.update_copy(
  530. PGDialect.colspecs,
  531. {
  532. sqltypes.Numeric: _PGNumeric,
  533. ENUM: _PGEnum, # needs force_unicode
  534. sqltypes.Enum: _PGEnum, # needs force_unicode
  535. HSTORE: _PGHStore,
  536. JSON: _PGJSON,
  537. sqltypes.JSON: _PGJSON,
  538. JSONB: _PGJSONB,
  539. UUID: _PGUUID,
  540. },
  541. )
  542. def __init__(
  543. self,
  544. use_native_unicode=True,
  545. client_encoding=None,
  546. use_native_hstore=True,
  547. use_native_uuid=True,
  548. executemany_mode="values_only",
  549. executemany_batch_page_size=100,
  550. executemany_values_page_size=1000,
  551. **kwargs
  552. ):
  553. PGDialect.__init__(self, **kwargs)
  554. self.use_native_unicode = use_native_unicode
  555. if not use_native_unicode and not util.py2k:
  556. raise exc.ArgumentError(
  557. "psycopg2 native_unicode mode is required under Python 3"
  558. )
  559. if not use_native_hstore:
  560. self._has_native_hstore = False
  561. self.use_native_hstore = use_native_hstore
  562. self.use_native_uuid = use_native_uuid
  563. self.supports_unicode_binds = use_native_unicode
  564. self.client_encoding = client_encoding
  565. # Parse executemany_mode argument, allowing it to be only one of the
  566. # symbol names
  567. self.executemany_mode = util.symbol.parse_user_argument(
  568. executemany_mode,
  569. {
  570. EXECUTEMANY_PLAIN: [None],
  571. EXECUTEMANY_BATCH: ["batch"],
  572. EXECUTEMANY_VALUES: ["values_only"],
  573. EXECUTEMANY_VALUES_PLUS_BATCH: ["values_plus_batch", "values"],
  574. },
  575. "executemany_mode",
  576. )
  577. if self.executemany_mode & EXECUTEMANY_VALUES:
  578. self.insert_executemany_returning = True
  579. self.executemany_batch_page_size = executemany_batch_page_size
  580. self.executemany_values_page_size = executemany_values_page_size
  581. if self.dbapi and hasattr(self.dbapi, "__version__"):
  582. m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", self.dbapi.__version__)
  583. if m:
  584. self.psycopg2_version = tuple(
  585. int(x) for x in m.group(1, 2, 3) if x is not None
  586. )
  587. if self.psycopg2_version < (2, 7):
  588. raise ImportError(
  589. "psycopg2 version 2.7 or higher is required."
  590. )
  591. def initialize(self, connection):
  592. super(PGDialect_psycopg2, self).initialize(connection)
  593. self._has_native_hstore = (
  594. self.use_native_hstore
  595. and self._hstore_oids(connection.connection) is not None
  596. )
  597. # PGDialect.initialize() checks server version for <= 8.2 and sets
  598. # this flag to False if so
  599. if not self.full_returning:
  600. self.insert_executemany_returning = False
  601. self.executemany_mode = EXECUTEMANY_PLAIN
  602. self.supports_sane_multi_rowcount = not (
  603. self.executemany_mode & EXECUTEMANY_BATCH
  604. )
  605. @classmethod
  606. def dbapi(cls):
  607. import psycopg2
  608. return psycopg2
  609. @classmethod
  610. def _psycopg2_extensions(cls):
  611. from psycopg2 import extensions
  612. return extensions
  613. @classmethod
  614. def _psycopg2_extras(cls):
  615. from psycopg2 import extras
  616. return extras
  617. @util.memoized_property
  618. def _isolation_lookup(self):
  619. extensions = self._psycopg2_extensions()
  620. return {
  621. "AUTOCOMMIT": extensions.ISOLATION_LEVEL_AUTOCOMMIT,
  622. "READ COMMITTED": extensions.ISOLATION_LEVEL_READ_COMMITTED,
  623. "READ UNCOMMITTED": extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
  624. "REPEATABLE READ": extensions.ISOLATION_LEVEL_REPEATABLE_READ,
  625. "SERIALIZABLE": extensions.ISOLATION_LEVEL_SERIALIZABLE,
  626. }
  627. def set_isolation_level(self, connection, level):
  628. try:
  629. level = self._isolation_lookup[level.replace("_", " ")]
  630. except KeyError as err:
  631. util.raise_(
  632. exc.ArgumentError(
  633. "Invalid value '%s' for isolation_level. "
  634. "Valid isolation levels for %s are %s"
  635. % (level, self.name, ", ".join(self._isolation_lookup))
  636. ),
  637. replace_context=err,
  638. )
  639. connection.set_isolation_level(level)
  640. def set_readonly(self, connection, value):
  641. connection.readonly = value
  642. def get_readonly(self, connection):
  643. return connection.readonly
  644. def set_deferrable(self, connection, value):
  645. connection.deferrable = value
  646. def get_deferrable(self, connection):
  647. return connection.deferrable
  648. def do_ping(self, dbapi_connection):
  649. cursor = None
  650. try:
  651. dbapi_connection.autocommit = True
  652. cursor = dbapi_connection.cursor()
  653. try:
  654. cursor.execute(self._dialect_specific_select_one)
  655. finally:
  656. cursor.close()
  657. if not dbapi_connection.closed:
  658. dbapi_connection.autocommit = False
  659. except self.dbapi.Error as err:
  660. if self.is_disconnect(err, dbapi_connection, cursor):
  661. return False
  662. else:
  663. raise
  664. else:
  665. return True
  666. def on_connect(self):
  667. extras = self._psycopg2_extras()
  668. extensions = self._psycopg2_extensions()
  669. fns = []
  670. if self.client_encoding is not None:
  671. def on_connect(conn):
  672. conn.set_client_encoding(self.client_encoding)
  673. fns.append(on_connect)
  674. if self.isolation_level is not None:
  675. def on_connect(conn):
  676. self.set_isolation_level(conn, self.isolation_level)
  677. fns.append(on_connect)
  678. if self.dbapi and self.use_native_uuid:
  679. def on_connect(conn):
  680. extras.register_uuid(None, conn)
  681. fns.append(on_connect)
  682. if util.py2k and self.dbapi and self.use_native_unicode:
  683. def on_connect(conn):
  684. extensions.register_type(extensions.UNICODE, conn)
  685. extensions.register_type(extensions.UNICODEARRAY, conn)
  686. fns.append(on_connect)
  687. if self.dbapi and self.use_native_hstore:
  688. def on_connect(conn):
  689. hstore_oids = self._hstore_oids(conn)
  690. if hstore_oids is not None:
  691. oid, array_oid = hstore_oids
  692. kw = {"oid": oid}
  693. if util.py2k:
  694. kw["unicode"] = True
  695. kw["array_oid"] = array_oid
  696. extras.register_hstore(conn, **kw)
  697. fns.append(on_connect)
  698. if self.dbapi and self._json_deserializer:
  699. def on_connect(conn):
  700. extras.register_default_json(
  701. conn, loads=self._json_deserializer
  702. )
  703. extras.register_default_jsonb(
  704. conn, loads=self._json_deserializer
  705. )
  706. fns.append(on_connect)
  707. if fns:
  708. def on_connect(conn):
  709. for fn in fns:
  710. fn(conn)
  711. return on_connect
  712. else:
  713. return None
  714. def do_executemany(self, cursor, statement, parameters, context=None):
  715. if (
  716. self.executemany_mode & EXECUTEMANY_VALUES
  717. and context
  718. and context.isinsert
  719. and context.compiled.insert_single_values_expr
  720. ):
  721. executemany_values = (
  722. "(%s)" % context.compiled.insert_single_values_expr
  723. )
  724. if not self.supports_unicode_statements:
  725. executemany_values = executemany_values.encode(self.encoding)
  726. # guard for statement that was altered via event hook or similar
  727. if executemany_values not in statement:
  728. executemany_values = None
  729. else:
  730. executemany_values = None
  731. if executemany_values:
  732. statement = statement.replace(executemany_values, "%s")
  733. if self.executemany_values_page_size:
  734. kwargs = {"page_size": self.executemany_values_page_size}
  735. else:
  736. kwargs = {}
  737. xtras = self._psycopg2_extras()
  738. context._psycopg2_fetched_rows = xtras.execute_values(
  739. cursor,
  740. statement,
  741. parameters,
  742. template=executemany_values,
  743. fetch=bool(context.compiled.returning),
  744. **kwargs
  745. )
  746. elif self.executemany_mode & EXECUTEMANY_BATCH:
  747. if self.executemany_batch_page_size:
  748. kwargs = {"page_size": self.executemany_batch_page_size}
  749. else:
  750. kwargs = {}
  751. self._psycopg2_extras().execute_batch(
  752. cursor, statement, parameters, **kwargs
  753. )
  754. else:
  755. cursor.executemany(statement, parameters)
  756. @util.memoized_instancemethod
  757. def _hstore_oids(self, conn):
  758. extras = self._psycopg2_extras()
  759. if hasattr(conn, "connection"):
  760. conn = conn.connection
  761. oids = extras.HstoreAdapter.get_oids(conn)
  762. if oids is not None and oids[0]:
  763. return oids[0:2]
  764. else:
  765. return None
  766. def create_connect_args(self, url):
  767. opts = url.translate_connect_args(username="user")
  768. is_multihost = False
  769. if "host" in url.query:
  770. is_multihost = isinstance(url.query["host"], (list, tuple))
  771. if opts:
  772. if "port" in opts:
  773. opts["port"] = int(opts["port"])
  774. opts.update(url.query)
  775. if is_multihost:
  776. opts["host"] = ",".join(url.query["host"])
  777. # send individual dbname, user, password, host, port
  778. # parameters to psycopg2.connect()
  779. return ([], opts)
  780. elif url.query:
  781. # any other connection arguments, pass directly
  782. opts.update(url.query)
  783. if is_multihost:
  784. opts["host"] = ",".join(url.query["host"])
  785. return ([], opts)
  786. else:
  787. # no connection arguments whatsoever; psycopg2.connect()
  788. # requires that "dsn" be present as a blank string.
  789. return ([""], opts)
  790. def is_disconnect(self, e, connection, cursor):
  791. if isinstance(e, self.dbapi.Error):
  792. # check the "closed" flag. this might not be
  793. # present on old psycopg2 versions. Also,
  794. # this flag doesn't actually help in a lot of disconnect
  795. # situations, so don't rely on it.
  796. if getattr(connection, "closed", False):
  797. return True
  798. # checks based on strings. in the case that .closed
  799. # didn't cut it, fall back onto these.
  800. str_e = str(e).partition("\n")[0]
  801. for msg in [
  802. # these error messages from libpq: interfaces/libpq/fe-misc.c
  803. # and interfaces/libpq/fe-secure.c.
  804. "terminating connection",
  805. "closed the connection",
  806. "connection not open",
  807. "could not receive data from server",
  808. "could not send data to server",
  809. # psycopg2 client errors, psycopg2/conenction.h,
  810. # psycopg2/cursor.h
  811. "connection already closed",
  812. "cursor already closed",
  813. # not sure where this path is originally from, it may
  814. # be obsolete. It really says "losed", not "closed".
  815. "losed the connection unexpectedly",
  816. # these can occur in newer SSL
  817. "connection has been closed unexpectedly",
  818. "SSL SYSCALL error: Bad file descriptor",
  819. "SSL SYSCALL error: EOF detected",
  820. "SSL error: decryption failed or bad record mac",
  821. "SSL SYSCALL error: Operation timed out",
  822. ]:
  823. idx = str_e.find(msg)
  824. if idx >= 0 and '"' not in str_e[:idx]:
  825. return True
  826. return False
  827. dialect = PGDialect_psycopg2