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())
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}
async with httpx.AsyncClient(timeout=60) as client:
response = await client.get(
client = httpx.AsyncClient(timeout=60)
request = client.build_request(
"GET",
f"{env_config.FILES_SERVER_URL}"
f"/api/v1/files/download_by_message/{chat_id}/{message_id}",
headers=headers,
)
response = await client.send(request, stream=True)
if response.status_code != 200:
return None
return response.content
return response, client

View File

@@ -1,8 +1,10 @@
import asyncio
import base64
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 asyncpg import exceptions
@@ -51,25 +53,30 @@ async def download_cached_file(object_id: int, object_type: str):
cache_data = cached_file.data
data = await download_file_from_cache(
cache_data["chat_id"], cache_data["message_id"]
data, filename, book = await asyncio.gather(
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:
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(
data,
return StreamingResponse(
response.aiter_bytes(),
headers={
"Content-Disposition": f"attachment; filename={filename}",
"X-Caption-B64": base64.b64encode(get_caption(book).encode("utf-8")).decode(
"latin-1"
),
},
background=BackgroundTask(close),
)