Update flask to version 3

pyproject.toml

@@ -54,13 +54,12 @@ dependencies = [
 
 [project.optional-dependencies]
 server = [
-    "Flask>=2.0",
-    "Werkzeug==2.0.3",
+    "Flask>=3.0",
     "flask-caching>=1.10",
     "flask-compress>=1.12",
     "flask-cors>=3",
     "flask-mail~=0.9.1",
-    "flask-restx==0.5",
+    "flask-restx>=1.0.0",
     "python-magic~=0.4",
     "simplejson~=3.0",
 ]

src/simdb/remote/apis/__init__.py

@@ -1,9 +1,10 @@
 import datetime
+import threading
 from pathlib import Path
 
 import appdirs
 import jwt
-from flask import Blueprint, Response, _app_ctx_stack, jsonify, request
+from flask import Blueprint, Response, jsonify, request
 from flask_restx import Resource
 
 from simdb import __version__
@@ -47,7 +48,7 @@ def setup_db(setup_state):
             args = config.get_section("database")
             setup_state.app.db = Database(
                 Database.DBMS.POSTGRESQL,
-                scopefunc=_app_ctx_stack.__ident_func__,
+                scopefunc=lambda: threading.get_ident(),
                 **args,
             )
         elif db_type == "sqlite":
@@ -57,7 +58,7 @@ def setup_db(setup_state):
             )
             file.parent.mkdir(parents=True, exist_ok=True)
             setup_state.app.db = Database(
-                Database.DBMS.SQLITE, scopefunc=_app_ctx_stack.__ident_func__, file=file
+                Database.DBMS.SQLITE, scopefunc=lambda: threading.get_ident(), file=file
             )
         else:
             raise RuntimeError(f"Unknown database type in configuration: {db_type}.")

src/simdb/remote/apis/files.py

@@ -1,11 +1,10 @@
 import gzip
-import json
 import uuid
 from pathlib import Path
 from typing import Dict, Iterable, List, Optional
 
 import magic
-from flask import Response, jsonify, request, send_file, stream_with_context
+from flask import Response, jsonify, request, send_file
 from flask_restx import Namespace, Resource
 from werkzeug.datastructures import FileStorage
 
@@ -14,13 +13,18 @@
 from simdb.database import DatabaseError, models
 from simdb.imas.checksum import checksum as imas_checksum
 from simdb.imas.utils import imas_files
-from simdb.json import CustomDecoder
 from simdb.remote.core.auth import User, requires_auth
 from simdb.remote.core.errors import error
 from simdb.remote.core.path import find_common_root, secure_path
 from simdb.remote.core.pydantic_utils import pydantic_validate
 from simdb.remote.core.typing import current_app
-from simdb.remote.models import FileDataList, FileGetDataResponse
+from simdb.remote.models import (
+    ChunkInfo,
+    FileDataList,
+    FileGetDataResponse,
+    FileRegistrationData,
+    FileUploadData,
+)
 from simdb.uri import URI
 
 api = Namespace("files", path="/")
@@ -68,10 +72,10 @@ def _verify_file(
 
 
 def _save_chunked_file(
-    file: FileStorage, chunk_info: Dict, path: Path, compressed: bool = True
+    file: FileStorage, chunk_info: ChunkInfo, path: Path, compressed: bool = True
 ):
     with path.open("r+b" if path.exists() else "wb") as file_out:
-        file_out.seek(chunk_info["chunk_size"] * chunk_info["chunk"])
+        file_out.seek(chunk_info.chunk_size * chunk_info.chunk)
         if compressed:
             with gzip.GzipFile(fileobj=file.stream, mode="rb") as gz_file:
                 file_out.write(gz_file.read())
@@ -81,7 +85,7 @@ def _save_chunked_file(
 
 def _stage_file_from_chunks(
     files: Iterable[FileStorage],
-    chunk_info: Dict,
+    chunk_info: Dict[str, ChunkInfo],
     sim_uuid: uuid.UUID,
     sim_files: List[models.File],
     common_root: Optional[Path],
@@ -107,7 +111,7 @@ def _stage_file_from_chunks(
         path = secure_path(sim_file.uri.path, common_root, staging_dir)
         path.parent.mkdir(parents=True, exist_ok=True)
         file_chunk_info = chunk_info.get(
-            sim_file.uuid.hex, {"chunk_size": 0, "chunk": 0, "num_chunks": 1}
+            sim_file.uuid.hex, ChunkInfo(chunk_size=0, chunk=0, num_chunks=1)
         )
         _save_chunked_file(file, file_chunk_info, path)
 
@@ -122,39 +126,33 @@ def _check_file_is_in_simulation(
     return sim_file
 
 
-def _process_simulation_data(data: dict) -> Response:
-    simulation = models.Simulation.from_data(data["simulation"])
+def _process_simulation_data(body: FileRegistrationData) -> Response:
+    simulation = models.Simulation.from_data_model(body.simulation)
     sim_file_paths = simulation.file_paths()
     common_root = find_common_root(sim_file_paths)
-    if DataObject.Type(data["obj_type"]) == DataObject.Type.FILE:
-        for file in data["files"]:
+    if body.obj_type == DataObject.Type.FILE:
+        for file in body.files:
             sim_file = _check_file_is_in_simulation(
-                simulation, uuid.UUID(file["file_uuid"]), file["file_type"]
+                simulation, file.file_uuid, file.file_type
             )
             _verify_file(simulation.uuid, sim_file, common_root)
-    elif DataObject.Type(data["obj_type"]) == DataObject.Type.IMAS:
-        file = data["files"][0]
+    elif body.obj_type == DataObject.Type.IMAS:
+        file = body.files[0]
         sim_files = (
-            simulation.inputs if file["file_type"] == "input" else simulation.outputs
+            simulation.inputs if file.file_type == "input" else simulation.outputs
         )
-        sim_file = next(f for f in sim_files if f.uuid == uuid.UUID(file["file_uuid"]))
-        _verify_file(simulation.uuid, sim_file, common_root, file["ids_list"])
+        sim_file = next(f for f in sim_files if f.uuid == file.file_uuid)
+        _verify_file(simulation.uuid, sim_file, common_root, file.ids_list)
     else:
-        raise ValueError("Unsupported object type {}".format(data["obj_type"]))
+        raise ValueError(f"Unsupported object type {body.obj_type}")
 
     return jsonify({})
 
 
 def _handle_file_upload() -> Response:
-    data: dict = json.load(request.files["data"].stream, cls=CustomDecoder)
+    body = FileUploadData.model_validate_json(request.files["data"].stream.read())
 
-    if "simulation" not in data:
-        return error("Simulation data not provided")
-
-    simulation = models.Simulation.from_data(data["simulation"])
-
-    chunk_info = data.get("chunk_info", {})
-    file_type = data["file_type"]
+    simulation = models.Simulation.from_data_model(body.simulation)
 
     files = request.files.getlist("files")
     if not files:
@@ -163,8 +161,10 @@ def _handle_file_upload() -> Response:
     sim_file_paths = simulation.file_paths()
     common_root = find_common_root(sim_file_paths)
 
-    sim_files = simulation.inputs if file_type == "input" else simulation.outputs
-    _stage_file_from_chunks(files, chunk_info, simulation.uuid, sim_files, common_root)
+    sim_files = simulation.inputs if body.file_type == "input" else simulation.outputs
+    _stage_file_from_chunks(
+        files, body.chunk_info or {}, simulation.uuid, sim_files, common_root
+    )
 
     return jsonify({})
 
@@ -180,9 +180,9 @@ def get(self, user: User) -> FileDataList:
     @requires_auth()
     def post(self, user: User):
         try:
-            data = request.get_json()
-            if data:
-                return _process_simulation_data(data)
+            if request.is_json:
+                body = FileRegistrationData.model_validate_json(request.get_data())
+                return _process_simulation_data(body)
             return _handle_file_upload()
 
         except ValueError as err:
@@ -209,11 +209,7 @@ def get(self, file_uuid: str, user: Optional[User] = None):
             if file.uri.path is None:
                 return error("File path is not set")
             mimetype = magic.from_file(file.uri.path, mime=True)
-            response = send_file(file.uri.path, mimetype=mimetype)
-            return Response(
-                stream_with_context(response.iter_content()),
-                content_type=response.headers["content-type"],
-            )
+            return send_file(file.uri.path, mimetype=mimetype)
         except DatabaseError as err:
             return error(str(err))
 

src/simdb/remote/app.py

@@ -13,7 +13,6 @@
 
 from simdb.config import Config
 from simdb.database.models import Base
-from simdb.json import CustomDecoder, CustomEncoder
 
 from .apis import blueprints
 from .core.auth._authenticator import Authenticator
@@ -98,8 +97,6 @@ def create_app(
     app.config["DEBUG"] = debug
     app.config["RESTX_INCLUDE_ALL_MODELS"] = True
     app.config["PROFILE"] = profile
-    app.json_encoder = CustomEncoder  # ty: ignore[invalid-assignment]
-    app.json_decoder = CustomDecoder  # ty: ignore[invalid-assignment]
     app.config.from_mapping(flask_options)
     app.simdb_config = config
     cache.init_app(app)

tests/remote/api/conftest.py

@@ -31,7 +31,7 @@
     SIMULATIONS.append(Simulation(Manifest()))
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")
 def client():
     if not has_flask:
         pytest.skip("Flask not installed")
@@ -49,22 +49,21 @@ def client():
     app = create_app(config=config, testing=True, debug=True)
     app.testing = True
 
-    with app.test_client() as client:
-        # with app.app_context():
-        for sim in SIMULATIONS:
-            app.db.insert_simulation(sim)
+    test_sims = [Simulation(Manifest()) for _ in range(100)]
 
-        app.db.session.commit()
-        app.db.session.close()
+    with app.app_context():
+        for sim in test_sims:
+            app.db.insert_simulation(sim)
 
+    with app.test_client() as client:
         yield client
 
     os.close(db_fd)
     Path(app.simdb_config.get_option("database.file")).unlink()
     shutil.rmtree(upload_dir)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")
 def client_copy_files():
     if not has_flask:
         pytest.skip("Flask not installed")
@@ -82,14 +81,11 @@ def client_copy_files():
     app = create_app(config=config, testing=True, debug=True)
     app.testing = True
 
-    with app.test_client() as client:
-        # with app.app_context():
+    with app.app_context():
         for sim in SIMULATIONS:
             app.db.insert_simulation(sim)
 
-        app.db.session.commit()
-        app.db.session.close()
-
+    with app.test_client() as client:
         yield client
 
     os.close(db_fd)