persist metadata on append (#575)

Co-authored-by: Ruslan Kuprieiev <[email protected]>

Author: Platob

.gitignore

@@ -7,3 +7,4 @@ __pycache__
 dist/
 *.egg-info
 build/
+venv/

s3fs/core.py

@@ -52,6 +52,25 @@ def setup_logging(level=None):
 
 _VALID_FILE_MODES = {"r", "w", "a", "rb", "wb", "ab"}
 
+_PRESERVE_KWARGS = [
+    "CacheControl",
+    "ContentDisposition",
+    "ContentEncoding",
+    "ContentLanguage",
+    "ContentLength",
+    "ContentType",
+    "Expires",
+    "WebsiteRedirectLocation",
+    "ServerSideEncryption",
+    "SSECustomerAlgorithm",
+    "SSEKMSKeyId",
+    "BucketKeyEnabled",
+    "StorageClass",
+    "ObjectLockMode",
+    "ObjectLockRetainUntilDate",
+    "ObjectLockLegalHoldStatus",
+    "Metadata",
+]
 
 key_acls = {
     "private",
@@ -1047,7 +1066,7 @@ async def _info(self, path, bucket=None, key=None, refresh=False, version_id=Non
                     "size": out["ContentLength"],
                     "name": "/".join([bucket, key]),
                     "type": "file",
-                    "StorageClass": "STANDARD",
+                    "StorageClass": out.get("StorageClass", "STANDARD"),
                     "VersionId": out.get("VersionId"),
                     "ContentType": out.get("ContentType"),
                 }
@@ -1817,14 +1836,38 @@ def __init__(
         self.append_block = False
 
         if "a" in mode and s3.exists(path):
-            loc = s3.info(path)["size"]
+            # See:
+            # put: https://boto3.amazonaws.com/v1/documentation/api/latest
+            # /reference/services/s3.html#S3.Client.put_object
+            #
+            # head: https://boto3.amazonaws.com/v1/documentation/api/latest
+            # /reference/services/s3.html#S3.Client.head_object
+            head = self._call_s3(
+                "head_object",
+                self.kwargs,
+                Bucket=bucket,
+                Key=key,
+                **version_id_kw(version_id),
+                **self.req_kw,
+            )
+
+            head = {
+                key: value
+                for key, value in head.items()
+                if key in _PRESERVE_KWARGS and key not in self.s3_additional_kwargs
+            }
+
+            loc = head.pop("ContentLength")
             if loc < 5 * 2 ** 20:
                 # existing file too small for multi-upload: download
                 self.write(self.fs.cat(self.path))
             else:
                 self.append_block = True
             self.loc = loc
 
+            # Reflect head
+            self.s3_additional_kwargs.update(head)
+
         if "r" in mode and "ETag" in self.details:
             self.req_kw["IfMatch"] = self.details["ETag"]
 

s3fs/tests/test_s3fs.py

@@ -15,6 +15,8 @@
 import moto
 from itertools import chain
 import fsspec.core
+from dateutil.tz import tzutc
+
 import s3fs.core
 from s3fs.core import S3FileSystem
 from s3fs.utils import ignoring, SSEParams
@@ -1290,6 +1292,36 @@ def test_append(s3):
         assert f.tell() == 20 * 2 ** 20 + 5
     assert s3.cat(a) == b"a" * 10 * 2 ** 20 + b"extra" + b"b" * 10 * 2 ** 20
 
+    # Keep Head Metadata
+    head = dict(
+        CacheControl="public",
+        ContentDisposition="string",
+        ContentEncoding="gzip",
+        ContentLanguage="ru-RU",
+        ContentType="text/csv",
+        Expires=datetime.datetime(2015, 1, 1, 0, 0, tzinfo=tzutc()),
+        Metadata={"string": "string"},
+        ServerSideEncryption="AES256",
+        StorageClass="REDUCED_REDUNDANCY",
+        WebsiteRedirectLocation="https://www.example.com/",
+        BucketKeyEnabled=False,
+    )
+    with s3.open(a, "wb", **head) as f:
+        f.write(b"data")
+
+    with s3.open(a, "ab") as f:
+        f.write(b"other")
+
+    with s3.open(a) as f:
+        filehead = {
+            k: v
+            for k, v in f._call_s3(
+                "head_object", f.kwargs, Bucket=f.bucket, Key=f.key
+            ).items()
+            if k in head
+        }
+        assert filehead == head
+
 
 def test_bigger_than_block_read(s3):
     with s3.open(test_bucket_name + "/2014-01-01.csv", "rb", block_size=3) as f: