terricain / aioboto3

Wrapper to use boto3 resources with the aiobotocore async backend
Apache License 2.0
719 stars 74 forks source link

Feature to support awaitable Fileobj in download_fileobj just like upload_fileobj #264

Closed shenchucheng closed 3 months ago

shenchucheng commented 2 years ago

Description

I find that it doesn't work as expected to use upload_fileobj method by passing an aiofiles obj in version 8.2.0, and the issue is already mentioned in #222 and fixed in the latest release. Thanks for your work! And the idea come up very naturally to support awaitable Fileobj in download_fileobj method just like upload_fileobj.

What I Did


async def download_fileobj(
    self, Bucket, Key, Fileobj, ExtraArgs=None, Callback=None, Config=None
):

    try:
        if ExtraArgs is None:
            ExtraArgs = {}
        resp = await self.get_object(Bucket=Bucket, Key=Key, **ExtraArgs)
    except ClientError as err:
        if err.response['Error']['Code'] == 'NoSuchKey':
            # Convert to 404 so it looks the same when boto3.download_file fails
            raise ClientError({'Error': {'Code': '404', 'Message': 'Not Found'}}, 'HeadObject')
        raise

    body = resp['Body']

    while True:
        data = await body.read(4096)

        if data == b'':
            break

        if Callback:
            try:
                Callback(len(data))
            except:  # noqa: E722
                pass

        write_result = Fileobj.write(data)
        if inspect.isawaitable(write_result):
            await write_result
        else:
            await asyncio.sleep(0.0)

It's worked in this way

import os
import aiofiles
import asyncio
import aioboto3

async def main():
    async with aiofiles.open('test.txt', 'w') as fd:
        await fd.write('hello world\n')
    async with aioboto3.client(
        service_name='s3',
        aws_access_key_id='...',
        aws_secret_access_key='...',
        endpoint_url='...',
    ) as s3:
        async with aiofiles.open('test.txt', 'rb') as fd:
            await s3.upload_fileobj(fd, 'bucket', 'test.txt')
        async with aiofiles.open('test1.txt', 'wb') as fd:
            await s3.download_fileobj('bucket', 'test.txt', fd)
        async with aiofiles.open('test1.txt', 'r') as fd:
            assert await fd.read() == 'hello world\n'
    os.remove('test.txt')
    os.remove('test1.txt')

asyncio.run(main())
terricain commented 2 years ago

Hey, feel free to raise a PR :)

kyboi commented 3 months ago

Hey, feel free to raise a PR :)

Opened PR for this at #332

terricain commented 3 months ago

This is done 😄