aws lib

Author: sanderegg

packages/aws-library/src/aws_library/s3/_client.py

@@ -28,19 +28,28 @@
 from settings_library.s3 import S3Settings
 from types_aiobotocore_s3 import S3Client
 from types_aiobotocore_s3.literals import BucketLocationConstraintType
-from types_aiobotocore_s3.type_defs import ObjectIdentifierTypeDef
+from types_aiobotocore_s3.type_defs import (
+    ListObjectsV2RequestRequestTypeDef,
+    ObjectIdentifierTypeDef,
+)
 
-from ._constants import MULTIPART_COPY_THRESHOLD, MULTIPART_UPLOADS_MIN_TOTAL_SIZE
+from ._constants import (
+    MULTIPART_COPY_THRESHOLD,
+    MULTIPART_UPLOADS_MIN_TOTAL_SIZE,
+    S3_OBJECT_DELIMITER,
+)
 from ._error_handler import s3_exception_handler, s3_exception_handler_async_gen
 from ._errors import S3DestinationNotEmptyError, S3KeyNotFoundError
 from ._models import (
     MultiPartUploadLinks,
+    PathCursor,
     S3DirectoryMetaData,
     S3MetaData,
     S3ObjectKey,
+    S3ObjectPrefix,
     UploadID,
 )
-from ._utils import compute_num_file_chunks
+from ._utils import compute_num_file_chunks, create_final_prefix
 
 _logger = logging.getLogger(__name__)
 
@@ -167,7 +176,99 @@ async def get_directory_metadata(
         size = 0
         async for s3_object in self._list_all_objects(bucket=bucket, prefix=prefix):
             size += s3_object.size
-        return S3DirectoryMetaData(size=size)
+        return S3DirectoryMetaData(prefix=S3ObjectPrefix(prefix), size=ByteSize(size))
+
+    @s3_exception_handler(_logger)
+    async def count_objects(
+        self,
+        *,
+        bucket: S3BucketName,
+        prefix: S3ObjectPrefix | None,
+        start_after: S3ObjectKey | None,
+        is_partial_prefix: bool = False,
+        use_delimiter: bool = True,
+    ) -> int:
+        """returns the number of entries in the bucket, defined
+        by prefix and start_after same as list_objects
+        """
+        paginator = self._client.get_paginator("list_objects_v2")
+        total_count = 0
+        async for page in paginator.paginate(
+            Bucket=bucket,
+            Prefix=create_final_prefix(prefix, is_partial_prefix=is_partial_prefix),
+            StartAfter=start_after or "",
+            Delimiter=S3_OBJECT_DELIMITER if use_delimiter else "",
+        ):
+            total_count += page.get("KeyCount", 0)
+        return total_count
+
+    @s3_exception_handler(_logger)
+    async def list_objects(
+        self,
+        *,
+        bucket: S3BucketName,
+        prefix: S3ObjectPrefix | None,
+        start_after: S3ObjectKey | None,
+        limit: int = _MAX_ITEMS_PER_PAGE,
+        next_cursor: PathCursor | None = None,
+        is_partial_prefix: bool = False,
+    ) -> tuple[list[S3MetaData | S3DirectoryMetaData], PathCursor | None]:
+        """returns a number of entries in the bucket, defined by limit
+        the entries are sorted alphabetically by key. If a cursor is returned
+        then the client can call the function again with the cursor to get the
+        next entries.
+
+        the first entry is defined by start_after
+        if start_after is None, the first entry is the first one in the bucket
+        if prefix is not None, only entries with the given prefix are returned
+        if prefix is None, all entries in the bucket are returned
+        if next_cursor is set, then the call will return the next entries after the cursor
+        if is_partial_prefix is set then the prefix is not auto-delimited
+        (if False equivalent to `ls /home/user/`
+        if True equivalent to `ls /home/user*`)
+        limit must be >= 1 and <= _AWS_MAX_ITEMS_PER_PAGE
+
+        Raises:
+            ValueError: in case of invalid limit
+        """
+        if limit < 1:
+            msg = "num_objects must be >= 1"
+            raise ValueError(msg)
+        if limit > _AWS_MAX_ITEMS_PER_PAGE:
+            msg = f"num_objects must be <= {_AWS_MAX_ITEMS_PER_PAGE}"
+            raise ValueError(msg)
+
+        list_config: ListObjectsV2RequestRequestTypeDef = {
+            "Bucket": bucket,
+            "Prefix": create_final_prefix(prefix, is_partial_prefix=is_partial_prefix),
+            "MaxKeys": limit,
+            "Delimiter": S3_OBJECT_DELIMITER,
+        }
+        if start_after:
+            list_config["StartAfter"] = start_after
+        if next_cursor:
+            list_config["ContinuationToken"] = next_cursor
+        listed_objects = await self._client.list_objects_v2(**list_config)
+        found_objects: list[S3MetaData | S3DirectoryMetaData] = []
+        if "CommonPrefixes" in listed_objects:
+            # we have folders here
+            list_subfolders = listed_objects["CommonPrefixes"]
+            found_objects.extend(
+                S3DirectoryMetaData.model_construct(
+                    prefix=S3ObjectPrefix(subfolder["Prefix"], size=None)
+                )
+                for subfolder in list_subfolders
+                if "Prefix" in subfolder
+            )
+        if "Contents" in listed_objects:
+            found_objects.extend(
+                S3MetaData.from_botocore_list_objects(obj)
+                for obj in listed_objects["Contents"]
+            )
+        next_cursor = None
+        if listed_objects["IsTruncated"]:
+            next_cursor = listed_objects["NextContinuationToken"]
+        return found_objects, next_cursor
 
     @s3_exception_handler_async_gen(_logger)
     async def list_objects_paginated(
@@ -459,7 +560,7 @@ async def copy_objects_recursively(
         dst_metadata = await self.get_directory_metadata(
             bucket=bucket, prefix=dst_prefix
         )
-        if dst_metadata.size > 0:
+        if dst_metadata.size and dst_metadata.size > 0:
             raise S3DestinationNotEmptyError(dst_prefix=dst_prefix)
         await limited_gather(
             *[

packages/aws-library/src/aws_library/s3/_constants.py

@@ -12,3 +12,4 @@
 
 PRESIGNED_LINK_MAX_SIZE: Final[ByteSize] = TypeAdapter(ByteSize).validate_python("5GiB")
 S3_MAX_FILE_SIZE: Final[ByteSize] = TypeAdapter(ByteSize).validate_python("5TiB")
+S3_OBJECT_DELIMITER: Final[str] = "/"

packages/aws-library/src/aws_library/s3/_models.py

@@ -1,23 +1,24 @@
 import datetime
-from dataclasses import dataclass
-from typing import TypeAlias
+from pathlib import Path
+from typing import TypeAlias, cast
 
 from models_library.api_schemas_storage.storage_schemas import ETag
 from models_library.basic_types import SHA256Str
-from pydantic import AnyUrl, BaseModel, ByteSize
+from pydantic import AnyUrl, BaseModel, ByteSize, Field
 from types_aiobotocore_s3.type_defs import HeadObjectOutputTypeDef, ObjectTypeDef
 
 S3ObjectKey: TypeAlias = str
+S3ObjectPrefix: TypeAlias = Path
 UploadID: TypeAlias = str
+PathCursor: TypeAlias = str
 
 
-@dataclass(frozen=True, slots=True, kw_only=True)
-class S3MetaData:
+class S3MetaData(BaseModel, frozen=True):
     object_key: S3ObjectKey
     last_modified: datetime.datetime
     e_tag: ETag
     sha256_checksum: SHA256Str | None
-    size: int
+    size: ByteSize
 
     @staticmethod
     def from_botocore_head_object(
@@ -27,12 +28,8 @@ def from_botocore_head_object(
             object_key=object_key,
             last_modified=obj["LastModified"],
             e_tag=obj["ETag"].strip('"'),
-            sha256_checksum=(
-                SHA256Str(obj.get("ChecksumSHA256"))
-                if obj.get("ChecksumSHA256")
-                else None
-            ),
-            size=obj["ContentLength"],
+            sha256_checksum=obj.get("ChecksumSHA256"),
+            size=ByteSize(obj["ContentLength"]),
         )
 
     @staticmethod
@@ -47,18 +44,22 @@ def from_botocore_list_objects(
             object_key=obj["Key"],
             last_modified=obj["LastModified"],
             e_tag=obj["ETag"].strip('"'),
-            sha256_checksum=(
-                SHA256Str(obj.get("ChecksumSHA256"))
-                if obj.get("ChecksumSHA256")
-                else None
-            ),
-            size=obj["Size"],
+            sha256_checksum=cast(SHA256Str | None, obj.get("ChecksumSHA256")),
+            size=ByteSize(obj["Size"]),
         )
 
+    def as_path(self) -> Path:
+        return Path(self.object_key)
 
-@dataclass(frozen=True)
-class S3DirectoryMetaData:
-    size: int
+
+class S3DirectoryMetaData(BaseModel, frozen=True):
+    prefix: S3ObjectPrefix
+    size: ByteSize | None = Field(
+        ..., description="Size of the directory if computed, None if unknown"
+    )
+
+    def as_path(self) -> Path:
+        return self.prefix
 
 
 class MultiPartUploadLinks(BaseModel):

packages/aws-library/src/aws_library/s3/_utils.py

@@ -2,6 +2,9 @@
 
 from pydantic import ByteSize, TypeAdapter
 
+from ._constants import S3_OBJECT_DELIMITER
+from ._models import S3ObjectPrefix
+
 _MULTIPART_MAX_NUMBER_OF_PARTS: Final[int] = 10000
 
 # this is artifically defined, if possible we keep a maximum number of requests for parallel
@@ -34,3 +37,15 @@ def compute_num_file_chunks(file_size: ByteSize) -> tuple[int, ByteSize]:
     raise ValueError(
         msg,
     )
+
+
+def create_final_prefix(
+    prefix: S3ObjectPrefix | None, *, is_partial_prefix: bool
+) -> str:
+    final_prefix = f"{prefix}" if prefix else ""
+    if prefix and not is_partial_prefix:
+        final_prefix = (
+            f"{final_prefix.rstrip(S3_OBJECT_DELIMITER)}{S3_OBJECT_DELIMITER}"
+        )
+
+    return final_prefix