refactor(api): replace json.loads with Pydantic validation in retention and website layers- #43588

api/services/data_source/website_service.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 import datetime
-import json
 from dataclasses import dataclass
 from typing import Any, NotRequired, TypedDict, cast
 
 import httpx
 from flask_login import current_user
+from pydantic import TypeAdapter, ValidationError
 
 from core.helper import encrypter
 from core.helper.http_client_pooling import get_pooled_http_client
@@ -36,6 +36,9 @@
 )
 
 
+_FIRECRAWL_DOCUMENTS_ADAPTER = TypeAdapter(list[FirecrawlDocumentData])
+
+
 @dataclass
 class CrawlOptions:
     """Options for crawling operations."""
@@ -378,7 +381,10 @@ def _get_firecrawl_url_data(
         if storage.exists(file_key):
             stored_data = storage.load_once(file_key)
             if stored_data:
-                crawl_data = json.loads(stored_data.decode("utf-8"))
+                try:
+                    crawl_data = _FIRECRAWL_DOCUMENTS_ADAPTER.validate_json(stored_data)
+                except ValidationError as e:
+                    raise ValueError(f"cached crawl data is not valid: {file_key}: {e}") from e
         else:
             firecrawl_app = FirecrawlApp(api_key=api_key, base_url=config.get("base_url"))
             result = firecrawl_app.check_crawl_status(job_id)

api/services/retention/conversation/message_export_service.py

@@ -8,9 +8,9 @@
 Does NOT touch Message.inputs / Message.user_feedback properties.
 """
 
+import contextlib
 import datetime
 import gzip
-import json
 import logging
 import tempfile
 from collections import defaultdict
@@ -20,7 +20,7 @@
 
 import orjson
 import sqlalchemy as sa
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError
 from sqlalchemy import select, tuple_
 from sqlalchemy.orm import Session
 
@@ -71,6 +71,13 @@ class AppMessageExportStats(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class _MessageMetadataModel(BaseModel):
+    """The part of `Message.message_metadata` an export reads."""
+
+    model_config = ConfigDict(extra="ignore")
+    retriever_resources: list[dict[str, JsonValue]] = Field(default_factory=list)
+
+
 class AppMessageExportService:
     @staticmethod
     def validate_export_filename(filename: str) -> str:
@@ -284,13 +291,10 @@ def _fetch_feedbacks(self, message_ids: list[str]) -> dict[str, list[AppMessageE
     def _build_record(row: Any, feedbacks_map: dict[str, list[AppMessageExportFeedback]]) -> AppMessageExportRecord:
         retriever_resources: list[Any] = []
         if row.message_metadata:
-            try:
-                metadata = json.loads(row.message_metadata)
-                value = metadata.get("retriever_resources", [])
-                if isinstance(value, list):
-                    retriever_resources = value
-            except (json.JSONDecodeError, TypeError):
-                pass
+            with contextlib.suppress(ValidationError, TypeError):
+                retriever_resources = _MessageMetadataModel.model_validate_json(
+                    row.message_metadata
+                ).retriever_resources
 
         message_id = str(row.id)
         return AppMessageExportRecord(

api/services/retention/workflow_run/archive_paid_plan_workflow_run.py

@@ -37,11 +37,12 @@
 from dataclasses import dataclass, field
 from enum import Enum
 from threading import Lock
-from typing import Any, NotRequired, TypedDict, TypeVar, cast
+from typing import Any, NotRequired, TypedDict, TypeVar
 
 import click
 import pyarrow as pa
 import pyarrow.parquet as pq
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from sqlalchemy import inspect, select
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -127,6 +128,19 @@ class ArchiveBundleIndexDict(TypedDict):
     campaign_ids: NotRequired[list[str]]
 
 
+_ARCHIVE_BUNDLE_INDEX_ADAPTER = TypeAdapter(ArchiveBundleIndexDict)
+
+
+class ShardManifestEntryModel(BaseModel):
+    """Fields `_build_bundle_index` reads from a V2 bundle manifest."""
+
+    model_config = ConfigDict(extra="allow")
+    schema_version: str
+    archive_format: str
+    run_ids: list[str]
+    campaign_id: str | None = None
+
+
 @dataclass(frozen=True)
 class ArchiveBundleIdentity:
     """Stable identity and object prefix for one V2 archive bundle."""
@@ -1002,10 +1016,10 @@ def _load_bundle_index(
     ) -> ArchiveBundleIndexDict:
         index_key = self._get_index_object_key(identity)
         payload = storage.get_object(index_key)
-        loaded = json.loads(payload)
-        if not isinstance(loaded, dict):
-            raise ValueError(f"archive index must be an object: {index_key}")
-        index = cast(ArchiveBundleIndexDict, loaded)
+        try:
+            index = _ARCHIVE_BUNDLE_INDEX_ADAPTER.validate_json(payload)
+        except ValidationError as e:
+            raise ValueError(f"archive index is not valid: {index_key}: {e}") from e
         expected_prefix = self._get_shard_object_prefix(identity)
         if index["schema_version"] != ARCHIVE_BUNDLE_SCHEMA_VERSION:
             raise ValueError(f"unsupported archive index schema_version: {index['schema_version']}")
@@ -1080,24 +1094,17 @@ def _build_bundle_index(
         campaign_ids: set[str] = set()
         for manifest_key in manifest_keys:
             manifest_payload = storage.get_object(manifest_key)
-            manifest = json.loads(manifest_payload)
-            if not isinstance(manifest, dict):
-                raise ValueError(f"archive manifest must be an object: {manifest_key}")
-            if manifest.get("schema_version") != ARCHIVE_BUNDLE_SCHEMA_VERSION:
-                raise ValueError(
-                    f"unsupported bundle schema_version in {manifest_key}: {manifest.get('schema_version')}"
-                )
-            if manifest.get("archive_format") != ARCHIVE_BUNDLE_FORMAT:
-                raise ValueError(
-                    f"unsupported bundle archive_format in {manifest_key}: {manifest.get('archive_format')}"
-                )
-            manifest_run_ids = manifest.get("run_ids")
-            if not isinstance(manifest_run_ids, list):
-                raise ValueError(f"manifest run_ids must be a list: {manifest_key}")
-            run_ids.update(str(run_id) for run_id in manifest_run_ids)
-            campaign_id = manifest.get("campaign_id")
-            if isinstance(campaign_id, str):
-                campaign_ids.add(campaign_id)
+            try:
+                manifest = ShardManifestEntryModel.model_validate_json(manifest_payload)
+            except ValidationError as e:
+                raise ValueError(f"archive manifest is not valid: {manifest_key}: {e}") from e
+            if manifest.schema_version != ARCHIVE_BUNDLE_SCHEMA_VERSION:
+                raise ValueError(f"unsupported bundle schema_version in {manifest_key}: {manifest.schema_version}")
+            if manifest.archive_format != ARCHIVE_BUNDLE_FORMAT:
+                raise ValueError(f"unsupported bundle archive_format in {manifest_key}: {manifest.archive_format}")
+            run_ids.update(str(run_id) for run_id in manifest.run_ids)
+            if manifest.campaign_id is not None:
+                campaign_ids.add(manifest.campaign_id)
         return ArchiveBundleIndexDict(
             schema_version=ARCHIVE_BUNDLE_SCHEMA_VERSION,
             archive_format=ARCHIVE_BUNDLE_FORMAT,

api/services/retention/workflow_run/bundle_archive_maintenance.py

@@ -21,12 +21,13 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, TypedDict, cast
+from typing import Any, NotRequired, TypedDict, cast
 
 import pyarrow.parquet as pq
 import sqlalchemy as sa
 from botocore.exceptions import ClientError, HTTPClientError
 from botocore.exceptions import ConnectionError as BotoCoreConnectionError
+from pydantic import TypeAdapter, ValidationError
 from sqlalchemy import delete, func, inspect, select
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.engine import CursorResult
@@ -83,15 +84,20 @@ class BundleManifest(TypedDict):
     object_prefix: str
     workflow_run_count: int
     workflow_node_execution_count: int
-    min_created_at: str
-    max_created_at: str
-    min_run_id: str
-    max_run_id: str
-    archived_at: str
+    # Optional so validation stays at the strength of the previous handwritten check, which
+    # listed only the first thirteen fields as required.
+    min_created_at: NotRequired[str]
+    max_created_at: NotRequired[str]
+    min_run_id: NotRequired[str]
+    max_run_id: NotRequired[str]
+    archived_at: NotRequired[str]
     tables: dict[str, TableManifestEntry]
     run_ids: list[str]
 
 
+_BUNDLE_MANIFEST_ADAPTER = TypeAdapter(BundleManifest)
+
+
 @dataclass(frozen=True)
 class ArchiveBundleCatalogEntry:
     """Immutable catalog identity and manifest-derived metrics for one V2 archive bundle."""
@@ -759,28 +765,10 @@ def _load_and_validate_manifest(
         *,
         object_prefix: str,
     ) -> BundleManifest:
-        loaded = json.loads(manifest_data)
-        if not isinstance(loaded, dict):
-            raise ValueError("manifest.json must be an object")
-        required_fields = {
-            "schema_version",
-            "archive_format",
-            "tenant_id",
-            "tenant_prefix",
-            "year",
-            "month",
-            "shard",
-            "bundle_id",
-            "object_prefix",
-            "workflow_run_count",
-            "workflow_node_execution_count",
-            "tables",
-            "run_ids",
-        }
-        missing_fields = sorted(required_fields - set(loaded))
-        if missing_fields:
-            raise ValueError(f"manifest missing required fields: {', '.join(missing_fields)}")
-        manifest = cast(BundleManifest, loaded)
+        try:
+            manifest = _BUNDLE_MANIFEST_ADAPTER.validate_json(manifest_data)
+        except ValidationError as e:
+            raise ValueError(f"manifest.json is not a valid bundle manifest: {e}") from e
         if manifest["schema_version"] != ARCHIVE_BUNDLE_SCHEMA_VERSION:
             raise ValueError(f"unsupported bundle schema_version: {manifest['schema_version']}")
         if manifest["archive_format"] != ARCHIVE_BUNDLE_FORMAT:
@@ -793,15 +781,10 @@ def _load_and_validate_manifest(
             raise ValueError("manifest run_ids count does not match workflow_run_count")
 
         tables = manifest["tables"]
-        if not isinstance(tables, dict):
-            raise ValueError("manifest tables must be an object")
         for table_name in ARCHIVED_TABLES:
             if table_name not in tables:
                 raise ValueError(f"manifest missing table: {table_name}")
             info = tables[table_name]
-            for key in ("row_count", "checksum", "size_bytes", "object_key"):
-                if key not in info:
-                    raise ValueError(f"manifest table {table_name} missing {key}")
             expected_key = f"{object_prefix}/{table_name}.parquet"
             if info["object_key"] != expected_key:
                 raise ValueError(

api/services/retention/workflow_run/delete_archived_workflow_run.py

@@ -9,7 +9,6 @@
 """
 
 import io
-import json
 import logging
 import time
 import zipfile
@@ -18,6 +17,7 @@
 from datetime import datetime
 from typing import TypedDict
 
+from pydantic import TypeAdapter, ValidationError
 from sqlalchemy.orm import Session, sessionmaker
 
 from extensions.ext_database import db
@@ -45,6 +45,8 @@ class _ArchiveManifest(TypedDict):
     tables: dict[str, _TableManifestEntry]
 
 
+_ARCHIVE_MANIFEST_ADAPTER = TypeAdapter(_ArchiveManifest)
+
 _ARCHIVED_TABLES = [
     "workflow_runs",
     "workflow_app_logs",
@@ -276,22 +278,10 @@ def _validate_archive_bundle(
                     manifest_data = archive.read("manifest.json")
                 except KeyError as e:
                     raise ValueError("manifest.json missing from archive bundle") from e
-                loaded = json.loads(manifest_data)
-                if not isinstance(loaded, dict):
-                    raise ValueError("manifest.json must be an object")
-                manifest = loaded
-
-                required_fields = {
-                    "schema_version",
-                    "workflow_run_id",
-                    "tenant_id",
-                    "app_id",
-                    "workflow_id",
-                    "tables",
-                }
-                missing_fields = sorted(required_fields - set(manifest))
-                if missing_fields:
-                    raise ValueError(f"manifest missing required fields: {', '.join(missing_fields)}")
+                try:
+                    manifest = _ARCHIVE_MANIFEST_ADAPTER.validate_json(manifest_data)
+                except ValidationError as e:
+                    raise ValueError(f"manifest.json is not a valid archive manifest: {e}") from e
                 if manifest["schema_version"] != ARCHIVE_SCHEMA_VERSION:
                     raise ValueError(
                         f"unsupported archive schema_version: {manifest['schema_version']} "
@@ -307,19 +297,12 @@ def _validate_archive_bundle(
                     raise ValueError("manifest workflow_id does not match delete target")
 
                 tables = manifest["tables"]
-                if not isinstance(tables, dict):
-                    raise ValueError("manifest tables must be an object")
                 missing_tables = [table_name for table_name in _ARCHIVED_TABLES if table_name not in tables]
                 if missing_tables:
                     raise ValueError(f"manifest missing tables: {', '.join(missing_tables)}")
 
                 for table_name in _ARCHIVED_TABLES:
                     info = tables[table_name]
-                    if not isinstance(info, dict):
-                        raise ValueError(f"manifest table entry must be an object: {table_name}")
-                    for key in ("row_count", "checksum", "size_bytes"):
-                        if key not in info:
-                            raise ValueError(f"manifest table {table_name} missing {key}")
                     member_path = f"{table_name}.jsonl"
                     try:
                         payload = archive.read(member_path)
@@ -343,7 +326,7 @@ def _validate_archive_bundle(
                             f"expected={info['row_count']}, actual={row_count}"
                         )
 
-                return manifest  # type: ignore[return-value]
+                return manifest
         except zipfile.BadZipFile as e:
             raise ValueError("archive bundle is not a valid zip file") from e