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.

133 lines
4.7 KiB

4 years ago
  1. import types
  2. import functools
  3. import zlib
  4. from requests.adapters import HTTPAdapter
  5. from .controller import CacheController
  6. from .cache import DictCache
  7. from .filewrapper import CallbackFileWrapper
  8. class CacheControlAdapter(HTTPAdapter):
  9. invalidating_methods = {"PUT", "DELETE"}
  10. def __init__(
  11. self,
  12. cache=None,
  13. cache_etags=True,
  14. controller_class=None,
  15. serializer=None,
  16. heuristic=None,
  17. cacheable_methods=None,
  18. *args,
  19. **kw
  20. ):
  21. super(CacheControlAdapter, self).__init__(*args, **kw)
  22. self.cache = cache or DictCache()
  23. self.heuristic = heuristic
  24. self.cacheable_methods = cacheable_methods or ("GET",)
  25. controller_factory = controller_class or CacheController
  26. self.controller = controller_factory(
  27. self.cache, cache_etags=cache_etags, serializer=serializer
  28. )
  29. def send(self, request, cacheable_methods=None, **kw):
  30. """
  31. Send a request. Use the request information to see if it
  32. exists in the cache and cache the response if we need to and can.
  33. """
  34. cacheable = cacheable_methods or self.cacheable_methods
  35. if request.method in cacheable:
  36. try:
  37. cached_response = self.controller.cached_request(request)
  38. except zlib.error:
  39. cached_response = None
  40. if cached_response:
  41. return self.build_response(request, cached_response, from_cache=True)
  42. # check for etags and add headers if appropriate
  43. request.headers.update(self.controller.conditional_headers(request))
  44. resp = super(CacheControlAdapter, self).send(request, **kw)
  45. return resp
  46. def build_response(
  47. self, request, response, from_cache=False, cacheable_methods=None
  48. ):
  49. """
  50. Build a response by making a request or using the cache.
  51. This will end up calling send and returning a potentially
  52. cached response
  53. """
  54. cacheable = cacheable_methods or self.cacheable_methods
  55. if not from_cache and request.method in cacheable:
  56. # Check for any heuristics that might update headers
  57. # before trying to cache.
  58. if self.heuristic:
  59. response = self.heuristic.apply(response)
  60. # apply any expiration heuristics
  61. if response.status == 304:
  62. # We must have sent an ETag request. This could mean
  63. # that we've been expired already or that we simply
  64. # have an etag. In either case, we want to try and
  65. # update the cache if that is the case.
  66. cached_response = self.controller.update_cached_response(
  67. request, response
  68. )
  69. if cached_response is not response:
  70. from_cache = True
  71. # We are done with the server response, read a
  72. # possible response body (compliant servers will
  73. # not return one, but we cannot be 100% sure) and
  74. # release the connection back to the pool.
  75. response.read(decode_content=False)
  76. response.release_conn()
  77. response = cached_response
  78. # We always cache the 301 responses
  79. elif response.status == 301:
  80. self.controller.cache_response(request, response)
  81. else:
  82. # Wrap the response file with a wrapper that will cache the
  83. # response when the stream has been consumed.
  84. response._fp = CallbackFileWrapper(
  85. response._fp,
  86. functools.partial(
  87. self.controller.cache_response, request, response
  88. ),
  89. )
  90. if response.chunked:
  91. super_update_chunk_length = response._update_chunk_length
  92. def _update_chunk_length(self):
  93. super_update_chunk_length()
  94. if self.chunk_left == 0:
  95. self._fp._close()
  96. response._update_chunk_length = types.MethodType(
  97. _update_chunk_length, response
  98. )
  99. resp = super(CacheControlAdapter, self).build_response(request, response)
  100. # See if we should invalidate the cache.
  101. if request.method in self.invalidating_methods and resp.ok:
  102. cache_url = self.controller.cache_url(request.url)
  103. self.cache.delete(cache_url)
  104. # Give the request a from_cache attr to let people use it
  105. resp.from_cache = from_cache
  106. return resp
  107. def close(self):
  108. self.cache.close()
  109. super(CacheControlAdapter, self).close()