Add warp function for pydantic model json dump

Signed-off-by: Wei-Chun, Chang <[email protected]>

Author: wcchang1115

recce/apis/check_api.py

@@ -10,6 +10,7 @@
 from recce.apis.run_func import submit_run
 from recce.exceptions import RecceException
 from recce.models import RunType, RunDAO, Check, CheckDAO, Run
+from recce.models.util import pydantic_model_json_dump
 
 check_router = APIRouter(tags=['check'])
 
@@ -240,7 +241,7 @@ async def export_handler():
     from ..models.state import recce_state
 
     try:
-        return recce_state.model_dump_json()
+        return pydantic_model_json_dump(recce_state)
     except RecceException as e:
         raise HTTPException(status_code=400, detail=e.message)
 

recce/models/state.py

@@ -8,6 +8,7 @@
 from pydantic import BaseModel, Field
 
 from .types import Run, Check
+from .util import pydantic_model_json_dump
 from .. import get_version
 
 logger = logging.getLogger('uvicorn')
@@ -28,7 +29,7 @@ def store(self, file_path):
         self.metadata = RecceStateMetadata()
         start_time = time.time()
         logger.info(f"Store recce state to '{file_path}'")
-        json_data = self.model_dump_json()
+        json_data = pydantic_model_json_dump(self)
         with open(file_path, 'w') as f:
             f.write(json_data)
         end_time = time.time()

recce/models/util.py

@@ -0,0 +1,16 @@
+import pydantic.version
+from pydantic import BaseModel
+
+from recce.exceptions import RecceException
+
+
+def pydantic_model_json_dump(model: BaseModel):
+    pydantic_version = pydantic.version.VERSION
+    pydantic_major = pydantic_version.split(".")[0]
+
+    if pydantic_major == "1":
+        return model.json()
+    elif pydantic_major == "2":
+        return model.model_dump_json()
+    else:
+        raise RecceException("Currently only support pydantic version 1 and 2.")

tests/models/test_state.py

@@ -2,6 +2,7 @@
 
 from recce.models import RunType, Run, RunDAO, CheckDAO, Check
 from recce.models.state import RecceState
+from recce.models.util import pydantic_model_json_dump
 
 
 def test_load():
@@ -13,7 +14,7 @@ def test_load():
     RunDAO(state).create(run)
     CheckDAO(state).create(check)
 
-    json_data = state.json()
+    json_data = pydantic_model_json_dump(state)
     new_state = RecceState(**json.loads(json_data))
 
     run_loaded = new_state.runs[0]