Source code for shapelib.core.cloud.providers.s3_provider

"""AWS S3 cloud storage provider."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any

from shapelib.factories.cloud_registry import register_cloud_provider

if TYPE_CHECKING:
    from shapelib.core.cloud.cloud_path import CloudPath




@register_cloud_provider
class S3Provider:
    """
    Cloud provider for AWS S3 storage.

    Handles file operations using the boto3 library for s3:// protocol URLs.

    Examples
    --------
    >>> provider = S3Provider()
    >>> provider.supports(CloudPath.from_string("s3://bucket/file.fgb"))
    True

    Big-O
    -----
    All operations are O(n) where n is the file size.
    """



    def supports(self, cloud_path: CloudPath) -> bool:
        """
        Check if this provider handles S3 protocol paths.

        Parameters
        ----------
        cloud_path : CloudPath
            The parsed cloud path to check.

        Returns
        -------
        bool
            True if the path uses s3:// protocol.
        """
        return cloud_path.is_s3_protocol




    def upload(self, source: Path, cloud_path: CloudPath) -> None:
        """
        Upload a local file to S3.

        Parameters
        ----------
        source : Path
            The local file to upload.
        cloud_path : CloudPath
            The S3 destination path (s3://bucket/key).

        Raises
        ------
        ImportError
            If boto3 is not installed.
        FileNotFoundError
            If the source file does not exist.
        ConnectionError
            If the upload fails.
        PermissionError
            If access is denied.
        """
        client = self._get_client()

        try:
            client.upload_file(str(source), cloud_path.s3_bucket, cloud_path.s3_key)
        except Exception as error:
            if "AccessDenied" in str(error) or "403" in str(error):
                raise PermissionError(f"Access denied: {cloud_path.path}") from error
            raise ConnectionError(f"S3 upload failed: {error}") from error




    def download(self, cloud_path: CloudPath, destination: Path) -> Path:
        """
        Download a file from S3.

        Parameters
        ----------
        cloud_path : CloudPath
            The S3 source path (s3://bucket/key).
        destination : Path
            The local destination path.

        Returns
        -------
        Path
            The path to the downloaded file.

        Raises
        ------
        ImportError
            If boto3 is not installed.
        FileNotFoundError
            If the S3 object does not exist.
        ConnectionError
            If the download fails.
        PermissionError
            If access is denied.
        """
        client = self._get_client()

        try:
            destination.parent.mkdir(parents=True, exist_ok=True)
            client.download_file(cloud_path.s3_bucket, cloud_path.s3_key, str(destination))
            return destination
        except Exception as error:
            error_str = str(error)
            if "NoSuchKey" in error_str or "404" in error_str:
                raise FileNotFoundError(f"Not found: {cloud_path.path}") from error
            if "AccessDenied" in error_str or "403" in error_str:
                raise PermissionError(f"Access denied: {cloud_path.path}") from error
            raise ConnectionError(f"S3 download failed: {error}") from error




    def file_exists(self, cloud_path: CloudPath) -> bool:
        """
        Check if an S3 object exists using head_object.

        Parameters
        ----------
        cloud_path : CloudPath
            The S3 path to check.

        Returns
        -------
        bool
            True if the object exists, False otherwise.
        """
        try:
            from botocore.exceptions import ClientError
        except ImportError:
            return False

        try:
            client = self._get_client()
            client.head_object(Bucket=cloud_path.s3_bucket, Key=cloud_path.s3_key)
            return True
        except ClientError:
            return False
        except Exception:
            return False


    def _get_client(self) -> Any:
        """
        Get or create the S3 client.

        Returns
        -------
        Any
            The S3 client instance (boto3.client).

        Raises
        ------
        ImportError
            If boto3 is not installed.
        """
        try:
            import boto3
        except ImportError as error:
            raise ImportError("boto3 is required for s3:// protocol. Install: pip install boto3") from error

        return boto3.client("s3", region_name="sa-east-1")