refactor: use pydantic instead of marshmalow

Author: bolinocroustibat

CHANGELOG.md

@@ -52,6 +52,7 @@
 - Simplify getting Sentry info by loading pyproject.toml info in config [#138](https://github.com/datagouv/hydra/pull/138)
 - Add a `POST` `/api/checks/` route for force crawling [#118](https://github.com/datagouv/hydra/pull/118)
 - Update `csv-detective` to 0.7.2 which doesn't include yanked version of `requests` anymore [#142](https://github.com/datagouv/hydra/pull/142) and [#144](https://github.com/datagouv/hydra/pull/144)
+- Use [Pydantic](https://docs.pydantic.dev/) instead of [Marshmallow](https://marshmallow.readthedocs.io/en/stable/) for API validation [#149](https://github.com/datagouv/hydra/pull/149)
 
 ## 1.0.1 (2023-01-04)
 

poetry.lock

@@ -24,10 +24,10 @@ coloredlogs = "^15.0.1"
 csv-detective = "0.7.2"
 dateparser = "^1.1.7"
 humanfriendly = "^10.0"
-marshmallow = "^3.14.1"
 minicli = "^0.5.3"
 minio = "7.2.7"
 pyarrow = "16.1.0"
+pydantic = "^2.8.2"
 python-dateutil = "^2.8.2"
 python-magic = "^0.4.25"
 progressist = "^0.1.0"

pyproject.toml

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import toml
+import tomllib
 
 log = logging.getLogger("udata-hydra")
 
@@ -24,7 +25,8 @@ def configure(self) -> None:
         # override with local settings
         local_settings = os.environ.get("HYDRA_SETTINGS", Path.cwd() / "config.toml")
         if Path(local_settings).exists():
-            configuration.update(toml.load(local_settings))
+            with open(Path(local_settings), "rb") as f:
+                configuration.update(tomllib.load(f))
 
         self.configuration = configuration
         self.check()

udata_hydra/__init__.py

@@ -40,7 +40,7 @@ async def process_resource(check_id: int, is_first_check: bool) -> None:
 
     Will call udata if first check or changes found, and update check with optional infos
     """
-    check: dict = await Check.get(check_id)
+    check: Optional[dict] = await Check.get(check_id)
     if not check:
         log.error(f"Check not found by id {check_id}")
         return

udata_hydra/analysis/resource.py

@@ -20,7 +20,10 @@ async def get(cls, check_id: int) -> Optional[dict]:
                 ON catalog.last_check = checks.id
                 WHERE checks.id = $1;
             """
-            return await connection.fetchrow(q, check_id)
+            record = await connection.fetchrow(q, check_id)
+            if record:
+                return dict(record)
+            return None
 
     @classmethod
     async def get_latest(
@@ -36,7 +39,10 @@ async def get_latest(
             WHERE checks.id = catalog.last_check
             AND catalog.{column} = $1
             """
-            return await connection.fetchrow(q, url or resource_id)
+            record = await connection.fetchrow(q, url or resource_id)
+            if record:
+                return dict(record)
+            return None
 
     @classmethod
     async def get_all(

udata_hydra/db/check.py

@@ -1,16 +1,20 @@
+from typing import Optional
+
 from udata_hydra import context
 
 
 class Resource:
     """Represents a resource in the "catalog" DB table"""
 
     @classmethod
-    async def get(cls, resource_id: str, column_name: str = "*") -> dict:
+    async def get(cls, resource_id: str, column_name: str = "*") -> Optional[dict]:
         pool = await context.pool()
         async with pool.acquire() as connection:
             q = f"""SELECT {column_name} FROM catalog WHERE resource_id = '{resource_id}';"""
-            resource = await connection.fetchrow(q)
-        return resource
+            record = await connection.fetchrow(q)
+            if record:
+                return dict(record)
+            return None
 
     @classmethod
     async def insert(

udata_hydra/db/resource.py

@@ -3,7 +3,7 @@
 
 import aiohttp
 from aiohttp import web
-from marshmallow import ValidationError
+from pydantic import ValidationError
 
 from udata_hydra import config, context
 from udata_hydra.crawl import check_url
@@ -16,20 +16,22 @@
 async def get_latest_check(request: web.Request) -> web.Response:
     """Get the latest check for a given URL or resource_id"""
     url, resource_id = get_request_params(request, params_names=["url", "resource_id"])
-    data: Optional[dict] = await Check.get_latest(url, resource_id)
-    if not data:
+
+    check: Optional[dict] = await Check.get_latest(url, resource_id)
+    if not check:
         raise web.HTTPNotFound()
-    if data["deleted"]:
+    if check["deleted"]:
         raise web.HTTPGone()
-    return web.json_response(CheckSchema().dump(dict(data)))
+
+    return web.Response(text=json.dumps(check, default=str), content_type="application/json")
 
 
 async def get_all_checks(request: web.Request) -> web.Response:
     url, resource_id = get_request_params(request, params_names=["url", "resource_id"])
     data: Optional[list] = await Check.get_all(url, resource_id)
     if not data:
         raise web.HTTPNotFound()
-    return web.json_response([CheckSchema().dump(dict(r)) for r in data])
+    return web.json_response([r for r in data])
 
 
 async def create_check(request: web.Request) -> web.Response:
@@ -40,7 +42,7 @@ async def create_check(request: web.Request) -> web.Response:
         payload: dict = await request.json()
         resource_id: str = payload["resource_id"]
     except ValidationError as err:
-        raise web.HTTPBadRequest(text=json.dumps(err.messages))
+        raise web.HTTPBadRequest(text=err.json())
     except KeyError as e:
         raise web.HTTPBadRequest(text=f"Missing key: {e}")
 
@@ -65,4 +67,4 @@ async def create_check(request: web.Request) -> web.Response:
     if not check:
         raise web.HTTPBadRequest(text=f"Check not created, status: {status}")
 
-    return web.json_response(CheckSchema().dump(dict(check)))
+    return web.Response(text=json.dumps(check, default=str), content_type="application/json")

udata_hydra/routes/checks.py

@@ -1,7 +1,8 @@
 import json
+from typing import Optional
 
 from aiohttp import web
-from marshmallow import ValidationError
+from pydantic import ValidationError
 
 from udata_hydra.db.resource import Resource
 from udata_hydra.schemas import ResourceSchema
@@ -14,11 +15,11 @@ async def get_resource(request: web.Request) -> web.Response:
     If resource is not found, respond with a 404 status code
     """
     [resource_id] = get_request_params(request, params_names=["resource_id"])
-    resource: dict = await Resource.get(resource_id)
+    resource: Optional[dict] = await Resource.get(resource_id)
     if not resource:
         raise web.HTTPNotFound()
 
-    return web.json_response(ResourceSchema().dump(dict(resource)))
+    return web.Response(text=json.dumps(resource, default=str), content_type="application/json")
 
 
 async def create_resource(request: web.Request) -> web.Response:
@@ -29,21 +30,21 @@ async def create_resource(request: web.Request) -> web.Response:
     """
     try:
         payload = await request.json()
-        valid_payload: dict = ResourceSchema().load(payload)
+        valid_payload = ResourceSchema.model_validate(payload)
     except ValidationError as err:
-        raise web.HTTPBadRequest(text=json.dumps(err.messages))
+        raise web.HTTPBadRequest(text=err.json())
 
-    resource: dict = valid_payload["document"]
-    if not resource:
+    document = valid_payload.document
+    if not document:
         raise web.HTTPBadRequest(text="Missing document body")
 
-    dataset_id = valid_payload["dataset_id"]
-    resource_id = valid_payload["resource_id"]
+    dataset_id = valid_payload.dataset_id
+    resource_id = valid_payload.resource_id
 
     await Resource.insert(
         dataset_id=dataset_id,
-        resource_id=resource_id,
-        url=resource["url"],
+        resource_id=str(resource_id),
+        url=document.url,
         priority=True,
     )
 
@@ -58,30 +59,30 @@ async def update_resource(request: web.Request) -> web.Response:
     """
     try:
         payload = await request.json()
-        valid_payload: dict = ResourceSchema().load(payload)
+        valid_payload = ResourceSchema.model_validate(payload)
     except ValidationError as err:
-        raise web.HTTPBadRequest(text=json.dumps(err.messages))
+        raise web.HTTPBadRequest(text=err.json())
 
-    resource: dict = valid_payload["document"]
-    if not resource:
+    document = valid_payload.document
+    if not document:
         raise web.HTTPBadRequest(text="Missing document body")
 
-    dataset_id: str = valid_payload["dataset_id"]
-    resource_id: str = valid_payload["resource_id"]
+    dataset_id: str = valid_payload.dataset_id
+    resource_id: str = str(valid_payload.resource_id)
 
-    await Resource.update_or_insert(dataset_id, resource_id, resource["url"])
+    await Resource.update_or_insert(dataset_id, resource_id, document.url)
 
     return web.json_response({"message": "updated"})
 
 
 async def delete_resource(request: web.Request) -> web.Response:
     try:
         payload = await request.json()
-        valid_payload: dict = ResourceSchema().load(payload)
+        valid_payload = ResourceSchema.model_validate(payload)
     except ValidationError as err:
-        raise web.HTTPBadRequest(text=json.dumps(err.messages))
+        raise web.HTTPBadRequest(text=err.json())
 
-    resource_id: str = valid_payload["resource_id"]
+    resource_id: str = str(valid_payload.resource_id)
 
     pool = request.app["pool"]
     async with pool.acquire() as connection:

udata_hydra/routes/resources.py

@@ -1,2 +1,2 @@
 from .check import CheckSchema
-from .resource import ResourceSchema
+from .resource import ResourceDocumentSchema, ResourceSchema