grung-db/src/grung/types.py

from __future__ import annotations

from collections import namedtuple
from dataclasses import dataclass, field
from typing import Dict, List

import nanoid
from tinydb import TinyDB, where

from grung.exceptions import PointerReferenceError

Metadata = namedtuple("Metadata", ["table", "fields", "backrefs"])


@dataclass
class Field:
    """
    Represents a single field in a Record.
    """

    name: str
    value_type: type = str
    default: str = None
    unique: bool = False

    def before_insert(self, value: value_type, db: TinyDB, record: Record) -> None:
        pass

    def after_insert(self, db: TinyDB, record: Record) -> None:
        pass

    def serialize(self, value: value_type) -> str:
        if value is not None:
            return str(value)

    def deserialize(self, value: value_type, db: TinyDB, recurse: bool = True) -> value_type:
        return value


@dataclass
class Integer(Field):
    value_type = int
    default: int = 0

    def deserialize(self, value: str, db: TinyDB, recurse: bool = True) -> value_type:
        return int(value)


class Record(Dict[(str, Field)]):
    """
    Base type for a single database record.
    """

    def __init__(self, raw_doc: dict = {}, doc_id: int = None, **params):
        self.doc_id = doc_id
        fields = self.__class__.fields()
        self._metadata = Metadata(
            table=self.__class__.__name__,
            fields={f.name: f for f in fields},
            backrefs=lambda value_type: (
                field for field in fields if type(field) == BackReference and field.value_type == value_type
            ),
        )
        super().__init__(dict({field.name: field.default for field in fields}, **raw_doc, **params))

    @classmethod
    def fields(self):
        return [
            # 1% collision rate at ~2M records
            Field("uid", default=nanoid.generate(size=8), unique=True)
        ]

    def serialize(self):
        """
        Serialie every field on the record
        """
        rec = {}
        for name, _field in self._metadata.fields.items():
            rec[name] = _field.serialize(self[name])
        return self.__class__(rec, doc_id=self.doc_id)

    def deserialize(self, db, recurse: bool = True):
        """
        Deserialize every field on the record
        """
        rec = {}
        for name, _field in self._metadata.fields.items():
            rec[name] = _field.deserialize(self[name], db, recurse=recurse)
        return self.__class__(rec, doc_id=self.doc_id)

    def before_insert(self, db: TinyDB) -> None:
        for name, _field in self._metadata.fields.items():
            _field.before_insert(self[name], db, self)

    def after_insert(self, db: TinyDB) -> None:
        for name, _field in self._metadata.fields.items():
            _field.after_insert(db, self)

    def __setattr__(self, key, value):
        if key in self:
            self[key] = value
        super().__setattr__(key, value)

    def __getattr__(self, attr_name):
        if attr_name in self:
            return self.get(attr_name)
        return super().__getattr__(attr_name)

    def __hash__(self):
        return hash(str(dict(self)))

    def __repr__(self):
        return (
            f"{self.__class__.__name__}[{self.doc_id}]("
            + ", ".join([f"{key}={val}" for (key, val) in self.items()])
            + ")"
        )


@dataclass
class Pointer(Field):
    """
    Store a string reference to a record.
    """

    name: str = ""
    value_type: type = Record

    def serialize(self, value: value_type) -> str:
        return Pointer.reference(value)

    def deserialize(self, value: str, db: TinyDB, recurse: bool = True) -> value_type:
        return Pointer.dereference(value, db, recurse)

    @classmethod
    def reference(cls, value: Record):
        if value:
            if not value.doc_id:
                raise PointerReferenceError(value)
            return f"{value._metadata.table}::{value.uid}"
        return None

    @classmethod
    def dereference(cls, value: str, db: TinyDB, recurse: bool = True):
        if not value:
            return
        elif type(value) == str:
            pt, puid = value.split("::")
            if puid:
                return db.table(pt).search(where("uid") == puid, recurse=recurse)[0]
        return value


@dataclass
class BackReference(Pointer):
    pass


@dataclass
class Collection(Field):
    """
    A collection of pointers.
    """

    value_type: type = Record
    default: List[value_type] = field(default_factory=lambda: [])

    def _pointer(self, rec):
        return Pointer(value_type=type(rec))

    def serialize(self, values: List[value_type]) -> List[str]:
        return [self._pointer(val).serialize(val) for val in values]

    def deserialize(self, values: List[str], db: TinyDB, recurse: bool = True) -> List[value_type]:
        """
        Recursively deserialize the objects in this collection
        """
        recs = []
        if not recurse:
            return values
        for val in values:
            recs.append(self._pointer(val).deserialize(val, db=db, recurse=recurse))
        return recs

    def after_insert(self, db: TinyDB, record: Record) -> None:
        """
        Populate any backreferences in the members of this collection with the parent record's uid.
        """
        if not record[self.name]:
            return

        for member in record[self.name]:
            target = Pointer.dereference(member, db=db)
            for backref in target._metadata.backrefs(type(record)):
                target[backref.name] = Pointer.reference(record)
                db.table(target._metadata.table).update({backref.name: record}, where("uid") == target.uid)