codekingpro/portable-devtools
114k
1# dialects/postgresql/psycopg.py
2# Copyright (C) 2005-2024 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+psycopg
11 :name: psycopg (a.k.a. psycopg 3)
12 :dbapi: psycopg
13 :connectstring: postgresql+psycopg://user:password@host:port/dbname[?key=value&key=value...]
14 :url: https://pypi.org/project/psycopg/
15
16``psycopg`` is the package and module name for version 3 of the ``psycopg``
17database driver, formerly known as ``psycopg2``. This driver is different
18enough from its ``psycopg2`` predecessor that SQLAlchemy supports it
19via a totally separate dialect; support for ``psycopg2`` is expected to remain
20for as long as that package continues to function for modern Python versions,
21and also remains the default dialect for the ``postgresql://`` dialect
22series.
23
24The SQLAlchemy ``psycopg`` dialect provides both a sync and an async
25implementation under the same dialect name. The proper version is
26selected depending on how the engine is created:
27
28* calling :func:`_sa.create_engine` with ``postgresql+psycopg://...`` will
29 automatically select the sync version, e.g.::
30
31 from sqlalchemy import create_engine
32 sync_engine = create_engine("postgresql+psycopg://scott:tiger@localhost/test")
33
34* calling :func:`_asyncio.create_async_engine` with
35 ``postgresql+psycopg://...`` will automatically select the async version,
36 e.g.::
37
38 from sqlalchemy.ext.asyncio import create_async_engine
39 asyncio_engine = create_async_engine("postgresql+psycopg://scott:tiger@localhost/test")
40
41The asyncio version of the dialect may also be specified explicitly using the
42``psycopg_async`` suffix, as::
43
44 from sqlalchemy.ext.asyncio import create_async_engine
45 asyncio_engine = create_async_engine("postgresql+psycopg_async://scott:tiger@localhost/test")
46
47.. seealso::
48
49 :ref:`postgresql_psycopg2` - The SQLAlchemy ``psycopg``
50 dialect shares most of its behavior with the ``psycopg2`` dialect.
51 Further documentation is available there.
52
53Using a different Cursor class
54------------------------------
55
56One of the differences between ``psycopg`` and the older ``psycopg2``
57is how bound parameters are handled: ``psycopg2`` would bind them
58client side, while ``psycopg`` by default will bind them server side.
59
60It's possible to configure ``psycopg`` to do client side binding by
61specifying the ``cursor_factory`` to be ``ClientCursor`` when creating
62the engine::
63
64 from psycopg import ClientCursor
65
66 client_side_engine = create_engine(
67 "postgresql+psycopg://...",
68 connect_args={"cursor_factory": ClientCursor},
69 )
70
71Similarly when using an async engine the ``AsyncClientCursor`` can be
72specified::
73
74 from psycopg import AsyncClientCursor
75
76 client_side_engine = create_async_engine(
77 "postgresql+psycopg://...",
78 connect_args={"cursor_factory": AsyncClientCursor},
79 )
80
81.. seealso::
82
83 `Client-side-binding cursors <https://www.psycopg.org/psycopg3/docs/advanced/cursors.html#client-side-binding-cursors>`_
84
85""" # noqa
86from __future__ import annotations
87
88import logging
89import re
90from typing import cast
91from typing import TYPE_CHECKING
92
93from . import ranges
94from ._psycopg_common import _PGDialect_common_psycopg
95from ._psycopg_common import _PGExecutionContext_common_psycopg
96from .base import INTERVAL
97from .base import PGCompiler
98from .base import PGIdentifierPreparer
99from .base import REGCONFIG
100from .json import JSON
101from .json import JSONB
102from .json import JSONPathType
103from .types import CITEXT
104from ... import pool
105from ... import util
106from ...engine import AdaptedConnection
107from ...sql import sqltypes
108from ...util.concurrency import await_fallback
109from ...util.concurrency import await_only
110
111if TYPE_CHECKING:
112 from typing import Iterable
113
114 from psycopg import AsyncConnection
115
116logger = logging.getLogger("sqlalchemy.dialects.postgresql")
117
118
119class _PGString(sqltypes.String):
120 render_bind_cast = True
121
122
123class _PGREGCONFIG(REGCONFIG):
124 render_bind_cast = True
125
126
127class _PGJSON(JSON):
128 render_bind_cast = True
129
130 def bind_processor(self, dialect):
131 return self._make_bind_processor(None, dialect._psycopg_Json)
132
133 def result_processor(self, dialect, coltype):
134 return None
135
136
137class _PGJSONB(JSONB):
138 render_bind_cast = True
139
140 def bind_processor(self, dialect):
141 return self._make_bind_processor(None, dialect._psycopg_Jsonb)
142
143 def result_processor(self, dialect, coltype):
144 return None
145
146
147class _PGJSONIntIndexType(sqltypes.JSON.JSONIntIndexType):
148 __visit_name__ = "json_int_index"
149
150 render_bind_cast = True
151
152
153class _PGJSONStrIndexType(sqltypes.JSON.JSONStrIndexType):
154 __visit_name__ = "json_str_index"
155
156 render_bind_cast = True
157
158
159class _PGJSONPathType(JSONPathType):
160 pass
161
162
163class _PGInterval(INTERVAL):
164 render_bind_cast = True
165
166
167class _PGTimeStamp(sqltypes.DateTime):
168 render_bind_cast = True
169
170
171class _PGDate(sqltypes.Date):
172 render_bind_cast = True
173
174
175class _PGTime(sqltypes.Time):
176 render_bind_cast = True
177
178
179class _PGInteger(sqltypes.Integer):
180 render_bind_cast = True
181
182
183class _PGSmallInteger(sqltypes.SmallInteger):
184 render_bind_cast = True
185
186
187class _PGNullType(sqltypes.NullType):
188 render_bind_cast = True
189
190
191class _PGBigInteger(sqltypes.BigInteger):
192 render_bind_cast = True
193
194
195class _PGBoolean(sqltypes.Boolean):
196 render_bind_cast = True
197
198
199class _PsycopgRange(ranges.AbstractSingleRangeImpl):
200 def bind_processor(self, dialect):
201 psycopg_Range = cast(PGDialect_psycopg, dialect)._psycopg_Range
202
203 def to_range(value):
204 if isinstance(value, ranges.Range):
205 value = psycopg_Range(
206 value.lower, value.upper, value.bounds, value.empty
207 )
208 return value
209
210 return to_range
211
212 def result_processor(self, dialect, coltype):
213 def to_range(value):
214 if value is not None:
215 value = ranges.Range(
216 value._lower,
217 value._upper,
218 bounds=value._bounds if value._bounds else "[)",
219 empty=not value._bounds,
220 )
221 return value
222
223 return to_range
224
225
226class _PsycopgMultiRange(ranges.AbstractMultiRangeImpl):
227 def bind_processor(self, dialect):
228 psycopg_Range = cast(PGDialect_psycopg, dialect)._psycopg_Range
229 psycopg_Multirange = cast(
230 PGDialect_psycopg, dialect
231 )._psycopg_Multirange
232
233 NoneType = type(None)
234
235 def to_range(value):
236 if isinstance(value, (str, NoneType, psycopg_Multirange)):
237 return value
238
239 return psycopg_Multirange(
240 [
241 psycopg_Range(
242 element.lower,
243 element.upper,
244 element.bounds,
245 element.empty,
246 )
247 for element in cast("Iterable[ranges.Range]", value)
248 ]
249 )
250
251 return to_range
252
253 def result_processor(self, dialect, coltype):
254 def to_range(value):
255 if value is None:
256 return None
257 else:
258 return ranges.MultiRange(
259 ranges.Range(
260 elem._lower,
261 elem._upper,
262 bounds=elem._bounds if elem._bounds else "[)",
263 empty=not elem._bounds,
264 )
265 for elem in value
266 )
267
268 return to_range
269
270
271class PGExecutionContext_psycopg(_PGExecutionContext_common_psycopg):
272 pass
273
274
275class PGCompiler_psycopg(PGCompiler):
276 pass
277
278
279class PGIdentifierPreparer_psycopg(PGIdentifierPreparer):
280 pass
281
282
283def _log_notices(diagnostic):
284 logger.info("%s: %s", diagnostic.severity, diagnostic.message_primary)
285
286
287class PGDialect_psycopg(_PGDialect_common_psycopg):
288 driver = "psycopg"
289
290 supports_statement_cache = True
291 supports_server_side_cursors = True
292 default_paramstyle = "pyformat"
293 supports_sane_multi_rowcount = True
294
295 execution_ctx_cls = PGExecutionContext_psycopg
296 statement_compiler = PGCompiler_psycopg
297 preparer = PGIdentifierPreparer_psycopg
298 psycopg_version = (0, 0)
299
300 _has_native_hstore = True
301 _psycopg_adapters_map = None
302
303 colspecs = util.update_copy(
304 _PGDialect_common_psycopg.colspecs,
305 {
306 sqltypes.String: _PGString,
307 REGCONFIG: _PGREGCONFIG,
308 JSON: _PGJSON,
309 CITEXT: CITEXT,
310 sqltypes.JSON: _PGJSON,
311 JSONB: _PGJSONB,
312 sqltypes.JSON.JSONPathType: _PGJSONPathType,
313 sqltypes.JSON.JSONIntIndexType: _PGJSONIntIndexType,
314 sqltypes.JSON.JSONStrIndexType: _PGJSONStrIndexType,
315 sqltypes.Interval: _PGInterval,
316 INTERVAL: _PGInterval,
317 sqltypes.Date: _PGDate,
318 sqltypes.DateTime: _PGTimeStamp,
319 sqltypes.Time: _PGTime,
320 sqltypes.Integer: _PGInteger,
321 sqltypes.SmallInteger: _PGSmallInteger,
322 sqltypes.BigInteger: _PGBigInteger,
323 ranges.AbstractSingleRange: _PsycopgRange,
324 ranges.AbstractMultiRange: _PsycopgMultiRange,
325 },
326 )
327
328 def __init__(self, **kwargs):
329 super().__init__(**kwargs)
330
331 if self.dbapi:
332 m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", self.dbapi.__version__)
333 if m:
334 self.psycopg_version = tuple(
335 int(x) for x in m.group(1, 2, 3) if x is not None
336 )
337
338 if self.psycopg_version < (3, 0, 2):
339 raise ImportError(
340 "psycopg version 3.0.2 or higher is required."
341 )
342
343 from psycopg.adapt import AdaptersMap
344
345 self._psycopg_adapters_map = adapters_map = AdaptersMap(
346 self.dbapi.adapters
347 )
348
349 if self._native_inet_types is False:
350 import psycopg.types.string
351
352 adapters_map.register_loader(
353 "inet", psycopg.types.string.TextLoader
354 )
355 adapters_map.register_loader(
356 "cidr", psycopg.types.string.TextLoader
357 )
358
359 if self._json_deserializer:
360 from psycopg.types.json import set_json_loads
361
362 set_json_loads(self._json_deserializer, adapters_map)
363
364 if self._json_serializer:
365 from psycopg.types.json import set_json_dumps
366
367 set_json_dumps(self._json_serializer, adapters_map)
368
369 def create_connect_args(self, url):
370 # see https://github.com/psycopg/psycopg/issues/83
371 cargs, cparams = super().create_connect_args(url)
372
373 if self._psycopg_adapters_map:
374 cparams["context"] = self._psycopg_adapters_map
375 if self.client_encoding is not None:
376 cparams["client_encoding"] = self.client_encoding
377 return cargs, cparams
378
379 def _type_info_fetch(self, connection, name):
380 from psycopg.types import TypeInfo
381
382 return TypeInfo.fetch(connection.connection.driver_connection, name)
383
384 def initialize(self, connection):
385 super().initialize(connection)
386
387 # PGDialect.initialize() checks server version for <= 8.2 and sets
388 # this flag to False if so
389 if not self.insert_returning:
390 self.insert_executemany_returning = False
391
392 # HSTORE can't be registered until we have a connection so that
393 # we can look up its OID, so we set up this adapter in
394 # initialize()
395 if self.use_native_hstore:
396 info = self._type_info_fetch(connection, "hstore")
397 self._has_native_hstore = info is not None
398 if self._has_native_hstore:
399 from psycopg.types.hstore import register_hstore
400
401 # register the adapter for connections made subsequent to
402 # this one
403 register_hstore(info, self._psycopg_adapters_map)
404
405 # register the adapter for this connection
406 register_hstore(info, connection.connection)
407
408 @classmethod
409 def import_dbapi(cls):
410 import psycopg
411
412 return psycopg
413
414 @classmethod
415 def get_async_dialect_cls(cls, url):
416 return PGDialectAsync_psycopg
417
418 @util.memoized_property
419 def _isolation_lookup(self):
420 return {
421 "READ COMMITTED": self.dbapi.IsolationLevel.READ_COMMITTED,
422 "READ UNCOMMITTED": self.dbapi.IsolationLevel.READ_UNCOMMITTED,
423 "REPEATABLE READ": self.dbapi.IsolationLevel.REPEATABLE_READ,
424 "SERIALIZABLE": self.dbapi.IsolationLevel.SERIALIZABLE,
425 }
426
427 @util.memoized_property
428 def _psycopg_Json(self):
429 from psycopg.types import json
430
431 return json.Json
432
433 @util.memoized_property
434 def _psycopg_Jsonb(self):
435 from psycopg.types import json
436
437 return json.Jsonb
438
439 @util.memoized_property
440 def _psycopg_TransactionStatus(self):
441 from psycopg.pq import TransactionStatus
442
443 return TransactionStatus
444
445 @util.memoized_property
446 def _psycopg_Range(self):
447 from psycopg.types.range import Range
448
449 return Range
450
451 @util.memoized_property
452 def _psycopg_Multirange(self):
453 from psycopg.types.multirange import Multirange
454
455 return Multirange
456
457 def _do_isolation_level(self, connection, autocommit, isolation_level):
458 connection.autocommit = autocommit
459 connection.isolation_level = isolation_level
460
461 def get_isolation_level(self, dbapi_connection):
462 status_before = dbapi_connection.info.transaction_status
463 value = super().get_isolation_level(dbapi_connection)
464
465 # don't rely on psycopg providing enum symbols, compare with
466 # eq/ne
467 if status_before == self._psycopg_TransactionStatus.IDLE:
468 dbapi_connection.rollback()
469 return value
470
471 def set_isolation_level(self, dbapi_connection, level):
472 if level == "AUTOCOMMIT":
473 self._do_isolation_level(
474 dbapi_connection, autocommit=True, isolation_level=None
475 )
476 else:
477 self._do_isolation_level(
478 dbapi_connection,
479 autocommit=False,
480 isolation_level=self._isolation_lookup[level],
481 )
482
483 def set_readonly(self, connection, value):
484 connection.read_only = value
485
486 def get_readonly(self, connection):
487 return connection.read_only
488
489 def on_connect(self):
490 def notices(conn):
491 conn.add_notice_handler(_log_notices)
492
493 fns = [notices]
494
495 if self.isolation_level is not None:
496
497 def on_connect(conn):
498 self.set_isolation_level(conn, self.isolation_level)
499
500 fns.append(on_connect)
501
502 # fns always has the notices function
503 def on_connect(conn):
504 for fn in fns:
505 fn(conn)
506
507 return on_connect
508
509 def is_disconnect(self, e, connection, cursor):
510 if isinstance(e, self.dbapi.Error) and connection is not None:
511 if connection.closed or connection.broken:
512 return True
513 return False
514
515 def _do_prepared_twophase(self, connection, command, recover=False):
516 dbapi_conn = connection.connection.dbapi_connection
517 if (
518 recover
519 # don't rely on psycopg providing enum symbols, compare with
520 # eq/ne
521 or dbapi_conn.info.transaction_status
522 != self._psycopg_TransactionStatus.IDLE
523 ):
524 dbapi_conn.rollback()
525 before_autocommit = dbapi_conn.autocommit
526 try:
527 if not before_autocommit:
528 self._do_autocommit(dbapi_conn, True)
529 dbapi_conn.execute(command)
530 finally:
531 if not before_autocommit:
532 self._do_autocommit(dbapi_conn, before_autocommit)
533
534 def do_rollback_twophase(
535 self, connection, xid, is_prepared=True, recover=False
536 ):
537 if is_prepared:
538 self._do_prepared_twophase(
539 connection, f"ROLLBACK PREPARED '{xid}'", recover=recover
540 )
541 else:
542 self.do_rollback(connection.connection)
543
544 def do_commit_twophase(
545 self, connection, xid, is_prepared=True, recover=False
546 ):
547 if is_prepared:
548 self._do_prepared_twophase(
549 connection, f"COMMIT PREPARED '{xid}'", recover=recover
550 )
551 else:
552 self.do_commit(connection.connection)
553
554 @util.memoized_property
555 def _dialect_specific_select_one(self):
556 return ";"
557
558
559class AsyncAdapt_psycopg_cursor:
560 __slots__ = ("_cursor", "await_", "_rows")
561
562 _psycopg_ExecStatus = None
563
564 def __init__(self, cursor, await_) -> None:
565 self._cursor = cursor
566 self.await_ = await_
567 self._rows = []
568
569 def __getattr__(self, name):
570 return getattr(self._cursor, name)
571
572 @property
573 def arraysize(self):
574 return self._cursor.arraysize
575
576 @arraysize.setter
577 def arraysize(self, value):
578 self._cursor.arraysize = value
579
580 def close(self):
581 self._rows.clear()
582 # Normal cursor just call _close() in a non-sync way.
583 self._cursor._close()
584
585 def execute(self, query, params=None, **kw):
586 result = self.await_(self._cursor.execute(query, params, **kw))
587 # sqlalchemy result is not async, so need to pull all rows here
588 res = self._cursor.pgresult
589
590 # don't rely on psycopg providing enum symbols, compare with
591 # eq/ne
592 if res and res.status == self._psycopg_ExecStatus.TUPLES_OK:
593 rows = self.await_(self._cursor.fetchall())
594 if not isinstance(rows, list):
595 self._rows = list(rows)
596 else:
597 self._rows = rows
598 return result
599
600 def executemany(self, query, params_seq):
601 return self.await_(self._cursor.executemany(query, params_seq))
602
603 def __iter__(self):
604 # TODO: try to avoid pop(0) on a list
605 while self._rows:
606 yield self._rows.pop(0)
607
608 def fetchone(self):
609 if self._rows:
610 # TODO: try to avoid pop(0) on a list
611 return self._rows.pop(0)
612 else:
613 return None
614
615 def fetchmany(self, size=None):
616 if size is None:
617 size = self._cursor.arraysize
618
619 retval = self._rows[0:size]
620 self._rows = self._rows[size:]
621 return retval
622
623 def fetchall(self):
624 retval = self._rows
625 self._rows = []
626 return retval
627
628
629class AsyncAdapt_psycopg_ss_cursor(AsyncAdapt_psycopg_cursor):
630 def execute(self, query, params=None, **kw):
631 self.await_(self._cursor.execute(query, params, **kw))
632 return self
633
634 def close(self):
635 self.await_(self._cursor.close())
636
637 def fetchone(self):
638 return self.await_(self._cursor.fetchone())
639
640 def fetchmany(self, size=0):
641 return self.await_(self._cursor.fetchmany(size))
642
643 def fetchall(self):
644 return self.await_(self._cursor.fetchall())
645
646 def __iter__(self):
647 iterator = self._cursor.__aiter__()
648 while True:
649 try:
650 yield self.await_(iterator.__anext__())
651 except StopAsyncIteration:
652 break
653
654
655class AsyncAdapt_psycopg_connection(AdaptedConnection):
656 _connection: AsyncConnection
657 __slots__ = ()
658 await_ = staticmethod(await_only)
659
660 def __init__(self, connection) -> None:
661 self._connection = connection
662
663 def __getattr__(self, name):
664 return getattr(self._connection, name)
665
666 def execute(self, query, params=None, **kw):
667 cursor = self.await_(self._connection.execute(query, params, **kw))
668 return AsyncAdapt_psycopg_cursor(cursor, self.await_)
669
670 def cursor(self, *args, **kw):
671 cursor = self._connection.cursor(*args, **kw)
672 if hasattr(cursor, "name"):
673 return AsyncAdapt_psycopg_ss_cursor(cursor, self.await_)
674 else:
675 return AsyncAdapt_psycopg_cursor(cursor, self.await_)
676
677 def commit(self):
678 self.await_(self._connection.commit())
679
680 def rollback(self):
681 self.await_(self._connection.rollback())
682
683 def close(self):
684 self.await_(self._connection.close())
685
686 @property
687 def autocommit(self):
688 return self._connection.autocommit
689
690 @autocommit.setter
691 def autocommit(self, value):
692 self.set_autocommit(value)
693
694 def set_autocommit(self, value):
695 self.await_(self._connection.set_autocommit(value))
696
697 def set_isolation_level(self, value):
698 self.await_(self._connection.set_isolation_level(value))
699
700 def set_read_only(self, value):
701 self.await_(self._connection.set_read_only(value))
702
703 def set_deferrable(self, value):
704 self.await_(self._connection.set_deferrable(value))
705
706
707class AsyncAdaptFallback_psycopg_connection(AsyncAdapt_psycopg_connection):
708 __slots__ = ()
709 await_ = staticmethod(await_fallback)
710
711
712class PsycopgAdaptDBAPI:
713 def __init__(self, psycopg) -> None:
714 self.psycopg = psycopg
715
716 for k, v in self.psycopg.__dict__.items():
717 if k != "connect":
718 self.__dict__[k] = v
719
720 def connect(self, *arg, **kw):
721 async_fallback = kw.pop("async_fallback", False)
722 creator_fn = kw.pop(
723 "async_creator_fn", self.psycopg.AsyncConnection.connect
724 )
725 if util.asbool(async_fallback):
726 return AsyncAdaptFallback_psycopg_connection(
727 await_fallback(creator_fn(*arg, **kw))
728 )
729 else:
730 return AsyncAdapt_psycopg_connection(
731 await_only(creator_fn(*arg, **kw))
732 )
733
734
735class PGDialectAsync_psycopg(PGDialect_psycopg):
736 is_async = True
737 supports_statement_cache = True
738
739 @classmethod
740 def import_dbapi(cls):
741 import psycopg
742 from psycopg.pq import ExecStatus
743
744 AsyncAdapt_psycopg_cursor._psycopg_ExecStatus = ExecStatus
745
746 return PsycopgAdaptDBAPI(psycopg)
747
748 @classmethod
749 def get_pool_class(cls, url):
750 async_fallback = url.query.get("async_fallback", False)
751
752 if util.asbool(async_fallback):
753 return pool.FallbackAsyncAdaptedQueuePool
754 else:
755 return pool.AsyncAdaptedQueuePool
756
757 def _type_info_fetch(self, connection, name):
758 from psycopg.types import TypeInfo
759
760 adapted = connection.connection
761 return adapted.await_(TypeInfo.fetch(adapted.driver_connection, name))
762
763 def _do_isolation_level(self, connection, autocommit, isolation_level):
764 connection.set_autocommit(autocommit)
765 connection.set_isolation_level(isolation_level)
766
767 def _do_autocommit(self, connection, value):
768 connection.set_autocommit(value)
769
770 def set_readonly(self, connection, value):
771 connection.set_read_only(value)
772
773 def set_deferrable(self, connection, value):
774 connection.set_deferrable(value)
775
776 def get_driver_connection(self, connection):
777 return connection._connection
778
779
780dialect = PGDialect_psycopg
781dialect_async = PGDialectAsync_psycopg
782