refactor: Add users_to_add and users_to_remove fields to replace user update mode

changes/1705.fix.md

@@ -0,0 +1 @@
+Refactor group mutation by adding `added_users` and `removed_users` fields to replace `user_update_mode`.

src/ai/backend/manager/models/group.py

@@ -46,7 +46,7 @@
     simple_db_mutate_returning_item,
 )
 from .storage import StorageSessionManager
-from .user import ModifyUserInput, UserRole
+from .user import UserRole
 from .utils import ExtendedAsyncSAEngine, execute_with_retry
 
 if TYPE_CHECKING:
@@ -393,8 +393,27 @@ class ModifyGroupInput(graphene.InputObjectType):
     is_active = graphene.Boolean(required=False)
     domain_name = graphene.String(required=False)
     total_resource_slots = graphene.JSONString(required=False)
-    user_update_mode = graphene.String(required=False)
-    user_uuids = graphene.List(lambda: graphene.String, required=False)
+    user_update_mode = graphene.String(
+        deprecation_reason=(
+            "Deprecated since 24.03.0. Recommend to use `added_users` and `removed_users` fields"
+        )
+    )
+    user_uuids = graphene.List(
+        lambda: graphene.String,
+        deprecation_reason=(
+            "Deprecated since 24.03.0. Recommend to use `added_users` and `removed_users` fields"
+        ),
+    )
+    added_users = graphene.List(
+        lambda: graphene.String,
+        required=False,
+        description="Added since 24.03.0. ID array of the users to be added to the group.",
+    )
+    removed_users = graphene.List(
+        lambda: graphene.String,
+        required=False,
+        description="Added since 24.03.0. ID array of the users to be removed from the group.",
+    )
     allowed_vfolder_hosts = graphene.JSONString(required=False)
     integration_id = graphene.String(required=False)
     resource_policy = graphene.String(required=False)
@@ -467,10 +486,11 @@ async def mutate(
         root,
         info: graphene.ResolveInfo,
         gid: uuid.UUID,
-        props: ModifyUserInput,
+        props: ModifyGroupInput,
     ) -> ModifyGroup:
         graph_ctx: GraphQueryContext = info.context
         data: Dict[str, Any] = {}
+        user_data: dict[str, set] = {}
         set_if_set(props, data, "name")
         set_if_set(props, data, "description")
         set_if_set(props, data, "is_active")
@@ -485,33 +505,54 @@ async def mutate(
         set_if_set(props, data, "integration_id")
         set_if_set(props, data, "resource_policy")
 
+        set_if_set(props, user_data, "added_users", clean_func=set)
+        set_if_set(props, user_data, "removed_users", clean_func=set)
+
         if "name" in data and _rx_slug.search(data["name"]) is None:
             raise ValueError("invalid name format. slug format required.")
         if props.user_update_mode not in (None, Undefined, "add", "remove"):
             raise ValueError("invalid user_update_mode")
         if not props.user_uuids:
             props.user_update_mode = None
-        if not data and props.user_update_mode is None:
+        if not data and not user_data and props.user_update_mode in (None, Undefined):
             return cls(ok=False, msg="nothing to update", group=None)
 
         async def _do_mutate() -> ModifyGroup:
-            async with graph_ctx.db.begin() as conn:
-                # TODO: refactor user addition/removal in groups as separate mutations
-                #       (to apply since 21.09)
+            async with graph_ctx.db.begin_session() as db_session:
+                # Using `user_update_mode` and `user_uuids` is deprecated
                 if props.user_update_mode == "add":
                     values = [{"user_id": uuid, "group_id": gid} for uuid in props.user_uuids]
-                    await conn.execute(
+                    await db_session.execute(
                         sa.insert(association_groups_users).values(values),
                     )
                 elif props.user_update_mode == "remove":
-                    await conn.execute(
+                    await db_session.execute(
                         sa.delete(association_groups_users).where(
                             (association_groups_users.c.user_id.in_(props.user_uuids))
                             & (association_groups_users.c.group_id == gid),
                         ),
                     )
+
+                added_users = user_data.get("added_users") or set()
+                removed_users = user_data.get("removed_users") or set()
+                if union := (added_users & removed_users):
+                    raise ValueError(
+                        "Should be no duplicate user id in `added_users` and `removed_users`."
+                        f" (ids: {list(union)})"
+                    )
+                if added_users:
+                    values = [{"user_id": uuid, "group_id": gid} for uuid in added_users]
+                    await db_session.execute(
+                        sa.insert(association_groups_users).values(values),
+                    )
+                if removed_users:
+                    values = [{"user_id": uuid, "group_id": gid} for uuid in removed_users]
+                    await db_session.execute(
+                        sa.insert(association_groups_users).values(values),
+                    )
+
                 if data:
-                    result = await conn.execute(
+                    result = await db_session.execute(
                         sa.update(groups).values(data).where(groups.c.id == gid).returning(groups),
                     )
                     if result.rowcount > 0:

src/ai/backend/manager/models/resource_policy.py

@@ -28,7 +28,6 @@
 )
 from .keypair import keypairs
 from .user import UserRole
-from .utils import deprecation_reason_msg, description_msg
 
 if TYPE_CHECKING:
     from .gql import GraphQueryContext
@@ -58,34 +57,6 @@
 )
 
 
-def user_max_vfolder_count(required: bool = False):
-    return graphene.Int(
-        required=required,
-        description=description_msg("24.03.1", "Limitation of the number of user vfolders."),
-    )
-
-
-def user_max_quota_scope_size(required: bool = False):
-    return BigInt(
-        required=required,
-        description=description_msg("24.03.1", "Limitation of the quota size of user vfolders."),
-    )
-
-
-def project_max_vfolder_count(required: bool = False):
-    return graphene.Int(
-        required=required,
-        description=description_msg("24.03.1", "Limitation of the number of project vfolders."),
-    )
-
-
-def project_max_quota_scope_size(required: bool = False):
-    return BigInt(
-        required=required,
-        description=description_msg("24.03.1", "Limitation of the quota size of project vfolders."),
-    )
-
-
 keypair_resource_policies = sa.Table(
     "keypair_resource_policies",
     mapper_registry.metadata,
@@ -172,9 +143,9 @@ class KeyPairResourcePolicy(graphene.ObjectType):
     idle_timeout = BigInt()
     allowed_vfolder_hosts = graphene.JSONString()
 
-    max_vfolder_count = graphene.Int(deprecation_reason=deprecation_reason_msg("23.09.4"))
-    max_vfolder_size = BigInt(deprecation_reason=deprecation_reason_msg("23.09.4"))
-    max_quota_scope_size = BigInt(deprecation_reason=deprecation_reason_msg("23.09.4"))
+    max_vfolder_count = graphene.Int(deprecation_reason="Deprecated since 23.09.4")
+    max_vfolder_size = BigInt(deprecation_reason="Deprecated since 23.09.4")
+    max_quota_scope_size = BigInt(deprecation_reason="Deprecated since 23.09.4")
 
     @classmethod
     def from_row(
@@ -326,15 +297,15 @@ class CreateKeyPairResourcePolicyInput(graphene.InputObjectType):
     allowed_vfolder_hosts = graphene.JSONString(required=False)
     max_vfolder_count = graphene.Int(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
     max_vfolder_size = BigInt(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
     max_quota_scope_size = BigInt(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
 
 
@@ -349,15 +320,15 @@ class ModifyKeyPairResourcePolicyInput(graphene.InputObjectType):
     allowed_vfolder_hosts = graphene.JSONString(required=False)
     max_vfolder_count = graphene.Int(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
     max_vfolder_size = BigInt(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
     max_quota_scope_size = BigInt(
         required=False,
-        deprecation_reason=deprecation_reason_msg("23.09.4"),
+        deprecation_reason="Deprecated since 23.09.4",
     )
 
 
@@ -471,9 +442,15 @@ class UserResourcePolicy(graphene.ObjectType):
     id = graphene.ID(required=True)
     name = graphene.String(required=True)
     created_at = GQLDateTime(required=True)
-    max_vfolder_count = user_max_vfolder_count()
-    max_quota_scope_size = user_max_quota_scope_size()
-    max_vfolder_size = BigInt(deprecation_reason=deprecation_reason_msg("23.09.1"))
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of user vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of user vfolders.",
+    )
+    max_vfolder_size = BigInt(deprecation_reason="Deprecated since 23.09.1")
 
     @classmethod
     def from_row(
@@ -545,13 +522,25 @@ async def batch_load_by_user(
 
 
 class CreateUserResourcePolicyInput(graphene.InputObjectType):
-    max_vfolder_count = user_max_vfolder_count()
-    max_quota_scope_size = user_max_quota_scope_size()
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of user vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of user vfolders.",
+    )
 
 
 class ModifyUserResourcePolicyInput(graphene.InputObjectType):
-    max_vfolder_count = user_max_vfolder_count()
-    max_quota_scope_size = user_max_quota_scope_size()
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of user vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of user vfolders.",
+    )
 
 
 class CreateUserResourcePolicy(graphene.Mutation):
@@ -648,9 +637,15 @@ class ProjectResourcePolicy(graphene.ObjectType):
     id = graphene.ID(required=True)
     name = graphene.String(required=True)
     created_at = GQLDateTime(required=True)
-    max_vfolder_count = project_max_vfolder_count()
-    max_quota_scope_size = project_max_quota_scope_size()
-    max_vfolder_size = BigInt(deprecation_reason=deprecation_reason_msg("23.09.1"))
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of project vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of project vfolders.",
+    )
+    max_vfolder_size = BigInt(deprecation_reason="Deprecated since 23.09.1")
 
     @classmethod
     def from_row(
@@ -723,13 +718,25 @@ async def batch_load_by_project(
 
 
 class CreateProjectResourcePolicyInput(graphene.InputObjectType):
-    max_vfolder_count = project_max_vfolder_count()
-    max_quota_scope_size = project_max_quota_scope_size()
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of project vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of project vfolders.",
+    )
 
 
 class ModifyProjectResourcePolicyInput(graphene.InputObjectType):
-    max_vfolder_count = project_max_vfolder_count()
-    max_quota_scope_size = project_max_quota_scope_size()
+    max_vfolder_count = graphene.Int(
+        required=False,
+        description="Added since 24.03.0. Limitation of the number of project vfolders.",
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        description="Added since 24.03.0. Limitation of the quota size of project vfolders.",
+    )
 
 
 class CreateProjectResourcePolicy(graphene.Mutation):

src/ai/backend/manager/models/storage.py

@@ -30,7 +30,6 @@
 from ..api.exceptions import InvalidAPIParameters, VFolderOperationFailed
 from ..exceptions import InvalidArgument
 from .base import Item, PaginatedList
-from .utils import description_msg
 
 if TYPE_CHECKING:
     from .gql import GraphQueryContext
@@ -209,9 +208,9 @@ class Meta:
     performance_metric = graphene.JSONString()
     usage = graphene.JSONString()
     proxy = graphene.String(
-        description=description_msg("24.03.0", "Name of the proxy which this volume belongs to.")
+        description="Added since 24.03.0. Name of the proxy which this volume belongs to."
     )
-    name = graphene.String(description=description_msg("24.03.0", "Name of the storage."))
+    name = graphene.String(description="Added since 24.03.0. Name of the storage.")
 
     async def resolve_hardware_metadata(self, info: graphene.ResolveInfo) -> HardwareMetadata:
         ctx: GraphQueryContext = info.context

src/ai/backend/manager/models/utils.py

@@ -362,17 +362,3 @@ def agg_to_array(column: sa.Column) -> sa.sql.functions.Function:
 
 def is_db_retry_error(e: Exception) -> bool:
     return isinstance(e, DBAPIError) and getattr(e.orig, "pgcode", None) == "40001"
-
-
-def description_msg(version: str, detail: str | None = None) -> str:
-    val = f"Added since {version}."
-    if detail:
-        val = f"{val} {detail}"
-    return val
-
-
-def deprecation_reason_msg(version: str, detail: str | None = None) -> str:
-    val = f"Deprecated since {version}."
-    if detail:
-        val = f"{val} {detail}"
-    return val