grung-db/src/grung/db.py

import inspect

from tinydb import TinyDB, table
from tinydb.table import Document

from grung.types import Record


class RecordTable(table.Table):
    """
    Wrapper around tinydb Tables that handles Records instead of dicts.
    """

    def __init__(self, storage, name, document_class: Document = Record, **kwargs):
        self.document_class = document_class
        super().__init__(storage, name, **kwargs)

    def insert(self, document):
        self._satisfy_constraints(document)
        if document.doc_id:
            last_insert_id = super().upsert(document)[0]
        else:
            last_insert_id = super().insert(dict(document))
        return self.get(doc_id=last_insert_id)

    def remove(self, document):
        if document.doc_id:
            super().remove(doc_ids=[document.doc_id])

    def _satisfy_constraints(self, document):
        # check for uniqueness, etc.
        pass


class GrungDB(TinyDB):
    """
    A TinyDB database instance that uses RecordTable instances for each table
    and Record instances for each document in the table.
    """

    default_table_name = "Record"
    _tables = {}

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.create_table(Record)

    def table(self, name: str) -> RecordTable:
        if name not in self._tables:
            raise RuntimeError(f"No such table: {name}")
        return self._tables[name]

    def create_table(self, table_class):
        name = table_class.__name__
        if name not in self._tables:
            self._tables[name] = RecordTable(self.storage, name, document_class=table_class)
        return self.table(name)

    def save(self, record):
        """
        Create or update a record in its table.
        """
        return self.table(record._metadata.table).insert(record)

    def delete(self, record):
        return self.table(record._metadata.table).remove(record)

    def __getattr__(self, attr_name):
        """
        Make tables attributes of the instance.
        """
        if attr_name in self._tables:
            return self.table(attr_name)
        return super().__getattr__(attr_name)

    @classmethod
    def with_schema(cls, schema_module, *args, **kwargs):
        db = GrungDB(*args, **kwargs)
        for name, obj in inspect.getmembers(schema_module):
            if type(obj) == type and issubclass(obj, Record):
                db.create_table(obj)
        return db