hyperboria/nexus/pylon/drivers/direct.py

from typing import Dict

import aiohttp.client_exceptions
from aiohttp_socks import ProxyError
from nexus.pylon.drivers.base import BaseDriver
from nexus.pylon.exceptions import (
    BadResponseError,
    NotFoundError,
)
from nexus.pylon.prepared_request import PreparedRequest
from nexus.pylon.proto.file_pb2 import Chunk as ChunkPb
from nexus.pylon.proto.file_pb2 import FileResponse as FileResponsePb
from python_socks import ProxyTimeoutError
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_random,
)


class DirectDriver(BaseDriver):
    allowed_content_type = None

    @retry(
        reraise=True,
        wait=wait_random(min=1, max=2),
        stop=stop_after_attempt(7),
        retry=retry_if_exception_type((ProxyError, aiohttp.client_exceptions.ClientPayloadError, ProxyTimeoutError)),
    )
    async def execute_prepared_file_request(self, prepared_file_request: PreparedRequest, params: Dict):
        async with self.get_session() as session:
            async with prepared_file_request.execute_with(session=session) as resp:
                if resp.status == 404:
                    raise NotFoundError(url=prepared_file_request.url)
                elif (
                    resp.status != 200
                    or (
                        self.allowed_content_type
                        and resp.headers.get('Content-Type', '').lower() not in self.allowed_content_type
                    )
                ):
                    raise BadResponseError(
                        request_headers=prepared_file_request.headers,
                        url=prepared_file_request.url,
                        status=resp.status,
                        headers=str(resp.headers),
                    )
                file_validator = self.validator(params)
                yield FileResponsePb(status=FileResponsePb.Status.BEGIN_TRANSMISSION, source=prepared_file_request.url)
                async for content, _ in resp.content.iter_chunks():
                    file_validator.update(content)
                    yield FileResponsePb(chunk=ChunkPb(content=content), source=prepared_file_request.url)
                file_validator.validate()