fix: improve new minio implementation (#546)

* feat: remove usage of use_minio feature flag

this rollout is done and we are 100% enabled on prod so we are safe to
remove this and simplify the logic

starting off by removing all references to use_minio in shared

* 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

* fix: address feedback

Author: joseph-sentry

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