extras.py 43 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325
  1. """Miscellaneous goodies for psycopg2
  2. This module is a generic place used to hold little helper functions
  3. and classes until a better place in the distribution is found.
  4. """
  5. # psycopg/extras.py - miscellaneous extra goodies for psycopg
  6. #
  7. # Copyright (C) 2003-2019 Federico Di Gregorio <fog@debian.org>
  8. # Copyright (C) 2020 The Psycopg Team
  9. #
  10. # psycopg2 is free software: you can redistribute it and/or modify it
  11. # under the terms of the GNU Lesser General Public License as published
  12. # by the Free Software Foundation, either version 3 of the License, or
  13. # (at your option) any later version.
  14. #
  15. # In addition, as a special exception, the copyright holders give
  16. # permission to link this program with the OpenSSL library (or with
  17. # modified versions of OpenSSL that use the same license as OpenSSL),
  18. # and distribute linked combinations including the two.
  19. #
  20. # You must obey the GNU Lesser General Public License in all respects for
  21. # all of the code used other than OpenSSL.
  22. #
  23. # psycopg2 is distributed in the hope that it will be useful, but WITHOUT
  24. # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  25. # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
  26. # License for more details.
  27. import os as _os
  28. import time as _time
  29. import re as _re
  30. from collections import namedtuple, OrderedDict
  31. import logging as _logging
  32. import psycopg2
  33. from psycopg2 import extensions as _ext
  34. from .extensions import cursor as _cursor
  35. from .extensions import connection as _connection
  36. from .extensions import adapt as _A, quote_ident
  37. from .compat import PY2, PY3, lru_cache
  38. from psycopg2._psycopg import ( # noqa
  39. REPLICATION_PHYSICAL, REPLICATION_LOGICAL,
  40. ReplicationConnection as _replicationConnection,
  41. ReplicationCursor as _replicationCursor,
  42. ReplicationMessage)
  43. # expose the json adaptation stuff into the module
  44. from psycopg2._json import ( # noqa
  45. json, Json, register_json, register_default_json, register_default_jsonb)
  46. # Expose range-related objects
  47. from psycopg2._range import ( # noqa
  48. Range, NumericRange, DateRange, DateTimeRange, DateTimeTZRange,
  49. register_range, RangeAdapter, RangeCaster)
  50. # Expose ipaddress-related objects
  51. from psycopg2._ipaddress import register_ipaddress # noqa
  52. class DictCursorBase(_cursor):
  53. """Base class for all dict-like cursors."""
  54. def __init__(self, *args, **kwargs):
  55. if 'row_factory' in kwargs:
  56. row_factory = kwargs['row_factory']
  57. del kwargs['row_factory']
  58. else:
  59. raise NotImplementedError(
  60. "DictCursorBase can't be instantiated without a row factory.")
  61. super(DictCursorBase, self).__init__(*args, **kwargs)
  62. self._query_executed = False
  63. self._prefetch = False
  64. self.row_factory = row_factory
  65. def fetchone(self):
  66. if self._prefetch:
  67. res = super(DictCursorBase, self).fetchone()
  68. if self._query_executed:
  69. self._build_index()
  70. if not self._prefetch:
  71. res = super(DictCursorBase, self).fetchone()
  72. return res
  73. def fetchmany(self, size=None):
  74. if self._prefetch:
  75. res = super(DictCursorBase, self).fetchmany(size)
  76. if self._query_executed:
  77. self._build_index()
  78. if not self._prefetch:
  79. res = super(DictCursorBase, self).fetchmany(size)
  80. return res
  81. def fetchall(self):
  82. if self._prefetch:
  83. res = super(DictCursorBase, self).fetchall()
  84. if self._query_executed:
  85. self._build_index()
  86. if not self._prefetch:
  87. res = super(DictCursorBase, self).fetchall()
  88. return res
  89. def __iter__(self):
  90. try:
  91. if self._prefetch:
  92. res = super(DictCursorBase, self).__iter__()
  93. first = next(res)
  94. if self._query_executed:
  95. self._build_index()
  96. if not self._prefetch:
  97. res = super(DictCursorBase, self).__iter__()
  98. first = next(res)
  99. yield first
  100. while True:
  101. yield next(res)
  102. except StopIteration:
  103. return
  104. class DictConnection(_connection):
  105. """A connection that uses `DictCursor` automatically."""
  106. def cursor(self, *args, **kwargs):
  107. kwargs.setdefault('cursor_factory', self.cursor_factory or DictCursor)
  108. return super(DictConnection, self).cursor(*args, **kwargs)
  109. class DictCursor(DictCursorBase):
  110. """A cursor that keeps a list of column name -> index mappings."""
  111. def __init__(self, *args, **kwargs):
  112. kwargs['row_factory'] = DictRow
  113. super(DictCursor, self).__init__(*args, **kwargs)
  114. self._prefetch = True
  115. def execute(self, query, vars=None):
  116. self.index = OrderedDict()
  117. self._query_executed = True
  118. return super(DictCursor, self).execute(query, vars)
  119. def callproc(self, procname, vars=None):
  120. self.index = OrderedDict()
  121. self._query_executed = True
  122. return super(DictCursor, self).callproc(procname, vars)
  123. def _build_index(self):
  124. if self._query_executed and self.description:
  125. for i in range(len(self.description)):
  126. self.index[self.description[i][0]] = i
  127. self._query_executed = False
  128. class DictRow(list):
  129. """A row object that allow by-column-name access to data."""
  130. __slots__ = ('_index',)
  131. def __init__(self, cursor):
  132. self._index = cursor.index
  133. self[:] = [None] * len(cursor.description)
  134. def __getitem__(self, x):
  135. if not isinstance(x, (int, slice)):
  136. x = self._index[x]
  137. return super(DictRow, self).__getitem__(x)
  138. def __setitem__(self, x, v):
  139. if not isinstance(x, (int, slice)):
  140. x = self._index[x]
  141. super(DictRow, self).__setitem__(x, v)
  142. def items(self):
  143. g = super(DictRow, self).__getitem__
  144. return ((n, g(self._index[n])) for n in self._index)
  145. def keys(self):
  146. return iter(self._index)
  147. def values(self):
  148. g = super(DictRow, self).__getitem__
  149. return (g(self._index[n]) for n in self._index)
  150. def get(self, x, default=None):
  151. try:
  152. return self[x]
  153. except Exception:
  154. return default
  155. def copy(self):
  156. return OrderedDict(self.items())
  157. def __contains__(self, x):
  158. return x in self._index
  159. def __reduce__(self):
  160. # this is apparently useless, but it fixes #1073
  161. return super(DictRow, self).__reduce__()
  162. def __getstate__(self):
  163. return self[:], self._index.copy()
  164. def __setstate__(self, data):
  165. self[:] = data[0]
  166. self._index = data[1]
  167. if PY2:
  168. iterkeys = keys
  169. itervalues = values
  170. iteritems = items
  171. has_key = __contains__
  172. def keys(self):
  173. return list(self.iterkeys())
  174. def values(self):
  175. return tuple(self.itervalues())
  176. def items(self):
  177. return list(self.iteritems())
  178. class RealDictConnection(_connection):
  179. """A connection that uses `RealDictCursor` automatically."""
  180. def cursor(self, *args, **kwargs):
  181. kwargs.setdefault('cursor_factory', self.cursor_factory or RealDictCursor)
  182. return super(RealDictConnection, self).cursor(*args, **kwargs)
  183. class RealDictCursor(DictCursorBase):
  184. """A cursor that uses a real dict as the base type for rows.
  185. Note that this cursor is extremely specialized and does not allow
  186. the normal access (using integer indices) to fetched data. If you need
  187. to access database rows both as a dictionary and a list, then use
  188. the generic `DictCursor` instead of `!RealDictCursor`.
  189. """
  190. def __init__(self, *args, **kwargs):
  191. kwargs['row_factory'] = RealDictRow
  192. super(RealDictCursor, self).__init__(*args, **kwargs)
  193. def execute(self, query, vars=None):
  194. self.column_mapping = []
  195. self._query_executed = True
  196. return super(RealDictCursor, self).execute(query, vars)
  197. def callproc(self, procname, vars=None):
  198. self.column_mapping = []
  199. self._query_executed = True
  200. return super(RealDictCursor, self).callproc(procname, vars)
  201. def _build_index(self):
  202. if self._query_executed and self.description:
  203. self.column_mapping = [d[0] for d in self.description]
  204. self._query_executed = False
  205. class RealDictRow(OrderedDict):
  206. """A `!dict` subclass representing a data record."""
  207. def __init__(self, *args, **kwargs):
  208. if args and isinstance(args[0], _cursor):
  209. cursor = args[0]
  210. args = args[1:]
  211. else:
  212. cursor = None
  213. super(RealDictRow, self).__init__(*args, **kwargs)
  214. if cursor is not None:
  215. # Required for named cursors
  216. if cursor.description and not cursor.column_mapping:
  217. cursor._build_index()
  218. # Store the cols mapping in the dict itself until the row is fully
  219. # populated, so we don't need to add attributes to the class
  220. # (hence keeping its maintenance, special pickle support, etc.)
  221. self[RealDictRow] = cursor.column_mapping
  222. def __setitem__(self, key, value):
  223. if RealDictRow in self:
  224. # We are in the row building phase
  225. mapping = self[RealDictRow]
  226. super(RealDictRow, self).__setitem__(mapping[key], value)
  227. if key == len(mapping) - 1:
  228. # Row building finished
  229. del self[RealDictRow]
  230. return
  231. super(RealDictRow, self).__setitem__(key, value)
  232. class NamedTupleConnection(_connection):
  233. """A connection that uses `NamedTupleCursor` automatically."""
  234. def cursor(self, *args, **kwargs):
  235. kwargs.setdefault('cursor_factory', self.cursor_factory or NamedTupleCursor)
  236. return super(NamedTupleConnection, self).cursor(*args, **kwargs)
  237. class NamedTupleCursor(_cursor):
  238. """A cursor that generates results as `~collections.namedtuple`.
  239. `!fetch*()` methods will return named tuples instead of regular tuples, so
  240. their elements can be accessed both as regular numeric items as well as
  241. attributes.
  242. >>> nt_cur = conn.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor)
  243. >>> rec = nt_cur.fetchone()
  244. >>> rec
  245. Record(id=1, num=100, data="abc'def")
  246. >>> rec[1]
  247. 100
  248. >>> rec.data
  249. "abc'def"
  250. """
  251. Record = None
  252. MAX_CACHE = 1024
  253. def execute(self, query, vars=None):
  254. self.Record = None
  255. return super(NamedTupleCursor, self).execute(query, vars)
  256. def executemany(self, query, vars):
  257. self.Record = None
  258. return super(NamedTupleCursor, self).executemany(query, vars)
  259. def callproc(self, procname, vars=None):
  260. self.Record = None
  261. return super(NamedTupleCursor, self).callproc(procname, vars)
  262. def fetchone(self):
  263. t = super(NamedTupleCursor, self).fetchone()
  264. if t is not None:
  265. nt = self.Record
  266. if nt is None:
  267. nt = self.Record = self._make_nt()
  268. return nt._make(t)
  269. def fetchmany(self, size=None):
  270. ts = super(NamedTupleCursor, self).fetchmany(size)
  271. nt = self.Record
  272. if nt is None:
  273. nt = self.Record = self._make_nt()
  274. return list(map(nt._make, ts))
  275. def fetchall(self):
  276. ts = super(NamedTupleCursor, self).fetchall()
  277. nt = self.Record
  278. if nt is None:
  279. nt = self.Record = self._make_nt()
  280. return list(map(nt._make, ts))
  281. def __iter__(self):
  282. try:
  283. it = super(NamedTupleCursor, self).__iter__()
  284. t = next(it)
  285. nt = self.Record
  286. if nt is None:
  287. nt = self.Record = self._make_nt()
  288. yield nt._make(t)
  289. while True:
  290. yield nt._make(next(it))
  291. except StopIteration:
  292. return
  293. # ascii except alnum and underscore
  294. _re_clean = _re.compile(
  295. '[' + _re.escape(' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~') + ']')
  296. def _make_nt(self):
  297. key = tuple(d[0] for d in self.description) if self.description else ()
  298. return self._cached_make_nt(key)
  299. @classmethod
  300. def _do_make_nt(cls, key):
  301. fields = []
  302. for s in key:
  303. s = cls._re_clean.sub('_', s)
  304. # Python identifier cannot start with numbers, namedtuple fields
  305. # cannot start with underscore. So...
  306. if s[0] == '_' or '0' <= s[0] <= '9':
  307. s = 'f' + s
  308. fields.append(s)
  309. nt = namedtuple("Record", fields)
  310. return nt
  311. @lru_cache(512)
  312. def _cached_make_nt(cls, key):
  313. return cls._do_make_nt(key)
  314. # Exposed for testability, and if someone wants to monkeypatch to tweak
  315. # the cache size.
  316. NamedTupleCursor._cached_make_nt = classmethod(_cached_make_nt)
  317. class LoggingConnection(_connection):
  318. """A connection that logs all queries to a file or logger__ object.
  319. .. __: https://docs.python.org/library/logging.html
  320. """
  321. def initialize(self, logobj):
  322. """Initialize the connection to log to `!logobj`.
  323. The `!logobj` parameter can be an open file object or a Logger/LoggerAdapter
  324. instance from the standard logging module.
  325. """
  326. self._logobj = logobj
  327. if _logging and isinstance(
  328. logobj, (_logging.Logger, _logging.LoggerAdapter)):
  329. self.log = self._logtologger
  330. else:
  331. self.log = self._logtofile
  332. def filter(self, msg, curs):
  333. """Filter the query before logging it.
  334. This is the method to overwrite to filter unwanted queries out of the
  335. log or to add some extra data to the output. The default implementation
  336. just does nothing.
  337. """
  338. return msg
  339. def _logtofile(self, msg, curs):
  340. msg = self.filter(msg, curs)
  341. if msg:
  342. if PY3 and isinstance(msg, bytes):
  343. msg = msg.decode(_ext.encodings[self.encoding], 'replace')
  344. self._logobj.write(msg + _os.linesep)
  345. def _logtologger(self, msg, curs):
  346. msg = self.filter(msg, curs)
  347. if msg:
  348. self._logobj.debug(msg)
  349. def _check(self):
  350. if not hasattr(self, '_logobj'):
  351. raise self.ProgrammingError(
  352. "LoggingConnection object has not been initialize()d")
  353. def cursor(self, *args, **kwargs):
  354. self._check()
  355. kwargs.setdefault('cursor_factory', self.cursor_factory or LoggingCursor)
  356. return super(LoggingConnection, self).cursor(*args, **kwargs)
  357. class LoggingCursor(_cursor):
  358. """A cursor that logs queries using its connection logging facilities."""
  359. def execute(self, query, vars=None):
  360. try:
  361. return super(LoggingCursor, self).execute(query, vars)
  362. finally:
  363. self.connection.log(self.query, self)
  364. def callproc(self, procname, vars=None):
  365. try:
  366. return super(LoggingCursor, self).callproc(procname, vars)
  367. finally:
  368. self.connection.log(self.query, self)
  369. class MinTimeLoggingConnection(LoggingConnection):
  370. """A connection that logs queries based on execution time.
  371. This is just an example of how to sub-class `LoggingConnection` to
  372. provide some extra filtering for the logged queries. Both the
  373. `initialize()` and `filter()` methods are overwritten to make sure
  374. that only queries executing for more than ``mintime`` ms are logged.
  375. Note that this connection uses the specialized cursor
  376. `MinTimeLoggingCursor`.
  377. """
  378. def initialize(self, logobj, mintime=0):
  379. LoggingConnection.initialize(self, logobj)
  380. self._mintime = mintime
  381. def filter(self, msg, curs):
  382. t = (_time.time() - curs.timestamp) * 1000
  383. if t > self._mintime:
  384. if PY3 and isinstance(msg, bytes):
  385. msg = msg.decode(_ext.encodings[self.encoding], 'replace')
  386. return msg + _os.linesep + " (execution time: %d ms)" % t
  387. def cursor(self, *args, **kwargs):
  388. kwargs.setdefault('cursor_factory',
  389. self.cursor_factory or MinTimeLoggingCursor)
  390. return LoggingConnection.cursor(self, *args, **kwargs)
  391. class MinTimeLoggingCursor(LoggingCursor):
  392. """The cursor sub-class companion to `MinTimeLoggingConnection`."""
  393. def execute(self, query, vars=None):
  394. self.timestamp = _time.time()
  395. return LoggingCursor.execute(self, query, vars)
  396. def callproc(self, procname, vars=None):
  397. self.timestamp = _time.time()
  398. return LoggingCursor.callproc(self, procname, vars)
  399. class LogicalReplicationConnection(_replicationConnection):
  400. def __init__(self, *args, **kwargs):
  401. kwargs['replication_type'] = REPLICATION_LOGICAL
  402. super(LogicalReplicationConnection, self).__init__(*args, **kwargs)
  403. class PhysicalReplicationConnection(_replicationConnection):
  404. def __init__(self, *args, **kwargs):
  405. kwargs['replication_type'] = REPLICATION_PHYSICAL
  406. super(PhysicalReplicationConnection, self).__init__(*args, **kwargs)
  407. class StopReplication(Exception):
  408. """
  409. Exception used to break out of the endless loop in
  410. `~ReplicationCursor.consume_stream()`.
  411. Subclass of `~exceptions.Exception`. Intentionally *not* inherited from
  412. `~psycopg2.Error` as occurrence of this exception does not indicate an
  413. error.
  414. """
  415. pass
  416. class ReplicationCursor(_replicationCursor):
  417. """A cursor used for communication on replication connections."""
  418. def create_replication_slot(self, slot_name, slot_type=None, output_plugin=None):
  419. """Create streaming replication slot."""
  420. command = "CREATE_REPLICATION_SLOT %s " % quote_ident(slot_name, self)
  421. if slot_type is None:
  422. slot_type = self.connection.replication_type
  423. if slot_type == REPLICATION_LOGICAL:
  424. if output_plugin is None:
  425. raise psycopg2.ProgrammingError(
  426. "output plugin name is required to create "
  427. "logical replication slot")
  428. command += "LOGICAL %s" % quote_ident(output_plugin, self)
  429. elif slot_type == REPLICATION_PHYSICAL:
  430. if output_plugin is not None:
  431. raise psycopg2.ProgrammingError(
  432. "cannot specify output plugin name when creating "
  433. "physical replication slot")
  434. command += "PHYSICAL"
  435. else:
  436. raise psycopg2.ProgrammingError(
  437. "unrecognized replication type: %s" % repr(slot_type))
  438. self.execute(command)
  439. def drop_replication_slot(self, slot_name):
  440. """Drop streaming replication slot."""
  441. command = "DROP_REPLICATION_SLOT %s" % quote_ident(slot_name, self)
  442. self.execute(command)
  443. def start_replication(
  444. self, slot_name=None, slot_type=None, start_lsn=0,
  445. timeline=0, options=None, decode=False, status_interval=10):
  446. """Start replication stream."""
  447. command = "START_REPLICATION "
  448. if slot_type is None:
  449. slot_type = self.connection.replication_type
  450. if slot_type == REPLICATION_LOGICAL:
  451. if slot_name:
  452. command += "SLOT %s " % quote_ident(slot_name, self)
  453. else:
  454. raise psycopg2.ProgrammingError(
  455. "slot name is required for logical replication")
  456. command += "LOGICAL "
  457. elif slot_type == REPLICATION_PHYSICAL:
  458. if slot_name:
  459. command += "SLOT %s " % quote_ident(slot_name, self)
  460. # don't add "PHYSICAL", before 9.4 it was just START_REPLICATION XXX/XXX
  461. else:
  462. raise psycopg2.ProgrammingError(
  463. "unrecognized replication type: %s" % repr(slot_type))
  464. if type(start_lsn) is str:
  465. lsn = start_lsn.split('/')
  466. lsn = "%X/%08X" % (int(lsn[0], 16), int(lsn[1], 16))
  467. else:
  468. lsn = "%X/%08X" % ((start_lsn >> 32) & 0xFFFFFFFF,
  469. start_lsn & 0xFFFFFFFF)
  470. command += lsn
  471. if timeline != 0:
  472. if slot_type == REPLICATION_LOGICAL:
  473. raise psycopg2.ProgrammingError(
  474. "cannot specify timeline for logical replication")
  475. command += " TIMELINE %d" % timeline
  476. if options:
  477. if slot_type == REPLICATION_PHYSICAL:
  478. raise psycopg2.ProgrammingError(
  479. "cannot specify output plugin options for physical replication")
  480. command += " ("
  481. for k, v in options.items():
  482. if not command.endswith('('):
  483. command += ", "
  484. command += "%s %s" % (quote_ident(k, self), _A(str(v)))
  485. command += ")"
  486. self.start_replication_expert(
  487. command, decode=decode, status_interval=status_interval)
  488. # allows replication cursors to be used in select.select() directly
  489. def fileno(self):
  490. return self.connection.fileno()
  491. # a dbtype and adapter for Python UUID type
  492. class UUID_adapter(object):
  493. """Adapt Python's uuid.UUID__ type to PostgreSQL's uuid__.
  494. .. __: https://docs.python.org/library/uuid.html
  495. .. __: https://www.postgresql.org/docs/current/static/datatype-uuid.html
  496. """
  497. def __init__(self, uuid):
  498. self._uuid = uuid
  499. def __conform__(self, proto):
  500. if proto is _ext.ISQLQuote:
  501. return self
  502. def getquoted(self):
  503. return ("'%s'::uuid" % self._uuid).encode('utf8')
  504. def __str__(self):
  505. return "'%s'::uuid" % self._uuid
  506. def register_uuid(oids=None, conn_or_curs=None):
  507. """Create the UUID type and an uuid.UUID adapter.
  508. :param oids: oid for the PostgreSQL :sql:`uuid` type, or 2-items sequence
  509. with oids of the type and the array. If not specified, use PostgreSQL
  510. standard oids.
  511. :param conn_or_curs: where to register the typecaster. If not specified,
  512. register it globally.
  513. """
  514. import uuid
  515. if not oids:
  516. oid1 = 2950
  517. oid2 = 2951
  518. elif isinstance(oids, (list, tuple)):
  519. oid1, oid2 = oids
  520. else:
  521. oid1 = oids
  522. oid2 = 2951
  523. _ext.UUID = _ext.new_type((oid1, ), "UUID",
  524. lambda data, cursor: data and uuid.UUID(data) or None)
  525. _ext.UUIDARRAY = _ext.new_array_type((oid2,), "UUID[]", _ext.UUID)
  526. _ext.register_type(_ext.UUID, conn_or_curs)
  527. _ext.register_type(_ext.UUIDARRAY, conn_or_curs)
  528. _ext.register_adapter(uuid.UUID, UUID_adapter)
  529. return _ext.UUID
  530. # a type, dbtype and adapter for PostgreSQL inet type
  531. class Inet(object):
  532. """Wrap a string to allow for correct SQL-quoting of inet values.
  533. Note that this adapter does NOT check the passed value to make
  534. sure it really is an inet-compatible address but DOES call adapt()
  535. on it to make sure it is impossible to execute an SQL-injection
  536. by passing an evil value to the initializer.
  537. """
  538. def __init__(self, addr):
  539. self.addr = addr
  540. def __repr__(self):
  541. return "%s(%r)" % (self.__class__.__name__, self.addr)
  542. def prepare(self, conn):
  543. self._conn = conn
  544. def getquoted(self):
  545. obj = _A(self.addr)
  546. if hasattr(obj, 'prepare'):
  547. obj.prepare(self._conn)
  548. return obj.getquoted() + b"::inet"
  549. def __conform__(self, proto):
  550. if proto is _ext.ISQLQuote:
  551. return self
  552. def __str__(self):
  553. return str(self.addr)
  554. def register_inet(oid=None, conn_or_curs=None):
  555. """Create the INET type and an Inet adapter.
  556. :param oid: oid for the PostgreSQL :sql:`inet` type, or 2-items sequence
  557. with oids of the type and the array. If not specified, use PostgreSQL
  558. standard oids.
  559. :param conn_or_curs: where to register the typecaster. If not specified,
  560. register it globally.
  561. """
  562. import warnings
  563. warnings.warn(
  564. "the inet adapter is deprecated, it's not very useful",
  565. DeprecationWarning)
  566. if not oid:
  567. oid1 = 869
  568. oid2 = 1041
  569. elif isinstance(oid, (list, tuple)):
  570. oid1, oid2 = oid
  571. else:
  572. oid1 = oid
  573. oid2 = 1041
  574. _ext.INET = _ext.new_type((oid1, ), "INET",
  575. lambda data, cursor: data and Inet(data) or None)
  576. _ext.INETARRAY = _ext.new_array_type((oid2, ), "INETARRAY", _ext.INET)
  577. _ext.register_type(_ext.INET, conn_or_curs)
  578. _ext.register_type(_ext.INETARRAY, conn_or_curs)
  579. return _ext.INET
  580. def wait_select(conn):
  581. """Wait until a connection or cursor has data available.
  582. The function is an example of a wait callback to be registered with
  583. `~psycopg2.extensions.set_wait_callback()`. This function uses
  584. :py:func:`~select.select()` to wait for data to become available, and
  585. therefore is able to handle/receive SIGINT/KeyboardInterrupt.
  586. """
  587. import select
  588. from psycopg2.extensions import POLL_OK, POLL_READ, POLL_WRITE
  589. while True:
  590. try:
  591. state = conn.poll()
  592. if state == POLL_OK:
  593. break
  594. elif state == POLL_READ:
  595. select.select([conn.fileno()], [], [])
  596. elif state == POLL_WRITE:
  597. select.select([], [conn.fileno()], [])
  598. else:
  599. raise conn.OperationalError("bad state from poll: %s" % state)
  600. except KeyboardInterrupt:
  601. conn.cancel()
  602. # the loop will be broken by a server error
  603. continue
  604. def _solve_conn_curs(conn_or_curs):
  605. """Return the connection and a DBAPI cursor from a connection or cursor."""
  606. if conn_or_curs is None:
  607. raise psycopg2.ProgrammingError("no connection or cursor provided")
  608. if hasattr(conn_or_curs, 'execute'):
  609. conn = conn_or_curs.connection
  610. curs = conn.cursor(cursor_factory=_cursor)
  611. else:
  612. conn = conn_or_curs
  613. curs = conn.cursor(cursor_factory=_cursor)
  614. return conn, curs
  615. class HstoreAdapter(object):
  616. """Adapt a Python dict to the hstore syntax."""
  617. def __init__(self, wrapped):
  618. self.wrapped = wrapped
  619. def prepare(self, conn):
  620. self.conn = conn
  621. # use an old-style getquoted implementation if required
  622. if conn.info.server_version < 90000:
  623. self.getquoted = self._getquoted_8
  624. def _getquoted_8(self):
  625. """Use the operators available in PG pre-9.0."""
  626. if not self.wrapped:
  627. return b"''::hstore"
  628. adapt = _ext.adapt
  629. rv = []
  630. for k, v in self.wrapped.items():
  631. k = adapt(k)
  632. k.prepare(self.conn)
  633. k = k.getquoted()
  634. if v is not None:
  635. v = adapt(v)
  636. v.prepare(self.conn)
  637. v = v.getquoted()
  638. else:
  639. v = b'NULL'
  640. # XXX this b'ing is painfully inefficient!
  641. rv.append(b"(" + k + b" => " + v + b")")
  642. return b"(" + b'||'.join(rv) + b")"
  643. def _getquoted_9(self):
  644. """Use the hstore(text[], text[]) function."""
  645. if not self.wrapped:
  646. return b"''::hstore"
  647. k = _ext.adapt(list(self.wrapped.keys()))
  648. k.prepare(self.conn)
  649. v = _ext.adapt(list(self.wrapped.values()))
  650. v.prepare(self.conn)
  651. return b"hstore(" + k.getquoted() + b", " + v.getquoted() + b")"
  652. getquoted = _getquoted_9
  653. _re_hstore = _re.compile(r"""
  654. # hstore key:
  655. # a string of normal or escaped chars
  656. "((?: [^"\\] | \\. )*)"
  657. \s*=>\s* # hstore value
  658. (?:
  659. NULL # the value can be null - not catched
  660. # or a quoted string like the key
  661. | "((?: [^"\\] | \\. )*)"
  662. )
  663. (?:\s*,\s*|$) # pairs separated by comma or end of string.
  664. """, _re.VERBOSE)
  665. @classmethod
  666. def parse(self, s, cur, _bsdec=_re.compile(r"\\(.)")):
  667. """Parse an hstore representation in a Python string.
  668. The hstore is represented as something like::
  669. "a"=>"1", "b"=>"2"
  670. with backslash-escaped strings.
  671. """
  672. if s is None:
  673. return None
  674. rv = {}
  675. start = 0
  676. for m in self._re_hstore.finditer(s):
  677. if m is None or m.start() != start:
  678. raise psycopg2.InterfaceError(
  679. "error parsing hstore pair at char %d" % start)
  680. k = _bsdec.sub(r'\1', m.group(1))
  681. v = m.group(2)
  682. if v is not None:
  683. v = _bsdec.sub(r'\1', v)
  684. rv[k] = v
  685. start = m.end()
  686. if start < len(s):
  687. raise psycopg2.InterfaceError(
  688. "error parsing hstore: unparsed data after char %d" % start)
  689. return rv
  690. @classmethod
  691. def parse_unicode(self, s, cur):
  692. """Parse an hstore returning unicode keys and values."""
  693. if s is None:
  694. return None
  695. s = s.decode(_ext.encodings[cur.connection.encoding])
  696. return self.parse(s, cur)
  697. @classmethod
  698. def get_oids(self, conn_or_curs):
  699. """Return the lists of OID of the hstore and hstore[] types.
  700. """
  701. conn, curs = _solve_conn_curs(conn_or_curs)
  702. # Store the transaction status of the connection to revert it after use
  703. conn_status = conn.status
  704. # column typarray not available before PG 8.3
  705. typarray = conn.info.server_version >= 80300 and "typarray" or "NULL"
  706. rv0, rv1 = [], []
  707. # get the oid for the hstore
  708. curs.execute("""\
  709. SELECT t.oid, %s
  710. FROM pg_type t JOIN pg_namespace ns
  711. ON typnamespace = ns.oid
  712. WHERE typname = 'hstore';
  713. """ % typarray)
  714. for oids in curs:
  715. rv0.append(oids[0])
  716. rv1.append(oids[1])
  717. # revert the status of the connection as before the command
  718. if (conn_status != _ext.STATUS_IN_TRANSACTION
  719. and not conn.autocommit):
  720. conn.rollback()
  721. return tuple(rv0), tuple(rv1)
  722. def register_hstore(conn_or_curs, globally=False, unicode=False,
  723. oid=None, array_oid=None):
  724. r"""Register adapter and typecaster for `!dict`\-\ |hstore| conversions.
  725. :param conn_or_curs: a connection or cursor: the typecaster will be
  726. registered only on this object unless *globally* is set to `!True`
  727. :param globally: register the adapter globally, not only on *conn_or_curs*
  728. :param unicode: if `!True`, keys and values returned from the database
  729. will be `!unicode` instead of `!str`. The option is not available on
  730. Python 3
  731. :param oid: the OID of the |hstore| type if known. If not, it will be
  732. queried on *conn_or_curs*.
  733. :param array_oid: the OID of the |hstore| array type if known. If not, it
  734. will be queried on *conn_or_curs*.
  735. The connection or cursor passed to the function will be used to query the
  736. database and look for the OID of the |hstore| type (which may be different
  737. across databases). If querying is not desirable (e.g. with
  738. :ref:`asynchronous connections <async-support>`) you may specify it in the
  739. *oid* parameter, which can be found using a query such as :sql:`SELECT
  740. 'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid*
  741. using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`.
  742. Note that, when passing a dictionary from Python to the database, both
  743. strings and unicode keys and values are supported. Dictionaries returned
  744. from the database have keys/values according to the *unicode* parameter.
  745. The |hstore| contrib module must be already installed in the database
  746. (executing the ``hstore.sql`` script in your ``contrib`` directory).
  747. Raise `~psycopg2.ProgrammingError` if the type is not found.
  748. """
  749. if oid is None:
  750. oid = HstoreAdapter.get_oids(conn_or_curs)
  751. if oid is None or not oid[0]:
  752. raise psycopg2.ProgrammingError(
  753. "hstore type not found in the database. "
  754. "please install it from your 'contrib/hstore.sql' file")
  755. else:
  756. array_oid = oid[1]
  757. oid = oid[0]
  758. if isinstance(oid, int):
  759. oid = (oid,)
  760. if array_oid is not None:
  761. if isinstance(array_oid, int):
  762. array_oid = (array_oid,)
  763. else:
  764. array_oid = tuple([x for x in array_oid if x])
  765. # create and register the typecaster
  766. if PY2 and unicode:
  767. cast = HstoreAdapter.parse_unicode
  768. else:
  769. cast = HstoreAdapter.parse
  770. HSTORE = _ext.new_type(oid, "HSTORE", cast)
  771. _ext.register_type(HSTORE, not globally and conn_or_curs or None)
  772. _ext.register_adapter(dict, HstoreAdapter)
  773. if array_oid:
  774. HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE)
  775. _ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)
  776. class CompositeCaster(object):
  777. """Helps conversion of a PostgreSQL composite type into a Python object.
  778. The class is usually created by the `register_composite()` function.
  779. You may want to create and register manually instances of the class if
  780. querying the database at registration time is not desirable (such as when
  781. using an :ref:`asynchronous connections <async-support>`).
  782. """
  783. def __init__(self, name, oid, attrs, array_oid=None, schema=None):
  784. self.name = name
  785. self.schema = schema
  786. self.oid = oid
  787. self.array_oid = array_oid
  788. self.attnames = [a[0] for a in attrs]
  789. self.atttypes = [a[1] for a in attrs]
  790. self._create_type(name, self.attnames)
  791. self.typecaster = _ext.new_type((oid,), name, self.parse)
  792. if array_oid:
  793. self.array_typecaster = _ext.new_array_type(
  794. (array_oid,), "%sARRAY" % name, self.typecaster)
  795. else:
  796. self.array_typecaster = None
  797. def parse(self, s, curs):
  798. if s is None:
  799. return None
  800. tokens = self.tokenize(s)
  801. if len(tokens) != len(self.atttypes):
  802. raise psycopg2.DataError(
  803. "expecting %d components for the type %s, %d found instead" %
  804. (len(self.atttypes), self.name, len(tokens)))
  805. values = [curs.cast(oid, token)
  806. for oid, token in zip(self.atttypes, tokens)]
  807. return self.make(values)
  808. def make(self, values):
  809. """Return a new Python object representing the data being casted.
  810. *values* is the list of attributes, already casted into their Python
  811. representation.
  812. You can subclass this method to :ref:`customize the composite cast
  813. <custom-composite>`.
  814. """
  815. return self._ctor(values)
  816. _re_tokenize = _re.compile(r"""
  817. \(? ([,)]) # an empty token, representing NULL
  818. | \(? " ((?: [^"] | "")*) " [,)] # or a quoted string
  819. | \(? ([^",)]+) [,)] # or an unquoted string
  820. """, _re.VERBOSE)
  821. _re_undouble = _re.compile(r'(["\\])\1')
  822. @classmethod
  823. def tokenize(self, s):
  824. rv = []
  825. for m in self._re_tokenize.finditer(s):
  826. if m is None:
  827. raise psycopg2.InterfaceError("can't parse type: %r" % s)
  828. if m.group(1) is not None:
  829. rv.append(None)
  830. elif m.group(2) is not None:
  831. rv.append(self._re_undouble.sub(r"\1", m.group(2)))
  832. else:
  833. rv.append(m.group(3))
  834. return rv
  835. def _create_type(self, name, attnames):
  836. self.type = namedtuple(name, attnames)
  837. self._ctor = self.type._make
  838. @classmethod
  839. def _from_db(self, name, conn_or_curs):
  840. """Return a `CompositeCaster` instance for the type *name*.
  841. Raise `ProgrammingError` if the type is not found.
  842. """
  843. conn, curs = _solve_conn_curs(conn_or_curs)
  844. # Store the transaction status of the connection to revert it after use
  845. conn_status = conn.status
  846. # Use the correct schema
  847. if '.' in name:
  848. schema, tname = name.split('.', 1)
  849. else:
  850. tname = name
  851. schema = 'public'
  852. # column typarray not available before PG 8.3
  853. typarray = conn.info.server_version >= 80300 and "typarray" or "NULL"
  854. # get the type oid and attributes
  855. curs.execute("""\
  856. SELECT t.oid, %s, attname, atttypid
  857. FROM pg_type t
  858. JOIN pg_namespace ns ON typnamespace = ns.oid
  859. JOIN pg_attribute a ON attrelid = typrelid
  860. WHERE typname = %%s AND nspname = %%s
  861. AND attnum > 0 AND NOT attisdropped
  862. ORDER BY attnum;
  863. """ % typarray, (tname, schema))
  864. recs = curs.fetchall()
  865. # revert the status of the connection as before the command
  866. if (conn_status != _ext.STATUS_IN_TRANSACTION
  867. and not conn.autocommit):
  868. conn.rollback()
  869. if not recs:
  870. raise psycopg2.ProgrammingError(
  871. "PostgreSQL type '%s' not found" % name)
  872. type_oid = recs[0][0]
  873. array_oid = recs[0][1]
  874. type_attrs = [(r[2], r[3]) for r in recs]
  875. return self(tname, type_oid, type_attrs,
  876. array_oid=array_oid, schema=schema)
  877. def register_composite(name, conn_or_curs, globally=False, factory=None):
  878. """Register a typecaster to convert a composite type into a tuple.
  879. :param name: the name of a PostgreSQL composite type, e.g. created using
  880. the |CREATE TYPE|_ command
  881. :param conn_or_curs: a connection or cursor used to find the type oid and
  882. components; the typecaster is registered in a scope limited to this
  883. object, unless *globally* is set to `!True`
  884. :param globally: if `!False` (default) register the typecaster only on
  885. *conn_or_curs*, otherwise register it globally
  886. :param factory: if specified it should be a `CompositeCaster` subclass: use
  887. it to :ref:`customize how to cast composite types <custom-composite>`
  888. :return: the registered `CompositeCaster` or *factory* instance
  889. responsible for the conversion
  890. """
  891. if factory is None:
  892. factory = CompositeCaster
  893. caster = factory._from_db(name, conn_or_curs)
  894. _ext.register_type(caster.typecaster, not globally and conn_or_curs or None)
  895. if caster.array_typecaster is not None:
  896. _ext.register_type(
  897. caster.array_typecaster, not globally and conn_or_curs or None)
  898. return caster
  899. def _paginate(seq, page_size):
  900. """Consume an iterable and return it in chunks.
  901. Every chunk is at most `page_size`. Never return an empty chunk.
  902. """
  903. page = []
  904. it = iter(seq)
  905. while True:
  906. try:
  907. for i in range(page_size):
  908. page.append(next(it))
  909. yield page
  910. page = []
  911. except StopIteration:
  912. if page:
  913. yield page
  914. return
  915. def execute_batch(cur, sql, argslist, page_size=100):
  916. r"""Execute groups of statements in fewer server roundtrips.
  917. Execute *sql* several times, against all parameters set (sequences or
  918. mappings) found in *argslist*.
  919. The function is semantically similar to
  920. .. parsed-literal::
  921. *cur*\.\ `~cursor.executemany`\ (\ *sql*\ , *argslist*\ )
  922. but has a different implementation: Psycopg will join the statements into
  923. fewer multi-statement commands, each one containing at most *page_size*
  924. statements, resulting in a reduced number of server roundtrips.
  925. After the execution of the function the `cursor.rowcount` property will
  926. **not** contain a total result.
  927. """
  928. for page in _paginate(argslist, page_size=page_size):
  929. sqls = [cur.mogrify(sql, args) for args in page]
  930. cur.execute(b";".join(sqls))
  931. def execute_values(cur, sql, argslist, template=None, page_size=100, fetch=False):
  932. '''Execute a statement using :sql:`VALUES` with a sequence of parameters.
  933. :param cur: the cursor to use to execute the query.
  934. :param sql: the query to execute. It must contain a single ``%s``
  935. placeholder, which will be replaced by a `VALUES list`__.
  936. Example: ``"INSERT INTO mytable (id, f1, f2) VALUES %s"``.
  937. :param argslist: sequence of sequences or dictionaries with the arguments
  938. to send to the query. The type and content must be consistent with
  939. *template*.
  940. :param template: the snippet to merge to every item in *argslist* to
  941. compose the query.
  942. - If the *argslist* items are sequences it should contain positional
  943. placeholders (e.g. ``"(%s, %s, %s)"``, or ``"(%s, %s, 42)``" if there
  944. are constants value...).
  945. - If the *argslist* items are mappings it should contain named
  946. placeholders (e.g. ``"(%(id)s, %(f1)s, 42)"``).
  947. If not specified, assume the arguments are sequence and use a simple
  948. positional template (i.e. ``(%s, %s, ...)``), with the number of
  949. placeholders sniffed by the first element in *argslist*.
  950. :param page_size: maximum number of *argslist* items to include in every
  951. statement. If there are more items the function will execute more than
  952. one statement.
  953. :param fetch: if `!True` return the query results into a list (like in a
  954. `~cursor.fetchall()`). Useful for queries with :sql:`RETURNING`
  955. clause.
  956. .. __: https://www.postgresql.org/docs/current/static/queries-values.html
  957. After the execution of the function the `cursor.rowcount` property will
  958. **not** contain a total result.
  959. While :sql:`INSERT` is an obvious candidate for this function it is
  960. possible to use it with other statements, for example::
  961. >>> cur.execute(
  962. ... "create table test (id int primary key, v1 int, v2 int)")
  963. >>> execute_values(cur,
  964. ... "INSERT INTO test (id, v1, v2) VALUES %s",
  965. ... [(1, 2, 3), (4, 5, 6), (7, 8, 9)])
  966. >>> execute_values(cur,
  967. ... """UPDATE test SET v1 = data.v1 FROM (VALUES %s) AS data (id, v1)
  968. ... WHERE test.id = data.id""",
  969. ... [(1, 20), (4, 50)])
  970. >>> cur.execute("select * from test order by id")
  971. >>> cur.fetchall()
  972. [(1, 20, 3), (4, 50, 6), (7, 8, 9)])
  973. '''
  974. from psycopg2.sql import Composable
  975. if isinstance(sql, Composable):
  976. sql = sql.as_string(cur)
  977. # we can't just use sql % vals because vals is bytes: if sql is bytes
  978. # there will be some decoding error because of stupid codec used, and Py3
  979. # doesn't implement % on bytes.
  980. if not isinstance(sql, bytes):
  981. sql = sql.encode(_ext.encodings[cur.connection.encoding])
  982. pre, post = _split_sql(sql)
  983. result = [] if fetch else None
  984. for page in _paginate(argslist, page_size=page_size):
  985. if template is None:
  986. template = b'(' + b','.join([b'%s'] * len(page[0])) + b')'
  987. parts = pre[:]
  988. for args in page:
  989. parts.append(cur.mogrify(template, args))
  990. parts.append(b',')
  991. parts[-1:] = post
  992. cur.execute(b''.join(parts))
  993. if fetch:
  994. result.extend(cur.fetchall())
  995. return result
  996. def _split_sql(sql):
  997. """Split *sql* on a single ``%s`` placeholder.
  998. Split on the %s, perform %% replacement and return pre, post lists of
  999. snippets.
  1000. """
  1001. curr = pre = []
  1002. post = []
  1003. tokens = _re.split(br'(%.)', sql)
  1004. for token in tokens:
  1005. if len(token) != 2 or token[:1] != b'%':
  1006. curr.append(token)
  1007. continue
  1008. if token[1:] == b's':
  1009. if curr is pre:
  1010. curr = post
  1011. else:
  1012. raise ValueError(
  1013. "the query contains more than one '%s' placeholder")
  1014. elif token[1:] == b'%':
  1015. curr.append(b'%')
  1016. else:
  1017. raise ValueError("unsupported format character: '%s'"
  1018. % token[1:].decode('ascii', 'replace'))
  1019. if curr is pre:
  1020. raise ValueError("the query doesn't contain any '%s' placeholder")
  1021. return pre, post