diff --git a/requirements.txt b/requirements.txt index 1f2a8398..cceba64b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -66,4 +66,3 @@ werkzeug==3.1.8 # flask # flask-cors -opengeodeweb-microservice==1.*,>=1.2.2 diff --git a/src/opengeodeweb_back/routes/blueprint_routes.py b/src/opengeodeweb_back/routes/blueprint_routes.py index e96e2f4f..2e986aee 100644 --- a/src/opengeodeweb_back/routes/blueprint_routes.py +++ b/src/opengeodeweb_back/routes/blueprint_routes.py @@ -4,6 +4,7 @@ import time import shutil import math +import typing from threading import Timer # Third party imports @@ -63,6 +64,18 @@ def allowed_files() -> flask.Response: return flask.make_response({"extensions": list(extensions)}, 200) +def _write_stream(path: str, stream: typing.IO[bytes], mode: str = "wb") -> None: + read_buffer_size = 1024 * 1024 + with open(path, mode) as destination: + while chunk := stream.read(read_buffer_size): + destination.write(chunk) + + +def _finalize_upload(filename: str) -> flask.Response: + print(f"{filename=}", flush=True) + return flask.make_response({"message": "File uploaded"}, 201) + + @routes.route( schemas_dict["upload_file"]["route"], methods=schemas_dict["upload_file"]["methods"], @@ -73,8 +86,12 @@ def upload_file() -> flask.Response: if not os.path.exists(UPLOAD_FOLDER_PATH): os.makedirs(UPLOAD_FOLDER_PATH, exist_ok=True) - # Multipart callers (e.g. Vease) still send the file as a "file" form part; - # streaming callers PUT the raw bytes as the body with ?filename= as a query param. + # Multipart callers (e.g. Vease) still send the whole file as a "file" form + # part. Everything else PUTs raw bytes with ?filename= as a query param: + # either the whole file in one request, or one of several chunks (when + # ?chunk_index=/?total_chunks= are also present) that get appended in + # order and assembled into the final file once the last one arrives. This + # keeps every request under cloud hosting's hard request-size limit. if flask.request.mimetype == "multipart/form-data": file = flask.request.files["file"] if file.filename is None: @@ -82,22 +99,30 @@ def upload_file() -> flask.Response: filename = werkzeug.utils.secure_filename(os.path.basename(file.filename)) file_path = os.path.join(UPLOAD_FOLDER_PATH, filename) file.save(file_path) - else: - raw_filename = flask.request.args.get("filename") - if not raw_filename: - flask.abort(400, "Filename is required") - filename = werkzeug.utils.secure_filename(os.path.basename(raw_filename)) - file_path = os.path.join(UPLOAD_FOLDER_PATH, filename) - chunk_size = 1024 * 1024 - with open(file_path, "wb") as destination: - while chunk := flask.request.stream.read(chunk_size): - destination.write(chunk) - print(f"{filename=}", flush=True) - if filename.lower().endswith(".csv.json"): - shutil.copyfile( - file_path, os.path.join(UPLOAD_FOLDER_PATH, filename[:-9] + ".json") - ) - return flask.make_response({"message": "File uploaded"}, 201) + return _finalize_upload(filename) + + raw_filename = flask.request.args.get("filename") + if not raw_filename: + flask.abort(400, "Filename is required") + filename = werkzeug.utils.secure_filename(os.path.basename(raw_filename)) + file_path = os.path.join(UPLOAD_FOLDER_PATH, filename) + + total_chunks = flask.request.args.get("total_chunks", type=int) + if total_chunks is None: + _write_stream(file_path, flask.request.stream) + return _finalize_upload(filename) + + chunk_index = flask.request.args.get("chunk_index", type=int) + if chunk_index is None or not 0 <= chunk_index < total_chunks: + flask.abort(400, "Invalid chunk_index") + + part_path = f"{file_path}.part" + _write_stream(part_path, flask.request.stream, "wb" if chunk_index == 0 else "ab") + if chunk_index < total_chunks - 1: + return flask.make_response({"message": "Chunk received"}, 200) + + os.replace(part_path, file_path) + return _finalize_upload(filename) @routes.route( diff --git a/tests/test_routes.py b/tests/test_routes.py index f8fffb25..3db46707 100644 --- a/tests/test_routes.py +++ b/tests/test_routes.py @@ -111,6 +111,60 @@ def test_upload_file_raw_missing_filename(client: FlaskClient) -> None: assert response.status_code == 400 +def test_upload_file_chunked( + client: FlaskClient, filename: str = "test.og_brep" +) -> None: + file = os.path.join(data_dir, filename) + with open(file, "rb") as opened_file: + file_bytes = opened_file.read() + + chunk_filename = "chunked_upload_test.og_brep" + chunk_size = max(1, len(file_bytes) // 3) + chunks = [ + file_bytes[index : index + chunk_size] + for index in range(0, len(file_bytes), chunk_size) + ] + total_chunks = len(chunks) + + uploaded_path = os.path.join(data_dir, chunk_filename) + try: + for chunk_index, chunk in enumerate(chunks): + response = client.put( + f"/opengeodeweb_back/upload_file" + f"?filename={chunk_filename}" + f"&chunk_index={chunk_index}" + f"&total_chunks={total_chunks}", + data=chunk, + content_type="application/octet-stream", + ) + if chunk_index < total_chunks - 1: + assert response.status_code == 200 + assert not os.path.exists(uploaded_path) + else: + assert response.status_code == 201 + + with open(uploaded_path, "rb") as uploaded_file: + assert uploaded_file.read() == file_bytes + finally: + if os.path.exists(uploaded_path): + os.remove(uploaded_path) + part_path = f"{uploaded_path}.part" + if os.path.exists(part_path): + os.remove(part_path) + + +def test_upload_file_chunked_invalid_chunk_index(client: FlaskClient) -> None: + response = client.put( + "/opengeodeweb_back/upload_file" + "?filename=invalid_chunk_index.og_brep" + "&chunk_index=2" + "&total_chunks=2", + data=b"some raw bytes", + content_type="application/octet-stream", + ) + assert response.status_code == 400 + + def test_missing_files(client: FlaskClient) -> None: route = f"/opengeodeweb_back/missing_files"