database.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. # coding=utf-8
  2. #
  3. # This file is part of Hypothesis, which may be found at
  4. # https://github.com/HypothesisWorks/hypothesis-python
  5. #
  6. # Most of this work is copyright (C) 2013-2018 David R. MacIver
  7. # (david@drmaciver.com), but it contains contributions by others. See
  8. # CONTRIBUTING.rst for a full list of people who may hold copyright, and
  9. # consult the git log if you need to determine who owns an individual
  10. # contribution.
  11. #
  12. # This Source Code Form is subject to the terms of the Mozilla Public License,
  13. # v. 2.0. If a copy of the MPL was not distributed with this file, You can
  14. # obtain one at http://mozilla.org/MPL/2.0/.
  15. #
  16. # END HEADER
  17. from __future__ import division, print_function, absolute_import
  18. import os
  19. import re
  20. import binascii
  21. import threading
  22. from hashlib import sha1
  23. from contextlib import contextmanager
  24. from hypothesis._settings import note_deprecation
  25. from hypothesis.internal.compat import FileNotFoundError, hbytes, \
  26. b64decode, b64encode
  27. sqlite3 = None
  28. SQLITE_PATH = re.compile(r"\.\(db|sqlite|sqlite3\)$")
  29. def _db_for_path(path=None):
  30. if path in (None, ':memory:'):
  31. return InMemoryExampleDatabase()
  32. path = str(path)
  33. if os.path.isdir(path):
  34. return DirectoryBasedExampleDatabase(path)
  35. if os.path.exists(path):
  36. return SQLiteExampleDatabase(path)
  37. if SQLITE_PATH.search(path):
  38. return SQLiteExampleDatabase(path)
  39. else:
  40. return DirectoryBasedExampleDatabase(path)
  41. class EDMeta(type):
  42. def __call__(self, *args, **kwargs):
  43. if self is ExampleDatabase:
  44. return _db_for_path(*args, **kwargs)
  45. return super(EDMeta, self).__call__(*args, **kwargs)
  46. class ExampleDatabase(EDMeta('ExampleDatabase', (object,), {})):
  47. """Interface class for storage systems.
  48. A key -> multiple distinct values mapping.
  49. Keys and values are binary data.
  50. """
  51. def save(self, key, value):
  52. """save this value under this key.
  53. If this value is already present for this key, silently do
  54. nothing
  55. """
  56. raise NotImplementedError('%s.save' % (type(self).__name__))
  57. def delete(self, key, value):
  58. """Remove this value from this key.
  59. If this value is not present, silently do nothing.
  60. """
  61. raise NotImplementedError('%s.delete' % (type(self).__name__))
  62. def move(self, src, dest, value):
  63. """Move value from key src to key dest. Equivalent to delete(src,
  64. value) followed by save(src, value) but may have a more efficient
  65. implementation.
  66. Note that value will be inserted at dest regardless of whether
  67. it is currently present at src.
  68. """
  69. if src == dest:
  70. self.save(src, value)
  71. return
  72. self.delete(src, value)
  73. self.save(dest, value)
  74. def fetch(self, key):
  75. """Return all values matching this key."""
  76. raise NotImplementedError('%s.fetch' % (type(self).__name__))
  77. def close(self):
  78. """Clear up any resources associated with this database."""
  79. raise NotImplementedError('%s.close' % (type(self).__name__))
  80. class InMemoryExampleDatabase(ExampleDatabase):
  81. def __init__(self):
  82. self.data = {}
  83. def __repr__(self):
  84. return 'InMemoryExampleDatabase(%r)' % (self.data,)
  85. def fetch(self, key):
  86. for v in self.data.get(key, ()):
  87. yield v
  88. def save(self, key, value):
  89. self.data.setdefault(key, set()).add(hbytes(value))
  90. def delete(self, key, value):
  91. self.data.get(key, set()).discard(hbytes(value))
  92. def close(self):
  93. pass
  94. class SQLiteExampleDatabase(ExampleDatabase):
  95. def __init__(self, path=u':memory:'):
  96. self.path = path
  97. self.db_created = False
  98. self.current_connection = threading.local()
  99. global sqlite3
  100. import sqlite3
  101. if path == u':memory:':
  102. note_deprecation(
  103. 'The SQLite database backend has been deprecated. '
  104. 'Use InMemoryExampleDatabase or set database_file=":memory:" '
  105. 'instead.'
  106. )
  107. else:
  108. note_deprecation(
  109. 'The SQLite database backend has been deprecated. '
  110. 'Set database_file to some path name not ending in .db, '
  111. '.sqlite or .sqlite3 to get the new directory based database '
  112. 'backend instead.'
  113. )
  114. def connection(self):
  115. if not hasattr(self.current_connection, 'connection'):
  116. self.current_connection.connection = sqlite3.connect(self.path)
  117. return self.current_connection.connection
  118. def close(self):
  119. if hasattr(self.current_connection, 'connection'):
  120. try:
  121. self.connection().close()
  122. finally:
  123. del self.current_connection.connection
  124. def __repr__(self):
  125. return u'%s(%s)' % (self.__class__.__name__, self.path)
  126. @contextmanager
  127. def cursor(self):
  128. conn = self.connection()
  129. cursor = conn.cursor()
  130. try:
  131. try:
  132. yield cursor
  133. finally:
  134. cursor.close()
  135. except BaseException:
  136. conn.rollback()
  137. raise
  138. else:
  139. conn.commit()
  140. def save(self, key, value):
  141. self.create_db_if_needed()
  142. with self.cursor() as cursor:
  143. try:
  144. cursor.execute("""
  145. insert into hypothesis_data_mapping(key, value)
  146. values(?, ?)
  147. """, (b64encode(key), b64encode(value)))
  148. except sqlite3.IntegrityError:
  149. pass
  150. def delete(self, key, value):
  151. self.create_db_if_needed()
  152. with self.cursor() as cursor:
  153. cursor.execute("""
  154. delete from hypothesis_data_mapping
  155. where key = ? and value = ?
  156. """, (b64encode(key), b64encode(value)))
  157. def fetch(self, key):
  158. self.create_db_if_needed()
  159. with self.cursor() as cursor:
  160. cursor.execute("""
  161. select value from hypothesis_data_mapping
  162. where key = ?
  163. """, (b64encode(key),))
  164. for (value,) in cursor:
  165. try:
  166. yield b64decode(value)
  167. except (binascii.Error, TypeError):
  168. pass
  169. def create_db_if_needed(self):
  170. if self.db_created:
  171. return
  172. with self.cursor() as cursor:
  173. cursor.execute("""
  174. create table if not exists hypothesis_data_mapping(
  175. key text,
  176. value text,
  177. unique(key, value)
  178. )
  179. """)
  180. self.db_created = True
  181. def mkdirp(path):
  182. try:
  183. os.makedirs(path)
  184. except OSError:
  185. pass
  186. return path
  187. def _hash(key):
  188. return sha1(key).hexdigest()[:16]
  189. class DirectoryBasedExampleDatabase(ExampleDatabase):
  190. def __init__(self, path):
  191. self.path = path
  192. self.keypaths = {}
  193. def __repr__(self):
  194. return 'DirectoryBasedExampleDatabase(%r)' % (self.path,)
  195. def close(self):
  196. pass
  197. def _key_path(self, key):
  198. try:
  199. return self.keypaths[key]
  200. except KeyError:
  201. pass
  202. directory = os.path.join(self.path, _hash(key))
  203. mkdirp(directory)
  204. self.keypaths[key] = directory
  205. return directory
  206. def _value_path(self, key, value):
  207. return os.path.join(
  208. self._key_path(key),
  209. sha1(value).hexdigest()[:16]
  210. )
  211. def fetch(self, key):
  212. kp = self._key_path(key)
  213. for path in os.listdir(kp):
  214. try:
  215. with open(os.path.join(kp, path), 'rb') as i:
  216. yield hbytes(i.read())
  217. except FileNotFoundError:
  218. pass
  219. def save(self, key, value):
  220. path = self._value_path(key, value)
  221. if not os.path.exists(path):
  222. suffix = binascii.hexlify(os.urandom(16))
  223. if not isinstance(suffix, str): # pragma: no branch
  224. # On Python 3, binascii.hexlify returns bytes
  225. suffix = suffix.decode('ascii')
  226. tmpname = path + '.' + suffix
  227. with open(tmpname, 'wb') as o:
  228. o.write(value)
  229. try:
  230. os.rename(tmpname, path)
  231. except OSError: # pragma: no cover
  232. os.unlink(tmpname)
  233. assert not os.path.exists(tmpname)
  234. def move(self, src, dest, value):
  235. if src == dest:
  236. self.save(src, value)
  237. return
  238. try:
  239. os.rename(
  240. self._value_path(src, value), self._value_path(dest, value))
  241. except OSError:
  242. self.delete(src, value)
  243. self.save(dest, value)
  244. def delete(self, key, value):
  245. try:
  246. os.unlink(self._value_path(key, value))
  247. except OSError:
  248. pass