141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
import gzip
|
|
import json
|
|
from functools import singledispatch
|
|
|
|
from kn_utils.logging import logger
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from pyinfra.storage.proto_data_loader import ProtoDataLoader
|
|
from pyinfra.storage.storages.storage import Storage
|
|
|
|
|
|
class DossierIdFileIdDownloadPayload(BaseModel):
|
|
dossierId: str
|
|
fileId: str
|
|
targetFileExtension: str
|
|
|
|
@property
|
|
def targetFilePath(self):
|
|
return f"{self.dossierId}/{self.fileId}.{self.targetFileExtension}"
|
|
|
|
|
|
class TenantIdDossierIdFileIdDownloadPayload(BaseModel):
|
|
tenantId: str
|
|
dossierId: str
|
|
fileId: str
|
|
targetFileExtension: str
|
|
|
|
@property
|
|
def targetFilePath(self):
|
|
return f"{self.tenantId}/{self.dossierId}/{self.fileId}.{self.targetFileExtension}"
|
|
|
|
|
|
class DossierIdFileIdUploadPayload(BaseModel):
|
|
dossierId: str
|
|
fileId: str
|
|
responseFileExtension: str
|
|
|
|
@property
|
|
def responseFilePath(self):
|
|
return f"{self.dossierId}/{self.fileId}.{self.responseFileExtension}"
|
|
|
|
|
|
class TenantIdDossierIdFileIdUploadPayload(BaseModel):
|
|
tenantId: str
|
|
dossierId: str
|
|
fileId: str
|
|
responseFileExtension: str
|
|
|
|
@property
|
|
def responseFilePath(self):
|
|
return f"{self.tenantId}/{self.dossierId}/{self.fileId}.{self.responseFileExtension}"
|
|
|
|
|
|
class TargetResponseFilePathDownloadPayload(BaseModel):
|
|
targetFilePath: str | dict[str, str]
|
|
|
|
|
|
class TargetResponseFilePathUploadPayload(BaseModel):
|
|
responseFilePath: str
|
|
|
|
|
|
def download_data_bytes_as_specified_in_message(storage: Storage, raw_payload: dict) -> dict[str, bytes] | bytes:
|
|
"""Convenience function to download a file specified in a message payload.
|
|
Supports both legacy and new payload formats. Also supports downloading multiple files at once, which should
|
|
be specified in a dictionary under the 'targetFilePath' key with the file path as value.
|
|
|
|
In all cases, the content will be returned as is (-> bytes).
|
|
"""
|
|
|
|
try:
|
|
if "tenantId" in raw_payload and "dossierId" in raw_payload:
|
|
payload = TenantIdDossierIdFileIdDownloadPayload(**raw_payload)
|
|
elif "tenantId" not in raw_payload and "dossierId" in raw_payload:
|
|
payload = DossierIdFileIdDownloadPayload(**raw_payload)
|
|
else:
|
|
payload = TargetResponseFilePathDownloadPayload(**raw_payload)
|
|
except ValidationError:
|
|
raise ValueError("No download file path found in payload, nothing to download.")
|
|
|
|
data = _download(payload.targetFilePath, storage)
|
|
|
|
return data
|
|
|
|
|
|
@singledispatch
|
|
def _download(file_path_or_file_path_dict: str | dict[str, str], storage: Storage) -> dict[str, bytes] | bytes:
|
|
pass
|
|
|
|
|
|
@_download.register(str)
|
|
def _download_single_file(file_path: str, storage: Storage) -> bytes:
|
|
if not storage.exists(file_path):
|
|
raise FileNotFoundError(f"File '{file_path}' does not exist in storage.")
|
|
|
|
data = storage.get_object(file_path)
|
|
logger.info(f"Downloaded {file_path} from storage.")
|
|
|
|
return data
|
|
|
|
|
|
@_download.register(dict)
|
|
def _download_multiple_files(file_path_dict: dict, storage: Storage) -> dict[str, bytes]:
|
|
return {key: _download(value, storage) for key, value in file_path_dict.items()}
|
|
|
|
|
|
def upload_data_as_specified_in_message(storage: Storage, raw_payload: dict, data):
|
|
"""Convenience function to upload a file specified in a message payload. For now, only json serializable data is
|
|
supported. The storage json consists of the raw_payload, which is extended with a 'data' key, containing the
|
|
data to be uploaded.
|
|
|
|
If the content is not a json serializable object, an exception will be raised.
|
|
If the result file identifier specifies compression with gzip (.gz), it will be compressed before upload.
|
|
|
|
This function can be extended in the future as needed (e.g. if we need to upload images), but since further
|
|
requirements are not specified at this point in time, and it is unclear what these would entail, the code is kept
|
|
simple for now to improve readability, maintainability and avoid refactoring efforts of generic solutions that
|
|
weren't as generic as they seemed.
|
|
"""
|
|
|
|
try:
|
|
if "tenantId" in raw_payload and "dossierId" in raw_payload:
|
|
payload = TenantIdDossierIdFileIdUploadPayload(**raw_payload)
|
|
elif "tenantId" not in raw_payload and "dossierId" in raw_payload:
|
|
payload = DossierIdFileIdUploadPayload(**raw_payload)
|
|
else:
|
|
payload = TargetResponseFilePathUploadPayload(**raw_payload)
|
|
except ValidationError:
|
|
raise ValueError("No upload file path found in payload, nothing to upload.")
|
|
|
|
if ".json" not in payload.responseFilePath:
|
|
raise ValueError("Only json serializable data can be uploaded.")
|
|
|
|
data = {**raw_payload, "data": data}
|
|
|
|
data = json.dumps(data).encode("utf-8")
|
|
data = gzip.compress(data) if ".gz" in payload.responseFilePath else data
|
|
|
|
storage.put_object(payload.responseFilePath, data)
|
|
|
|
logger.info(f"Uploaded {payload.responseFilePath} to storage.")
|