71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import logging
|
|
from typing import Dict
|
|
|
|
from pyinfra.config import CONFIG
|
|
from pyinfra.exceptions import InvalidStorageItemFormat
|
|
from pyinfra.server.packing import string_to_bytes
|
|
from pyinfra.visitor.strategies.download.multi import MultiDownloadStrategy
|
|
from pyinfra.visitor.strategies.download.single import SingleDownloadStrategy
|
|
|
|
logger = logging.getLogger()
|
|
|
|
|
|
def build_storage_upload_info(analysis_payload, request_metadata):
|
|
storage_upload_info = {**request_metadata, "id": analysis_payload["metadata"].get("id", 0)}
|
|
storage_upload_info["fileId"] = build_file_path(
|
|
storage_upload_info, storage_upload_info.get("operation", CONFIG.service.response_folder)
|
|
)
|
|
return storage_upload_info
|
|
|
|
|
|
def build_file_path(storage_upload_info, folder):
|
|
return f"{storage_upload_info['fileId']}" + (f"/{folder}" if folder else "")
|
|
|
|
|
|
def validate(data):
|
|
if not ("data" in data and "metadata" in data):
|
|
raise InvalidStorageItemFormat(f"Expected a mapping with keys 'data' and 'metadata', got {data}.")
|
|
|
|
|
|
def wrap(data):
|
|
return {"data": data, "metadata": {}}
|
|
|
|
|
|
def standardize(data) -> Dict:
|
|
"""Storage items can be a blob or a blob with metadata. Standardizes to the latter.
|
|
|
|
Cases:
|
|
1) backend upload: data as bytes
|
|
2) Some Python service's upload: data as bytes of a json string "{'data': <str>, 'metadata': <dict>}",
|
|
where value of key 'data' was encoded with bytes_to_string(...)
|
|
|
|
Returns:
|
|
{"data": bytes, "metadata": dict}
|
|
"""
|
|
|
|
def is_blob_without_metadata(data):
|
|
return isinstance(data, bytes)
|
|
|
|
def is_blob_with_metadata(data: Dict):
|
|
return isinstance(data, dict)
|
|
|
|
if is_blob_without_metadata(data):
|
|
return wrap(data)
|
|
|
|
elif is_blob_with_metadata(data):
|
|
validate(data)
|
|
return data
|
|
|
|
else: # Fallback / used for testing with simple data
|
|
logger.warning("Encountered storage data in unexpected format.")
|
|
assert isinstance(data, str)
|
|
return wrap(string_to_bytes(data))
|
|
|
|
|
|
def get_download_strategy(download_strategy_type=None):
|
|
download_strategies = {
|
|
"single": SingleDownloadStrategy(),
|
|
"multi": MultiDownloadStrategy(),
|
|
}
|
|
return download_strategies.get(download_strategy_type or CONFIG.service.download_strategy, SingleDownloadStrategy())
|