st-ten-1/src/lib/db/crud_db.py

from peewee import TextField
from playhouse.shortcuts import model_to_dict

from . import db, models_reference


class Crud_DB:
    def __init__(self, table_name, select=None, filters=None, pagination=250):
        self.table_model = models_reference[table_name]
        self.table_pk = self.table_model._meta.primary_key
        self.table_fields = list(self.table_model._meta.fields.keys())
        self.pagination = pagination
        self.default_filters = {}
        self.sorting = {}
        self.filters = {}
        if filters is None:
            filters = {}
        for column_name, filter in filters.items():
            self.filter(column_name, filter, filter_storage=self.default_filters)

    @db.connection_context()
    @db.atomic()
    def commit(self, data, deleted_rows=None):
        if deleted_rows is not None:
            if hasattr(self.table_model, "crud_delete"):
                deleted = self.table_model.crud_delete(deleted_rows)
            else:
                deleted = self.table_model.delete().where(self.table_pk << deleted_rows).execute()
            if deleted != len(deleted_rows):
                raise AssertionError(f"deleted {deleted} rows instead of the expected {len(deleted_rows)}")
        # SQLITE DOES NOT SUPPORT UPDATE, ONLY REPLACE
        complete_data = []
        for rn, row in enumerate(data):
            pk = row[self.table_pk.name]
            try:
                if pk is not None:
                    raise self.table_model.DoesNotExist()
                db_row.update(model_to_dict(self.table_model.get_by_id(pk)))
            except self.table_model.DoesNotExist:
                db_row = row
            complete_data.append(db_row)
        if hasattr(self.table_model, "crud_update"):
            self.table_model.crud_update(complete_data)
        else:
            self.table_model.insert_many(complete_data).on_conflict_replace().execute()

    def revert(self):
        self.sorting.clear()
        self.filters.clear()

    @db.connection_context()
    def get(self, page=0):
        if self.table_model._meta.composite_key:
            raise NotImplementedError("need to implement composite_key")
        if hasattr(self.table_model, "crud_select"):
            query = self.table_model.crud_select()
        else:
            query = self.table_model.select()
        if len(self.default_filters) + len(self.filters) > 0:
            query = query.where(*self.default_filters.values(), *self.filters.values())
        if len(self.sorting) > 0:
            query = query.order_by(*self.sorting.values())
        total_count = query.count()
        last_page = int(total_count / self.pagination)
        if page < 0:
            page = last_page - page + 1
        page = min(max(page, 0), last_page)
        self.page = page
        if self.pagination is not False:
            query = query.paginate(self.page + 1, self.pagination)
        self.data = query.dicts().execute()
        self.shown_keys = [r[self.table_pk.name] for r in self.data]
        return self.data, total_count, self.page, last_page

    def sort(self, column_name, is_ascending=None, sorting_storage=None):
        if sorting_storage is None:
            sorting_storage = self.sorting
        if column_name is not None:
            if is_ascending is None:
                sorting_storage.pop(column_name, None)
            else:
                field = self.table_model._meta.fields[column_name]
                if not is_ascending:
                    field = -field
                sorting_storage[column_name] = field
        else:
            sorting_storage.clear()

    def filter(self, column_name, filter, filter_storage=None):  # only contains/equality filter
        if filter_storage is None:
            filter_storage = self.filters
        if filter is None or filter == "":
            if column_name in self.filters:
                filter_storage.pop(column_name, None)
        else:
            field = self.table_model._meta.fields[column_name]
            if type(field) is TextField:
                filter_storage[column_name] = field.contains(filter)
            else:
                filter_storage[column_name] = field == filter