change user in new context

Author: mishushakov

template/server/api/models/context.py

@@ -6,6 +6,7 @@ class Context(BaseModel):
     id: StrictStr = Field(description="Context ID")
     language: StrictStr = Field(description="Language of the context")
     cwd: StrictStr = Field(description="Current working directory of the context")
+    user: StrictStr = Field(description="User of the context")
 
     def __hash__(self):
         return hash(self.id)

template/server/api/models/create_context.py

@@ -4,6 +4,10 @@
 
 
 class CreateContext(BaseModel):
+    user: Optional[StrictStr] = Field(
+        default="root",
+        description="User to run the context",
+    )
     cwd: Optional[StrictStr] = Field(
         default="/home/user",
         description="Current working directory",

template/server/contexts.py

@@ -24,10 +24,16 @@ def normalize_language(language: Optional[str]) -> str:
     return language
 
 
-async def create_context(client, websockets: dict, language: str, cwd: str) -> Context:
+def get_kernel_name(language: str, user: str) -> str:
+    if user == "root":
+        return language+"_root"
+    return language
+
+
+async def create_context(client, websockets: dict, language: str, cwd: str, user: str) -> Context:
     data = {
         "path": str(uuid.uuid4()),
-        "kernel": {"name": language},
+        "kernel": {"name": get_kernel_name(language, user)}, # replace with root kernel when user is root
         "type": "notebook",
         "name": str(uuid.uuid4()),
     }
@@ -59,4 +65,4 @@ async def create_context(client, websockets: dict, language: str, cwd: str) -> C
             status_code=500,
         )
 
-    return Context(language=language, id=context_id, cwd=cwd)
+    return Context(language=language, id=context_id, cwd=cwd, user=user)

template/server/main.py

@@ -91,7 +91,7 @@ async def post_execute(request: ExecutionRequest):
             if not context_id:
                 try:
                     context = await create_context(
-                        client, websockets, language, "/home/user"
+                        client, websockets, language, "/home/user", "root"
                     )
                 except Exception as e:
                     return PlainTextResponse(str(e), status_code=500)
@@ -127,9 +127,10 @@ async def post_contexts(request: CreateContext) -> Context:
 
     language = normalize_language(request.language)
     cwd = request.cwd or "/home/user"
+    user = request.user or "root"
 
     try:
-        return await create_context(client, websockets, language, cwd)
+        return await create_context(client, websockets, language, cwd, user)
     except Exception as e:
         return PlainTextResponse(str(e), status_code=500)