python_arango_ogm.db.pao_model

  1from __future__ import annotations
  2
  3import copy
  4import sys
  5from abc import ABC, abstractmethod
  6from enum import StrEnum, auto
  7from functools import partialmethod
  8from typing import Any, Dict, Sequence, Type, Union
  9
 10from python_arango_ogm.db.pao_edges import PAOEdgeDef
 11from python_arango_ogm.db.pao_fields import Field, FloatField
 12from python_arango_ogm.utils import str_util
 13from python_arango_ogm.db.pao_db_base import PAODBBase
 14
 15class LevelEnum(StrEnum):
 16    """
 17    Level Enum, used to specify when schema validation is applied
 18    """
 19
 20    NONE = auto()
 21    NEW = auto()
 22    MODERATE = auto()
 23    STRICT = auto()
 24
 25class PAOModel(ABC):
 26    LEVEL = LevelEnum.STRICT
 27    ADDITIONAL_PROPERTIES = False
 28    SCHEMA_NAME = None
 29    db:PAODBBase = None
 30
 31    def __init__(self):
 32        super().__init__()
 33
 34    @classmethod
 35    def is_field(cls, attribute_name: str) -> bool:
 36        attr = getattr(cls, attribute_name)
 37        return ((not attribute_name.startswith('_')) and
 38                (issubclass(type(attr), Field) or issubclass(type(attr), PAOEdgeDef)))
 39
 40    @classmethod
 41    def all(cls, sort_fields:Dict[str, str], marshall=True) -> str:
 42        records = cls.db.get_by_attributes(cls.collection_name(), sort_key_dict=sort_fields)
 43        return cls.marshall_rows(records) if marshall else records
 44
 45    @classmethod
 46    def insert(cls, attributes:Dict[str, Any]):
 47        # TODO: See if we have timestamp fields (created_at, updated_at) and add accordingly
 48        doc = cls.add_timestamps(attributes, created=True, updated=True)
 49        return cls.db.insert_doc(cls.collection_name(), doc)
 50    @classmethod
 51    def upsert(cls, attributes:Dict[str, Any], insert_attrs:Dict[str, Any], update_attrs:Dict[str, Any]):
 52        insert_doc = cls.add_timestamps(insert_attrs, created=True, updated=True)
 53        update_doc = cls.add_timestamps(update_attrs, updated=True)
 54        return cls.db.upsert_doc(cls.collection_name(), attributes, insert_doc, update_doc)
 55
 56    @classmethod
 57    def find_by_key(cls, key, marshall:bool=True) -> Union[Dict[str, Any], Type[PAOModel]]:
 58        """ Find a single record by given key and return """
 59        record = cls.db.find_by_key(cls.collection_name(), key)
 60        return cls.marshall_row(record) if marshall and record else record
 61
 62    @classmethod
 63    def find_by_attributes(cls, attributes:Dict[str, Any], marshall:bool=True):
 64        """ Find a single record by given attributes and return """
 65        records = cls.db.find_by_attributes(cls.collection_name(), {'_key': key})
 66        return cls.marshall_rows(records) if marshall else records
 67
 68    @classmethod
 69    def get_by_attributes(cls, attributes:Dict[str, Any], sort_keys: Dict[str, str], marshall:bool=True):
 70        """ Find records by given attributes, sorting by sort_keys and return """
 71        cls.db.get_by_attributes(cls.collection_name(), {'_key': key})
 72
 73    @classmethod
 74    def remove_by_key(cls, key):
 75        cls.db.remove_by_key(cls.collection_name(), key)
 76
 77    @classmethod
 78    def get_fields(cls) -> Dict[str, Union[Field, PAOEdgeDef]]:
 79        model_fields = {f:getattr(cls, f) for f in dir(cls) if cls.is_field(f)}
 80        return model_fields
 81
 82    @classmethod
 83    def get_edge_defs(cls) -> Dict[str, PAOEdgeDef]:
 84        return {e:getattr(cls, e) for e in dir(cls) if isinstance(getattr(cls, e), PAOEdgeDef)}
 85
 86    def get_edges(self, edge_def:PAOEdgeDef) -> Sequence[PAOEdgeDef]:
 87        if not self._key:
 88            raise AttributeError("_key field is not present (this model instance hasn't been obtained from a marshalling method).")
 89
 90        return edge_def.associated_edges({"key": self._key})
 91
 92    def insert_edge(self, edge_def:PAOEdgeDef, to: Union[str, PAOModel]) -> Sequence[PAOEdgeDef]:
 93        if not self._key:
 94            raise AttributeError("_key field is not present (this model instance hasn't been obtained from a marshalling method).")
 95
 96        return edge_def.insert_edge({"from": self._key, "to": to})
 97
 98    @classmethod
 99    def collection_name(cls) -> str:
100        if cls.SCHEMA_NAME:
101            coll_name = cls.SCHEMA_NAME
102        else:
103            coll_name = str_util.snake_text(cls.__name__.split('Model')[0])
104        return coll_name
105
106    @classmethod
107    def add_timestamps(cls, field_dict:Dict[str, Any], created:bool=False, updated:bool=False) -> Dict[str, Any]:
108        fields = cls.get_fields()
109        doc = copy.copy(field_dict)
110        if created and 'created_at' in fields:
111            doc['created_at'] = '`DATE_NOW()`'
112        if updated and 'updated_at' in fields:
113            doc['updated_at'] = '`DATE_NOW()`'
114        return doc
115
116    @classmethod
117    def marshall_row(cls, field_dict:Dict[str, Any]) -> Type[PAOModel]:
118        model = PAOModel()
119        for k, v in field_dict.items():
120            setattr(model, k, v)
121
122        # edge_defs = cls.get_edge_defs()
123        # for edge_def in edge_defs.values():
124        #     to_model = edge_def.model_class_to()
125        #     def edge_accessor(self):
126        #         self.get_edges(self.key, edge)
127        #
128        #     # edge_accessor = partialmethod(get_edges, edge_def=edge_def)
129        #
130        #     setattr(model, f"{to_model.collection_name()}_edges", edge_accessor)
131        return model
132
133    @classmethod
134    def marshall_rows(cls, rows:Sequence[Dict[str, Any]]) -> Sequence[Type[PAOModel]]:
135        return [cls.marshall_row(row) for row in rows]