diff --git a/swh/scheduler/model.py b/swh/scheduler/model.py index 5de117d..f225117 100644 --- a/swh/scheduler/model.py +++ b/swh/scheduler/model.py @@ -1,90 +1,125 @@ # 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 from uuid import UUID from typing import Any, Dict, List, Optional, Tuple import attr import attr.converters from attrs_strict import type_validator @attr.s class BaseSchedulerModel: + """Base class for database-backed objects. + + These database-backed objects are defined through attrs-based attributes + that match the columns of the database 1:1. This is a (very) lightweight + ORM. + + These attrs-based attributes have metadata specific to the functionality + expected from these fields in the database: + + - `primary_key`: the column is a primary key; it should be filtered out + when doing an `update` of the object + - `auto_primary_key`: the column is a primary key, which is automatically handled + by the database. It will not be inserted to. This must be matched with a + database-side default value. + - `auto_now_add`: the column is a timestamp that is set to the current time when + the object is inserted, and never updated afterwards. This must be matched with + a database-side default value. + - `auto_now`: the column is a timestamp that is set to the current time when + the object is inserted or updated. + + """ + + _pk_cols: Optional[Tuple[str, ...]] = None _select_cols: Optional[Tuple[str, ...]] = None _insert_cols_and_metavars: Optional[Tuple[Tuple[str, ...], Tuple[str, ...]]] = None + @classmethod + def primary_key_columns(cls) -> Tuple[str, ...]: + """Get the primary key columns for this object type""" + if cls._pk_cols is None: + columns: List[str] = [] + for field in attr.fields(cls): + if any( + field.metadata.get(flag) + for flag in ("auto_primary_key", "primary_key") + ): + columns.append(field.name) + cls._pk_cols = tuple(sorted(columns)) + + return cls._pk_cols + @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 + This implements support for the `auto_*` field metadata attributes. """ 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") + field.metadata.get(flag) + for flag in ("auto_now_add", "auto_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 @attr.s class Lister(BaseSchedulerModel): name = attr.ib(type=str, validator=[type_validator()]) instance_name = attr.ib(type=str, validator=[type_validator()], factory=str) # Populated by database id = attr.ib( type=Optional[UUID], validator=type_validator(), default=None, - metadata={"primary_key": True}, + metadata={"auto_primary_key": True}, ) current_state = attr.ib( type=Dict[str, Any], validator=[type_validator()], factory=dict ) created = attr.ib( type=Optional[datetime.datetime], validator=[type_validator()], default=None, metadata={"auto_now_add": True}, ) updated = attr.ib( type=Optional[datetime.datetime], validator=[type_validator()], default=None, metadata={"auto_now": True}, ) diff --git a/swh/scheduler/tests/test_model.py b/swh/scheduler/tests/test_model.py index 26a0129..47bb618 100644 --- a/swh/scheduler/tests/test_model.py +++ b/swh/scheduler/tests/test_model.py @@ -1,77 +1,94 @@ # 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}) + id = attr.ib(type=str, metadata={"auto_primary_key": True}) test1 = attr.ib(type=str) assert TestModel.insert_columns_and_metavars() == (("test1",), ("%(test1)s",)) + + +def test_insert_primary_key(): + @attr.s + class TestModel(model.BaseSchedulerModel): + id = attr.ib(type=str, metadata={"auto_primary_key": True}) + test1 = attr.ib(type=str) + + assert TestModel.primary_key_columns() == ("id",) + + @attr.s + class TestModel2(model.BaseSchedulerModel): + col1 = attr.ib(type=str, metadata={"primary_key": True}) + col2 = attr.ib(type=str, metadata={"primary_key": True}) + test1 = attr.ib(type=str) + + assert TestModel2.primary_key_columns() == ("col1", "col2")