Consolidate request JSON validation

Author: adamtheturtle

src/mock_vws/_flask_server/target_manager.py

@@ -7,7 +7,7 @@
 from collections.abc import Callable, Mapping, Sequence
 from enum import Enum, StrEnum, auto
 from http import HTTPMethod, HTTPStatus
-from typing import Annotated, TypeIs, assert_never
+from typing import Annotated, NotRequired, TypedDict, TypeIs, assert_never
 from zoneinfo import ZoneInfo
 
 from beartype import beartype
@@ -16,6 +16,7 @@
     BaseModel,
     BeforeValidator,
     ConfigDict,
+    TypeAdapter,
     ValidationError,
     model_validator,
 )
@@ -244,6 +245,59 @@ def to_vumark_database(self) -> VuMarkDatabase:
         )
 
 
+class OAuth2ClientCredentialBody(BaseModel):
+    """A request to register an OAuth2 client credential."""
+
+    model_config = ConfigDict(strict=True, extra="forbid")
+
+    client_id: str
+    client_secret: str
+    scopes: tuple[str, ...]
+
+
+class ImageTargetBody(TypedDict):
+    """An image target sent to the storage service."""
+
+    name: str
+    width: float
+    image_base64: str
+    active_flag: bool
+    processing_time_seconds: float
+    application_metadata: str | None
+    target_id: str
+    last_modified_date: NotRequired[str]
+    delete_date_optional: NotRequired[str | None]
+    upload_date: NotRequired[str]
+    tracking_rating: NotRequired[int]
+    current_month_recos: NotRequired[int]
+    previous_month_recos: NotRequired[int]
+    total_recos: NotRequired[int]
+    reco_rating: NotRequired[str]
+
+
+class ImageTargetUpdateBody(TypedDict):
+    """Fields which may update a stored image target."""
+
+    name: NotRequired[str]
+    width: NotRequired[float]
+    active_flag: NotRequired[bool]
+    application_metadata: NotRequired[str | None]
+    image: NotRequired[str]
+
+
+class RecognitionCountsBody(TypedDict):
+    """Recognition counts to update on a stored target."""
+
+    current_month_recos: NotRequired[int]
+    previous_month_recos: NotRequired[int]
+    total_recos: NotRequired[int]
+
+
+_IMAGE_TARGET_ADAPTER = TypeAdapter(type=ImageTargetBody)
+_IMAGE_TARGET_UPDATE_ADAPTER = TypeAdapter(type=ImageTargetUpdateBody)
+_RECOGNITION_COUNTS_ADAPTER = TypeAdapter(type=RecognitionCountsBody)
+
+
 @beartype
 class _InvalidRequestBodyError(Exception):
     """A request body which cannot be used to create a resource.
@@ -685,11 +739,13 @@ def get_oauth2_client_credentials() -> Response:
 @beartype
 def put_oauth2_client_credential() -> Response:
     """Add or replace an OAuth2 client credential."""
-    value = json.loads(s=request.data)
+    value = OAuth2ClientCredentialBody.model_validate_json(
+        json_data=request.data,
+    )
     credential = OAuth2ClientCredential(
-        client_id=value["client_id"],  # pyrefly: ignore [unknown-argument-type]
-        client_secret=value["client_secret"],  # pyrefly: ignore [unknown-argument-type]
-        scopes=tuple(value["scopes"]),  # pyrefly: ignore [unknown-argument-type]
+        client_id=value.client_id,
+        client_secret=value.client_secret,
+        scopes=value.scopes,
     )
     TARGET_MANAGER.add_oauth2_client_credential(credential=credential)
     return Response(response="", status=HTTPStatus.NO_CONTENT)
@@ -719,19 +775,22 @@ def create_target(database_name: str) -> Response:
     :status 201: The target has been created.
     :status 404: There is no cloud database with the given name.
     """
-    request_json = json.loads(s=request.data)
+    request_body = _IMAGE_TARGET_ADAPTER.validate_json(
+        request.data,
+        strict=True,
+    )
     settings = TargetManagerSettings.model_validate(obj={})
 
-    image_bytes = base64.b64decode(s=request_json["image_base64"])  # pyrefly: ignore [unknown-argument-type]
+    image_bytes = base64.b64decode(s=request_body["image_base64"])
     target_tracking_rater = settings.target_rater.to_target_rater()
     target = ImageTarget(
-        name=request_json["name"],  # pyrefly: ignore [unknown-argument-type]
-        width=request_json["width"],  # pyrefly: ignore [unknown-argument-type]
+        name=request_body["name"],
+        width=request_body["width"],
         image_value=image_bytes,
-        active_flag=request_json["active_flag"],  # pyrefly: ignore [unknown-argument-type]
-        processing_time_seconds=request_json["processing_time_seconds"],  # pyrefly: ignore [unknown-argument-type]
-        application_metadata=request_json["application_metadata"],  # pyrefly: ignore [unknown-argument-type]
-        target_id=request_json["target_id"],  # pyrefly: ignore [unknown-argument-type]
+        active_flag=request_body["active_flag"],
+        processing_time_seconds=request_body["processing_time_seconds"],
+        application_metadata=request_body["application_metadata"],
+        target_id=request_body["target_id"],
         target_tracking_rater=target_tracking_rater,
     )
     with TARGET_MANAGER.lock:
@@ -807,7 +866,10 @@ def delete_target(database_name: str, target_id: str) -> Response:
 @beartype
 def update_target(database_name: str, target_id: str) -> Response:
     """Update a target."""
-    request_json = json.loads(s=request.data)
+    request_body = _IMAGE_TARGET_UPDATE_ADAPTER.validate_json(
+        request.data,
+        strict=True,
+    )
 
     with TARGET_MANAGER.lock:
         database = _find_cloud_database(database_name=database_name)
@@ -816,26 +878,25 @@ def update_target(database_name: str, target_id: str) -> Response:
 
         target = database.get_target(target_id=target_id)
 
-        name = request_json.get("name", target.name)  # pyrefly: ignore [unknown-variable-type]
-        active_flag = request_json.get("active_flag", target.active_flag)  # pyrefly: ignore [unknown-variable-type]
+        name = request_body.get("name", target.name)
+        active_flag = request_body.get("active_flag", target.active_flag)
 
         gmt = ZoneInfo(key="GMT")
         last_modified_date = datetime.datetime.now(tz=gmt)
 
-        width = request_json.get("width", target.width)  # pyrefly: ignore [unknown-variable-type]
-        application_metadata = request_json.get(  # pyrefly: ignore [unknown-variable-type]
-            "application_metadata",
-            target.application_metadata,
+        width = request_body.get("width", target.width)
+        application_metadata = request_body.get(
+            "application_metadata", target.application_metadata
         )
         image_value = target.image_value
-        if "image" in request_json:
-            image_value = base64.b64decode(s=request_json["image"])  # pyrefly: ignore [unknown-argument-type]
+        if "image" in request_body:
+            image_value = base64.b64decode(s=request_body["image"])
         new_target = copy.replace(
             target,
-            name=name,  # pyrefly: ignore [unknown-argument-type]
-            width=width,  # pyrefly: ignore [unknown-argument-type]
-            active_flag=active_flag,  # pyrefly: ignore [unknown-argument-type]
-            application_metadata=application_metadata,  # pyrefly: ignore [unknown-argument-type]
+            name=name,
+            width=width,
+            active_flag=active_flag,
+            application_metadata=application_metadata,
             image_value=image_value,
             last_modified_date=last_modified_date,
         )
@@ -880,7 +941,10 @@ def set_target_recognition_counts(
 
     :status 200: The recognition counts have been set.
     """
-    request_json = json.loads(s=request.data)
+    request_body = _RECOGNITION_COUNTS_ADAPTER.validate_json(
+        request.data,
+        strict=True,
+    )
 
     with TARGET_MANAGER.lock:
         database = _find_cloud_database(database_name=database_name)
@@ -891,15 +955,13 @@ def set_target_recognition_counts(
 
         new_target = copy.replace(
             target,
-            current_month_recos=request_json.get(  # pyrefly: ignore [unknown-argument-type]
-                "current_month_recos",
-                target.current_month_recos,
+            current_month_recos=request_body.get(
+                "current_month_recos", target.current_month_recos
             ),
-            previous_month_recos=request_json.get(  # pyrefly: ignore [unknown-argument-type]
-                "previous_month_recos",
-                target.previous_month_recos,
+            previous_month_recos=request_body.get(
+                "previous_month_recos", target.previous_month_recos
             ),
-            total_recos=request_json.get("total_recos", target.total_recos),  # pyrefly: ignore [unknown-argument-type]
+            total_recos=request_body.get("total_recos", target.total_recos),
         )
 
         database.targets.remove(target)

src/mock_vws/_flask_server/vws.py

@@ -8,18 +8,18 @@
 import email.utils
 import gzip
 import html
-import json
 import logging
 import threading
 import time
 import uuid
 from enum import StrEnum, auto
 from http import HTTPMethod, HTTPStatus
-from typing import assert_never
+from typing import NotRequired, TypedDict, assert_never
 
 import requests
 from beartype import beartype
 from flask import Flask, Response, request
+from pydantic import TypeAdapter
 from pydantic_settings import BaseSettings
 from werkzeug.exceptions import MethodNotAllowed, NotFound
 
@@ -100,6 +100,30 @@
 _LOGGER = logging.getLogger(name=__name__)
 
 
+class _AddTargetBody(TypedDict):
+    """A validated add-target request body."""
+
+    name: str
+    width: float
+    image: str
+    active_flag: NotRequired[bool | None]
+    application_metadata: NotRequired[str | None]
+
+
+class _UpdateTargetBody(TypedDict):
+    """A validated update-target request body."""
+
+    width: NotRequired[float]
+    active_flag: NotRequired[bool | None]
+    application_metadata: NotRequired[str | None]
+    name: NotRequired[str]
+    image: NotRequired[str]
+
+
+_ADD_TARGET_BODY_ADAPTER = TypeAdapter(type=_AddTargetBody)
+_UPDATE_TARGET_BODY_ADAPTER = TypeAdapter(type=_UpdateTargetBody)
+
+
 @beartype
 class _ImageMatcherChoice(StrEnum):
     """Image matcher choices."""
@@ -750,22 +774,25 @@ def add_target() -> Response:
 
     # We do not use ``request.get_json(force=True)`` because this only works
     # when the content type is given as ``application/json``.
-    request_json = json.loads(s=request.data)
-    name = request_json["name"]  # pyrefly: ignore [unknown-variable-type]
-    active_flag = request_json.get("active_flag")  # pyrefly: ignore [unknown-variable-type]
+    request_json = _ADD_TARGET_BODY_ADAPTER.validate_json(
+        request.data,
+        strict=True,
+    )
+    name = request_json["name"]
+    active_flag = request_json.get("active_flag")
     if active_flag is None:
         active_flag = True
 
     # This rater is not used.
     target_tracking_rater = HardcodedTargetTrackingRater(rating=1)
 
     new_target = ImageTarget(
-        name=name,  # pyrefly: ignore [unknown-argument-type]
-        width=request_json["width"],  # pyrefly: ignore [unknown-argument-type]
-        image_value=base64.b64decode(s=request_json["image"]),  # pyrefly: ignore [unknown-argument-type]
-        active_flag=active_flag,  # pyrefly: ignore [unknown-argument-type]
+        name=name,
+        width=request_json["width"],
+        image_value=base64.b64decode(s=request_json["image"]),
+        active_flag=active_flag,
         processing_time_seconds=settings.processing_time_seconds,
-        application_metadata=request_json.get("application_metadata"),  # pyrefly: ignore [unknown-argument-type]
+        application_metadata=request_json.get("application_metadata"),
         target_tracking_rater=target_tracking_rater,
     )
 
@@ -1220,7 +1247,10 @@ def update_target(target_id: str) -> Response:
     settings = VWSSettings.model_validate(obj={})
     # We do not use ``request.get_json(force=True)`` because this only works
     # when the content type is given as ``application/json``.
-    request_json = json.loads(s=request.data)
+    request_json = _UPDATE_TARGET_BODY_ADAPTER.validate_json(
+        request.data,
+        strict=True,
+    )
     databases = get_all_cloud_databases()
     database = get_database_matching_server_keys(
         request_headers=dict(request.headers),
@@ -1239,10 +1269,10 @@ def update_target(target_id: str) -> Response:
 
     update_values: dict[str, str | int | float | bool | None] = {}
     if "width" in request_json:
-        update_values["width"] = request_json["width"]  # pyrefly: ignore [unknown-argument-type]
+        update_values["width"] = request_json["width"]
 
     if "active_flag" in request_json:
-        active_flag = request_json["active_flag"]  # pyrefly: ignore [unknown-variable-type]
+        active_flag = request_json["active_flag"]
         if active_flag is None:
             _LOGGER.warning(
                 msg=(
@@ -1251,10 +1281,10 @@ def update_target(target_id: str) -> Response:
                 ),
             )
             raise FailError(status_code=HTTPStatus.BAD_REQUEST)
-        update_values["active_flag"] = active_flag  # pyrefly: ignore [unknown-argument-type]
+        update_values["active_flag"] = active_flag
 
     if "application_metadata" in request_json:
-        application_metadata = request_json["application_metadata"]  # pyrefly: ignore [unknown-variable-type]
+        application_metadata = request_json["application_metadata"]
         if application_metadata is None:
             _LOGGER.warning(
                 msg=(
@@ -1263,15 +1293,15 @@ def update_target(target_id: str) -> Response:
                 ),
             )
             raise FailError(status_code=HTTPStatus.BAD_REQUEST)
-        update_values["application_metadata"] = application_metadata  # pyrefly: ignore [unknown-argument-type]
+        update_values["application_metadata"] = application_metadata
 
     if "name" in request_json:
-        name = request_json["name"]  # pyrefly: ignore [unknown-variable-type]
-        update_values["name"] = name  # pyrefly: ignore [unknown-argument-type]
+        name = request_json["name"]
+        update_values["name"] = name
 
     if "image" in request_json:
-        image = request_json["image"]  # pyrefly: ignore [unknown-variable-type]
-        update_values["image"] = image  # pyrefly: ignore [unknown-argument-type]
+        image = request_json["image"]
+        update_values["image"] = image
 
     put_url = (
         f"{settings.target_manager_base_url}/cloud_databases/"