notebook-snapshot-extension/snapshot/storage/s3.py

"""S3-compatible storage adapter backed by boto3.

Works with AWS S3, MinIO, and any S3-API-compatible service.
"""

from .base import StorageAdapter, validate_key


class S3ConnectionError(Exception):
    """Raised when the S3 connectivity / credentials check fails at startup.

    Wraps the underlying boto3/botocore exception with a friendly, actionable
    message naming the endpoint, bucket, and the most likely fix.
    """


class S3Storage(StorageAdapter):
    """Store objects in an S3-compatible bucket via boto3.

    Works with AWS S3 and MinIO (just point ``endpoint`` at the MinIO host).
    """

    def __init__(
        self,
        endpoint: str,
        access_key: str,
        secret_key: str,
        bucket: str,
        secure: bool = False,
        prefix: str = "",
    ) -> None:
        """Initialize the adapter and verify connectivity.

        boto3 is imported lazily inside this constructor so that importing
        ``snapshot.storage.s3`` does not require boto3 to be installed.

        On construction, the adapter performs a HEAD on the configured bucket
        to validate network reachability and credentials. If the bucket does
        not exist, it is created automatically. Any connectivity, credential,
        or permission problem is surfaced as :class:`S3ConnectionError` with
        a message that names the endpoint, bucket, and likely fix.

        Parameters
        ----------
        endpoint: str
            S3 endpoint host (without scheme), e.g. ``localhost:9000`` for
            MinIO or ``s3.us-east-1.amazonaws.com`` for AWS S3.
        access_key: str
            S3 access key ID.
        secret_key: str
            S3 secret access key.
        bucket: str
            Target bucket name. Created if it does not exist.
        secure: bool, optional
            Use HTTPS instead of HTTP. Defaults to ``False``.
        prefix: str, optional
            Prefix prepended to every stored key. Defaults to ``""``.

        Raises
        ------
        S3ConnectionError
            If the endpoint is unreachable, credentials are missing/invalid,
            or the bucket cannot be created due to permissions.
        """
        try:
            import boto3
            from botocore.client import Config as BotoConfig
            from botocore.exceptions import ClientError
        except ImportError as exc:
            raise ImportError(
                "boto3 is required for S3 storage. "
                "Install it with: pip install snapshot[s3]"
            ) from exc

        self._bucket = bucket
        self._prefix = prefix
        self._endpoint = endpoint
        scheme = "https" if secure else "http"
        self._client = boto3.client(
            "s3",
            endpoint_url=f"{scheme}://{endpoint}",
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name="us-east-1",
            config=BotoConfig(signature_version="s3v4"),
        )
        self._ClientError = ClientError
        self._connect_and_ensure_bucket()

    def _connect_and_ensure_bucket(self) -> None:
        """Verify connectivity/credentials, then ensure the bucket exists.

        Raises :class:`S3ConnectionError` with a friendly message on failure.
        """
        try:
            self._client.head_bucket(Bucket=self._bucket)
            return  # bucket exists, connectivity and credentials confirmed
        except self._ClientError as exc:
            code = exc.response.get("Error", {}).get("Code", "")
            if code in ("404", "NoSuchBucket"):
                # Bucket missing — try to create it. Permission failure here
                # surfaces as S3ConnectionError rather than the raw boto3 error.
                try:
                    self._client.create_bucket(Bucket=self._bucket)
                except self._ClientError as create_exc:
                    ccode = create_exc.response.get("Error", {}).get("Code", "")
                    raise S3ConnectionError(
                        f"S3 bucket '{self._bucket}' does not exist and could "
                        f"not be created (code={ccode}). Check that the "
                        f"credentials have s3:CreateBucket permission on "
                        f"endpoint {self._endpoint}."
                    ) from create_exc
                return
            # Other ClientError: permission denied, account disabled, etc.
            raise S3ConnectionError(
                f"S3 access denied (endpoint={self._endpoint}, "
                f"bucket={self._bucket}, code={code}). Check that the "
                f"credentials have s3:ListBucket and s3:GetBucketLocation "
                f"permissions."
            ) from exc
        except Exception as exc:  # noqa: BLE001 - catch any boto3/botocore error
            # EndpointConnectionError, NoCredentialsError, InvalidAccessKeyId,
            # SSLError, ConnectTimeoutError, etc. all land here.
            cls = type(exc).__name__
            hint = ""
            if "Credentials" in cls or "AccessKey" in cls or "Auth" in cls:
                hint = (
                    " Check that S3_ACCESS_KEY / S3_SECRET_KEY env vars (or "
                    "c.SnapshotConfig.s3_access_key / s3_secret_key) are set "
                    "and correct."
                )
            elif "Endpoint" in cls or "Connection" in cls or "Connect" in cls or "SSLError" in cls or "timeout" in cls.lower():
                hint = (
                    f" Check that the endpoint '{self._endpoint}' is reachable, "
                    f"the URL is correct, and s3_secure matches the protocol."
                )
            raise S3ConnectionError(
                f"Cannot connect to S3 endpoint {self._endpoint} "
                f"(bucket={self._bucket}): {cls}: {exc}.{hint}"
            ) from exc

    def _object_key(self, key: str) -> str:
        """Return the full S3 object key, including the configured prefix."""
        return f"{self._prefix}{key}"

    def put(self, key: str, data: bytes) -> None:
        """Store ``data`` under ``key`` in the configured bucket."""
        validate_key(key)
        self._client.put_object(
            Bucket=self._bucket,
            Key=self._object_key(key),
            Body=data,
        )

    def get(self, key: str) -> bytes:
        """Return the bytes stored under ``key``.

        Raises
        ------
        KeyError
            If the object does not exist.
        """
        validate_key(key)
        try:
            response = self._client.get_object(
                Bucket=self._bucket,
                Key=self._object_key(key),
            )
        except self._ClientError as exc:
            error_code = exc.response.get("Error", {}).get("Code", "")
            if error_code in ("NoSuchKey", "404"):
                raise KeyError(key) from exc
            raise
        return response["Body"].read()

    def list_keys(self, prefix: str = "") -> list[str]:
        """Return a sorted list of keys starting with ``prefix``.

        The configured bucket prefix is stripped from returned keys.
        """
        if prefix:
            validate_key(prefix)
        full_prefix = self._object_key(prefix)
        keys = []
        continuation_token = None
        while True:
            kwargs: dict = {
                "Bucket": self._bucket,
                "Prefix": full_prefix,
            }
            if continuation_token:
                kwargs["ContinuationToken"] = continuation_token
            response = self._client.list_objects_v2(**kwargs)
            for obj in response.get("Contents", []):
                full_key = obj["Key"]
                if full_key.startswith(self._prefix):
                    stripped = full_key[len(self._prefix):]
                    keys.append(stripped)
            if not response.get("IsTruncated"):
                break
            continuation_token = response.get("NextContinuationToken")
        return sorted(keys)

    def delete(self, key: str) -> None:
        """Delete the object at ``key``.

        Idempotent: S3 ``delete_object`` succeeds even when the key is absent.
        """
        validate_key(key)
        self._client.delete_object(
            Bucket=self._bucket,
            Key=self._object_key(key),
        )

    def exists(self, key: str) -> bool:
        """Return whether ``key`` exists in the bucket."""
        validate_key(key)
        try:
            self._client.head_object(
                Bucket=self._bucket,
                Key=self._object_key(key),
            )
            return True
        except self._ClientError as exc:
            error_code = exc.response.get("Error", {}).get("Code", "")
            if error_code in ("404", "NoSuchKey"):
                return False
            raise