The pinnacle in listing technology. Goes fast

fast-s3-lister.py

import asyncio
import datetime
import queue
import threading

from typing import Iterator

import aioboto3
from aiobotocore.config import AioConfig

class S3Lister:
    """
    Asynchronous divide-and-conquer S3 object key lister with optional bucket URI formatting.
    """

    def __init__(
        self,
        bucket: str,
        prefix: str = "",
        prune=None,
        as_bucket_uri: bool = True,
        concurrency: int = 128,
        buffer: int = 50_000,
    ):
        """
        :param bucket: S3 Bucket Name
        :param prefix: S3 Prefix to list
        :param prune: Function that takes a dict of partition keys and returns True if the partition should be skipped
        :param as_bucket_uri: If True, the returned keys will be formatted as s3://bucket/key
        :param concurrency: Number of concurrent requests to S3
        :param buffer: Maximum number of keys to buffer in the lister before waiting on consumption to continue
        """

        self._bucket = bucket
        self._prefix = prefix
        self._prune = prune
        self.as_bucket_uri = as_bucket_uri
        self._buffer = buffer
        self._sem = asyncio.Semaphore(concurrency)
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._thread.start()
        self._aio_config = AioConfig(
            max_pool_connections=concurrency,
            retries={"mode": "adaptive", "max_attempts": 10, "total_max_attempts": 10},
            connect_timeout=5,
            read_timeout=30,
        )
        self.requests = 0

    def close(self):
        self._loop.call_soon_threadsafe(self._loop.stop)
        self._thread.join()

    async def _list_prefix(self, prefix: str):
        try:
            paginator = self._s3.get_paginator("list_objects_v2")
            await self._sem.acquire()
            async for page in paginator.paginate(
                Bucket=self._bucket,
                Prefix=prefix,
                Delimiter="/",
            ):
                self.requests += 1
                for cp in page.get("CommonPrefixes", []):
                    if self._prune and self._prune(
                        {
                            part.split("=", 1)[0]: part.split("=", 1)[1]
                            for part in cp["Prefix"].strip("/").split("/")
                            if "=" in part
                        }
                    ):
                        logger.info("Skipping %s", cp["Prefix"])
                        continue
                    self._tg.create_task(self._list_prefix(cp["Prefix"]))

                for obj in page.get("Contents", []):
                    k = obj["Key"]
                    if self.as_bucket_uri:
                        k = f"s3://{self._bucket}/{k}"
                    self._q.put(k)
        finally:
            self._sem.release()

    async def _crawl(self):
        try:
            async with aioboto3.Session().client("s3", config=self._aio_config) as s3:
                self._s3 = s3
                async with asyncio.TaskGroup() as tg:
                    self._tg = tg
                    await self._sem.acquire()
                    tg.create_task(self._list_prefix(self._prefix))
            self._q.shutdown()
        except:
            self._q.put(None)
            raise

    def __iter__(self) -> Iterator[str]:
        self._q = queue.Queue(maxsize=self._buffer)
        future = asyncio.run_coroutine_threadsafe(self._crawl(), self._loop)

        while True:
            try:
                item = self._q.get()
                if item is None:
                    break
                yield item
            except queue.ShutDown:
                break

        future.result()
        self.close()

    @staticmethod
    def newer_than_pruner(dt: datetime.date):
        """Partition pruner that only includes partitions equal or newer than the given date."""
        return lambda partitions: not (
            "ingest_dt" not in partitions
            or datetime.date.fromisoformat(partitions["ingest_dt"]) >= dt
        )