You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

utils.py 11 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. import inspect
  2. import os
  3. import pkgutil
  4. import warnings
  5. from importlib import import_module
  6. from threading import local
  7. from django.conf import settings
  8. from django.core.exceptions import ImproperlyConfigured
  9. from django.utils import six
  10. from django.utils._os import npath, upath
  11. from django.utils.deprecation import RemovedInDjango110Warning
  12. from django.utils.functional import cached_property
  13. from django.utils.module_loading import import_string
  14. DEFAULT_DB_ALIAS = 'default'
  15. DJANGO_VERSION_PICKLE_KEY = '_django_version'
  16. class Error(Exception if six.PY3 else StandardError): # NOQA: StandardError undefined on PY3
  17. pass
  18. class InterfaceError(Error):
  19. pass
  20. class DatabaseError(Error):
  21. pass
  22. class DataError(DatabaseError):
  23. pass
  24. class OperationalError(DatabaseError):
  25. pass
  26. class IntegrityError(DatabaseError):
  27. pass
  28. class InternalError(DatabaseError):
  29. pass
  30. class ProgrammingError(DatabaseError):
  31. pass
  32. class NotSupportedError(DatabaseError):
  33. pass
  34. class DatabaseErrorWrapper(object):
  35. """
  36. Context manager and decorator that re-throws backend-specific database
  37. exceptions using Django's common wrappers.
  38. """
  39. def __init__(self, wrapper):
  40. """
  41. wrapper is a database wrapper.
  42. It must have a Database attribute defining PEP-249 exceptions.
  43. """
  44. self.wrapper = wrapper
  45. def __enter__(self):
  46. pass
  47. def __exit__(self, exc_type, exc_value, traceback):
  48. if exc_type is None:
  49. return
  50. for dj_exc_type in (
  51. DataError,
  52. OperationalError,
  53. IntegrityError,
  54. InternalError,
  55. ProgrammingError,
  56. NotSupportedError,
  57. DatabaseError,
  58. InterfaceError,
  59. Error,
  60. ):
  61. db_exc_type = getattr(self.wrapper.Database, dj_exc_type.__name__)
  62. if issubclass(exc_type, db_exc_type):
  63. dj_exc_value = dj_exc_type(*exc_value.args)
  64. dj_exc_value.__cause__ = exc_value
  65. # Only set the 'errors_occurred' flag for errors that may make
  66. # the connection unusable.
  67. if dj_exc_type not in (DataError, IntegrityError):
  68. self.wrapper.errors_occurred = True
  69. six.reraise(dj_exc_type, dj_exc_value, traceback)
  70. def __call__(self, func):
  71. # Note that we are intentionally not using @wraps here for performance
  72. # reasons. Refs #21109.
  73. def inner(*args, **kwargs):
  74. with self:
  75. return func(*args, **kwargs)
  76. return inner
  77. def load_backend(backend_name):
  78. """
  79. Return a database backend's "base" module given a fully qualified database
  80. backend name, or raise an error if it doesn't exist.
  81. """
  82. # This backend was renamed in Django 1.9.
  83. if backend_name == 'django.db.backends.postgresql_psycopg2':
  84. backend_name = 'django.db.backends.postgresql'
  85. try:
  86. return import_module('%s.base' % backend_name)
  87. except ImportError as e_user:
  88. # The database backend wasn't found. Display a helpful error message
  89. # listing all possible (built-in) database backends.
  90. backend_dir = os.path.join(os.path.dirname(upath(__file__)), 'backends')
  91. try:
  92. builtin_backends = [
  93. name for _, name, ispkg in pkgutil.iter_modules([npath(backend_dir)])
  94. if ispkg and name not in {'base', 'dummy', 'postgresql_psycopg2'}
  95. ]
  96. except EnvironmentError:
  97. builtin_backends = []
  98. if backend_name not in ['django.db.backends.%s' % b for b in
  99. builtin_backends]:
  100. backend_reprs = map(repr, sorted(builtin_backends))
  101. error_msg = ("%r isn't an available database backend.\n"
  102. "Try using 'django.db.backends.XXX', where XXX "
  103. "is one of:\n %s\nError was: %s" %
  104. (backend_name, ", ".join(backend_reprs), e_user))
  105. raise ImproperlyConfigured(error_msg)
  106. else:
  107. # If there's some other error, this must be an error in Django
  108. raise
  109. class ConnectionDoesNotExist(Exception):
  110. pass
  111. class ConnectionHandler(object):
  112. def __init__(self, databases=None):
  113. """
  114. databases is an optional dictionary of database definitions (structured
  115. like settings.DATABASES).
  116. """
  117. self._databases = databases
  118. self._connections = local()
  119. @cached_property
  120. def databases(self):
  121. if self._databases is None:
  122. self._databases = settings.DATABASES
  123. if self._databases == {}:
  124. self._databases = {
  125. DEFAULT_DB_ALIAS: {
  126. 'ENGINE': 'django.db.backends.dummy',
  127. },
  128. }
  129. if self._databases[DEFAULT_DB_ALIAS] == {}:
  130. self._databases[DEFAULT_DB_ALIAS]['ENGINE'] = 'django.db.backends.dummy'
  131. if DEFAULT_DB_ALIAS not in self._databases:
  132. raise ImproperlyConfigured("You must define a '%s' database" % DEFAULT_DB_ALIAS)
  133. return self._databases
  134. def ensure_defaults(self, alias):
  135. """
  136. Puts the defaults into the settings dictionary for a given connection
  137. where no settings is provided.
  138. """
  139. try:
  140. conn = self.databases[alias]
  141. except KeyError:
  142. raise ConnectionDoesNotExist("The connection %s doesn't exist" % alias)
  143. conn.setdefault('ATOMIC_REQUESTS', False)
  144. conn.setdefault('AUTOCOMMIT', True)
  145. conn.setdefault('ENGINE', 'django.db.backends.dummy')
  146. if conn['ENGINE'] == 'django.db.backends.' or not conn['ENGINE']:
  147. conn['ENGINE'] = 'django.db.backends.dummy'
  148. conn.setdefault('CONN_MAX_AGE', 0)
  149. conn.setdefault('OPTIONS', {})
  150. conn.setdefault('TIME_ZONE', None)
  151. for setting in ['NAME', 'USER', 'PASSWORD', 'HOST', 'PORT']:
  152. conn.setdefault(setting, '')
  153. def prepare_test_settings(self, alias):
  154. """
  155. Makes sure the test settings are available in the 'TEST' sub-dictionary.
  156. """
  157. try:
  158. conn = self.databases[alias]
  159. except KeyError:
  160. raise ConnectionDoesNotExist("The connection %s doesn't exist" % alias)
  161. test_settings = conn.setdefault('TEST', {})
  162. for key in ['CHARSET', 'COLLATION', 'NAME', 'MIRROR']:
  163. test_settings.setdefault(key, None)
  164. def __getitem__(self, alias):
  165. if hasattr(self._connections, alias):
  166. return getattr(self._connections, alias)
  167. self.ensure_defaults(alias)
  168. self.prepare_test_settings(alias)
  169. db = self.databases[alias]
  170. backend = load_backend(db['ENGINE'])
  171. conn = backend.DatabaseWrapper(db, alias)
  172. setattr(self._connections, alias, conn)
  173. return conn
  174. def __setitem__(self, key, value):
  175. setattr(self._connections, key, value)
  176. def __delitem__(self, key):
  177. delattr(self._connections, key)
  178. def __iter__(self):
  179. return iter(self.databases)
  180. def all(self):
  181. return [self[alias] for alias in self]
  182. def close_all(self):
  183. for alias in self:
  184. try:
  185. connection = getattr(self._connections, alias)
  186. except AttributeError:
  187. continue
  188. connection.close()
  189. class ConnectionRouter(object):
  190. def __init__(self, routers=None):
  191. """
  192. If routers is not specified, will default to settings.DATABASE_ROUTERS.
  193. """
  194. self._routers = routers
  195. @cached_property
  196. def routers(self):
  197. if self._routers is None:
  198. self._routers = settings.DATABASE_ROUTERS
  199. routers = []
  200. for r in self._routers:
  201. if isinstance(r, six.string_types):
  202. router = import_string(r)()
  203. else:
  204. router = r
  205. routers.append(router)
  206. return routers
  207. def _router_func(action):
  208. def _route_db(self, model, **hints):
  209. chosen_db = None
  210. for router in self.routers:
  211. try:
  212. method = getattr(router, action)
  213. except AttributeError:
  214. # If the router doesn't have a method, skip to the next one.
  215. pass
  216. else:
  217. chosen_db = method(model, **hints)
  218. if chosen_db:
  219. return chosen_db
  220. instance = hints.get('instance')
  221. if instance is not None and instance._state.db:
  222. return instance._state.db
  223. return DEFAULT_DB_ALIAS
  224. return _route_db
  225. db_for_read = _router_func('db_for_read')
  226. db_for_write = _router_func('db_for_write')
  227. def allow_relation(self, obj1, obj2, **hints):
  228. for router in self.routers:
  229. try:
  230. method = router.allow_relation
  231. except AttributeError:
  232. # If the router doesn't have a method, skip to the next one.
  233. pass
  234. else:
  235. allow = method(obj1, obj2, **hints)
  236. if allow is not None:
  237. return allow
  238. return obj1._state.db == obj2._state.db
  239. def allow_migrate(self, db, app_label, **hints):
  240. for router in self.routers:
  241. try:
  242. method = router.allow_migrate
  243. except AttributeError:
  244. # If the router doesn't have a method, skip to the next one.
  245. continue
  246. if six.PY3:
  247. sig = inspect.signature(router.allow_migrate)
  248. has_deprecated_signature = not any(
  249. p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
  250. )
  251. else:
  252. argspec = inspect.getargspec(router.allow_migrate)
  253. has_deprecated_signature = len(argspec.args) == 3 and not argspec.keywords
  254. if has_deprecated_signature:
  255. warnings.warn(
  256. "The signature of allow_migrate has changed from "
  257. "allow_migrate(self, db, model) to "
  258. "allow_migrate(self, db, app_label, model_name=None, **hints). "
  259. "Support for the old signature will be removed in Django 1.10.",
  260. RemovedInDjango110Warning)
  261. model = hints.get('model')
  262. allow = None if model is None else method(db, model)
  263. else:
  264. allow = method(db, app_label, **hints)
  265. if allow is not None:
  266. return allow
  267. return True
  268. def allow_migrate_model(self, db, model):
  269. return self.allow_migrate(
  270. db,
  271. model._meta.app_label,
  272. model_name=model._meta.model_name,
  273. model=model,
  274. )
  275. def get_migratable_models(self, app_config, db, include_auto_created=False):
  276. """
  277. Return app models allowed to be synchronized on provided db.
  278. """
  279. models = app_config.get_models(include_auto_created=include_auto_created)
  280. return [model for model in models if self.allow_migrate_model(db, model)]