LobotomyLabs/classquiz-ai
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

✨ Added streaming up- and downloads, fixed file sizes and added a (not working) dashboard

Browse Source
This commit is contained in:
Mawoka
2023-05-29 16:24:38 +02:00
parent ef73234606
commit a8b59ce6f5
17 changed files with 414 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classquiz/storage/__init__.py
+20 -43
View File
@@ -2,20 +2,17 @@
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from io import BytesIO
from typing import Optional
from typing import Optional, BinaryIO
from .deta_storage import DetaStorage
from .local_storage import LocalStorage
from .s3_storage import S3Storage
from typing import Generator
class Storage:
def __init__(
self,
backend: str,
deta_key: Optional[str],
deta_id: Optional[str],
storage_path: Optional[str],
access_key: str | None = None,
secret_key: str | None = None,
@@ -23,65 +20,45 @@ class Storage:
base_url: str | None = None,
):
self.backend = backend
self.deta_key: str | None = deta_key
self.deta_id: str | None = deta_id
# self.deta_base_url = f"https://drive.deta.sh/v1/{deta_id}/classquiz1"
self.access_key = access_key
self.secret_key = secret_key
self.bucket_name = bucket_name
self.base_url = base_url
self.deta_instance = None
if backend == "deta":
if deta_key is None or deta_id is None:
raise ValueError("deta_key and deta_id must be provided")
if self.base_url is None:
self.base_url = f"https://drive.deta.sh/v1/{deta_id}/classquiz1"
self.deta_instance = DetaStorage(
deta_base_url=self.base_url,
deta_key=self.deta_key,
deta_id=self.deta_id,
)
self.instance: LocalStorage | S3Storage | None = None
elif backend == "local":
if backend == "local":
if storage_path is None:
raise ValueError("storage_path must be provided")
else:
self.local_instance = LocalStorage(base_path=storage_path)
self.instance = LocalStorage(base_path=storage_path)
elif backend == "s3":
if access_key is None or secret_key is None or bucket_name is None or base_url is None:
raise ValueError("Not all parameters given")
self.s3_instance = S3Storage(
self.instance = S3Storage(
base_url=base_url, access_key=access_key, secret_key=secret_key, bucket_name=bucket_name
)
else:
raise NotImplementedError(f"Backend {backend} not implemented")
async def download(self, file_name: str) -> BytesIO | None:
if self.backend == "deta":
return await self.deta_instance.download(file_name)
elif self.backend == "local":
return await self.local_instance.get_file(file_name)
elif self.backend == "s3":
return await self.s3_instance.download(file_name)
async def download(self, file_name: str) -> Generator | None:
"""
No support for s3 since it doesn't make sense relaying the traffic through this backend
:param file_name:
:return:
"""
yield self.instance.download(file_name)
async def upload(self, file_name: str, file_data: bytes) -> None:
if self.backend == "deta":
return await self.deta_instance.upload(file=file_data, file_name=file_name)
elif self.backend == "local":
return await self.local_instance.write_file(file_name=file_name, data=file_data)
elif self.backend == "s3":
return await self.s3_instance.upload(file=file_data, file_name=file_name)
async def upload(self, file_name: str, file_data: BinaryIO) -> None:
return await self.instance.upload(file=file_data, file_name=file_name)
async def delete(self, file_names: [str]) -> None:
if self.backend == "deta":
return await self.deta_instance.delete(file_names=file_names)
elif self.backend == "local":
return await self.local_instance.delete_file(file_names=file_names)
elif self.backend == "s3":
return await self.s3_instance.delete(file_names=file_names)
return await self.instance.delete(file_names=file_names)
async def get_url(self, file_name: str, expiry: int) -> str:
if self.backend == "s3":
return self.s3_instance.get_url(file_name=file_name, expire=expiry)
return self.instance.get_url(file_name=file_name, expire=expiry)
async def get_file_size(self, file_name: str) -> int | None:
return self.instance.size(file_name)
classquiz/storage/deta_storage.py
-58
View File
@@ -1,58 +0,0 @@
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from io import BytesIO
from classquiz.storage.errors import DeletionFailedError, SavingFailedError, DownloadingFailedError
from aiohttp import ClientSession
class DetaStorage:
def __init__(self, deta_base_url: str, deta_id: str, deta_key: str):
self.deta_url = deta_base_url
self.deta_id = deta_id
self.deta_key = deta_key
self.headers = {
"X-Api-Key": self.deta_key,
}
async def download(self, file_name: str) -> BytesIO | None:
"""
:param file_name: The name of the file to be downloaded
:return: Either bytes if successful download or None if failed
"""
async with ClientSession(headers=self.headers) as session, session.get(
f"{self.deta_url}/files/download?name={file_name}"
) as response:
if response.status == 200:
return BytesIO(await response.read())
elif response.status == 404:
return None
else:
raise DownloadingFailedError
async def upload(self, file: bytes, file_name: str) -> None:
"""
:param file: The file in bytes
:param file_name: The name of the file
:return:
"""
async with ClientSession(headers=self.headers) as session, session.post(
f"{self.deta_url}/files?name={file_name}", data=file
) as response:
if response.status == 201:
return None
else:
print(response.status, await response.json())
raise SavingFailedError
async def delete(self, file_names: [str]) -> None:
async with ClientSession(headers=self.headers) as session, session.delete(
f"{self.deta_url}/files", json={"names": file_names}
) as response:
if response.status == 200:
return None
else:
raise DeletionFailedError
classquiz/storage/local_storage.py
+20 -7
View File
@@ -2,32 +2,45 @@
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
import io
import os
from typing import BinaryIO, Generator
import aiofiles
import aiofiles.os
_DEFAULT_CHUNK_SIZE = 32768 # bytes; arbitrary
async def aioshutil_copyfileobj(async_fsrc, async_fdst, *, chunksize: int = _DEFAULT_CHUNK_SIZE) -> None:
while (chunk := await async_fsrc.read(chunksize)) != b"":
await async_fdst.write(chunk)
class LocalStorage:
def __init__(self, base_path: str):
self.base_path = base_path
async def get_file(self, file_name: str) -> io.BytesIO | None:
async def download(self, file_name: str) -> Generator | None:
try:
async with aiofiles.open(file=os.path.join(self.base_path, file_name), mode="rb") as f:
return io.BytesIO(await f.read())
yield f.read()
except FileNotFoundError:
return None
yield None
async def write_file(self, file_name: str, data: bytes) -> None:
async def upload(self, file_name: str, data: BinaryIO) -> None:
async with aiofiles.open(file=os.path.join(self.base_path, file_name), mode="wb") as f:
await f.write(data)
await aioshutil_copyfileobj(data, f)
async def delete_file(self, file_names: [str]) -> None:
async def delete(self, file_names: [str]) -> None:
for i in file_names:
try:
await aiofiles.os.remove(os.path.join(self.base_path, i))
except FileNotFoundError:
pass
return None
def size(self, file_name: str) -> int | None:
try:
return os.stat(os.path.join(self.base_path, file_name))
except FileNotFoundError:
return None
classquiz/storage/s3_storage.py
+26 -13
View File
@@ -8,8 +8,7 @@
import hashlib
import hmac
from datetime import datetime, timedelta
from io import BytesIO
from typing import Tuple
from typing import Tuple, BinaryIO, Generator
from aiohttp import ClientSession
import minio
@@ -113,17 +112,9 @@ class S3Storage:
return headers, request_url
async def download(self, file_name: str):
headers, url = self._generate_aws_signature_v4(method="GET", path=f"/{file_name}")
async with ClientSession() as session, session.get(url, headers=headers) as resp:
if resp.status == 200:
return BytesIO(await resp.read())
elif resp.status == 404:
return None
else:
raise DownloadingFailedError
async def upload(self, file: bytes, file_name: str, content_type: str | None = "application/octet-stream") -> None:
async def upload(
self, file: BinaryIO, file_name: str, content_type: str | None = "application/octet-stream"
) -> None:
headers, url = self._generate_aws_signature_v4(method="PUT", path=f"/{file_name}")
async with ClientSession() as session, session.put(url, headers=headers, data=file) as resp:
if resp.status == 200:
@@ -145,3 +136,25 @@ class S3Storage:
return self.client.presigned_get_object(
object_name=file_name, bucket_name=self.bucket_name, expires=timedelta(seconds=expire)
)
def size(self, file_name: str) -> int | None:
res = self.client.stat_object(bucket_name=self.bucket_name, object_name=file_name)
if res is None:
return None
return res.size
async def download(self, file_name: str) -> Generator:
headers, url = self._generate_aws_signature_v4(method="GET", path=f"/{file_name}")
async with ClientSession() as session, session.get(url, headers=headers) as resp:
if resp.status == 200:
async for i in resp.content.iter_chunked(1024):
yield i
elif resp.status == 404:
yield None
else:
raise DownloadingFailedError
# client = httpx.AsyncClient()
# async with client.stream("GET", url, headers=headers) as resp:
# if resp.status == 200:
# yield resp.aiter_bytes()