fix: improve new minio implementation

shared/rollouts/features.py

@@ -2,5 +2,4 @@
 
 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")
-USE_MINIO = Feature("use_minio")
+NEW_MINIO = Feature("new_minio")

shared/storage/__init__.py

@@ -1,56 +1,24 @@
+from typing import Literal, cast
+
 from shared.config import get_config
-from shared.rollouts.features import USE_MINIO, USE_NEW_MINIO
-from shared.storage.aws import AWSStorageService
+from shared.rollouts.features import 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] = {}
 
 
 def get_appropriate_storage_service(
     repoid: int | None = None,
     force_minio=False,
 ) -> BaseStorageService:
-    chosen_storage: str = get_config("services", "chosen_storage", default="minio")  # type: ignore
-    if force_minio:
-        chosen_storage = "minio"
-
+    minio_config = get_config("services", "minio", default={})
     if repoid:
-        if USE_MINIO.check_value(repoid, default=False):
-            chosen_storage = "minio"
-
-        if 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)
+        new_minio_mode = cast(
+            Literal["read", "write"] | None,
+            NEW_MINIO.check_value(repoid, default=None),  # type: ignore
+        )
+        return MinioStorageService(
+            minio_config,
+            new_mode=new_minio_mode,
         )
-
-    return _storage_service_cache[chosen_storage]
-
-
-def _get_appropriate_storage_service_given_storage(
-    chosen_storage: str,
-) -> BaseStorageService:
-    if chosen_storage == "gcp":
-        gcp_config = get_config("services", "gcp", default={})
-        return GCPStorageService(gcp_config)
-    elif chosen_storage == "aws":
-        aws_config = get_config("services", "aws", default={})
-        return AWSStorageService(aws_config)
-    elif chosen_storage == "gcp_with_fallback":
-        gcp_config = get_config("services", "gcp", default={})
-        gcp_service = GCPStorageService(gcp_config)
-        aws_config = get_config("services", "aws", default={})
-        aws_service = AWSStorageService(aws_config)
-        return StorageWithFallbackService(gcp_service, aws_service)
-    elif chosen_storage == "new_minio":
-        minio_config = get_config("services", "minio", default={})
-        return NewMinioStorageService(minio_config)
     else:
-        minio_config = get_config("services", "minio", default={})
         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