feat: Add django-guardin permissions

server/account/management/commands/ensuresuperuser.py

@@ -1,6 +1,7 @@
 from django.core.management.base import BaseCommand
 
 from account.models import User
+from account.services import UserService
 
 
 class Command(BaseCommand):
@@ -14,4 +15,4 @@ def add_arguments(self, parser):
 
     def handle(self, *args, **options):
         if not User.objects.exists():
-            User.objects.create_superuser(email=options["email"], password=options["password"])
+            UserService.create_superuser(email=options["email"], password=options["password"])

server/account/models.py

@@ -4,32 +4,22 @@
 
 
 class UserManager(BaseUserManager):
-    def _create_user(self, email, password=None, **extra_fields):
-        if not email:
-            raise ValueError("The email field must be set")
-
-        email = self.normalize_email(email)
-        user = self.model(email=email, **extra_fields)
-        if extra_fields.get("is_service_account", False):
-            user.set_unusable_password()
-        else:
-            user.set_password(password)
-        user.save(using=self._db)
-        return user
-
-    def create_user(self, username=None, email=None, password=None, **extra_fields):
-        extra_fields.setdefault("is_staff", False)
-        extra_fields.setdefault("is_superuser", False)
-        return self._create_user(email, password, **extra_fields)
-
-    def create_superuser(self, username=None, email=None, password=None, **extra_fields):
-        extra_fields.setdefault("is_staff", True)
-        extra_fields.setdefault("is_superuser", True)
-        return self._create_user(email, password, **extra_fields)
-
-    def create_service_account(self, email, **extra_fields):
-        extra_fields.setdefault("is_service_account", True)
-        return self._create_user(email, password=None, **extra_fields)
+    def create_user(self, username=None, email=None, password=None, role=None, **extra_fields):
+        from account.services.user import UserService
+
+        return UserService.create_user(username=username, email=email, password=password, role=role, **extra_fields)
+
+    def create_superuser(self, username=None, email=None, password=None, role=None, **extra_fields):
+        from account.services.user import UserService
+
+        return UserService.create_superuser(
+            username=username, email=email, password=password, role=role, **extra_fields
+        )
+
+    def create_service_account(self, email, role=None, **extra_fields):
+        from account.services.user import UserService
+
+        return UserService.create_service_account(email=email, role=role, **extra_fields)
 
 
 class User(AbstractUser):

server/account/pipelines.py

@@ -2,6 +2,7 @@
 
 from django.conf import settings
 from social_core.exceptions import AuthForbidden
+
 from utils.functions import import_from_string
 
 logger = logging.getLogger(__name__)

server/account/serializers.py

@@ -2,7 +2,7 @@
 
 from account.models import User
 
-from .services import ServiceAccountNameService
+from .services import ServiceAccountNameService, UserService
 
 
 class AccountSerializer(serializers.Serializer):
@@ -40,7 +40,7 @@ def create(self, validated_data):
         is_active = validated_data.get("is_active")
         is_staff = validated_data.get("is_staff")
         dataset_ids = validated_data.get("data_set_ids", [])
-        service_account = User.objects.create_service_account(email=email, is_active=is_active, is_staff=is_staff)
+        service_account = UserService.create_service_account(email=email, is_active=is_active, is_staff=is_staff)
         service_account.data_sets.add(*dataset_ids)
         return service_account
 

server/account/services/__init__.py

@@ -1,4 +1,5 @@
 from .service_accounts import ServiceAccountNameService
 from .sso_provider import SSOProviderService
+from .user import UserService
 
-__all__ = ["ServiceAccountNameService", "SSOProviderService"]
+__all__ = ["ServiceAccountNameService", "SSOProviderService", "UserService"]

server/account/services/user.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Optional, Type
+
+from django.contrib.auth import get_user_model
+from django.db import models
+from guardian.shortcuts import assign_perm, remove_perm
+
+from catalog.utils import AdminRole, BaseRole, UserRole
+
+User = get_user_model()
+logger = logging.getLogger(__name__)
+
+
+class UserService:
+    @staticmethod
+    def _create_user_base(email, password=None, **extra_fields):
+        if not email:
+            raise ValueError("The email field must be set")
+
+        from django.contrib.auth.base_user import BaseUserManager
+
+        manager = BaseUserManager()
+        email = manager.normalize_email(email)
+        user = User(email=email, **extra_fields)
+
+        if extra_fields.get("is_service_account", False):
+            user.set_unusable_password()
+        else:
+            if password:
+                user.set_password(password)
+
+        user.save()
+        return user
+
+    @staticmethod
+    def create_user(
+        username=None,
+        email=None,
+        password=None,
+        role: Type[BaseRole] = UserRole,
+        is_staff: bool = False,
+        is_superuser: bool = False,
+        **extra_fields,
+    ) -> User:
+        extra_fields.setdefault("is_staff", is_staff)
+        extra_fields.setdefault("is_superuser", is_superuser)
+
+        user = UserService._create_user_base(email, password, **extra_fields)
+        UserService.assign_role(user, role)
+
+        return user
+
+    @staticmethod
+    def create_superuser(
+        username=None, email=None, password=None, role: Type[BaseRole] = AdminRole, **extra_fields
+    ) -> User:
+        extra_fields.setdefault("is_staff", True)
+        extra_fields.setdefault("is_superuser", True)
+
+        user = UserService._create_user_base(email, password, **extra_fields)
+
+        UserService.assign_role(user, role)
+
+        return user
+
+    @staticmethod
+    def create_service_account(email, role: Optional[Type[BaseRole]] = None, **extra_fields) -> User:
+        extra_fields.setdefault("is_service_account", True)
+
+        user = UserService._create_user_base(email, password=None, **extra_fields)
+
+        if role is None:
+            if not extra_fields.get("is_staff", False):
+                role = AdminRole
+            else:
+                role = UserRole
+
+        if role is not None:
+            UserService.assign_role(user, role)
+
+        return user
+
+    @staticmethod
+    def assign_role(user: User, role: Type[BaseRole], obj: Optional[models.Model] = None) -> None:
+        try:
+            role.assign(user, obj)
+        except Exception as e:
+            logger.error(f"Failed to assign {role.__name__} role to user {user.id}: {str(e)}", exc_info=True)
+            raise
+
+    @staticmethod
+    def assign_permission(user: User, permission: str, obj: Optional[models.Model] = None) -> None:
+        try:
+            assign_perm(permission, user, obj)
+        except Exception as e:
+            logger.error(f"Failed to assign permission '{permission}' to user {user.id}: {str(e)}", exc_info=True)
+            raise
+
+    @staticmethod
+    def remove_permission(user: User, permission: str, obj: Optional[models.Model] = None) -> None:
+        try:
+            remove_perm(permission, user, obj)
+        except Exception as e:
+            logger.error(f"Failed to remove permission '{permission}' from user {user.id}: {str(e)}", exc_info=True)
+            raise
+
+    @staticmethod
+    def assign_role_to_object(user: User, role: Type[BaseRole], obj: models.Model) -> None:
+        UserService.assign_role(user, role, obj)