Optimize downloading

This commit is contained in:
2022-02-06 14:47:11 +03:00
parent 9251b47cbf
commit 2702283872
2 changed files with 29 additions and 17 deletions

View File

@@ -32,17 +32,22 @@ async def upload_file(content: bytes, filename: str, caption: str) -> UploadedFi
return UploadedFile.parse_obj(response.json()) return UploadedFile.parse_obj(response.json())
async def download_file(chat_id: int, message_id: int) -> Optional[bytes]: async def download_file(
chat_id: int, message_id: int
) -> Optional[tuple[httpx.Response, httpx.AsyncClient]]:
headers = {"Authorization": env_config.FILES_SERVER_API_KEY} headers = {"Authorization": env_config.FILES_SERVER_API_KEY}
async with httpx.AsyncClient(timeout=60) as client: client = httpx.AsyncClient(timeout=60)
response = await client.get( request = client.build_request(
f"{env_config.FILES_SERVER_URL}" "GET",
f"/api/v1/files/download_by_message/{chat_id}/{message_id}", f"{env_config.FILES_SERVER_URL}"
headers=headers, f"/api/v1/files/download_by_message/{chat_id}/{message_id}",
) headers=headers,
)
if response.status_code != 200: response = await client.send(request, stream=True)
return None
return response.content if response.status_code != 200:
return None
return response, client

View File

@@ -1,8 +1,10 @@
import asyncio
import base64 import base64
from fastapi import APIRouter, Depends, HTTPException, status, Request from fastapi import APIRouter, Depends, HTTPException, status, Request
from fastapi.responses import StreamingResponse
from starlette.responses import Response from starlette.background import BackgroundTask
from arq.connections import ArqRedis from arq.connections import ArqRedis
from asyncpg import exceptions from asyncpg import exceptions
@@ -51,25 +53,30 @@ async def download_cached_file(object_id: int, object_type: str):
cache_data = cached_file.data cache_data = cached_file.data
data = await download_file_from_cache( data, filename, book = await asyncio.gather(
cache_data["chat_id"], cache_data["message_id"] download_file_from_cache(cache_data["chat_id"], cache_data["message_id"]),
get_filename(object_id, object_type),
get_book(object_id),
) )
if data is None: if data is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND) raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
filename = await get_filename(object_id, object_type) response, client = data
book = await get_book(object_id) async def close():
await response.aclose()
await client.aclose()
return Response( return StreamingResponse(
data, response.aiter_bytes(),
headers={ headers={
"Content-Disposition": f"attachment; filename={filename}", "Content-Disposition": f"attachment; filename={filename}",
"X-Caption-B64": base64.b64encode(get_caption(book).encode("utf-8")).decode( "X-Caption-B64": base64.b64encode(get_caption(book).encode("utf-8")).decode(
"latin-1" "latin-1"
), ),
}, },
background=BackgroundTask(close),
) )