AdvLabDB/advlabdb/customClasses.py

from flask import flash, redirect, request, url_for
from flask_admin import AdminIndexView
from flask_admin.contrib.sqla import ModelView
from flask_security import current_user

from advlabdb.exceptions import DataBaseException, ModelViewException


def adminViewIsAccessible():
    return current_user.has_role("admin")


def assistantViewIsAccessible():
    return current_user.has_role("assistant")


class CustomIndexView(AdminIndexView):
    def inaccessible_callback(self, name, **kwargs):
        # Redirect to login page if user doesn't have access
        return redirect(url_for("security.login", next=request.url))


class SecureAdminIndexView(CustomIndexView):
    def is_accessible(self):
        return adminViewIsAccessible()


class SecureAssistantIndexView(CustomIndexView):
    def is_accessible(self):
        return assistantViewIsAccessible()


class CustomModelView(ModelView):
    can_export = True
    can_set_page_size = True

    create_modal = True
    edit_modal = True
    details_modal = True

    queryFilter = None
    customCreateModel = None

    def inaccessible_callback(self, name, **kwargs):
        # Redirect to login page if user doesn't have access
        return redirect(url_for("security.login", next=request.url))

    def get_query(self):
        if self.queryFilter:
            return super().get_query().filter(self.queryFilter())
        else:
            return super().get_query()

    def get_count_query(self):
        if self.queryFilter:
            return super().get_count_query().filter(self.queryFilter())
        else:
            return super().get_count_query()

    def handle_view_exception(self, exc):
        if type(exc) in (ModelViewException, DataBaseException):
            flash(str(exc), "error")
            return True

        return super().handle_view_exception(exc)

    def create_model(self, form):
        if not self.customCreateModel:
            return super().create_model(form)
        else:
            try:
                model = self.customCreateModel(form)

                self.session.add(model)

                self.on_model_change(form, model, True)

                self.session.commit()
            except Exception as ex:
                flash(str(ex), "error")

                self.session.rollback()
            else:
                self.after_model_change(form, model, True)

                return model


class SecureAdminModelView(CustomModelView):
    list_template = "admin_list.html"
    create_template = "admin_create.html"
    edit_template = "admin_edit.html"

    def is_accessible(self):
        return adminViewIsAccessible()


class SecureAssistantModelView(CustomModelView):
    list_template = "assistant_list.html"
    create_template = "assistant_create.html"
    edit_template = "assistant_edit.html"

    def is_accessible(self):
        return assistantViewIsAccessible()