Merge pull request #207 from Aiven-Open/joelynch/use-iterable

typing: allow passing Iterable of unknown size to delete_keys

Author: bogdanp05

rohmu/object_storage/azure.py

@@ -7,6 +7,7 @@
 
 from azure.core.exceptions import HttpResponseError, IncompleteReadError, ResourceExistsError
 from azure.storage.blob import BlobServiceClient, ContentSettings
+from collections.abc import Iterable, Sized
 from rohmu.common.statsd import StatsdConfig
 from rohmu.errors import (
     FileNotFoundFromStorageError,
@@ -34,7 +35,7 @@
 )
 from rohmu.typing import Metadata
 from rohmu.util import batched
-from typing import Any, BinaryIO, Collection, Iterator, Optional, Tuple, Union
+from typing import Any, BinaryIO, Iterator, Optional, Tuple, Union
 from typing_extensions import Self
 
 import azure.common
@@ -317,9 +318,12 @@ def delete_key(self, key: str, preserve_trailing_slash: bool = False) -> None:
 
         return result
 
-    def delete_keys(self, keys: Collection[str], preserve_trailing_slash: bool = False) -> None:
+    def delete_keys(self, keys: Iterable[str], preserve_trailing_slash: bool = False) -> None:
         container_client = self.get_blob_service_client().get_container_client(container=self.container_name)
-        self.log.debug("Deleting %i keys", len(keys))
+        if isinstance(keys, Sized):
+            self.log.debug("Deleting %i keys", len(keys))
+        else:
+            self.log.debug("Deleting keys")
         for keys_batch in batched(keys, 256):  # 256 is the maximum batch size for Azure
             paths_to_delete = []
             for key in keys_batch:

rohmu/object_storage/base.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterable
 from contextlib import suppress
 from dataclasses import dataclass, field
 from io import BytesIO
@@ -258,7 +259,7 @@ def format_key_from_backend(self, key: str) -> str:
     def delete_key(self, key: str, preserve_trailing_slash: bool = False) -> None:
         raise NotImplementedError
 
-    def delete_keys(self, keys: Collection[str], preserve_trailing_slash: bool = False) -> None:
+    def delete_keys(self, keys: Iterable[str], preserve_trailing_slash: bool = False) -> None:
         """Delete specified keys"""
         for key in keys:
             self.delete_key(key, preserve_trailing_slash=preserve_trailing_slash)

rohmu/object_storage/google.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Sized
 from contextlib import contextmanager
 from googleapiclient.discovery import build
 from googleapiclient.errors import BatchError, HttpError
@@ -52,7 +53,6 @@
     BinaryIO,
     Callable,
     cast,
-    Collection,
     Iterable,
     Iterator,
     Optional,
@@ -486,7 +486,7 @@ def delete_key(self, key: str, preserve_trailing_slash: bool = False) -> None:
             reporter.report(self.stats)
             self.notifier.object_deleted(key)
 
-    def delete_keys(self, keys: Collection[str], preserve_trailing_slash: bool = False) -> None:
+    def delete_keys(self, keys: Iterable[str], preserve_trailing_slash: bool = False) -> None:
         retry_deletion_keys: list[str] = []
 
         def _delete_keys_callback(key: str, response: HttpRequest, exception: HttpError | None) -> None:
@@ -500,7 +500,10 @@ def _delete_keys_callback(key: str, response: HttpRequest, exception: HttpError
             # Retry outsize of callback to avoid retry multiplication
             retry_deletion_keys.append(key)
 
-        self.log.debug("Deleting %i keys", len(keys))
+        if isinstance(keys, Sized):
+            self.log.debug("Deleting %i keys", len(keys))
+        else:
+            self.log.debug("Deleting keys")
         reporter = Reporter(StorageOperation.delete_key)
         with self._object_client() as object_resource:
             for keys_batch in batched(keys, 100):  # Docs ask to not batch more than 100 requests

rohmu/object_storage/s3.py

@@ -43,7 +43,7 @@
 from rohmu.typing import Metadata
 from rohmu.util import batched, ProgressStream
 from threading import RLock
-from typing import Any, BinaryIO, cast, Collection, Iterator, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, BinaryIO, cast, Iterable, Iterator, Optional, Tuple, TYPE_CHECKING, Union
 from typing_extensions import Self
 
 import botocore.client
@@ -344,9 +344,9 @@ def delete_key(self, key: str, preserve_trailing_slash: bool = False) -> None:
         self.get_client().delete_object(Bucket=self.bucket_name, Key=path)
         self.notifier.object_deleted(key=key)
 
-    def delete_keys(self, keys: Collection[str], preserve_trailing_slash: bool = False) -> None:
-        self.stats.operation(StorageOperation.delete_key, count=len(keys))
+    def delete_keys(self, keys: Iterable[str], preserve_trailing_slash: bool = False) -> None:
         for batch in batched(keys, 1000):  # Cannot delete more than 1000 objects at a time
+            self.stats.operation(StorageOperation.delete_key, count=len(batch))
             formatted_keys = [
                 self.format_key_for_backend(
                     k,