Use caching for the etag using stored hash if available

flask_pymongo/__init__.py

@@ -183,12 +183,22 @@ def get_upload(filename):
         response.headers["Content-Disposition"] = f"attachment; filename={filename}"
         response.content_length = fileobj.length
         response.last_modified = fileobj.upload_date
-        # Compute the sha1 sum of the file for the etag.
-        pos = fileobj.tell()
-        raw_data = fileobj.read()
-        fileobj.seek(pos)
-        digest = hashlib.sha1(raw_data).hexdigest()
-        response.set_etag(digest)
+
+        # 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
         response.make_conditional(request)
@@ -237,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

tests/util.py

@@ -33,5 +33,5 @@ def setUp(self):
     def tearDown(self):
         assert self.mongo.cx is not None
         self.mongo.cx.drop_database(self.dbname)
-
+        self.mongo.cx.close()
         super().tearDown()