AdvLabDB/advlabdb/customClasses.py

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

from advlabdb.exceptions import ModelViewException, DataBaseException


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


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

    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 SecureModelView(ModelView):
    can_export = True
    can_set_page_size = True

    create_modal = True
    edit_modal = True
    details_modal = True

    list_template = "admin_list.html"
    create_template = "admin_create.html"
    edit_template = "admin_edit.html"

    queryFilter = None

    def is_accessible(self):
        return adminViewIsAccessible()

    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)