db.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. import logging
  2. from asgiref.sync import sync_to_async
  3. from django.contrib.sessions.backends.base import CreateError, SessionBase, UpdateError
  4. from django.core.exceptions import SuspiciousOperation
  5. from django.db import DatabaseError, IntegrityError, router, transaction
  6. from django.utils import timezone
  7. from django.utils.functional import cached_property
  8. class SessionStore(SessionBase):
  9. """
  10. Implement database session store.
  11. """
  12. def __init__(self, session_key=None):
  13. super().__init__(session_key)
  14. @classmethod
  15. def get_model_class(cls):
  16. # Avoids a circular import and allows importing SessionStore when
  17. # django.contrib.sessions is not in INSTALLED_APPS.
  18. from django.contrib.sessions.models import Session
  19. return Session
  20. @cached_property
  21. def model(self):
  22. return self.get_model_class()
  23. def _get_session_from_db(self):
  24. try:
  25. return self.model.objects.get(
  26. session_key=self.session_key, expire_date__gt=timezone.now()
  27. )
  28. except (self.model.DoesNotExist, SuspiciousOperation) as e:
  29. if isinstance(e, SuspiciousOperation):
  30. logger = logging.getLogger("django.security.%s" % e.__class__.__name__)
  31. logger.warning(str(e))
  32. self._session_key = None
  33. async def _aget_session_from_db(self):
  34. try:
  35. return await self.model.objects.aget(
  36. session_key=self.session_key, expire_date__gt=timezone.now()
  37. )
  38. except (self.model.DoesNotExist, SuspiciousOperation) as e:
  39. if isinstance(e, SuspiciousOperation):
  40. logger = logging.getLogger("django.security.%s" % e.__class__.__name__)
  41. logger.warning(str(e))
  42. self._session_key = None
  43. def load(self):
  44. s = self._get_session_from_db()
  45. return self.decode(s.session_data) if s else {}
  46. async def aload(self):
  47. s = await self._aget_session_from_db()
  48. return self.decode(s.session_data) if s else {}
  49. def exists(self, session_key):
  50. return self.model.objects.filter(session_key=session_key).exists()
  51. async def aexists(self, session_key):
  52. return await self.model.objects.filter(session_key=session_key).aexists()
  53. def create(self):
  54. while True:
  55. self._session_key = self._get_new_session_key()
  56. try:
  57. # Save immediately to ensure we have a unique entry in the
  58. # database.
  59. self.save(must_create=True)
  60. except CreateError:
  61. # Key wasn't unique. Try again.
  62. continue
  63. self.modified = True
  64. return
  65. async def acreate(self):
  66. while True:
  67. self._session_key = await self._aget_new_session_key()
  68. try:
  69. # Save immediately to ensure we have a unique entry in the
  70. # database.
  71. await self.asave(must_create=True)
  72. except CreateError:
  73. # Key wasn't unique. Try again.
  74. continue
  75. self.modified = True
  76. return
  77. def create_model_instance(self, data):
  78. """
  79. Return a new instance of the session model object, which represents the
  80. current session state. Intended to be used for saving the session data
  81. to the database.
  82. """
  83. return self.model(
  84. session_key=self._get_or_create_session_key(),
  85. session_data=self.encode(data),
  86. expire_date=self.get_expiry_date(),
  87. )
  88. async def acreate_model_instance(self, data):
  89. """See create_model_instance()."""
  90. return self.model(
  91. session_key=await self._aget_or_create_session_key(),
  92. session_data=self.encode(data),
  93. expire_date=await self.aget_expiry_date(),
  94. )
  95. def save(self, must_create=False):
  96. """
  97. Save the current session data to the database. If 'must_create' is
  98. True, raise a database error if the saving operation doesn't create a
  99. new entry (as opposed to possibly updating an existing entry).
  100. """
  101. if self.session_key is None:
  102. return self.create()
  103. data = self._get_session(no_load=must_create)
  104. obj = self.create_model_instance(data)
  105. using = router.db_for_write(self.model, instance=obj)
  106. try:
  107. with transaction.atomic(using=using):
  108. obj.save(
  109. force_insert=must_create, force_update=not must_create, using=using
  110. )
  111. except IntegrityError:
  112. if must_create:
  113. raise CreateError
  114. raise
  115. except DatabaseError:
  116. if not must_create:
  117. raise UpdateError
  118. raise
  119. async def asave(self, must_create=False):
  120. """See save()."""
  121. if self.session_key is None:
  122. return await self.acreate()
  123. data = await self._aget_session(no_load=must_create)
  124. obj = await self.acreate_model_instance(data)
  125. using = router.db_for_write(self.model, instance=obj)
  126. try:
  127. # This code MOST run in a transaction, so it requires
  128. # @sync_to_async wrapping until transaction.atomic() supports
  129. # async.
  130. @sync_to_async
  131. def sync_transaction():
  132. with transaction.atomic(using=using):
  133. obj.save(
  134. force_insert=must_create,
  135. force_update=not must_create,
  136. using=using,
  137. )
  138. await sync_transaction()
  139. except IntegrityError:
  140. if must_create:
  141. raise CreateError
  142. raise
  143. except DatabaseError:
  144. if not must_create:
  145. raise UpdateError
  146. raise
  147. def delete(self, session_key=None):
  148. if session_key is None:
  149. if self.session_key is None:
  150. return
  151. session_key = self.session_key
  152. try:
  153. self.model.objects.get(session_key=session_key).delete()
  154. except self.model.DoesNotExist:
  155. pass
  156. async def adelete(self, session_key=None):
  157. if session_key is None:
  158. if self.session_key is None:
  159. return
  160. session_key = self.session_key
  161. try:
  162. obj = await self.model.objects.aget(session_key=session_key)
  163. await obj.adelete()
  164. except self.model.DoesNotExist:
  165. pass
  166. @classmethod
  167. def clear_expired(cls):
  168. cls.get_model_class().objects.filter(expire_date__lt=timezone.now()).delete()
  169. @classmethod
  170. async def aclear_expired(cls):
  171. await cls.get_model_class().objects.filter(
  172. expire_date__lt=timezone.now()
  173. ).adelete()