diff --git a/requirements.txt b/requirements.txt --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ # dependency lines, see https://pip.readthedocs.org/en/1.1/requirements.html arrow +attrs celery >= 4 # Not a direct dependency: missing from celery 4.4.4 future >= 0.18.0 diff --git a/swh/scheduler/api/client.py b/swh/scheduler/api/client.py --- a/swh/scheduler/api/client.py +++ b/swh/scheduler/api/client.py @@ -6,6 +6,7 @@ from swh.core.api import RPCClient +from .serializers import ENCODERS, DECODERS from ..interface import SchedulerInterface @@ -15,3 +16,6 @@ """ backend_class = SchedulerInterface + + extra_type_decoders = DECODERS + extra_type_encoders = ENCODERS diff --git a/swh/scheduler/api/serializers.py b/swh/scheduler/api/serializers.py new file mode 100644 --- /dev/null +++ b/swh/scheduler/api/serializers.py @@ -0,0 +1,28 @@ +# Copyright (C) 2020 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +"""Decoder and encoders for swh.scheduler.model objects.""" + +from typing import Callable, Dict, List, Tuple + +import attr + +import swh.scheduler.model as model + + +def _encode_model_object(obj): + d = attr.asdict(obj) + d["__type__"] = type(obj).__name__ + return d + + +ENCODERS: List[Tuple[type, str, Callable]] = [ + (model.BaseSchedulerModel, "scheduler_model", _encode_model_object), +] + + +DECODERS: Dict[str, Callable] = { + "scheduler_model": lambda d: getattr(model, d.pop("__type__"))(**d) +} diff --git a/swh/scheduler/api/server.py b/swh/scheduler/api/server.py --- a/swh/scheduler/api/server.py +++ b/swh/scheduler/api/server.py @@ -13,6 +13,7 @@ from swh.scheduler import get_scheduler from swh.scheduler.interface import SchedulerInterface +from .serializers import ENCODERS, DECODERS scheduler = None @@ -26,7 +27,8 @@ class SchedulerServerApp(RPCServerApp): - pass + extra_type_decoders = DECODERS + extra_type_encoders = ENCODERS app = SchedulerServerApp( diff --git a/swh/scheduler/model.py b/swh/scheduler/model.py new file mode 100644 --- /dev/null +++ b/swh/scheduler/model.py @@ -0,0 +1,26 @@ +# Copyright (C) 2020 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +from typing import List, Optional, Tuple + +import attr +import attr.converters + + +@attr.s +class BaseSchedulerModel: + _columns: Optional[List[str]] = None + + @classmethod + def columns(cls) -> List[str]: + if cls._columns is None: + columns: List[Tuple[str, int]] = [] + for field in attr.fields(cls): + columns.append( + (field.name, 0 if field.metadata.get("primary_key") else 1) + ) + cls._columns = [c[0] for c in sorted(columns, key=lambda c: (c[1], c[0]))] + + return cls._columns