address review

Author: blink1073

flask_pymongo/__init__.py

@@ -27,7 +27,6 @@
 __all__ = ("PyMongo", "ASCENDING", "DESCENDING", "BSONObjectIdConverter", "BSONProvider")
 
 import hashlib
-from collections import OrderedDict
 from mimetypes import guess_type
 from typing import Any
 
@@ -67,8 +66,6 @@ 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)
@@ -187,19 +184,20 @@ def get_upload(filename):
         response.content_length = fileobj.length
         response.last_modified = fileobj.upload_date
 
-        # 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)
+        # GridFS does not manage its own checksum.
+        # Try to use a sha1 sum that we have added during a save_file.
+        # Fall back to a legacy md5 sum if it exists.
+        # Otherwise, compute the sha1 sum directly.
+        try:
+            etag = fileobj.sha1
+        except AttributeError:
+            etag = fileobj.md5
+            if etag is None:
+                pos = fileobj.tell()
+                raw_data = fileobj.read()
+                fileobj.seek(pos)
+                etag = hashlib.sha1(raw_data).hexdigest()
+        response.set_etag(etag)
 
         response.cache_control.max_age = cache_for
         response.cache_control.public = True
@@ -249,5 +247,23 @@ 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)
-        id = storage.put(fileobj, filename=filename, content_type=content_type, **kwargs)
-        return id
+
+        # GridFS does not manage its own checksum, so we attach a sha1 to the file
+        # for use as an etag.
+        hashingfile = _Wrapper(fileobj)
+        with storage.new_file(filename=filename, content_type=content_type, **kwargs) as grid_file:
+            grid_file.write(hashingfile)
+            grid_file.sha1 = hashingfile.hash.hexdigest()
+            return grid_file._id
+
+
+class _Wrapper:
+    def __init__(self, file):
+        self.file = file
+        self.hash = hashlib.sha1()
+
+    def read(self, n):
+        data = self.file.read(n)
+        if data:
+            self.hash.update(data)
+        return data