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,57 @@ +# 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: + _select_cols: Optional[Tuple[str, ...]] = None + _insert_cols_and_metavars: Optional[Tuple[Tuple[str, ...], Tuple[str, ...]]] = None + + @classmethod + def select_columns(cls) -> Tuple[str, ...]: + """Get all the database columns needed for a `select` on this object type""" + if cls._select_cols is None: + columns: List[str] = [] + for field in attr.fields(cls): + columns.append(field.name) + cls._select_cols = tuple(sorted(columns)) + + return cls._select_cols + + @classmethod + def insert_columns_and_metavars(cls) -> Tuple[Tuple[str, ...], Tuple[str, ...]]: + """Get the database columns and metavars needed for an `insert` or `update` on + this object type. + + This supports the following attributes as booleans in the field's metadata: + - primary_key: handled by the database; never inserted or updated + - auto_now_add: handled by the database; set to now() on insertion, never + updated + - auto_now: handled by the client; set to now() on every insertion and update + """ + if cls._insert_cols_and_metavars is None: + zipped_cols_and_metavars: List[Tuple[str, str]] = [] + + for field in attr.fields(cls): + if any( + field.metadata.get(flag) for flag in ("auto_now_add", "primary_key") + ): + continue + elif field.metadata.get("auto_now"): + zipped_cols_and_metavars.append((field.name, "now()")) + else: + zipped_cols_and_metavars.append((field.name, f"%({field.name})s")) + + zipped_cols_and_metavars.sort() + + cols, metavars = zip(*zipped_cols_and_metavars) + cls._insert_cols_and_metavars = cols, metavars + + return cls._insert_cols_and_metavars diff --git a/swh/scheduler/tests/test_model.py b/swh/scheduler/tests/test_model.py new file mode 100644 --- /dev/null +++ b/swh/scheduler/tests/test_model.py @@ -0,0 +1,77 @@ +# 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 + +import datetime + +import attr + +from swh.scheduler import model + + +def test_select_columns(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str) + test1 = attr.ib(type=str) + a_first_attr = attr.ib(type=str) + + @property + def test2(self): + """This property should not show up in the extracted columns""" + return self.test1 + + assert TestModel.select_columns() == ("a_first_attr", "id", "test1") + + +def test_insert_columns(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str) + test1 = attr.ib(type=str) + + @property + def test2(self): + """This property should not show up in the extracted columns""" + return self.test1 + + assert TestModel.insert_columns_and_metavars() == ( + ("id", "test1"), + ("%(id)s", "%(test1)s"), + ) + + +def test_insert_columns_auto_now_add(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str) + test1 = attr.ib(type=str) + added = attr.ib(type=datetime.datetime, metadata={"auto_now_add": True}) + + assert TestModel.insert_columns_and_metavars() == ( + ("id", "test1"), + ("%(id)s", "%(test1)s"), + ) + + +def test_insert_columns_auto_now(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str) + test1 = attr.ib(type=str) + updated = attr.ib(type=datetime.datetime, metadata={"auto_now": True}) + + assert TestModel.insert_columns_and_metavars() == ( + ("id", "test1", "updated"), + ("%(id)s", "%(test1)s", "now()"), + ) + + +def test_insert_columns_primary_key(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str, metadata={"primary_key": True}) + test1 = attr.ib(type=str) + + assert TestModel.insert_columns_and_metavars() == (("test1",), ("%(test1)s",))