hstore.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420
  1. # postgresql/hstore.py
  2. # Copyright (C) 2005-2017 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: http://www.opensource.org/licenses/mit-license.php
  7. import re
  8. from .base import ischema_names
  9. from .array import ARRAY
  10. from ... import types as sqltypes
  11. from ...sql import functions as sqlfunc
  12. from ...sql import operators
  13. from ... import util
  14. __all__ = ('HSTORE', 'hstore')
  15. idx_precedence = operators._PRECEDENCE[operators.json_getitem_op]
  16. GETITEM = operators.custom_op(
  17. "->", precedence=idx_precedence, natural_self_precedent=True,
  18. eager_grouping=True
  19. )
  20. HAS_KEY = operators.custom_op(
  21. "?", precedence=idx_precedence, natural_self_precedent=True,
  22. eager_grouping=True
  23. )
  24. HAS_ALL = operators.custom_op(
  25. "?&", precedence=idx_precedence, natural_self_precedent=True,
  26. eager_grouping=True
  27. )
  28. HAS_ANY = operators.custom_op(
  29. "?|", precedence=idx_precedence, natural_self_precedent=True,
  30. eager_grouping=True
  31. )
  32. CONTAINS = operators.custom_op(
  33. "@>", precedence=idx_precedence, natural_self_precedent=True,
  34. eager_grouping=True
  35. )
  36. CONTAINED_BY = operators.custom_op(
  37. "<@", precedence=idx_precedence, natural_self_precedent=True,
  38. eager_grouping=True
  39. )
  40. class HSTORE(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
  41. """Represent the PostgreSQL HSTORE type.
  42. The :class:`.HSTORE` type stores dictionaries containing strings, e.g.::
  43. data_table = Table('data_table', metadata,
  44. Column('id', Integer, primary_key=True),
  45. Column('data', HSTORE)
  46. )
  47. with engine.connect() as conn:
  48. conn.execute(
  49. data_table.insert(),
  50. data = {"key1": "value1", "key2": "value2"}
  51. )
  52. :class:`.HSTORE` provides for a wide range of operations, including:
  53. * Index operations::
  54. data_table.c.data['some key'] == 'some value'
  55. * Containment operations::
  56. data_table.c.data.has_key('some key')
  57. data_table.c.data.has_all(['one', 'two', 'three'])
  58. * Concatenation::
  59. data_table.c.data + {"k1": "v1"}
  60. For a full list of special methods see
  61. :class:`.HSTORE.comparator_factory`.
  62. For usage with the SQLAlchemy ORM, it may be desirable to combine
  63. the usage of :class:`.HSTORE` with :class:`.MutableDict` dictionary
  64. now part of the :mod:`sqlalchemy.ext.mutable`
  65. extension. This extension will allow "in-place" changes to the
  66. dictionary, e.g. addition of new keys or replacement/removal of existing
  67. keys to/from the current dictionary, to produce events which will be
  68. detected by the unit of work::
  69. from sqlalchemy.ext.mutable import MutableDict
  70. class MyClass(Base):
  71. __tablename__ = 'data_table'
  72. id = Column(Integer, primary_key=True)
  73. data = Column(MutableDict.as_mutable(HSTORE))
  74. my_object = session.query(MyClass).one()
  75. # in-place mutation, requires Mutable extension
  76. # in order for the ORM to detect
  77. my_object.data['some_key'] = 'some value'
  78. session.commit()
  79. When the :mod:`sqlalchemy.ext.mutable` extension is not used, the ORM
  80. will not be alerted to any changes to the contents of an existing
  81. dictionary, unless that dictionary value is re-assigned to the
  82. HSTORE-attribute itself, thus generating a change event.
  83. .. versionadded:: 0.8
  84. .. seealso::
  85. :class:`.hstore` - render the PostgreSQL ``hstore()`` function.
  86. """
  87. __visit_name__ = 'HSTORE'
  88. hashable = False
  89. text_type = sqltypes.Text()
  90. def __init__(self, text_type=None):
  91. """Construct a new :class:`.HSTORE`.
  92. :param text_type: the type that should be used for indexed values.
  93. Defaults to :class:`.types.Text`.
  94. .. versionadded:: 1.1.0
  95. """
  96. if text_type is not None:
  97. self.text_type = text_type
  98. class Comparator(
  99. sqltypes.Indexable.Comparator, sqltypes.Concatenable.Comparator):
  100. """Define comparison operations for :class:`.HSTORE`."""
  101. def has_key(self, other):
  102. """Boolean expression. Test for presence of a key. Note that the
  103. key may be a SQLA expression.
  104. """
  105. return self.operate(HAS_KEY, other, result_type=sqltypes.Boolean)
  106. def has_all(self, other):
  107. """Boolean expression. Test for presence of all keys in jsonb
  108. """
  109. return self.operate(HAS_ALL, other, result_type=sqltypes.Boolean)
  110. def has_any(self, other):
  111. """Boolean expression. Test for presence of any key in jsonb
  112. """
  113. return self.operate(HAS_ANY, other, result_type=sqltypes.Boolean)
  114. def contains(self, other, **kwargs):
  115. """Boolean expression. Test if keys (or array) are a superset
  116. of/contained the keys of the argument jsonb expression.
  117. """
  118. return self.operate(CONTAINS, other, result_type=sqltypes.Boolean)
  119. def contained_by(self, other):
  120. """Boolean expression. Test if keys are a proper subset of the
  121. keys of the argument jsonb expression.
  122. """
  123. return self.operate(
  124. CONTAINED_BY, other, result_type=sqltypes.Boolean)
  125. def _setup_getitem(self, index):
  126. return GETITEM, index, self.type.text_type
  127. def defined(self, key):
  128. """Boolean expression. Test for presence of a non-NULL value for
  129. the key. Note that the key may be a SQLA expression.
  130. """
  131. return _HStoreDefinedFunction(self.expr, key)
  132. def delete(self, key):
  133. """HStore expression. Returns the contents of this hstore with the
  134. given key deleted. Note that the key may be a SQLA expression.
  135. """
  136. if isinstance(key, dict):
  137. key = _serialize_hstore(key)
  138. return _HStoreDeleteFunction(self.expr, key)
  139. def slice(self, array):
  140. """HStore expression. Returns a subset of an hstore defined by
  141. array of keys.
  142. """
  143. return _HStoreSliceFunction(self.expr, array)
  144. def keys(self):
  145. """Text array expression. Returns array of keys."""
  146. return _HStoreKeysFunction(self.expr)
  147. def vals(self):
  148. """Text array expression. Returns array of values."""
  149. return _HStoreValsFunction(self.expr)
  150. def array(self):
  151. """Text array expression. Returns array of alternating keys and
  152. values.
  153. """
  154. return _HStoreArrayFunction(self.expr)
  155. def matrix(self):
  156. """Text array expression. Returns array of [key, value] pairs."""
  157. return _HStoreMatrixFunction(self.expr)
  158. comparator_factory = Comparator
  159. def bind_processor(self, dialect):
  160. if util.py2k:
  161. encoding = dialect.encoding
  162. def process(value):
  163. if isinstance(value, dict):
  164. return _serialize_hstore(value).encode(encoding)
  165. else:
  166. return value
  167. else:
  168. def process(value):
  169. if isinstance(value, dict):
  170. return _serialize_hstore(value)
  171. else:
  172. return value
  173. return process
  174. def result_processor(self, dialect, coltype):
  175. if util.py2k:
  176. encoding = dialect.encoding
  177. def process(value):
  178. if value is not None:
  179. return _parse_hstore(value.decode(encoding))
  180. else:
  181. return value
  182. else:
  183. def process(value):
  184. if value is not None:
  185. return _parse_hstore(value)
  186. else:
  187. return value
  188. return process
  189. ischema_names['hstore'] = HSTORE
  190. class hstore(sqlfunc.GenericFunction):
  191. """Construct an hstore value within a SQL expression using the
  192. PostgreSQL ``hstore()`` function.
  193. The :class:`.hstore` function accepts one or two arguments as described
  194. in the PostgreSQL documentation.
  195. E.g.::
  196. from sqlalchemy.dialects.postgresql import array, hstore
  197. select([hstore('key1', 'value1')])
  198. select([
  199. hstore(
  200. array(['key1', 'key2', 'key3']),
  201. array(['value1', 'value2', 'value3'])
  202. )
  203. ])
  204. .. versionadded:: 0.8
  205. .. seealso::
  206. :class:`.HSTORE` - the PostgreSQL ``HSTORE`` datatype.
  207. """
  208. type = HSTORE
  209. name = 'hstore'
  210. class _HStoreDefinedFunction(sqlfunc.GenericFunction):
  211. type = sqltypes.Boolean
  212. name = 'defined'
  213. class _HStoreDeleteFunction(sqlfunc.GenericFunction):
  214. type = HSTORE
  215. name = 'delete'
  216. class _HStoreSliceFunction(sqlfunc.GenericFunction):
  217. type = HSTORE
  218. name = 'slice'
  219. class _HStoreKeysFunction(sqlfunc.GenericFunction):
  220. type = ARRAY(sqltypes.Text)
  221. name = 'akeys'
  222. class _HStoreValsFunction(sqlfunc.GenericFunction):
  223. type = ARRAY(sqltypes.Text)
  224. name = 'avals'
  225. class _HStoreArrayFunction(sqlfunc.GenericFunction):
  226. type = ARRAY(sqltypes.Text)
  227. name = 'hstore_to_array'
  228. class _HStoreMatrixFunction(sqlfunc.GenericFunction):
  229. type = ARRAY(sqltypes.Text)
  230. name = 'hstore_to_matrix'
  231. #
  232. # parsing. note that none of this is used with the psycopg2 backend,
  233. # which provides its own native extensions.
  234. #
  235. # My best guess at the parsing rules of hstore literals, since no formal
  236. # grammar is given. This is mostly reverse engineered from PG's input parser
  237. # behavior.
  238. HSTORE_PAIR_RE = re.compile(r"""
  239. (
  240. "(?P<key> (\\ . | [^"])* )" # Quoted key
  241. )
  242. [ ]* => [ ]* # Pair operator, optional adjoining whitespace
  243. (
  244. (?P<value_null> NULL ) # NULL value
  245. | "(?P<value> (\\ . | [^"])* )" # Quoted value
  246. )
  247. """, re.VERBOSE)
  248. HSTORE_DELIMITER_RE = re.compile(r"""
  249. [ ]* , [ ]*
  250. """, re.VERBOSE)
  251. def _parse_error(hstore_str, pos):
  252. """format an unmarshalling error."""
  253. ctx = 20
  254. hslen = len(hstore_str)
  255. parsed_tail = hstore_str[max(pos - ctx - 1, 0):min(pos, hslen)]
  256. residual = hstore_str[min(pos, hslen):min(pos + ctx + 1, hslen)]
  257. if len(parsed_tail) > ctx:
  258. parsed_tail = '[...]' + parsed_tail[1:]
  259. if len(residual) > ctx:
  260. residual = residual[:-1] + '[...]'
  261. return "After %r, could not parse residual at position %d: %r" % (
  262. parsed_tail, pos, residual)
  263. def _parse_hstore(hstore_str):
  264. """Parse an hstore from its literal string representation.
  265. Attempts to approximate PG's hstore input parsing rules as closely as
  266. possible. Although currently this is not strictly necessary, since the
  267. current implementation of hstore's output syntax is stricter than what it
  268. accepts as input, the documentation makes no guarantees that will always
  269. be the case.
  270. """
  271. result = {}
  272. pos = 0
  273. pair_match = HSTORE_PAIR_RE.match(hstore_str)
  274. while pair_match is not None:
  275. key = pair_match.group('key').replace(r'\"', '"').replace(
  276. "\\\\", "\\")
  277. if pair_match.group('value_null'):
  278. value = None
  279. else:
  280. value = pair_match.group('value').replace(
  281. r'\"', '"').replace("\\\\", "\\")
  282. result[key] = value
  283. pos += pair_match.end()
  284. delim_match = HSTORE_DELIMITER_RE.match(hstore_str[pos:])
  285. if delim_match is not None:
  286. pos += delim_match.end()
  287. pair_match = HSTORE_PAIR_RE.match(hstore_str[pos:])
  288. if pos != len(hstore_str):
  289. raise ValueError(_parse_error(hstore_str, pos))
  290. return result
  291. def _serialize_hstore(val):
  292. """Serialize a dictionary into an hstore literal. Keys and values must
  293. both be strings (except None for values).
  294. """
  295. def esc(s, position):
  296. if position == 'value' and s is None:
  297. return 'NULL'
  298. elif isinstance(s, util.string_types):
  299. return '"%s"' % s.replace("\\", "\\\\").replace('"', r'\"')
  300. else:
  301. raise ValueError("%r in %s position is not a string." %
  302. (s, position))
  303. return ', '.join('%s=>%s' % (esc(k, 'key'), esc(v, 'value'))
  304. for k, v in val.items())