PyCS/pycs/database/base.py

from __future__ import annotations

import typing as T

from flask import abort
from sqlalchemy_serializer import SerializerMixin

from pycs import db
from pycs.database.util import commit_on_return

class BaseModel(db.Model, SerializerMixin):
    """ Base model class """
    __abstract__ = True

    # setup of the SerializerMixin
    date_format = '%s'  # Unixtimestamp (seconds)
    datetime_format = '%d. %b. %Y %H:%M:%S'
    time_format = '%H:%M'


    id = db.Column(db.Integer, primary_key=True)

    serialize_only = ("id",)

    def __repr__(self):
        attrs = self.serialize()
        content = ", ".join([f"{attr}={value}" for attr, value in attrs.items()])
        return f"<{self.__class__.__name__}: {content}>"


    def serialize(self) -> dict:
        """ default model serialize method. adds identifier as alias for id """
        res = self.to_dict()
        res["identifier"] = self.id
        return res


    @commit_on_return
    def delete(self) -> dict:
        """
        delete this instance from the database

        :return: serialized self
        """
        db.session.delete(self)
        dump = self.serialize()

        return dump


    # do an alias
    remove = delete


    @classmethod
    def new(cls, commit: bool = True, **kwargs):
        """ creates a new object. optionally commits the created object. """
        obj = cls(**kwargs)
        db.session.add(obj)

        if commit:
            obj.commit()

        return obj

    @classmethod
    def get_or_create(cls, **kwargs) -> T.Tuple[BaseModel, bool]:
        """ get an object from the DB based on the kwargs, or create an object with these. """

        is_new = False

        obj = cls.query.filter_by(**kwargs).one_or_none()

        if obj is None:
            obj = cls.new(commit=False, **kwargs)
            is_new = True

        return obj, is_new


    @classmethod
    def get_or_404(cls, obj_id: int) -> BaseModel:
        """ get an object for the given id or raise 404 error if the object is not present """
        obj = cls.query.get(obj_id)

        if obj is None:
            abort(404, f"{cls.__name__} with ID {obj_id} could not be found!")

        return obj


    @staticmethod
    def commit():
        """ commit current session """
        db.session.commit()

    @staticmethod
    def flush():
        """ flush current session """
        db.session.flush()


class NamedBaseModel(BaseModel):
    """ Extends the base model with a name attribute. """
    __abstract__ = True

    name = db.Column(db.String, nullable=False)

    serialize_only = BaseModel.serialize_only + ("name",)

    @commit_on_return
    def set_name(self, name: str):
        """ set the name attribute """
        self.name = name