capabilities rework

api/kiwi_vpn_api/db_new/capabilities.py

@@ -13,9 +13,11 @@ class Capability(Enum):
     issue = "issue"
     renew = "renew"
 
+    def __repr__(self) -> str:
+        return self.value
+
 
 class UserCapabilityBase(SQLModel):
-    user_name: str = Field(primary_key=True, foreign_key="user.name")
     capability_name: str = Field(primary_key=True)
 
     @property
@@ -27,6 +29,8 @@ class UserCapabilityBase(SQLModel):
 
 
 class UserCapability(UserCapabilityBase, table=True):
+    user_name: str = Field(primary_key=True, foreign_key="user.name")
+
     user: "User" = Relationship(
-        back_populates="capabilities"
+        back_populates="capabilities",
     )

api/kiwi_vpn_api/db_new/user.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import Any
 
 from pydantic import root_validator
 from sqlalchemy.exc import IntegrityError
@@ -109,34 +109,19 @@ class User(UserBase, table=True):
             db.delete(self)
             db.commit()
 
-    def extend_capabilities(self, capabilities: Sequence[Capability]) -> None:
-        """
-        Extend this user's capabilities
-        """
+    def get_capabilities(self) -> set[Capability]:
+        return set(
+            capability._
+            for capability in self.capabilities
+        )
 
-        for capability in capabilities:
-            user_capability = UserCapability(
+    def set_capabilities(self, capabilities: set[Capability]) -> None:
+        self.capabilities = [
+            UserCapability(
                 user_name=self.name,
                 capability_name=capability.value,
-            )
-
-            if user_capability not in self.capabilities:
-                self.capabilities.append(user_capability)
-
-    def remove_capabilities(self, capabilities: Sequence[Capability]) -> None:
-        """
-        Remove from this user's capabilities
-        """
-
-        for capability in capabilities:
-            try:
-                self.capabilities.remove(UserCapability(
-                    user_name=self.name,
-                    capability_name=capability.value,
-                ))
-
-            except ValueError:
-                pass
+            ) for capability in capabilities
+        ]
 
 
 class UserCreate(UserBase):