Use a fixed size cache to store sha sums

Author: blink1073

flask_pymongo/__init__.py

@@ -27,6 +27,7 @@
 __all__ = ("PyMongo", "ASCENDING", "DESCENDING", "BSONObjectIdConverter", "BSONProvider")
 
 import hashlib
+from collections import OrderedDict
 from mimetypes import guess_type
 from typing import Any
 
@@ -66,6 +67,8 @@ def __init__(
     ) -> None:
         self.cx: MongoClient | None = None
         self.db: Database | None = None
+        self._hash_cache = OrderedDict()
+        self._hash_limit = 128
 
         if app is not None:
             self.init_app(app, uri, *args, **kwargs)
@@ -184,10 +187,18 @@ def get_upload(filename):
         response.content_length = fileobj.length
         response.last_modified = fileobj.upload_date
 
-        # Get or compute the sha1 sum for the etag.
-        metadata = fileobj.metadata
-        sha1_sum = metadata and metadata.get("sha1_sum")
-        sha1_sum = sha1_sum or self._compute_sha(fileobj)
+        # GridFS does not manage its own checksum, so we manage our own using its
+        # metadata storage, to be used for the etag.
+        sha1_sum = self._hash_cache.get(str(fileobj._id))
+        if sha1_sum is None:
+            # Compute the checksum of the file for the etag.
+            pos = fileobj.tell()
+            raw_data = fileobj.read()
+            fileobj.seek(pos)
+            sha1_sum = hashlib.sha1(raw_data).hexdigest()
+            while len(self._hash_cache) >= self._hash_limit:
+                self._hash_cache.popitem()
+            self._hash_cache[str(fileobj._id)] = sha1_sum
         response.set_etag(sha1_sum)
 
         response.cache_control.max_age = cache_for
@@ -238,19 +249,5 @@ def save_upload(filename):
             db_obj = self.db
         assert db_obj is not None, "Please initialize the app before calling save_file!"
         storage = GridFS(db_obj, base)
-
-        # GridFS does not manage its own hash, so we manage our own using its
-        # metadata storage, to be used for the etag.
-        sha1_sum = self._compute_sha(fileobj)
-        metadata = dict(sha1_sum=sha1_sum)
-        id = storage.put(
-            fileobj, filename=filename, content_type=content_type, metadata=metadata, **kwargs
-        )
+        id = storage.put(fileobj, filename=filename, content_type=content_type, **kwargs)
         return id
-
-    def _compute_sha(self, fileobj: Any) -> str:
-        """Compute the sha sum of a file object."""
-        pos = fileobj.tell()
-        raw_data = fileobj.read()
-        fileobj.seek(pos)
-        return hashlib.sha1(raw_data).hexdigest()