112 lines
4.2 KiB
Python
112 lines
4.2 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.adapter.logging_utils import logging_error, logging_debug
|
|
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:
|
|
if (
|
|
hasattr(message, "extra_data")
|
|
and isinstance(message.extra_data, dict)
|
|
and len(message.extra_data) > 0
|
|
):
|
|
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"Parsing JSON: {jsonerror}")
|
|
self.is_valid = False
|
|
else:
|
|
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"Parsing JSON: {jsonerror}")
|
|
self.is_valid = False
|
|
else:
|
|
self.form_data = message.body()
|
|
|
|
def on_field(self, field: Field):
|
|
logging_debug(field)
|
|
self.form_data[field.field_name.decode("utf-8")] = field.value
|
|
logging_debug(self.form_data)
|
|
|
|
def on_file(self, file: File):
|
|
logging_debug(file)
|
|
self.form_data[file.field_name.decode("utf-8")] = UploadFile(
|
|
file.file_object, size=file.size, filename=file.file_name
|
|
)
|
|
logging_debug(self.form_data)
|