Stanislav Hejny
87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
import io
|
|
import json
|
|
import logging
|
|
|
|
import python_multipart
|
|
from fastapi import UploadFile
|
|
from python_multipart import multipart
|
|
from python_multipart.multipart import Field, File
|
|
|
|
from amqp.model.model import DataMessage
|
|
|
|
|
|
def process_form_data(json_input) -> dict:
|
|
"""
|
|
Processes the input JSON to extract form data. The JSON has two top entries: 'files' and 'form'.
|
|
This is part of the file stream implementation in AMQ adapter to stream large files in chunks.
|
|
'files' contains information about the AMQ queue from which read the file, while 'form' contains
|
|
additional non-file form values.
|
|
This function combines the file and values at the top level in the returned dictionary.
|
|
Args:
|
|
json_input: JSON string or dict containing the form and files data
|
|
|
|
Returns:
|
|
dict: Processed form data with files data overriding non-empty values
|
|
"""
|
|
# Load JSON if input is string
|
|
if isinstance(json_input, str):
|
|
data = json.loads(json_input)
|
|
else:
|
|
data = json_input
|
|
|
|
result = {}
|
|
|
|
# Process form data - take first element of each list
|
|
for key, value in data['form'].items():
|
|
if value: # Only process if list is not empty
|
|
result[key] = value[0]
|
|
|
|
# Override with files data where available
|
|
for key, value in data['files'].items():
|
|
if value: # Only override if files data is non-empty
|
|
result[key] = value[0]
|
|
|
|
return result
|
|
|
|
|
|
class CleverMultiPartParser:
|
|
def __init__(self, message: DataMessage):
|
|
self.form_data = {}
|
|
self.is_multipart = False
|
|
self.is_json = False
|
|
self.is_valid = True
|
|
# Convert each list of strings to a single string joined by newline, then to bytes
|
|
headers = {key: "\n".join(value).encode("utf-8") for key, value in message.headers.items()}
|
|
content_type: str | bytes | None = headers.get("Content-Type")
|
|
content_type, params = multipart.parse_options_header(content_type)
|
|
if content_type.decode('utf-8') in ["application/octet-stream",
|
|
"multipart/form-data",
|
|
"application/x-www-form-urlencoded",
|
|
"application/x-url-encoded"]:
|
|
try:
|
|
python_multipart.parse_form(headers, io.BytesIO(message.body()), self.on_field, self.on_file)
|
|
self.is_multipart = True
|
|
except Exception as e:
|
|
logging.error(f"Error parsing multipart/form-data: {e}. Trying parsing as JSON.")
|
|
try:
|
|
self.form_data = json.loads(message.body())
|
|
self.is_json = True
|
|
if self.form_data['files'] and self.form_data['form']:
|
|
self.form_data = process_form_data(self.form_data)
|
|
except Exception as jsonerror:
|
|
logging.error(f"Error parsing JSON: {jsonerror}")
|
|
self.is_valid = False
|
|
else:
|
|
self.form_data = message.body()
|
|
|
|
def on_field(self, field: Field):
|
|
print(field)
|
|
self.form_data[field.field_name.decode('utf-8')] = field.value
|
|
print(self.form_data)
|
|
|
|
def on_file(self, file: File):
|
|
print(file)
|
|
self.form_data[file.field_name.decode('utf-8')] = \
|
|
UploadFile(file.file_object, size=file.size, filename=file.file_name)
|
|
print(self.form_data)
|