fix: improve new minio implementation

The team mentioned wanting to ship the new read implementation first
then the new write implementation so this commit enables that.

i'm folding the NewMinioStorageService functionality back into the
MinioStorageService and creating 2 new feature flags one for read and
one for write.

Based on the values of those feature flags the MinioStorageService
will choose an implementation of read/write.

I also removed the old USE_NEW_MINIO feature flag, and modified the
get_appropriate_storage_service function to make it so it no longer
caches the MinioStorageService and instead the MinioStorageService is
caching its internal minio_client since that is the expensive part
of creating a new MinioStorageService instance.

I also moved some code around

Author: joseph-sentry

shared/rollouts/features.py

@@ -2,4 +2,6 @@
 
 BUNDLE_THRESHOLD_FLAG = Feature("bundle_threshold_flag")
 INCLUDE_GITHUB_COMMENT_ACTIONS_BY_OWNER = Feature("include_github_comment_actions")
-USE_NEW_MINIO = Feature("use_new_minio")
+
+READ_NEW_MINIO = Feature("read_new_minio")
+WRITE_NEW_MINIO = Feature("write_new_minio")

shared/storage/__init__.py

@@ -1,13 +1,7 @@
-from typing import Literal
-
 from shared.config import get_config
-from shared.rollouts.features import USE_NEW_MINIO
-from shared.storage.aws import AWSStorageService
+from shared.rollouts.features import READ_NEW_MINIO, WRITE_NEW_MINIO
 from shared.storage.base import BaseStorageService
-from shared.storage.fallback import StorageWithFallbackService
-from shared.storage.gcp import GCPStorageService
 from shared.storage.minio import MinioStorageService
-from shared.storage.new_minio import NewMinioStorageService
 
 _storage_service_cache: dict[str, BaseStorageService] = {}
 
@@ -16,24 +10,20 @@ def get_appropriate_storage_service(
     repoid: int | None = None,
     force_minio=False,
 ) -> BaseStorageService:
-    chosen_storage = "minio"
-    if repoid and USE_NEW_MINIO.check_value(repoid, default=False):
-        chosen_storage = "new_minio"
-
-    if chosen_storage not in _storage_service_cache:
-        _storage_service_cache[chosen_storage] = (
-            _get_appropriate_storage_service_given_storage(chosen_storage)
+    return get_minio_storage_service(repoid)
+
+
+def get_minio_storage_service(
+    repo_id: int | None,
+) -> MinioStorageService:
+    minio_config = get_config("services", "minio", default={})
+    if repo_id:
+        new_read = READ_NEW_MINIO.check_value(repo_id, default=False)
+        new_write = WRITE_NEW_MINIO.check_value(repo_id, default=False)
+        return MinioStorageService(
+            minio_config,
+            new_read=new_read,
+            new_write=new_write,
         )
-
-    return _storage_service_cache[chosen_storage]
-
-
-def _get_appropriate_storage_service_given_storage(
-    chosen_storage: Literal["minio", "new_minio"],
-) -> BaseStorageService:
-    if chosen_storage == "new_minio":
-        minio_config = get_config("services", "minio", default={})
-        return NewMinioStorageService(minio_config)
-    elif chosen_storage == "minio":
-        minio_config = get_config("services", "minio", default={})
+    else:
         return MinioStorageService(minio_config)

shared/storage/compression.py

@@ -0,0 +1,34 @@
+import gzip
+import importlib.metadata
+from typing import IO
+
+
+class GZipStreamReader:
+    def __init__(self, fileobj: IO[bytes]):
+        self.data = fileobj
+
+    def read(self, size: int = -1, /) -> bytes:
+        curr_data = self.data.read(size)
+
+        if not curr_data:
+            return b""
+
+        return gzip.compress(curr_data)
+
+
+def zstd_decoded_by_default() -> bool:
+    try:
+        version = importlib.metadata.version("urllib3")
+    except importlib.metadata.PackageNotFoundError:
+        return False
+
+    if version < "2.0.0":
+        return False
+
+    distribution = importlib.metadata.metadata("urllib3")
+    if requires_dist := distribution.get_all("Requires-Dist"):
+        for req in requires_dist:
+            if "[zstd]" in req:
+                return True
+
+    return False