feat: Support gzip compression in server and clients

Cargo.toml

@@ -52,6 +52,7 @@ console = "0.16.3"
 ed25519-dalek = { version = "2.2.0" }
 elegant-departure = "0.3.2"
 figment = "0.10.19"
+flate2 = "1.1.9"
 futures = "0.3.32"
 futures-util = "0.3.32"
 gcp_auth = "0.12.6"

clients/python/README.md

@@ -92,9 +92,13 @@ session.put(b"payload", origin="203.0.113.42")
 ### Compression
 
 Uploads are compressed with Zstd by default, and downloads are transparently decompressed.
-Compression can be overridden per-upload:
+The supported algorithms are `"zstd"` and `"gzip"`, or `"none"` to disable compression.
+Compression can be set per-usecase or overridden per-upload:
 
 ```python
+# compress all uploads of this usecase with zstd (the default):
+usecase = Usecase("my-usecase", compression="zstd")
+
 # upload as-is and record no encoding:
 session.put(video_data, compress="none")
 

clients/python/src/objectstore_client/_compression.py

@@ -0,0 +1,175 @@
+"""Streaming compression and decompression for the supported algorithms."""
+
+from __future__ import annotations
+
+import gzip
+import io
+import shutil
+import zlib
+from typing import IO, Protocol, cast
+
+import zstandard
+
+from objectstore_client.metadata import Compression
+
+SUPPORTED_COMPRESSIONS: tuple[Compression, ...] = ("zstd", "gzip")
+"""Compression algorithms that the client can apply and transparently decode."""
+
+
+class _Compressor(Protocol):
+    """A streaming compression object, such as `zlib.compressobj()`."""
+
+    def compress(self, data: bytes, /) -> bytes: ...
+
+    def flush(self) -> bytes: ...
+
+
+class _CompressionReader(io.RawIOBase):
+    """
+    A readable stream that compresses `source` with `compressor` as it is read.
+
+    `tell()` reports the number of compressed bytes produced so far. `seek(0)` is
+    allowed as long as no data has been read yet, so that urllib3 can record the body
+    position. Closing the reader does not close `source`.
+    """
+
+    _CHUNK_SIZE = 64 * 1024
+
+    def __init__(self, source: IO[bytes], compressor: _Compressor) -> None:
+        super().__init__()
+        self._source = source
+        self._compressor = compressor
+        self._buffer = bytearray()
+        self._finished = False
+        self._position = 0
+
+    def readable(self) -> bool:
+        return True
+
+    def read(self, size: int | None = -1, /) -> bytes:
+        if size is None or size < 0:
+            size = -1
+        while not self._finished and (size < 0 or len(self._buffer) < size):
+            if chunk := self._source.read(self._CHUNK_SIZE):
+                self._buffer += self._compressor.compress(chunk)
+            else:
+                self._buffer += self._compressor.flush()
+                self._finished = True
+
+        if size < 0:
+            size = len(self._buffer)
+        out = bytes(self._buffer[:size])
+        del self._buffer[:size]
+        self._position += len(out)
+        return out
+
+    def readall(self) -> bytes:
+        return self.read(-1)
+
+    def readinto(self, buffer: memoryview | bytearray) -> int:  # type: ignore[override]
+        data = self.read(len(buffer))
+        buffer[: len(data)] = data
+        return len(data)
+
+    def tell(self) -> int:
+        return self._position
+
+    def seek(self, offset: int, whence: int = io.SEEK_SET, /) -> int:
+        if offset == 0 and whence == io.SEEK_SET and self._position == 0:
+            return 0
+        raise OSError("Cannot seek in a compressed stream after reading has started")
+
+
+class _GzipReader(gzip.GzipFile):
+    """Decode gzip from a stream, closing the source with the reader."""
+
+    def __init__(self, source: IO[bytes]) -> None:
+        super().__init__(fileobj=source, mode="rb")
+        self._source = source
+
+    def seekable(self) -> bool:
+        return False
+
+    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+        raise io.UnsupportedOperation("gzip download stream is not seekable")
+
+    def close(self) -> None:
+        if self.closed:
+            return
+        try:
+            super().close()
+        finally:
+            self._source.close()
+
+
+def _gzip_compressor() -> _Compressor:
+    """Return a zlib compressor that emits the gzip format.
+
+    By default, zlib wraps deflate data in the zlib format. Adding 16 to the window
+    size (`wbits`) selects the gzip header and trailer instead, which is what
+    `Content-Encoding: gzip` requires. See `zlib.compressobj`.
+    """
+    return zlib.compressobj(wbits=16 + zlib.MAX_WBITS)
+
+
+def validate(compression: str | None, *, allow_none: bool = True) -> None:
+    """Raise `ValueError` if `compression` is not a supported algorithm.
+
+    `None` is always accepted. `"none"` is accepted only if `allow_none` is set.
+    """
+    if compression is None or (allow_none and compression == "none"):
+        return
+    if compression not in SUPPORTED_COMPRESSIONS:
+        raise ValueError(f"Invalid compression: {compression}")
+
+
+def compressing_reader(source: IO[bytes], compression: Compression) -> IO[bytes]:
+    """Return a readable stream yielding `source` compressed with `compression`.
+
+    The returned stream compresses lazily as it is read, reports the number of
+    compressed bytes produced via `tell()`, and supports `seek(0)` only until
+    reading starts. Closing it does not close `source`.
+    """
+    if compression == "zstd":
+        compressor = zstandard.ZstdCompressor().compressobj()
+        return cast(IO[bytes], _CompressionReader(source, compressor))
+    if compression == "gzip":
+        return cast(IO[bytes], _CompressionReader(source, _gzip_compressor()))
+    raise ValueError(f"Invalid compression: {compression}")
+
+
+def compress_into(
+    source: IO[bytes], dest: IO[bytes], compression: Compression
+) -> tuple[int, int]:
+    """Compress `source` into `dest`, returning the bytes read and written.
+
+    For gzip, `dest` must support `tell()` to measure the compressed size.
+    """
+    if compression == "zstd":
+        read, written = zstandard.ZstdCompressor().copy_stream(source, dest)
+        return read, written
+    if compression == "gzip":
+        start = dest.tell()
+        with gzip.GzipFile(
+            filename="", fileobj=dest, mode="wb", compresslevel=6, mtime=0
+        ) as writer:
+            shutil.copyfileobj(source, writer)
+            read = writer.tell()
+        return read, dest.tell() - start
+    raise ValueError(f"Invalid compression: {compression}")
+
+
+def decompressing_reader(stream: IO[bytes], compression: str) -> IO[bytes]:
+    """Return a readable stream decoding `stream` compressed with `compression`.
+
+    Concatenated zstd frames and gzip members are decoded as a single payload.
+    The returned reader is non-seekable and closes the source when closed.
+    """
+    if compression == "zstd":
+        dctx = zstandard.ZstdDecompressor()
+        return cast(IO[bytes], dctx.stream_reader(stream, read_across_frames=True))
+    if compression == "gzip":
+        return cast(IO[bytes], _GzipReader(stream))
+    raise NotImplementedError(
+        f"Transparent decoding of `{compression}` is not implemented"
+    )

clients/python/src/objectstore_client/_resumable.py

@@ -12,8 +12,8 @@
 from urllib.parse import urlencode
 
 import urllib3
-import zstandard
 
+from objectstore_client import _compression
 from objectstore_client.errors import RequestError, raise_for_status
 from objectstore_client.metadata import Compression, ExpirationPolicy
 from objectstore_client.metrics import measure_storage_operation
@@ -236,7 +236,7 @@ def cancel(self) -> None:
 def prepare_body(
     contents: bytes | IO[bytes],
     source_size: int,
-    compress: bool,
+    compression: Compression | None,
     spool_dir: str | PathLike[str] | None,
 ) -> Iterator[tuple[IO[bytes], int, int] | None]:
     """Yield the body, source size, and compressed size (equal if uncompressed).
@@ -248,7 +248,7 @@ def prepare_body(
     with (
         BytesIO(contents) if isinstance(contents, bytes) else nullcontext(contents)
     ) as source:
-        if not compress:
+        if compression is None or compression == "none":
             yield source, source_size, source_size
             return
 
@@ -259,8 +259,8 @@ def prepare_body(
             return
 
         with temporary as body:
-            source_size, compressed_size = zstandard.ZstdCompressor().copy_stream(
-                source, body
+            source_size, compressed_size = _compression.compress_into(
+                source, body, compression
             )
             body.seek(0)
             yield body, source_size, compressed_size

clients/python/src/objectstore_client/client.py

@@ -12,10 +12,9 @@
 from urllib.parse import urlparse
 
 import urllib3
-import zstandard
 from urllib3.connectionpool import HTTPConnectionPool
 
-from objectstore_client import _resumable, presign, utils
+from objectstore_client import _compression, _resumable, presign, utils
 from objectstore_client.auth import Permission, TokenGenerator, TokenProvider
 from objectstore_client.errors import (
     ExpiryExtensionRejected,
@@ -91,7 +90,7 @@ def __init__(
         """Create a usecase.
 
         Raises ``ValueError`` if the name violates the naming rules on ``Usecase``,
-        or if the ``resumable_threshold_bytes`` is invalid.
+        or if either ``compression`` or ``resumable_threshold_bytes`` is invalid.
         """
         if not 3 <= len(name) <= 64:
             raise ValueError("usecase name must be 3–64 ASCII characters long")
@@ -108,6 +107,7 @@ def __init__(
                 'resumable_threshold_bytes must be a non-negative integer or "none"'
             )
         self._resumable_threshold_bytes = resumable_threshold_bytes
+        _compression.validate(compression)
         self._compression = compression
         self._expiration_policy = expiration_policy
 
@@ -437,10 +437,8 @@ def put(
         """
         if compress is not None and precompressed is not None:
             raise ValueError("Cannot pass both `compress` and `precompressed`")
-        if compress and compress not in ("none", "zstd"):
-            raise ValueError(f"Invalid compression: {compress}")
-        if precompressed and precompressed != "zstd":
-            raise ValueError(f"Invalid compression: {precompressed}")
+        _compression.validate(compress)
+        _compression.validate(precompressed, allow_none=False)
 
         if resumable_threshold_bytes is None:
             resumable_threshold_bytes = self._usecase._resumable_threshold_bytes
@@ -486,13 +484,13 @@ def put(
                 with _resumable.prepare_body(
                     contents,
                     source_size,
-                    compress_with != "none",
+                    compress_with,
                     self._spool_dir,
                 ) as prepared:
                     if prepared is None:
                         # Temporary-file creation failed before consuming the input.
                         result_key, source_size, compressed_size = self._put_direct(
-                            contents, key, headers, compress=compress_with != "none"
+                            contents, key, headers, compression=compress_with
                         )
                     else:
                         body, source_size, compressed_size = prepared
@@ -513,11 +511,11 @@ def put(
                             # Send the prepared body without compressing it again.
                             headers["Content-Length"] = str(compressed_size)
                             result_key, _, _ = self._put_direct(
-                                body, key, headers, compress=False
+                                body, key, headers, compression="none"
                             )
             else:
                 result_key, source_size, compressed_size = self._put_direct(
-                    contents, key, headers, compress=compress_with != "none"
+                    contents, key, headers, compression=compress_with
                 )
 
             if precompressed is None and source_size is not None:
@@ -543,7 +541,7 @@ def _put_direct(
         contents: bytes | IO[bytes],
         key: str | None,
         headers: dict[str, str],
-        compress: bool,
+        compression: Compression,
     ) -> tuple[str, int | None, int | None]:
         """Return the key and final source/body positions when available."""
         with (
@@ -552,12 +550,8 @@ def _put_direct(
             body = original_body
             retries = None  # by default use the pool's value, set by the Client
             try:
-                if compress:
-                    cctx = zstandard.ZstdCompressor()
-                    body = cctx.stream_reader(original_body, closefd=False)
-                    body = cast(IO[bytes], utils._ZstdCompressionReaderWrapper(body))
-
-                if compress:
+                if compression != "none":
+                    body = _compression.compressing_reader(original_body, compression)
                     # For on-the-fly compression, don't attempt read retries,
                     # as the stream cannot be rewound after data has been consumed.
                     pool_retries = self._pool.retries
@@ -645,15 +639,8 @@ def get(
             )
             decompressed = False
             if metadata.compression and decompress and not encoding_accepted:
-                if metadata.compression != "zstd":
-                    raise NotImplementedError(
-                        "Transparent decoding of anything but `zstd` is not "
-                        "implemented yet"
-                    )
-
+                stream = _compression.decompressing_reader(stream, metadata.compression)
                 metadata.compression = None
-                dctx = zstandard.ZstdDecompressor()
-                stream = dctx.stream_reader(stream, read_across_frames=True)
                 decompressed = True
 
             span.set_attribute("objectstore.decompressed", decompressed)
@@ -933,8 +920,7 @@ def _create_upload(
         """
         if object_length < 0:
             raise ValueError("Object length must not be negative")
-        if compression is not None and compression not in ("none", "zstd"):
-            raise ValueError(f"Invalid compression: {compression}")
+        _compression.validate(compression)
 
         headers = self._metadata_headers(
             compression=compression,

clients/python/src/objectstore_client/metadata.py

@@ -9,7 +9,7 @@
 
 from objectstore_client.utils import decode_header_value
 
-Compression = Literal["zstd"] | Literal["none"]
+Compression = Literal["zstd", "gzip"] | Literal["none"]
 
 HEADER_EXPIRATION = "x-sn-expiration"
 HEADER_TIME_CREATED = "x-sn-time-created"