fix: Use async_open in a few places (#4929)

Use async_open in a few places
This commit is contained in:
Christophe Bornet 2024-11-28 19:06:01 +01:00 • committed by GitHub
commit c816832b44
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 15 additions and 15 deletions

View file

@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import time import time
from http import HTTPStatus from http import HTTPStatus
from typing import TYPE_CHECKING, Annotated from typing import TYPE_CHECKING, Annotated
@ -556,7 +557,7 @@ async def create_upload_file(
""" """
try: try:
flow_id_str = str(flow_id) flow_id_str = str(flow_id)
file_path = save_uploaded_file(file, folder_name=flow_id_str) file_path = await asyncio.to_thread(save_uploaded_file, file, folder_name=flow_id_str)
return UploadFileResponse( return UploadFileResponse(
flow_id=flow_id_str, flow_id=flow_id_str,

View file

@ -118,7 +118,7 @@ class DirectoryReader:
except UnicodeDecodeError: except UnicodeDecodeError:
# This is happening in Windows, so we need to open the file in binary mode # This is happening in Windows, so we need to open the file in binary mode
# The file is always just a python file, so we can safely read it as utf-8 # The file is always just a python file, so we can safely read it as utf-8
with async_open(_file_path, "rb") as f: async with async_open(_file_path, "rb") as f:
return (await f.read()).decode("utf-8") return (await f.read()).decode("utf-8")
def get_files(self): def get_files(self):

View file

@ -9,6 +9,7 @@ from pathlib import Path
from uuid import UUID from uuid import UUID
import orjson import orjson
from aiofile import async_open
from emoji import demojize, purely_emoji from emoji import demojize, purely_emoji
from loguru import logger from loguru import logger
from sqlalchemy.exc import NoResultFound from sqlalchemy.exc import NoResultFound
@ -545,13 +546,14 @@ async def load_flows_from_directory() -> None:
user_id = user.id user_id = user.id
_flows_path = Path(flows_path) _flows_path = Path(flows_path)
files = [f for f in _flows_path.iterdir() if f.is_file()] files = [f for f in _flows_path.iterdir() if f.is_file()]
for f in files: for file_path in files:
if f.suffix != ".json": if file_path.suffix != ".json":
continue continue
logger.info(f"Loading flow from file: {f.name}") logger.info(f"Loading flow from file: {file_path.name}")
content = f.read_text(encoding="utf-8") async with async_open(file_path, "r", encoding="utf-8") as f:
content = await f.read()
flow = orjson.loads(content) flow = orjson.loads(content)
no_json_name = f.stem no_json_name = file_path.stem
flow_endpoint_name = flow.get("endpoint_name") flow_endpoint_name = flow.get("endpoint_name")
if _is_valid_uuid(no_json_name): if _is_valid_uuid(no_json_name):
flow["id"] = no_json_name flow["id"] = no_json_name

View file

@ -1,6 +1,6 @@
import asyncio
from pathlib import Path from pathlib import Path
from aiofile import async_open
from loguru import logger from loguru import logger
from .service import StorageService from .service import StorageService
@ -33,11 +33,9 @@ class LocalStorageService(StorageService):
folder_path.mkdir(parents=True, exist_ok=True) folder_path.mkdir(parents=True, exist_ok=True)
file_path = folder_path / file_name file_path = folder_path / file_name
def write_file(file_path: Path, data: bytes) -> None:
file_path.write_bytes(data)
try: try:
await asyncio.to_thread(write_file, file_path, data) async with async_open(file_path, "wb") as f:
await f.write(data)
logger.info(f"File {file_name} saved successfully in flow {flow_id}.") logger.info(f"File {file_name} saved successfully in flow {flow_id}.")
except Exception: except Exception:
logger.exception(f"Error saving file {file_name} in flow {flow_id}") logger.exception(f"Error saving file {file_name} in flow {flow_id}")
@ -57,10 +55,9 @@ class LocalStorageService(StorageService):
msg = f"File {file_name} not found in flow {flow_id}" msg = f"File {file_name} not found in flow {flow_id}"
raise FileNotFoundError(msg) raise FileNotFoundError(msg)
def read_file(file_path: Path) -> bytes: async with async_open(file_path, "rb") as f:
return file_path.read_bytes() content = await f.read()
content = await asyncio.to_thread(read_file, file_path)
logger.debug(f"File {file_name} retrieved successfully from flow {flow_id}.") logger.debug(f"File {file_name} retrieved successfully from flow {flow_id}.")
return content return content