|
import os
|
|
import aiofiles
|
|
from pathlib import Path
|
|
from typing import AsyncGenerator
|
|
|
|
from .settings import settings
|
|
|
|
class StorageManager:
|
|
def __init__(self, base_path: str = settings.STORAGE_PATH):
|
|
self.base_path = Path(base_path)
|
|
self.base_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
async def _get_full_path(self, user_id: int, file_path: str) -> Path:
|
|
# Ensure file_path is relative and safe
|
|
relative_path = Path(file_path).relative_to('/') if str(file_path).startswith('/') else Path(file_path)
|
|
full_path = self.base_path / str(user_id) / relative_path
|
|
full_path.parent.mkdir(parents=True, exist_ok=True)
|
|
return full_path
|
|
|
|
async def save_file(self, user_id: int, file_path: str, file_content: bytes):
|
|
full_path = await self._get_full_path(user_id, file_path)
|
|
async with aiofiles.open(full_path, "wb") as f:
|
|
await f.write(file_content)
|
|
return str(full_path)
|
|
|
|
async def get_file(self, user_id: int, file_path: str) -> AsyncGenerator:
|
|
full_path = await self._get_full_path(user_id, file_path)
|
|
if not full_path.exists():
|
|
raise FileNotFoundError(f"File not found: {file_path}")
|
|
|
|
async with aiofiles.open(full_path, "rb") as f:
|
|
while chunk := await f.read(8192):
|
|
yield chunk
|
|
|
|
async def delete_file(self, user_id: int, file_path: str):
|
|
full_path = await self._get_full_path(user_id, file_path)
|
|
if full_path.exists():
|
|
os.remove(full_path)
|
|
|
|
parent_dir = full_path.parent
|
|
while parent_dir != self.base_path and parent_dir.exists():
|
|
try:
|
|
if not any(parent_dir.iterdir()):
|
|
parent_dir.rmdir()
|
|
parent_dir = parent_dir.parent
|
|
else:
|
|
break
|
|
except OSError:
|
|
break
|
|
|
|
async def file_exists(self, user_id: int, file_path: str) -> bool:
|
|
full_path = await self._get_full_path(user_id, file_path)
|
|
return full_path.exists()
|
|
|
|
storage_manager = StorageManager()
|