PageRenderTime 51ms CodeModel.GetById 17ms RepoModel.GetById 1ms app.codeStats 0ms

/hd-venv/lib/python2.7/site-packages/psycopg2/extras.py

https://github.com/Riegerb/GSWD_Tutorial
Python | 973 lines | 864 code | 57 blank | 52 comment | 32 complexity | e94ad63b67f5505bf2e38eaee9c085e2 MD5 | raw file
  1. """Miscellaneous goodies for psycopg2
  2. This module is a generic place used to hold little helper functions
  3. and classes untill a better place in the distribution is found.
  4. """
  5. # psycopg/extras.py - miscellaneous extra goodies for psycopg
  6. #
  7. # Copyright (C) 2003-2010 Federico Di Gregorio <fog@debian.org>
  8. #
  9. # psycopg2 is free software: you can redistribute it and/or modify it
  10. # under the terms of the GNU Lesser General Public License as published
  11. # by the Free Software Foundation, either version 3 of the License, or
  12. # (at your option) any later version.
  13. #
  14. # In addition, as a special exception, the copyright holders give
  15. # permission to link this program with the OpenSSL library (or with
  16. # modified versions of OpenSSL that use the same license as OpenSSL),
  17. # and distribute linked combinations including the two.
  18. #
  19. # You must obey the GNU Lesser General Public License in all respects for
  20. # all of the code used other than OpenSSL.
  21. #
  22. # psycopg2 is distributed in the hope that it will be useful, but WITHOUT
  23. # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  24. # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
  25. # License for more details.
  26. import os
  27. import sys
  28. import time
  29. import warnings
  30. import re as regex
  31. try:
  32. import logging
  33. except:
  34. logging = None
  35. import psycopg2
  36. from psycopg2 import extensions as _ext
  37. from psycopg2.extensions import cursor as _cursor
  38. from psycopg2.extensions import connection as _connection
  39. from psycopg2.extensions import adapt as _A
  40. from psycopg2.extensions import b
  41. class DictCursorBase(_cursor):
  42. """Base class for all dict-like cursors."""
  43. def __init__(self, *args, **kwargs):
  44. if 'row_factory' in kwargs:
  45. row_factory = kwargs['row_factory']
  46. del kwargs['row_factory']
  47. else:
  48. raise NotImplementedError(
  49. "DictCursorBase can't be instantiated without a row factory.")
  50. super(DictCursorBase, self).__init__(*args, **kwargs)
  51. self._query_executed = 0
  52. self._prefetch = 0
  53. self.row_factory = row_factory
  54. def fetchone(self):
  55. if self._prefetch:
  56. res = super(DictCursorBase, self).fetchone()
  57. if self._query_executed:
  58. self._build_index()
  59. if not self._prefetch:
  60. res = super(DictCursorBase, self).fetchone()
  61. return res
  62. def fetchmany(self, size=None):
  63. if self._prefetch:
  64. res = super(DictCursorBase, self).fetchmany(size)
  65. if self._query_executed:
  66. self._build_index()
  67. if not self._prefetch:
  68. res = super(DictCursorBase, self).fetchmany(size)
  69. return res
  70. def fetchall(self):
  71. if self._prefetch:
  72. res = super(DictCursorBase, self).fetchall()
  73. if self._query_executed:
  74. self._build_index()
  75. if not self._prefetch:
  76. res = super(DictCursorBase, self).fetchall()
  77. return res
  78. def __iter__(self):
  79. if self._prefetch:
  80. res = super(DictCursorBase, self).__iter__()
  81. first = res.next()
  82. if self._query_executed:
  83. self._build_index()
  84. if not self._prefetch:
  85. res = super(DictCursorBase, self).__iter__()
  86. first = res.next()
  87. yield first
  88. while 1:
  89. yield res.next()
  90. class DictConnection(_connection):
  91. """A connection that uses `DictCursor` automatically."""
  92. def cursor(self, *args, **kwargs):
  93. kwargs.setdefault('cursor_factory', DictCursor)
  94. return super(DictConnection, self).cursor(*args, **kwargs)
  95. class DictCursor(DictCursorBase):
  96. """A cursor that keeps a list of column name -> index mappings."""
  97. def __init__(self, *args, **kwargs):
  98. kwargs['row_factory'] = DictRow
  99. super(DictCursor, self).__init__(*args, **kwargs)
  100. self._prefetch = 1
  101. def execute(self, query, vars=None):
  102. self.index = {}
  103. self._query_executed = 1
  104. return super(DictCursor, self).execute(query, vars)
  105. def callproc(self, procname, vars=None):
  106. self.index = {}
  107. self._query_executed = 1
  108. return super(DictCursor, self).callproc(procname, vars)
  109. def _build_index(self):
  110. if self._query_executed == 1 and self.description:
  111. for i in range(len(self.description)):
  112. self.index[self.description[i][0]] = i
  113. self._query_executed = 0
  114. class DictRow(list):
  115. """A row object that allow by-colmun-name access to data."""
  116. __slots__ = ('_index',)
  117. def __init__(self, cursor):
  118. self._index = cursor.index
  119. self[:] = [None] * len(cursor.description)
  120. def __getitem__(self, x):
  121. if not isinstance(x, (int, slice)):
  122. x = self._index[x]
  123. return list.__getitem__(self, x)
  124. def __setitem__(self, x, v):
  125. if not isinstance(x, (int, slice)):
  126. x = self._index[x]
  127. list.__setitem__(self, x, v)
  128. def items(self):
  129. return list(self.iteritems())
  130. def keys(self):
  131. return self._index.keys()
  132. def values(self):
  133. return tuple(self[:])
  134. def has_key(self, x):
  135. return x in self._index
  136. def get(self, x, default=None):
  137. try:
  138. return self[x]
  139. except:
  140. return default
  141. def iteritems(self):
  142. for n, v in self._index.iteritems():
  143. yield n, list.__getitem__(self, v)
  144. def iterkeys(self):
  145. return self._index.iterkeys()
  146. def itervalues(self):
  147. return list.__iter__(self)
  148. def copy(self):
  149. return dict(self.iteritems())
  150. def __contains__(self, x):
  151. return x in self._index
  152. def __getstate__(self):
  153. return self[:], self._index.copy()
  154. def __setstate__(self, data):
  155. self[:] = data[0]
  156. self._index = data[1]
  157. # drop the crusty Py2 methods
  158. if sys.version_info[0] > 2:
  159. items = iteritems; del iteritems
  160. keys = iterkeys; del iterkeys
  161. values = itervalues; del itervalues
  162. del has_key
  163. class RealDictConnection(_connection):
  164. """A connection that uses `RealDictCursor` automatically."""
  165. def cursor(self, *args, **kwargs):
  166. kwargs.setdefault('cursor_factory', RealDictCursor)
  167. return super(RealDictConnection, self).cursor(*args, **kwargs)
  168. class RealDictCursor(DictCursorBase):
  169. """A cursor that uses a real dict as the base type for rows.
  170. Note that this cursor is extremely specialized and does not allow
  171. the normal access (using integer indices) to fetched data. If you need
  172. to access database rows both as a dictionary and a list, then use
  173. the generic `DictCursor` instead of `!RealDictCursor`.
  174. """
  175. def __init__(self, *args, **kwargs):
  176. kwargs['row_factory'] = RealDictRow
  177. super(RealDictCursor, self).__init__(*args, **kwargs)
  178. self._prefetch = 0
  179. def execute(self, query, vars=None):
  180. self.column_mapping = []
  181. self._query_executed = 1
  182. return super(RealDictCursor, self).execute(query, vars)
  183. def callproc(self, procname, vars=None):
  184. self.column_mapping = []
  185. self._query_executed = 1
  186. return super(RealDictCursor, self).callproc(procname, vars)
  187. def _build_index(self):
  188. if self._query_executed == 1 and self.description:
  189. for i in range(len(self.description)):
  190. self.column_mapping.append(self.description[i][0])
  191. self._query_executed = 0
  192. class RealDictRow(dict):
  193. """A `!dict` subclass representing a data record."""
  194. __slots__ = ('_column_mapping')
  195. def __init__(self, cursor):
  196. dict.__init__(self)
  197. # Required for named cursors
  198. if cursor.description and not cursor.column_mapping:
  199. cursor._build_index()
  200. self._column_mapping = cursor.column_mapping
  201. def __setitem__(self, name, value):
  202. if type(name) == int:
  203. name = self._column_mapping[name]
  204. return dict.__setitem__(self, name, value)
  205. def __getstate__(self):
  206. return (self.copy(), self._column_mapping[:])
  207. def __setstate__(self, data):
  208. self.update(data[0])
  209. self._column_mapping = data[1]
  210. class NamedTupleConnection(_connection):
  211. """A connection that uses `NamedTupleCursor` automatically."""
  212. def cursor(self, *args, **kwargs):
  213. kwargs.setdefault('cursor_factory', NamedTupleCursor)
  214. return super(NamedTupleConnection, self).cursor(*args, **kwargs)
  215. class NamedTupleCursor(_cursor):
  216. """A cursor that generates results as `~collections.namedtuple`.
  217. `!fetch*()` methods will return named tuples instead of regular tuples, so
  218. their elements can be accessed both as regular numeric items as well as
  219. attributes.
  220. >>> nt_cur = conn.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor)
  221. >>> rec = nt_cur.fetchone()
  222. >>> rec
  223. Record(id=1, num=100, data="abc'def")
  224. >>> rec[1]
  225. 100
  226. >>> rec.data
  227. "abc'def"
  228. """
  229. Record = None
  230. def execute(self, query, vars=None):
  231. self.Record = None
  232. return super(NamedTupleCursor, self).execute(query, vars)
  233. def executemany(self, query, vars):
  234. self.Record = None
  235. return super(NamedTupleCursor, self).executemany(query, vars)
  236. def callproc(self, procname, vars=None):
  237. self.Record = None
  238. return super(NamedTupleCursor, self).callproc(procname, vars)
  239. def fetchone(self):
  240. t = super(NamedTupleCursor, self).fetchone()
  241. if t is not None:
  242. nt = self.Record
  243. if nt is None:
  244. nt = self.Record = self._make_nt()
  245. return nt(*t)
  246. def fetchmany(self, size=None):
  247. ts = super(NamedTupleCursor, self).fetchmany(size)
  248. nt = self.Record
  249. if nt is None:
  250. nt = self.Record = self._make_nt()
  251. return [nt(*t) for t in ts]
  252. def fetchall(self):
  253. ts = super(NamedTupleCursor, self).fetchall()
  254. nt = self.Record
  255. if nt is None:
  256. nt = self.Record = self._make_nt()
  257. return [nt(*t) for t in ts]
  258. def __iter__(self):
  259. it = super(NamedTupleCursor, self).__iter__()
  260. t = it.next()
  261. nt = self.Record
  262. if nt is None:
  263. nt = self.Record = self._make_nt()
  264. yield nt(*t)
  265. while 1:
  266. yield nt(*it.next())
  267. try:
  268. from collections import namedtuple
  269. except ImportError, _exc:
  270. def _make_nt(self):
  271. raise self._exc
  272. else:
  273. def _make_nt(self, namedtuple=namedtuple):
  274. return namedtuple("Record", [d[0] for d in self.description or ()])
  275. class LoggingConnection(_connection):
  276. """A connection that logs all queries to a file or logger__ object.
  277. .. __: http://docs.python.org/library/logging.html
  278. """
  279. def initialize(self, logobj):
  280. """Initialize the connection to log to `!logobj`.
  281. The `!logobj` parameter can be an open file object or a Logger
  282. instance from the standard logging module.
  283. """
  284. self._logobj = logobj
  285. if logging and isinstance(logobj, logging.Logger):
  286. self.log = self._logtologger
  287. else:
  288. self.log = self._logtofile
  289. def filter(self, msg, curs):
  290. """Filter the query before logging it.
  291. This is the method to overwrite to filter unwanted queries out of the
  292. log or to add some extra data to the output. The default implementation
  293. just does nothing.
  294. """
  295. return msg
  296. def _logtofile(self, msg, curs):
  297. msg = self.filter(msg, curs)
  298. if msg: self._logobj.write(msg + os.linesep)
  299. def _logtologger(self, msg, curs):
  300. msg = self.filter(msg, curs)
  301. if msg: self._logobj.debug(msg)
  302. def _check(self):
  303. if not hasattr(self, '_logobj'):
  304. raise self.ProgrammingError(
  305. "LoggingConnection object has not been initialize()d")
  306. def cursor(self, *args, **kwargs):
  307. self._check()
  308. kwargs.setdefault('cursor_factory', LoggingCursor)
  309. return super(LoggingConnection, self).cursor(*args, **kwargs)
  310. class LoggingCursor(_cursor):
  311. """A cursor that logs queries using its connection logging facilities."""
  312. def execute(self, query, vars=None):
  313. try:
  314. return super(LoggingCursor, self).execute(query, vars)
  315. finally:
  316. self.connection.log(self.query, self)
  317. def callproc(self, procname, vars=None):
  318. try:
  319. return super(LoggingCursor, self).callproc(procname, vars)
  320. finally:
  321. self.connection.log(self.query, self)
  322. class MinTimeLoggingConnection(LoggingConnection):
  323. """A connection that logs queries based on execution time.
  324. This is just an example of how to sub-class `LoggingConnection` to
  325. provide some extra filtering for the logged queries. Both the
  326. `inizialize()` and `filter()` methods are overwritten to make sure
  327. that only queries executing for more than ``mintime`` ms are logged.
  328. Note that this connection uses the specialized cursor
  329. `MinTimeLoggingCursor`.
  330. """
  331. def initialize(self, logobj, mintime=0):
  332. LoggingConnection.initialize(self, logobj)
  333. self._mintime = mintime
  334. def filter(self, msg, curs):
  335. t = (time.time() - curs.timestamp) * 1000
  336. if t > self._mintime:
  337. return msg + os.linesep + " (execution time: %d ms)" % t
  338. def cursor(self, *args, **kwargs):
  339. kwargs.setdefault('cursor_factory', MinTimeLoggingCursor)
  340. return LoggingConnection.cursor(self, *args, **kwargs)
  341. class MinTimeLoggingCursor(LoggingCursor):
  342. """The cursor sub-class companion to `MinTimeLoggingConnection`."""
  343. def execute(self, query, vars=None):
  344. self.timestamp = time.time()
  345. return LoggingCursor.execute(self, query, vars)
  346. def callproc(self, procname, vars=None):
  347. self.timestamp = time.time()
  348. return LoggingCursor.execute(self, procname, vars)
  349. # a dbtype and adapter for Python UUID type
  350. class UUID_adapter(object):
  351. """Adapt Python's uuid.UUID__ type to PostgreSQL's uuid__.
  352. .. __: http://docs.python.org/library/uuid.html
  353. .. __: http://www.postgresql.org/docs/current/static/datatype-uuid.html
  354. """
  355. def __init__(self, uuid):
  356. self._uuid = uuid
  357. def prepare(self, conn):
  358. pass
  359. def getquoted(self):
  360. return "'"+str(self._uuid)+"'::uuid"
  361. __str__ = getquoted
  362. def register_uuid(oids=None, conn_or_curs=None):
  363. """Create the UUID type and an uuid.UUID adapter.
  364. :param oids: oid for the PostgreSQL :sql:`uuid` type, or 2-items sequence
  365. with oids of the type and the array. If not specified, use PostgreSQL
  366. standard oids.
  367. :param conn_or_curs: where to register the typecaster. If not specified,
  368. register it globally.
  369. """
  370. import uuid
  371. if not oids:
  372. oid1 = 2950
  373. oid2 = 2951
  374. elif isinstance(oids, (list, tuple)):
  375. oid1, oid2 = oids
  376. else:
  377. oid1 = oids
  378. oid2 = 2951
  379. _ext.UUID = _ext.new_type((oid1, ), "UUID",
  380. lambda data, cursor: data and uuid.UUID(data) or None)
  381. _ext.UUIDARRAY = _ext.new_array_type((oid2,), "UUID[]", _ext.UUID)
  382. _ext.register_type(_ext.UUID, conn_or_curs)
  383. _ext.register_type(_ext.UUIDARRAY, conn_or_curs)
  384. _ext.register_adapter(uuid.UUID, UUID_adapter)
  385. return _ext.UUID
  386. # a type, dbtype and adapter for PostgreSQL inet type
  387. class Inet(object):
  388. """Wrap a string to allow for correct SQL-quoting of inet values.
  389. Note that this adapter does NOT check the passed value to make
  390. sure it really is an inet-compatible address but DOES call adapt()
  391. on it to make sure it is impossible to execute an SQL-injection
  392. by passing an evil value to the initializer.
  393. """
  394. def __init__(self, addr):
  395. self.addr = addr
  396. def __repr__(self):
  397. return "%s(%r)" % (self.__class__.__name__, self.addr)
  398. def prepare(self, conn):
  399. self._conn = conn
  400. def getquoted(self):
  401. obj = _A(self.addr)
  402. if hasattr(obj, 'prepare'):
  403. obj.prepare(self._conn)
  404. return obj.getquoted() + b("::inet")
  405. def __conform__(self, foo):
  406. if foo is _ext.ISQLQuote:
  407. return self
  408. def __str__(self):
  409. return str(self.addr)
  410. def register_inet(oid=None, conn_or_curs=None):
  411. """Create the INET type and an Inet adapter.
  412. :param oid: oid for the PostgreSQL :sql:`inet` type, or 2-items sequence
  413. with oids of the type and the array. If not specified, use PostgreSQL
  414. standard oids.
  415. :param conn_or_curs: where to register the typecaster. If not specified,
  416. register it globally.
  417. """
  418. if not oid:
  419. oid1 = 869
  420. oid2 = 1041
  421. elif isinstance(oid, (list, tuple)):
  422. oid1, oid2 = oid
  423. else:
  424. oid1 = oid
  425. oid2 = 1041
  426. _ext.INET = _ext.new_type((oid1, ), "INET",
  427. lambda data, cursor: data and Inet(data) or None)
  428. _ext.INETARRAY = _ext.new_array_type((oid2, ), "INETARRAY", _ext.INET)
  429. _ext.register_type(_ext.INET, conn_or_curs)
  430. _ext.register_type(_ext.INETARRAY, conn_or_curs)
  431. return _ext.INET
  432. def register_tstz_w_secs(oids=None, conn_or_curs=None):
  433. """The function used to register an alternate type caster for
  434. :sql:`TIMESTAMP WITH TIME ZONE` to deal with historical time zones with
  435. seconds in the UTC offset.
  436. These are now correctly handled by the default type caster, so currently
  437. the function doesn't do anything.
  438. """
  439. warnings.warn("deprecated", DeprecationWarning)
  440. import select
  441. from psycopg2.extensions import POLL_OK, POLL_READ, POLL_WRITE
  442. from psycopg2 import OperationalError
  443. def wait_select(conn):
  444. """Wait until a connection or cursor has data available.
  445. The function is an example of a wait callback to be registered with
  446. `~psycopg2.extensions.set_wait_callback()`. This function uses `!select()`
  447. to wait for data available.
  448. """
  449. while 1:
  450. state = conn.poll()
  451. if state == POLL_OK:
  452. break
  453. elif state == POLL_READ:
  454. select.select([conn.fileno()], [], [])
  455. elif state == POLL_WRITE:
  456. select.select([], [conn.fileno()], [])
  457. else:
  458. raise OperationalError("bad state from poll: %s" % state)
  459. def _solve_conn_curs(conn_or_curs):
  460. """Return the connection and a DBAPI cursor from a connection or cursor."""
  461. if hasattr(conn_or_curs, 'execute'):
  462. conn = conn_or_curs.connection
  463. curs = conn.cursor(cursor_factory=_cursor)
  464. else:
  465. conn = conn_or_curs
  466. curs = conn.cursor(cursor_factory=_cursor)
  467. return conn, curs
  468. class HstoreAdapter(object):
  469. """Adapt a Python dict to the hstore syntax."""
  470. def __init__(self, wrapped):
  471. self.wrapped = wrapped
  472. def prepare(self, conn):
  473. self.conn = conn
  474. # use an old-style getquoted implementation if required
  475. if conn.server_version < 90000:
  476. self.getquoted = self._getquoted_8
  477. def _getquoted_8(self):
  478. """Use the operators available in PG pre-9.0."""
  479. if not self.wrapped:
  480. return b("''::hstore")
  481. adapt = _ext.adapt
  482. rv = []
  483. for k, v in self.wrapped.iteritems():
  484. k = adapt(k)
  485. k.prepare(self.conn)
  486. k = k.getquoted()
  487. if v is not None:
  488. v = adapt(v)
  489. v.prepare(self.conn)
  490. v = v.getquoted()
  491. else:
  492. v = b('NULL')
  493. # XXX this b'ing is painfully inefficient!
  494. rv.append(b("(") + k + b(" => ") + v + b(")"))
  495. return b("(") + b('||').join(rv) + b(")")
  496. def _getquoted_9(self):
  497. """Use the hstore(text[], text[]) function."""
  498. if not self.wrapped:
  499. return b("''::hstore")
  500. k = _ext.adapt(self.wrapped.keys())
  501. k.prepare(self.conn)
  502. v = _ext.adapt(self.wrapped.values())
  503. v.prepare(self.conn)
  504. return b("hstore(") + k.getquoted() + b(", ") + v.getquoted() + b(")")
  505. getquoted = _getquoted_9
  506. _re_hstore = regex.compile(r"""
  507. # hstore key:
  508. # a string of normal or escaped chars
  509. "((?: [^"\\] | \\. )*)"
  510. \s*=>\s* # hstore value
  511. (?:
  512. NULL # the value can be null - not catched
  513. # or a quoted string like the key
  514. | "((?: [^"\\] | \\. )*)"
  515. )
  516. (?:\s*,\s*|$) # pairs separated by comma or end of string.
  517. """, regex.VERBOSE)
  518. @classmethod
  519. def parse(self, s, cur, _bsdec=regex.compile(r"\\(.)")):
  520. """Parse an hstore representation in a Python string.
  521. The hstore is represented as something like::
  522. "a"=>"1", "b"=>"2"
  523. with backslash-escaped strings.
  524. """
  525. if s is None:
  526. return None
  527. rv = {}
  528. start = 0
  529. for m in self._re_hstore.finditer(s):
  530. if m is None or m.start() != start:
  531. raise psycopg2.InterfaceError(
  532. "error parsing hstore pair at char %d" % start)
  533. k = _bsdec.sub(r'\1', m.group(1))
  534. v = m.group(2)
  535. if v is not None:
  536. v = _bsdec.sub(r'\1', v)
  537. rv[k] = v
  538. start = m.end()
  539. if start < len(s):
  540. raise psycopg2.InterfaceError(
  541. "error parsing hstore: unparsed data after char %d" % start)
  542. return rv
  543. @classmethod
  544. def parse_unicode(self, s, cur):
  545. """Parse an hstore returning unicode keys and values."""
  546. if s is None:
  547. return None
  548. s = s.decode(_ext.encodings[cur.connection.encoding])
  549. return self.parse(s, cur)
  550. @classmethod
  551. def get_oids(self, conn_or_curs):
  552. """Return the lists of OID of the hstore and hstore[] types.
  553. """
  554. conn, curs = _solve_conn_curs(conn_or_curs)
  555. # Store the transaction status of the connection to revert it after use
  556. conn_status = conn.status
  557. # column typarray not available before PG 8.3
  558. typarray = conn.server_version >= 80300 and "typarray" or "NULL"
  559. rv0, rv1 = [], []
  560. # get the oid for the hstore
  561. curs.execute("""\
  562. SELECT t.oid, %s
  563. FROM pg_type t JOIN pg_namespace ns
  564. ON typnamespace = ns.oid
  565. WHERE typname = 'hstore';
  566. """ % typarray)
  567. for oids in curs:
  568. rv0.append(oids[0])
  569. rv1.append(oids[1])
  570. # revert the status of the connection as before the command
  571. if (conn_status != _ext.STATUS_IN_TRANSACTION
  572. and not conn.autocommit):
  573. conn.rollback()
  574. return tuple(rv0), tuple(rv1)
  575. def register_hstore(conn_or_curs, globally=False, unicode=False,
  576. oid=None, array_oid=None):
  577. """Register adapter and typecaster for `!dict`\-\ |hstore| conversions.
  578. :param conn_or_curs: a connection or cursor: the typecaster will be
  579. registered only on this object unless *globally* is set to `!True`
  580. :param globally: register the adapter globally, not only on *conn_or_curs*
  581. :param unicode: if `!True`, keys and values returned from the database
  582. will be `!unicode` instead of `!str`. The option is not available on
  583. Python 3
  584. :param oid: the OID of the |hstore| type if known. If not, it will be
  585. queried on *conn_or_curs*.
  586. :param array_oid: the OID of the |hstore| array type if known. If not, it
  587. will be queried on *conn_or_curs*.
  588. The connection or cursor passed to the function will be used to query the
  589. database and look for the OID of the |hstore| type (which may be different
  590. across databases). If querying is not desirable (e.g. with
  591. :ref:`asynchronous connections <async-support>`) you may specify it in the
  592. *oid* parameter, which can be found using a query such as :sql:`SELECT
  593. 'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid*
  594. using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`.
  595. Note that, when passing a dictionary from Python to the database, both
  596. strings and unicode keys and values are supported. Dictionaries returned
  597. from the database have keys/values according to the *unicode* parameter.
  598. The |hstore| contrib module must be already installed in the database
  599. (executing the ``hstore.sql`` script in your ``contrib`` directory).
  600. Raise `~psycopg2.ProgrammingError` if the type is not found.
  601. """
  602. if oid is None:
  603. oid = HstoreAdapter.get_oids(conn_or_curs)
  604. if oid is None or not oid[0]:
  605. raise psycopg2.ProgrammingError(
  606. "hstore type not found in the database. "
  607. "please install it from your 'contrib/hstore.sql' file")
  608. else:
  609. array_oid = oid[1]
  610. oid = oid[0]
  611. if isinstance(oid, int):
  612. oid = (oid,)
  613. if array_oid is not None:
  614. if isinstance(array_oid, int):
  615. array_oid = (array_oid,)
  616. else:
  617. array_oid = tuple([x for x in array_oid if x])
  618. # create and register the typecaster
  619. if sys.version_info[0] < 3 and unicode:
  620. cast = HstoreAdapter.parse_unicode
  621. else:
  622. cast = HstoreAdapter.parse
  623. HSTORE = _ext.new_type(oid, "HSTORE", cast)
  624. _ext.register_type(HSTORE, not globally and conn_or_curs or None)
  625. _ext.register_adapter(dict, HstoreAdapter)
  626. if array_oid:
  627. HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE)
  628. _ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)
  629. class CompositeCaster(object):
  630. """Helps conversion of a PostgreSQL composite type into a Python object.
  631. The class is usually created by the `register_composite()` function.
  632. You may want to create and register manually instances of the class if
  633. querying the database at registration time is not desirable (such as when
  634. using an :ref:`asynchronous connections <async-support>`).
  635. .. attribute:: name
  636. The name of the PostgreSQL type.
  637. .. attribute:: oid
  638. The oid of the PostgreSQL type.
  639. .. attribute:: array_oid
  640. The oid of the PostgreSQL array type, if available.
  641. .. attribute:: type
  642. The type of the Python objects returned. If :py:func:`collections.namedtuple()`
  643. is available, it is a named tuple with attributes equal to the type
  644. components. Otherwise it is just the `!tuple` object.
  645. .. attribute:: attnames
  646. List of component names of the type to be casted.
  647. .. attribute:: atttypes
  648. List of component type oids of the type to be casted.
  649. """
  650. def __init__(self, name, oid, attrs, array_oid=None):
  651. self.name = name
  652. self.oid = oid
  653. self.array_oid = array_oid
  654. self.attnames = [ a[0] for a in attrs ]
  655. self.atttypes = [ a[1] for a in attrs ]
  656. self._create_type(name, self.attnames)
  657. self.typecaster = _ext.new_type((oid,), name, self.parse)
  658. if array_oid:
  659. self.array_typecaster = _ext.new_array_type(
  660. (array_oid,), "%sARRAY" % name, self.typecaster)
  661. else:
  662. self.array_typecaster = None
  663. def parse(self, s, curs):
  664. if s is None:
  665. return None
  666. tokens = self.tokenize(s)
  667. if len(tokens) != len(self.atttypes):
  668. raise psycopg2.DataError(
  669. "expecting %d components for the type %s, %d found instead" %
  670. (len(self.atttypes), self.name, len(tokens)))
  671. attrs = [ curs.cast(oid, token)
  672. for oid, token in zip(self.atttypes, tokens) ]
  673. return self._ctor(*attrs)
  674. _re_tokenize = regex.compile(r"""
  675. \(? ([,)]) # an empty token, representing NULL
  676. | \(? " ((?: [^"] | "")*) " [,)] # or a quoted string
  677. | \(? ([^",)]+) [,)] # or an unquoted string
  678. """, regex.VERBOSE)
  679. _re_undouble = regex.compile(r'(["\\])\1')
  680. @classmethod
  681. def tokenize(self, s):
  682. rv = []
  683. for m in self._re_tokenize.finditer(s):
  684. if m is None:
  685. raise psycopg2.InterfaceError("can't parse type: %r" % s)
  686. if m.group(1) is not None:
  687. rv.append(None)
  688. elif m.group(2) is not None:
  689. rv.append(self._re_undouble.sub(r"\1", m.group(2)))
  690. else:
  691. rv.append(m.group(3))
  692. return rv
  693. def _create_type(self, name, attnames):
  694. try:
  695. from collections import namedtuple
  696. except ImportError:
  697. self.type = tuple
  698. self._ctor = lambda *args: tuple(args)
  699. else:
  700. self.type = namedtuple(name, attnames)
  701. self._ctor = self.type
  702. @classmethod
  703. def _from_db(self, name, conn_or_curs):
  704. """Return a `CompositeCaster` instance for the type *name*.
  705. Raise `ProgrammingError` if the type is not found.
  706. """
  707. conn, curs = _solve_conn_curs(conn_or_curs)
  708. # Store the transaction status of the connection to revert it after use
  709. conn_status = conn.status
  710. # Use the correct schema
  711. if '.' in name:
  712. schema, tname = name.split('.', 1)
  713. else:
  714. tname = name
  715. schema = 'public'
  716. # column typarray not available before PG 8.3
  717. typarray = conn.server_version >= 80300 and "typarray" or "NULL"
  718. # get the type oid and attributes
  719. curs.execute("""\
  720. SELECT t.oid, %s, attname, atttypid
  721. FROM pg_type t
  722. JOIN pg_namespace ns ON typnamespace = ns.oid
  723. JOIN pg_attribute a ON attrelid = typrelid
  724. WHERE typname = %%s AND nspname = %%s
  725. AND attnum > 0 AND NOT attisdropped
  726. ORDER BY attnum;
  727. """ % typarray, (tname, schema))
  728. recs = curs.fetchall()
  729. # revert the status of the connection as before the command
  730. if (conn_status != _ext.STATUS_IN_TRANSACTION
  731. and not conn.autocommit):
  732. conn.rollback()
  733. if not recs:
  734. raise psycopg2.ProgrammingError(
  735. "PostgreSQL type '%s' not found" % name)
  736. type_oid = recs[0][0]
  737. array_oid = recs[0][1]
  738. type_attrs = [ (r[2], r[3]) for r in recs ]
  739. return CompositeCaster(tname, type_oid, type_attrs,
  740. array_oid=array_oid)
  741. def register_composite(name, conn_or_curs, globally=False):
  742. """Register a typecaster to convert a composite type into a tuple.
  743. :param name: the name of a PostgreSQL composite type, e.g. created using
  744. the |CREATE TYPE|_ command
  745. :param conn_or_curs: a connection or cursor used to find the type oid and
  746. components; the typecaster is registered in a scope limited to this
  747. object, unless *globally* is set to `!True`
  748. :param globally: if `!False` (default) register the typecaster only on
  749. *conn_or_curs*, otherwise register it globally
  750. :return: the registered `CompositeCaster` instance responsible for the
  751. conversion
  752. .. versionchanged:: 2.4.3
  753. added support for array of composite types
  754. """
  755. caster = CompositeCaster._from_db(name, conn_or_curs)
  756. _ext.register_type(caster.typecaster, not globally and conn_or_curs or None)
  757. if caster.array_typecaster is not None:
  758. _ext.register_type(caster.array_typecaster, not globally and conn_or_curs or None)
  759. return caster
  760. __all__ = filter(lambda k: not k.startswith('_'), locals().keys())