fix: Rename "XXXInput" to "CreateXXInput"

- fix: Missing max_concurrent_sftp_sessions GraphQL fields when creating/modifying keypair resource policies
lablup · Nov 5, 2023 · 35047b4 · 35047b4
1 parent b8428c8
commit 35047b4
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 28 deletions.
diff --git a/src/ai/backend/manager/models/domain.py b/src/ai/backend/manager/models/domain.py
@@ -40,9 +40,9 @@
     "domains",
     "DomainRow",
     "Domain",
-    "DomainInput",
-    "ModifyDomainInput",
+    "CreateDomainInput",
     "CreateDomain",
+    "ModifyDomainInput",
     "ModifyDomain",
     "DeleteDomain",
     "DomainDotfile",
@@ -170,7 +170,7 @@ async def batch_load_by_name(
             )
 
 
-class DomainInput(graphene.InputObjectType):
+class CreateDomainInput(graphene.InputObjectType):
     description = graphene.String(required=False)
     is_active = graphene.Boolean(required=False, default=True)
     total_resource_slots = graphene.JSONString(required=False)
@@ -194,7 +194,7 @@ class CreateDomain(graphene.Mutation):
 
     class Arguments:
         name = graphene.String(required=True)
-        props = DomainInput(required=True)
+        props = CreateDomainInput(required=True)
 
     ok = graphene.Boolean()
     msg = graphene.String()
@@ -206,7 +206,7 @@ async def mutate(
         root,
         info: graphene.ResolveInfo,
         name: str,
-        props: DomainInput,
+        props: CreateDomainInput,
     ) -> CreateDomain:
         if _rx_slug.search(name) is None:
             return cls(False, "invalid name format. slug format required.", None)

diff --git a/src/ai/backend/manager/models/group.py b/src/ai/backend/manager/models/group.py
@@ -62,9 +62,9 @@
     "AssocGroupUserRow",
     "resolve_group_name_or_id",
     "Group",
-    "GroupInput",
-    "ModifyGroupInput",
+    "CreateGroupInput",
     "CreateGroup",
+    "ModifyGroupInput",
     "ModifyGroup",
     "DeleteGroup",
     "GroupDotfile",
@@ -376,7 +376,7 @@ async def get_groups_for_user(
             ]
 
 
-class GroupInput(graphene.InputObjectType):
+class CreateGroupInput(graphene.InputObjectType):
     description = graphene.String(required=False, default="")
     is_active = graphene.Boolean(required=False, default=True)
     domain_name = graphene.String(required=True)
@@ -404,7 +404,7 @@ class CreateGroup(graphene.Mutation):
 
     class Arguments:
         name = graphene.String(required=True)
-        props = GroupInput(required=True)
+        props = CreateGroupInput(required=True)
 
     ok = graphene.Boolean()
     msg = graphene.String()
@@ -420,7 +420,7 @@ async def mutate(
         root,
         info: graphene.ResolveInfo,
         name: str,
-        props: GroupInput,
+        props: CreateGroupInput,
     ) -> CreateGroup:
         if _rx_slug.search(name) is None:
             raise ValueError("invalid name format. slug format required.")

diff --git a/src/ai/backend/manager/models/keypair.py b/src/ai/backend/manager/models/keypair.py
@@ -49,8 +49,9 @@
     "KeyPair",
     "KeyPairList",
     "UserInfo",
-    "KeyPairInput",
+    "CreateKeyPairInput",
     "CreateKeyPair",
+    "ModifyKeyPairInput",
     "ModifyKeyPair",
     "DeleteKeyPair",
     "Dotfile",
@@ -500,7 +501,7 @@ class Meta:
     items = graphene.List(KeyPair, required=True)
 
 
-class KeyPairInput(graphene.InputObjectType):
+class CreateKeyPairInput(graphene.InputObjectType):
     is_active = graphene.Boolean(required=False, default=True)
     is_admin = graphene.Boolean(required=False, default=False)
     resource_policy = graphene.String(required=True)
@@ -531,7 +532,7 @@ class CreateKeyPair(graphene.Mutation):
 
     class Arguments:
         user_id = graphene.String(required=True)
-        props = KeyPairInput(required=True)
+        props = CreateKeyPairInput(required=True)
 
     ok = graphene.Boolean()
     msg = graphene.String()
@@ -543,7 +544,7 @@ async def mutate(
         root,
         info: graphene.ResolveInfo,
         user_id: str,
-        props: KeyPairInput,
+        props: CreateKeyPairInput,
     ) -> CreateKeyPair:
         from .user import users  # noqa
 

diff --git a/src/ai/backend/manager/models/resource_policy.py b/src/ai/backend/manager/models/resource_policy.py
@@ -317,28 +317,48 @@ async def batch_load_by_ak(
 
 class CreateKeyPairResourcePolicyInput(graphene.InputObjectType):
     default_for_unspecified = graphene.String(required=True)
-    total_resource_slots = graphene.JSONString(required=True)
-    max_session_lifetime = graphene.Int(required=True, default_value=0)
+    total_resource_slots = graphene.JSONString(required=False, default_value={})
+    max_session_lifetime = graphene.Int(required=False, default_value=0)
     max_concurrent_sessions = graphene.Int(required=True)
+    max_concurrent_sftp_sessions = graphene.Int(required=False, default_value=1)
     max_containers_per_session = graphene.Int(required=True)
     idle_timeout = BigInt(required=True)
     allowed_vfolder_hosts = graphene.JSONString(required=False)
-    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(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
+    max_vfolder_size = BigInt(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
 
 
 class ModifyKeyPairResourcePolicyInput(graphene.InputObjectType):
     default_for_unspecified = graphene.String(required=False)
     total_resource_slots = graphene.JSONString(required=False)
     max_session_lifetime = graphene.Int(required=False)
     max_concurrent_sessions = graphene.Int(required=False)
+    max_concurrent_sftp_sessions = graphene.Int(required=False)
     max_containers_per_session = graphene.Int(required=False)
     idle_timeout = BigInt(required=False)
     allowed_vfolder_hosts = graphene.JSONString(required=False)
-    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(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
+    max_vfolder_size = BigInt(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
+    max_quota_scope_size = BigInt(
+        required=False,
+        deprecation_reason=deprecation_reason_msg("23.09.4"),
+    )
 
 
 class CreateKeyPairResourcePolicy(graphene.Mutation):
@@ -366,6 +386,7 @@ async def mutate(
             "total_resource_slots": ResourceSlot.from_user_input(props.total_resource_slots, None),
             "max_session_lifetime": props.max_session_lifetime,
             "max_concurrent_sessions": props.max_concurrent_sessions,
+            "max_concurrent_sftp_sessions": props.max_concurrent_sessions,
             "max_containers_per_session": props.max_containers_per_session,
             "idle_timeout": props.idle_timeout,
             "allowed_vfolder_hosts": props.allowed_vfolder_hosts,
@@ -399,7 +420,10 @@ async def mutate(
     ) -> ModifyKeyPairResourcePolicy:
         data: Dict[str, Any] = {}
         set_if_set(
-            props, data, "default_for_unspecified", clean_func=lambda v: DefaultForUnspecified[v]
+            props,
+            data,
+            "default_for_unspecified",
+            clean_func=lambda v: DefaultForUnspecified[v],
         )
         set_if_set(
             props,
@@ -409,6 +433,7 @@ async def mutate(
         )
         set_if_set(props, data, "max_session_lifetime")
         set_if_set(props, data, "max_concurrent_sessions")
+        set_if_set(props, data, "max_concurrent_sftp_sessions")
         set_if_set(props, data, "max_containers_per_session")
         set_if_set(props, data, "idle_timeout")
         set_if_set(props, data, "allowed_vfolder_hosts")

diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py
@@ -57,9 +57,9 @@
     "UserList",
     "UserGroup",
     "UserRole",
-    "UserInput",
-    "ModifyUserInput",
+    "CreateUserInput",
     "CreateUser",
+    "ModifyUserInput",
     "ModifyUser",
     "DeleteUser",
     "UserStatus",
@@ -507,7 +507,7 @@ class Meta:
     items = graphene.List(User, required=True)
 
 
-class UserInput(graphene.InputObjectType):
+class CreateUserInput(graphene.InputObjectType):
     username = graphene.String(required=True)
     password = graphene.String(required=True)
     need_password_change = graphene.Boolean(required=True)
@@ -552,7 +552,7 @@ class CreateUser(graphene.Mutation):
 
     class Arguments:
         email = graphene.String(required=True)
-        props = UserInput(required=True)
+        props = CreateUserInput(required=True)
 
     ok = graphene.Boolean()
     msg = graphene.String()
@@ -564,7 +564,7 @@ async def mutate(
         root,
         info: graphene.ResolveInfo,
         email: str,
-        props: UserInput,
+        props: CreateUserInput,
     ) -> CreateUser:
         graph_ctx: GraphQueryContext = info.context
         username = props.username if props.username else email