elaunira-r2index 0.1.0__py3-none-any.whl

elaunira/r2index/async_uploader.py

@@ -0,0 +1,149 @@
+"""Asynchronous R2 uploader using aioboto3."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+import aioboto3
+from aiobotocore.config import AioConfig
+
+from .exceptions import UploadError
+from .uploader import R2Config
+
+# 100MB threshold and part size for multipart uploads
+MULTIPART_THRESHOLD = 100 * 1024 * 1024
+MULTIPART_PART_SIZE = 100 * 1024 * 1024
+
+
+class AsyncR2Uploader:
+    """Asynchronous R2 uploader using aioboto3."""
+
+    def __init__(self, config: R2Config) -> None:
+        """
+        Initialize the async R2 uploader.
+
+        Args:
+            config: R2 configuration with credentials and endpoint.
+        """
+        self.config = config
+        self._session = aioboto3.Session()
+
+    async def upload_file(
+        self,
+        file_path: str | Path,
+        object_key: str,
+        content_type: str | None = None,
+        progress_callback: Callable[[int], None] | None = None,
+    ) -> str:
+        """
+        Upload a file to R2 asynchronously.
+
+        Uses multipart upload for files larger than 100MB.
+
+        Args:
+            file_path: Path to the file to upload.
+            object_key: The key (path) to store the object under in R2.
+            content_type: Optional content type for the object.
+            progress_callback: Optional callback called with bytes uploaded so far.
+
+        Returns:
+            The object key of the uploaded file.
+
+        Raises:
+            UploadError: If the upload fails.
+        """
+        file_path = Path(file_path)
+
+        if not file_path.exists():
+            raise UploadError(f"File not found: {file_path}")
+
+        transfer_config = AioConfig(
+            max_pool_connections=10,
+        )
+
+        extra_args = {}
+        if content_type:
+            extra_args["ContentType"] = content_type
+
+        try:
+            async with self._session.client(
+                "s3",
+                aws_access_key_id=self.config.access_key_id,
+                aws_secret_access_key=self.config.secret_access_key,
+                endpoint_url=self.config.endpoint_url,
+                region_name=self.config.region,
+                config=transfer_config,
+            ) as client:
+                callback = None
+                if progress_callback:
+                    callback = _AsyncProgressCallback(progress_callback)
+
+                await client.upload_file(
+                    str(file_path),
+                    self.config.bucket,
+                    object_key,
+                    ExtraArgs=extra_args if extra_args else None,
+                    Callback=callback,
+                )
+        except Exception as e:
+            raise UploadError(f"Failed to upload file to R2: {e}") from e
+
+        return object_key
+
+    async def delete_object(self, object_key: str) -> None:
+        """
+        Delete an object from R2 asynchronously.
+
+        Args:
+            object_key: The key of the object to delete.
+
+        Raises:
+            UploadError: If the deletion fails.
+        """
+        try:
+            async with self._session.client(
+                "s3",
+                aws_access_key_id=self.config.access_key_id,
+                aws_secret_access_key=self.config.secret_access_key,
+                endpoint_url=self.config.endpoint_url,
+                region_name=self.config.region,
+            ) as client:
+                await client.delete_object(Bucket=self.config.bucket, Key=object_key)
+        except Exception as e:
+            raise UploadError(f"Failed to delete object from R2: {e}") from e
+
+    async def object_exists(self, object_key: str) -> bool:
+        """
+        Check if an object exists in R2 asynchronously.
+
+        Args:
+            object_key: The key of the object to check.
+
+        Returns:
+            True if the object exists, False otherwise.
+        """
+        try:
+            async with self._session.client(
+                "s3",
+                aws_access_key_id=self.config.access_key_id,
+                aws_secret_access_key=self.config.secret_access_key,
+                endpoint_url=self.config.endpoint_url,
+                region_name=self.config.region,
+            ) as client:
+                await client.head_object(Bucket=self.config.bucket, Key=object_key)
+                return True
+        except client.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                return False
+            raise UploadError(f"Failed to check object existence: {e}") from e
+
+
+class _AsyncProgressCallback:
+    """Wrapper to track cumulative progress for aioboto3 callback."""
+
+    def __init__(self, callback: Callable[[int], None]) -> None:
+        self._callback = callback
+        self._bytes_transferred = 0
+
+    def __call__(self, bytes_amount: int) -> None:
+        self._bytes_transferred += bytes_amount
+        self._callback(self._bytes_transferred)

elaunira/r2index/checksums.py

@@ -0,0 +1,127 @@
+"""Streaming checksum computation for large files."""
+
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO
+
+# 8MB chunk size for memory-efficient processing of large files
+CHUNK_SIZE = 8 * 1024 * 1024
+
+
+@dataclass
+class ChecksumResult:
+    """Result of checksum computation."""
+
+    md5: str
+    sha1: str
+    sha256: str
+    sha512: str
+    size: int
+
+
+def compute_checksums(file_path: str | Path) -> ChecksumResult:
+    """
+    Compute MD5, SHA1, SHA256, and SHA512 checksums for a file.
+
+    Reads the file in chunks for memory-efficient processing of large files.
+    All checksums are computed in a single pass through the file.
+
+    Args:
+        file_path: Path to the file to compute checksums for.
+
+    Returns:
+        ChecksumResult containing all checksums and file size.
+    """
+    file_path = Path(file_path)
+
+    md5_hash = hashlib.md5()
+    sha1_hash = hashlib.sha1()
+    sha256_hash = hashlib.sha256()
+    sha512_hash = hashlib.sha512()
+
+    size = 0
+
+    with open(file_path, "rb") as f:
+        size = _compute_from_file_object(f, md5_hash, sha1_hash, sha256_hash, sha512_hash)
+
+    return ChecksumResult(
+        md5=md5_hash.hexdigest(),
+        sha1=sha1_hash.hexdigest(),
+        sha256=sha256_hash.hexdigest(),
+        sha512=sha512_hash.hexdigest(),
+        size=size,
+    )
+
+
+def compute_checksums_from_file_object(file_obj: BinaryIO) -> ChecksumResult:
+    """
+    Compute checksums from a file-like object.
+
+    Args:
+        file_obj: Binary file-like object to read from.
+
+    Returns:
+        ChecksumResult containing all checksums and total bytes read.
+    """
+    md5_hash = hashlib.md5()
+    sha1_hash = hashlib.sha1()
+    sha256_hash = hashlib.sha256()
+    sha512_hash = hashlib.sha512()
+
+    size = _compute_from_file_object(file_obj, md5_hash, sha1_hash, sha256_hash, sha512_hash)
+
+    return ChecksumResult(
+        md5=md5_hash.hexdigest(),
+        sha1=sha1_hash.hexdigest(),
+        sha256=sha256_hash.hexdigest(),
+        sha512=sha512_hash.hexdigest(),
+        size=size,
+    )
+
+
+def _compute_from_file_object(
+    file_obj: BinaryIO,
+    md5_hash: "hashlib._Hash",
+    sha1_hash: "hashlib._Hash",
+    sha256_hash: "hashlib._Hash",
+    sha512_hash: "hashlib._Hash",
+) -> int:
+    """
+    Internal helper to compute checksums from a file object.
+
+    Returns the total number of bytes read.
+    """
+    size = 0
+
+    while True:
+        chunk = file_obj.read(CHUNK_SIZE)
+        if not chunk:
+            break
+
+        size += len(chunk)
+        md5_hash.update(chunk)
+        sha1_hash.update(chunk)
+        sha256_hash.update(chunk)
+        sha512_hash.update(chunk)
+
+    return size
+
+
+async def compute_checksums_async(file_path: str | Path) -> ChecksumResult:
+    """
+    Compute checksums asynchronously.
+
+    Note: This uses synchronous file I/O in a way that doesn't block the event loop
+    for too long by processing in chunks. For truly async file I/O, consider using
+    aiofiles, but for CPU-bound hashing, the benefit is minimal.
+
+    Args:
+        file_path: Path to the file to compute checksums for.
+
+    Returns:
+        ChecksumResult containing all checksums and file size.
+    """
+    import asyncio
+
+    return await asyncio.to_thread(compute_checksums, file_path)