From 85f49824cac3b6083bf50fe43ac08f27eecbc83c Mon Sep 17 00:00:00 2001
From: Mo8it <mo8it@proton.me>
Date: Mon, 9 May 2022 02:13:18 +0200
Subject: [PATCH] Add UserHasRole filter

---
 advlabdb/adminModelViews.py | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/advlabdb/adminModelViews.py b/advlabdb/adminModelViews.py
index 4ce372b..7eaef29 100644
--- a/advlabdb/adminModelViews.py
+++ b/advlabdb/adminModelViews.py
@@ -6,7 +6,7 @@ import numpy as np
 from flask import flash, has_request_context, redirect, request, url_for
 from flask_admin import expose
 from flask_admin.contrib.sqla.fields import QuerySelectField, QuerySelectMultipleField
-from flask_admin.contrib.sqla.filters import FilterEqual
+from flask_admin.contrib.sqla.filters import BooleanEqualFilter, FilterEqual
 from flask_admin.helpers import get_form_data
 from flask_admin.menu import MenuLink
 from flask_admin.model.template import EndpointLinkRowAction
@@ -14,7 +14,7 @@ from flask_security import admin_change_password, current_user, hash_password
 from flask_wtf import FlaskForm
 from flask_wtf.file import FileAllowed, FileField, FileRequired
 from matplotlib.figure import Figure
-from sqlalchemy import and_, func, or_
+from sqlalchemy import and_, not_, or_
 from werkzeug.utils import secure_filename
 from wtforms import Form
 from wtforms.fields import (
@@ -258,13 +258,14 @@ class UserView(SecureAdminModelView):
                     model, password, notify=False
                 )  # Password is automatically hashed with this function
 
-        if model.has_role("assistant") and not model.assistant:
+        if model.has_role("assistant") and model.assistant is None:
             semester_experiments = form.semester_experiments.data if form.semester_experiments else []
 
             assistant = Assistant(user=model, semester_experiments=semester_experiments)
             self.session.add(assistant)
 
-        if model.has_role("admin") and not model.admin:
+        if model.has_role("admin") and model.admin is None:
+            flash("Admin role was assigned!", "danger")
             admin = Admin(user=model)
             self.session.add(admin)
 
@@ -826,6 +827,17 @@ class SemesterExperimentView(SecureAdminModelView):
         return updateSuccessful
 
 
+def userHasRoleFilterFactory(column, role_name):
+    class UserHasRoleFilter(BooleanEqualFilter):
+        def apply(self, query, value, alias=None):
+            if bool(int(value)):  # value is string "0" or "1"
+                return query.filter(self.column.user.has(User.roles.any(Role.name == role_name)))
+            else:
+                return query.filter(self.column.user.has(not_(User.roles.any(Role.name == role_name))))
+
+    return UserHasRoleFilter(column, f"Currently has {role_name} role")
+
+
 class AssistantView(SecureAdminModelView):
     class UserEndpointLinkRowAction(CustomIdEndpointLinkRowAction):
         def customId(self, row):
@@ -854,7 +866,10 @@ class AssistantView(SecureAdminModelView):
         "user.last_name",
         "user.email",
     ]
-    column_filters = ("user.active",)
+    column_filters = (
+        userHasRoleFilterFactory(Assistant, "assistant"),
+        "user.active",
+    )
     form_excluded_columns = [
         "experiment_marks",
         "appointments",
@@ -900,7 +915,10 @@ class AdminView(SecureAdminModelView):
         "user.first_name",
         "user.last_name",
     ]
-    column_filters = ("user.active",)
+    column_filters = (
+        userHasRoleFilterFactory(Admin, "admin"),
+        "user.active",
+    )
 
     column_extra_row_actions = [
         UserEndpointLinkRowAction(