Merge branch 'morosi-feat' into 'master'

Add bucket context and improve e3.aws.s3.S3

See merge request it/e3-aws!54

Author: adanaja

src/e3/aws/s3/__init__.py

@@ -1,14 +1,30 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING
+import os
 import logging
 from botocore.exceptions import ClientError
+from contextlib import contextmanager
+import botocore
+import boto3
 
 if TYPE_CHECKING:
     from typing import Any
+    from collections.abc import Iterable, Iterator
 
 logger = logging.getLogger("e3.aws.s3")
 
 
+class BucketExistsError(Exception):
+    """Exception when a bucket already exists."""
+
+    def __init__(self, bucket: str) -> None:
+        """Initialize BucketExistsError.
+
+        :param bucket: name of the bucket
+        """
+        self.bucket = bucket
+
+
 class KeyExistsError(Exception):
     """Exception when a key already exists."""
 
@@ -47,6 +63,42 @@ def __init__(
         self.client = client
         self.bucket = bucket
 
+    def create_bucket(self, *, exist_ok: bool = True) -> None:
+        """Create the bucket.
+
+        :param exist_ok: don't raise an exception if the bucket already exists
+        :raises BucketExistsError: if the bucket already exists
+        """
+        try:
+            params: dict[str, Any] = {}
+
+            # us-east-1 is the default location
+            region = self.client.meta.region_name
+            if region != "us-east-1":
+                params["CreateBucketConfiguration"] = {"LocationConstraint": region}
+
+            self.client.create_bucket(Bucket=self.bucket, **params)
+        except ClientError as error:
+            # Raise any non already exists error
+            if error.response["Error"]["Code"] not in [
+                "BucketAlreadyExists",
+                "BucketAlreadyOwnedByYou",
+            ]:
+                raise
+
+            if not exist_ok:
+                raise BucketExistsError(self.bucket) from error
+
+    def clear_bucket(self) -> None:
+        """Clear objects from S3."""
+        for obj in list(self.iterate()):
+            self.client.delete_object(Bucket=self.bucket, Key=obj["Key"])
+
+    def delete_bucket(self) -> None:
+        """Clear and delete the bucket."""
+        self.clear_bucket()
+        self.client.delete_bucket(Bucket=self.bucket)
+
     def push(self, key: str, content: bytes, exist_ok: bool | None = None) -> None:
         """Push content to S3.
 
@@ -93,9 +145,84 @@ def get(self, key: str, default: bytes | None = None) -> bytes:
                 raise KeyNotFoundError(key) from e
             raise e
 
+    def iterate(self, *, prefix: str | None = None) -> Iterable[dict[str, Any]]:
+        """Iterate all objects from S3.
+
+        :param prefix: limit to objects with that prefix
+        :return: an iterator over objects from S3
+        """
+        params = {"Bucket": self.bucket}
+
+        if prefix is not None:
+            params["Prefix"] = prefix
+
+        paginator = self.client.get_paginator("list_objects_v2")
+        for page in paginator.paginate(**params):
+            for content in page.get("Contents", []):
+                yield content
+
     def delete(self, key: str) -> None:
         """Delete content from S3.
 
         :param key: object key
         """
         self.client.delete_object(Bucket=self.bucket, Key=key)
+
+    @property
+    def bucket_exists(self) -> bool:
+        """Return if the bucket exists."""
+        try:
+            self.client.head_bucket(Bucket=self.bucket)
+            return True
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                return False
+            raise
+
+    @property
+    def key_count(self) -> int:
+        """Return the number of keys from S3."""
+        return len(list(self.iterate()))
+
+
+@contextmanager
+def bucket(
+    name: str,
+    *,
+    client: botocore.client.S3 | None = None,
+    region: str | None = None,
+    auto_create: bool = True,
+    auto_delete: bool = False,
+    exist_ok: bool = True,
+) -> Iterator[S3]:
+    """Context manager to create and make AWS API calls on a bucket.
+
+    If auto_create is True, the bucket is created when entering the
+    context. If the bucket already exists and exist_ok is False, an
+    exception is raised.
+
+    If auto_delete is True, the bucket is cleared and deleted when
+    leaving the context.
+
+    :param name: name of the bucket
+    :param client: a client for the S3 API
+    :param region: region of the client (default AWS_DEFAULT_REGION)
+    :param auto_create: create the bucket when entering the context
+    :param auto_delete: delete the bucket when leaving the context
+    :param exist_ok: don't raise an exception if the bucket already exists
+    :raises BucketExistsError: if the bucket already exists
+    """
+    if client is None:
+        region = region if region is not None else os.environ["AWS_DEFAULT_REGION"]
+        client = boto3.client("s3", region_name=region)
+
+    s3 = S3(client=client, bucket=name)
+
+    if auto_create:
+        s3.create_bucket(exist_ok=exist_ok)
+
+    try:
+        yield s3
+    finally:
+        if auto_delete:
+            s3.delete_bucket()

tests/tests_e3_aws/s3/main_test.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING
 import pytest
+from unittest.mock import ANY
 from moto import mock_aws
 import boto3
-from e3.aws.s3 import S3, KeyExistsError, KeyNotFoundError
+from e3.aws import s3
+from e3.aws.s3 import S3, BucketExistsError, KeyExistsError, KeyNotFoundError
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -51,6 +53,46 @@ def test_get_not_found_error(client: S3) -> None:
         client.get("foo")
 
 
+def test_iterate(client: S3) -> None:
+    """Test iterating content."""
+    client.push("hello", b"hello")
+    client.push("prefix/world", b"world")
+    assert list(client.iterate()) == [
+        {
+            "ChecksumAlgorithm": ANY,
+            "ETag": ANY,
+            "Key": "hello",
+            "LastModified": ANY,
+            "Size": 5,
+            "StorageClass": "STANDARD",
+        },
+        {
+            "ChecksumAlgorithm": ANY,
+            "ETag": ANY,
+            "Key": "prefix/world",
+            "LastModified": ANY,
+            "Size": 5,
+            "StorageClass": ANY,
+        },
+    ]
+
+
+def test_iterate_prefix(client: S3) -> None:
+    """Test iterating content with a prefix."""
+    client.push("hello", b"hello")
+    client.push("prefix/world", b"world")
+    assert list(client.iterate(prefix="prefix")) == [
+        {
+            "ChecksumAlgorithm": ANY,
+            "ETag": ANY,
+            "Key": "prefix/world",
+            "LastModified": ANY,
+            "Size": 5,
+            "StorageClass": ANY,
+        }
+    ]
+
+
 def test_delete(client: S3) -> None:
     """Test deleting content."""
     test_push(client)
@@ -61,3 +103,39 @@ def test_delete(client: S3) -> None:
 def test_delete_missing(client: S3) -> None:
     """Test deleting content with a missing key."""
     client.delete("foo")
+
+
+@mock_aws
+def test_bucket() -> None:
+    """Test creating a bucket in a context."""
+    with s3.bucket("test", region="eu-west-1") as client:
+        # Bucket should have been created
+        assert client.bucket_exists
+        assert client.key_count == 0
+        client.push("foo", b"hello")
+        assert client.key_count == 1
+
+    # Bucket should still exist
+    assert client.bucket_exists
+    assert client.get("foo") == b"hello"
+
+
+@mock_aws
+def test_bucket_auto_delete() -> None:
+    """Test auto bucket deletion in a context."""
+    with s3.bucket("test", region="eu-west-1", auto_delete=True) as client:
+        # Bucket should have been created
+        assert client.bucket_exists
+        assert client.key_count == 0
+        client.push("foo", b"hello")
+        assert client.key_count == 1
+
+    # Bucket should no longer exist
+    assert not client.bucket_exists
+
+
+def test_bucket_already_exist_error(client: S3) -> None:
+    """Test creating an already existing bucket in a context."""
+    with pytest.raises(BucketExistsError):
+        with s3.bucket("test", region="eu-west-1", exist_ok=False) as client:
+            assert client.bucket_exists