extensions.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. """psycopg extensions to the DBAPI-2.0
  2. This module holds all the extensions to the DBAPI-2.0 provided by psycopg.
  3. - `connection` -- the new-type inheritable connection class
  4. - `cursor` -- the new-type inheritable cursor class
  5. - `lobject` -- the new-type inheritable large object class
  6. - `adapt()` -- exposes the PEP-246_ compatible adapting mechanism used
  7. by psycopg to adapt Python types to PostgreSQL ones
  8. .. _PEP-246: https://www.python.org/dev/peps/pep-0246/
  9. """
  10. # psycopg/extensions.py - DBAPI-2.0 extensions specific to psycopg
  11. #
  12. # Copyright (C) 2003-2019 Federico Di Gregorio <fog@debian.org>
  13. # Copyright (C) 2020 The Psycopg Team
  14. #
  15. # psycopg2 is free software: you can redistribute it and/or modify it
  16. # under the terms of the GNU Lesser General Public License as published
  17. # by the Free Software Foundation, either version 3 of the License, or
  18. # (at your option) any later version.
  19. #
  20. # In addition, as a special exception, the copyright holders give
  21. # permission to link this program with the OpenSSL library (or with
  22. # modified versions of OpenSSL that use the same license as OpenSSL),
  23. # and distribute linked combinations including the two.
  24. #
  25. # You must obey the GNU Lesser General Public License in all respects for
  26. # all of the code used other than OpenSSL.
  27. #
  28. # psycopg2 is distributed in the hope that it will be useful, but WITHOUT
  29. # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  30. # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
  31. # License for more details.
  32. import re as _re
  33. from psycopg2._psycopg import ( # noqa
  34. BINARYARRAY, BOOLEAN, BOOLEANARRAY, BYTES, BYTESARRAY, DATE, DATEARRAY,
  35. DATETIMEARRAY, DECIMAL, DECIMALARRAY, FLOAT, FLOATARRAY, INTEGER,
  36. INTEGERARRAY, INTERVAL, INTERVALARRAY, LONGINTEGER, LONGINTEGERARRAY,
  37. ROWIDARRAY, STRINGARRAY, TIME, TIMEARRAY, UNICODE, UNICODEARRAY,
  38. AsIs, Binary, Boolean, Float, Int, QuotedString, )
  39. try:
  40. from psycopg2._psycopg import ( # noqa
  41. MXDATE, MXDATETIME, MXDATETIMETZ, MXINTERVAL, MXTIME, MXDATEARRAY,
  42. MXDATETIMEARRAY, MXDATETIMETZARRAY, MXINTERVALARRAY, MXTIMEARRAY,
  43. DateFromMx, TimeFromMx, TimestampFromMx, IntervalFromMx, )
  44. except ImportError:
  45. pass
  46. from psycopg2._psycopg import ( # noqa
  47. PYDATE, PYDATETIME, PYDATETIMETZ, PYINTERVAL, PYTIME, PYDATEARRAY,
  48. PYDATETIMEARRAY, PYDATETIMETZARRAY, PYINTERVALARRAY, PYTIMEARRAY,
  49. DateFromPy, TimeFromPy, TimestampFromPy, IntervalFromPy, )
  50. from psycopg2._psycopg import ( # noqa
  51. adapt, adapters, encodings, connection, cursor,
  52. lobject, Xid, libpq_version, parse_dsn, quote_ident,
  53. string_types, binary_types, new_type, new_array_type, register_type,
  54. ISQLQuote, Notify, Diagnostics, Column, ConnectionInfo,
  55. QueryCanceledError, TransactionRollbackError,
  56. set_wait_callback, get_wait_callback, encrypt_password, )
  57. """Isolation level values."""
  58. ISOLATION_LEVEL_AUTOCOMMIT = 0
  59. ISOLATION_LEVEL_READ_UNCOMMITTED = 4
  60. ISOLATION_LEVEL_READ_COMMITTED = 1
  61. ISOLATION_LEVEL_REPEATABLE_READ = 2
  62. ISOLATION_LEVEL_SERIALIZABLE = 3
  63. ISOLATION_LEVEL_DEFAULT = None
  64. """psycopg connection status values."""
  65. STATUS_SETUP = 0
  66. STATUS_READY = 1
  67. STATUS_BEGIN = 2
  68. STATUS_SYNC = 3 # currently unused
  69. STATUS_ASYNC = 4 # currently unused
  70. STATUS_PREPARED = 5
  71. # This is a useful mnemonic to check if the connection is in a transaction
  72. STATUS_IN_TRANSACTION = STATUS_BEGIN
  73. """psycopg asynchronous connection polling values"""
  74. POLL_OK = 0
  75. POLL_READ = 1
  76. POLL_WRITE = 2
  77. POLL_ERROR = 3
  78. """Backend transaction status values."""
  79. TRANSACTION_STATUS_IDLE = 0
  80. TRANSACTION_STATUS_ACTIVE = 1
  81. TRANSACTION_STATUS_INTRANS = 2
  82. TRANSACTION_STATUS_INERROR = 3
  83. TRANSACTION_STATUS_UNKNOWN = 4
  84. def register_adapter(typ, callable):
  85. """Register 'callable' as an ISQLQuote adapter for type 'typ'."""
  86. adapters[(typ, ISQLQuote)] = callable
  87. # The SQL_IN class is the official adapter for tuples starting from 2.0.6.
  88. class SQL_IN(object):
  89. """Adapt any iterable to an SQL quotable object."""
  90. def __init__(self, seq):
  91. self._seq = seq
  92. self._conn = None
  93. def prepare(self, conn):
  94. self._conn = conn
  95. def getquoted(self):
  96. # this is the important line: note how every object in the
  97. # list is adapted and then how getquoted() is called on it
  98. pobjs = [adapt(o) for o in self._seq]
  99. if self._conn is not None:
  100. for obj in pobjs:
  101. if hasattr(obj, 'prepare'):
  102. obj.prepare(self._conn)
  103. qobjs = [o.getquoted() for o in pobjs]
  104. return b'(' + b', '.join(qobjs) + b')'
  105. def __str__(self):
  106. return str(self.getquoted())
  107. class NoneAdapter(object):
  108. """Adapt None to NULL.
  109. This adapter is not used normally as a fast path in mogrify uses NULL,
  110. but it makes easier to adapt composite types.
  111. """
  112. def __init__(self, obj):
  113. pass
  114. def getquoted(self, _null=b"NULL"):
  115. return _null
  116. def make_dsn(dsn=None, **kwargs):
  117. """Convert a set of keywords into a connection strings."""
  118. if dsn is None and not kwargs:
  119. return ''
  120. # If no kwarg is specified don't mung the dsn, but verify it
  121. if not kwargs:
  122. parse_dsn(dsn)
  123. return dsn
  124. # Override the dsn with the parameters
  125. if 'database' in kwargs:
  126. if 'dbname' in kwargs:
  127. raise TypeError(
  128. "you can't specify both 'database' and 'dbname' arguments")
  129. kwargs['dbname'] = kwargs.pop('database')
  130. # Drop the None arguments
  131. kwargs = {k: v for (k, v) in kwargs.items() if v is not None}
  132. if dsn is not None:
  133. tmp = parse_dsn(dsn)
  134. tmp.update(kwargs)
  135. kwargs = tmp
  136. dsn = " ".join(["%s=%s" % (k, _param_escape(str(v)))
  137. for (k, v) in kwargs.items()])
  138. # verify that the returned dsn is valid
  139. parse_dsn(dsn)
  140. return dsn
  141. def _param_escape(s,
  142. re_escape=_re.compile(r"([\\'])"),
  143. re_space=_re.compile(r'\s')):
  144. """
  145. Apply the escaping rule required by PQconnectdb
  146. """
  147. if not s:
  148. return "''"
  149. s = re_escape.sub(r'\\\1', s)
  150. if re_space.search(s):
  151. s = "'" + s + "'"
  152. return s
  153. # Create default json typecasters for PostgreSQL 9.2 oids
  154. from psycopg2._json import register_default_json, register_default_jsonb # noqa
  155. try:
  156. JSON, JSONARRAY = register_default_json()
  157. JSONB, JSONBARRAY = register_default_jsonb()
  158. except ImportError:
  159. pass
  160. del register_default_json, register_default_jsonb
  161. # Create default Range typecasters
  162. from psycopg2. _range import Range # noqa
  163. del Range
  164. # Add the "cleaned" version of the encodings to the key.
  165. # When the encoding is set its name is cleaned up from - and _ and turned
  166. # uppercase, so an encoding not respecting these rules wouldn't be found in the
  167. # encodings keys and would raise an exception with the unicode typecaster
  168. for k, v in list(encodings.items()):
  169. k = k.replace('_', '').replace('-', '').upper()
  170. encodings[k] = v
  171. del k, v