feature: fetcher and uploader

Author: FelixAbrahamsson

.gitignore

@@ -0,0 +1,19 @@
+.venv
+.cache
+__pycache__
+.pytest_cache
+*.egg-info
+.vscode
+dist
+.eggs/
+build/
+_build/
+*.pyc
+rclone.conf
+*.json
+scratch.py
+
+AUTHORS
+ChangeLog
+
+*.env*

.python-version

@@ -0,0 +1 @@
+3.8.8

README.md

@@ -0,0 +1,57 @@
+# Download/upload images from/to s3 very quickly
+
+## Setup
+
+```
+poetry install
+```
+
+## Usage
+
+Download image files
+
+```python
+from PIL import Image
+from fast_s3 import Fetcher
+
+
+large_list_of_image_paths = [...]
+
+fetcher = Fetcher(
+    paths=large_list_of_image_paths,
+    endpoint_url="https://s3.my-path-to-s3",
+    aws_access_key_id="my-key-id",
+    aws_secret_access_key="my-secret-key",
+    region_name="my-region",
+    bucket_name="my-bucket",
+    ordered=True,  # returns files in the same order as paths
+    buffer_size=1024,
+    n_workers=32,
+)
+
+for file in fetcher:
+    Image.open(file.buffer).save(file.path)
+```
+
+Upload files
+
+```python
+from fast_s3 import Uploader
+
+
+large_list_of_files = [...]
+large_list_of_paths = [...]
+
+uploader = Uploader(
+    endpoint_url="https://s3.my-path-to-s3",
+    aws_access_key_id="my-key-id",
+    aws_secret_access_key="my-secret-key",
+    region_name="my-region",
+    bucket_name="my-bucket",
+)
+
+uploader.upload_files(
+    large_list_of_files,
+    large_list_of_paths,
+)
+```

fast_s3/__init__.py

@@ -0,0 +1,2 @@
+from .fetcher import Fetcher
+from .uploader import Uploader

fast_s3/fetcher.py

@@ -0,0 +1,79 @@
+import io
+from pathlib import Path
+from typing import List, Union
+
+from .file import File
+from .transfer_manager import transfer_manager
+
+
+class Fetcher:
+    def __init__(
+        self,
+        paths: List[Union[str, Path]],
+        endpoint_url: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        region_name: str,
+        bucket_name: str,
+        ordered=True,
+        buffer_size=1024,
+        n_workers=32,
+    ):
+        self.paths = paths
+        self.ordered = ordered
+        self.buffer_size = buffer_size
+        self.transfer_manager = transfer_manager(
+            endpoint_url=endpoint_url,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            region_name=region_name,
+            n_workers=n_workers,
+        )
+        self.bucket_name = bucket_name
+        self.files: List[File] = []
+        self.current_path_index = 0
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __iter__(self):
+        for _ in range(self.buffer_size):
+            self.queue_download_()
+
+        if self.ordered:
+            for _ in range(len(self)):
+                file = self.files.pop(0)
+                file.future.result()
+                yield file
+                self.queue_download_()
+        else:
+            for _ in range(len(self)):
+                for index, file in enumerate(self.files):
+                    if file.future.done():
+                        break
+                else:
+                    index = 0
+                file = self.files.pop(index)
+                file.future.result()
+                yield file
+                self.queue_download_()
+
+    def queue_download_(self):
+        if self.current_path_index < len(self):
+            buffer = io.BytesIO()
+            path = self.paths[self.current_path_index]
+            self.files.append(
+                File(
+                    buffer=buffer,
+                    future=self.transfer_manager.download(
+                        fileobj=buffer,
+                        bucket=self.bucket_name,
+                        key=str(path),
+                    ),
+                    path=path,
+                )
+            )
+            self.current_path_index += 1
+
+    def close(self):
+        self.transfer_manager.shutdown()

fast_s3/file.py

@@ -0,0 +1,15 @@
+import io
+from pathlib import Path
+from typing import Union
+
+from pydantic import BaseModel
+from s3transfer.futures import TransferFuture
+
+
+class File(BaseModel):
+    buffer: io.BytesIO
+    future: TransferFuture
+    path: Union[str, Path]
+
+    class Config:
+        arbitrary_types_allowed = True

fast_s3/transfer_manager.py

@@ -0,0 +1,27 @@
+import boto3
+import boto3.s3.transfer as s3transfer
+import botocore
+from s3transfer.manager import TransferManager
+
+
+def transfer_manager(
+    endpoint_url: str,
+    aws_access_key_id: str,
+    aws_secret_access_key: str,
+    region_name: str,
+    n_workers=32,
+) -> TransferManager:
+    session = boto3.Session(
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
+        region_name=region_name,
+    )
+    client = session.client(
+        "s3",
+        endpoint_url=endpoint_url,
+        config=botocore.config.Config(max_pool_connections=n_workers),
+    )
+    transfer_config = s3transfer.TransferConfig(
+        use_threads=True, max_concurrency=n_workers
+    )
+    return s3transfer.create_transfer_manager(client, transfer_config)

fast_s3/uploader.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+from typing import List, Tuple, Union
+
+from s3transfer.futures import TransferFuture
+
+from .transfer_manager import transfer_manager
+
+
+class Uploader:
+    def __init__(
+        self,
+        endpoint_url: str,
+        aws_access_key_id: str,
+        aws_secret_access_key: str,
+        region_name: str,
+        bucket_name: str,
+        n_workers=32,
+    ):
+        self.transfer_manager = transfer_manager(
+            endpoint_url=endpoint_url,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            region_name=region_name,
+            n_workers=n_workers,
+        )
+        self.bucket_name = bucket_name
+        self.futures: List[TransferFuture] = []
+
+    def upload_files(
+        self,
+        source: List[Union[str, bytes]],
+        destination: List[Union[str, Path]],
+    ):
+        if len(source) != len(destination):
+            raise ValueError(
+                "The number of source files and destination paths must be equal."
+            )
+        for file, path in zip(source, destination):
+            self.futures.append(
+                self.transfer_manager.upload(
+                    fileobj=file,
+                    bucket=self.bucket_name,
+                    key=str(path),
+                )
+            )
+
+    def await_futures(self):
+        for future in self.futures:
+            future.result()
+        self.futures = []
+
+    def close(self):
+        self.transfer_manager.shutdown()

notebooks/.gitignore

@@ -0,0 +1,4 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore