backend/models/request_log.py

from datetime import datetime
from typing import Optional, Any
from pydantic import BaseModel, Field, GetJsonSchemaHandler, ConfigDict
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
from bson import ObjectId


class PyObjectId(ObjectId):
    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: Any
    ) -> core_schema.CoreSchema:
        def validate_from_str(input_value: str) -> ObjectId:
            if isinstance(input_value, ObjectId):
                return input_value
            if isinstance(input_value, str):
                if ObjectId.is_valid(input_value):
                    return ObjectId(input_value)
                raise ValueError("Invalid ObjectId string")
            raise ValueError("Invalid ObjectId type")

        return core_schema.no_info_plain_validator_function(validate_from_str)

    @classmethod
    def __get_pydantic_json_schema__(
        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
    ) -> JsonSchemaValue:
        return {"type": "string"}


class RequestLog(BaseModel):
    id: Optional[PyObjectId] = Field(default=None, alias="_id")
    endpoint: Optional[str] = None
    response_time: float
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    status_code: Optional[int] = None
    error_type: Optional[str] = None

    model_config = ConfigDict(
        populate_by_name=True,
        arbitrary_types_allowed=True,
        json_encoders={ObjectId: str}
    )

    def to_mongo(self) -> dict:
        data = self.model_dump(by_alias=True, exclude={"id"}, mode='python')
        if self.id:
            data["_id"] = self.id
        return data

    @classmethod
    def from_mongo(cls, data: dict) -> "RequestLog":
        return cls(**data)