[Feature]: Store user SSH key on the server #2053

Review: move `ssh_public_key` to `User` model (from `UserWithCreds`)

Author: peterschmidt85

src/dstack/_internal/core/models/users.py

@@ -30,6 +30,7 @@ class User(CoreModel):
     email: Optional[str]
     active: bool
     permissions: UserPermissions
+    ssh_public_key: Optional[str] = None
 
 
 class UserTokenCreds(CoreModel):
@@ -38,5 +39,4 @@ class UserTokenCreds(CoreModel):
 
 class UserWithCreds(User):
     creds: UserTokenCreds
-    ssh_public_key: Optional[str] = None
     ssh_private_key: Optional[str] = None

src/dstack/_internal/server/services/users.py

@@ -225,6 +225,7 @@ def user_model_to_user(user_model: UserModel) -> User:
         email=user_model.email,
         active=user_model.active,
         permissions=get_user_permissions(user_model),
+        ssh_public_key=user_model.ssh_public_key,
     )
 
 
@@ -237,8 +238,8 @@ def user_model_to_user_with_creds(user_model: UserModel) -> UserWithCreds:
         email=user_model.email,
         active=user_model.active,
         permissions=get_user_permissions(user_model),
-        creds=UserTokenCreds(token=user_model.token.get_plaintext_or_error()),
         ssh_public_key=user_model.ssh_public_key,
+        creds=UserTokenCreds(token=user_model.token.get_plaintext_or_error()),
         ssh_private_key=user_model.ssh_private_key,
     )
 

src/tests/_internal/server/routers/test_projects.py

@@ -77,6 +77,7 @@ async def test_returns_projects(self, test_db, session: AsyncSession, client: As
                     "permissions": {
                         "can_create_projects": True,
                     },
+                    "ssh_public_key": None,
                 },
                 "created_at": "2023-01-02T03:04:00+00:00",
                 "backends": [],
@@ -244,6 +245,7 @@ async def test_creates_project(self, test_db, session: AsyncSession, client: Asy
                 "permissions": {
                     "can_create_projects": True,
                 },
+                "ssh_public_key": user.ssh_public_key,
             },
             "created_at": "2023-01-02T03:04:00+00:00",
             "backends": [],
@@ -259,6 +261,7 @@ async def test_creates_project(self, test_db, session: AsyncSession, client: Asy
                         "permissions": {
                             "can_create_projects": True,
                         },
+                        "ssh_public_key": user.ssh_public_key,
                     },
                     "project_role": ProjectRole.ADMIN,
                     "permissions": {
@@ -693,6 +696,7 @@ async def test_returns_project(self, test_db, session: AsyncSession, client: Asy
                 "permissions": {
                     "can_create_projects": True,
                 },
+                "ssh_public_key": None,
             },
             "created_at": "2023-01-02T03:04:00+00:00",
             "backends": [],
@@ -708,6 +712,7 @@ async def test_returns_project(self, test_db, session: AsyncSession, client: Asy
                         "permissions": {
                             "can_create_projects": True,
                         },
+                        "ssh_public_key": None,
                     },
                     "project_role": ProjectRole.ADMIN,
                     "permissions": {
@@ -937,6 +942,7 @@ async def test_sets_project_members(self, test_db, session: AsyncSession, client
                     "permissions": {
                         "can_create_projects": True,
                     },
+                    "ssh_public_key": admin.ssh_public_key,
                 },
                 "project_role": ProjectRole.ADMIN,
                 "permissions": {
@@ -954,6 +960,7 @@ async def test_sets_project_members(self, test_db, session: AsyncSession, client
                     "permissions": {
                         "can_create_projects": True,
                     },
+                    "ssh_public_key": user1.ssh_public_key,
                 },
                 "project_role": ProjectRole.ADMIN,
                 "permissions": {
@@ -971,6 +978,7 @@ async def test_sets_project_members(self, test_db, session: AsyncSession, client
                     "permissions": {
                         "can_create_projects": True,
                     },
+                    "ssh_public_key": user2.ssh_public_key,
                 },
                 "project_role": ProjectRole.USER,
                 "permissions": {
@@ -1027,6 +1035,7 @@ async def test_sets_project_members_by_email(
                     "permissions": {
                         "can_create_projects": True,
                     },
+                    "ssh_public_key": user1.ssh_public_key,
                 },
                 "project_role": ProjectRole.ADMIN,
                 "permissions": {

src/tests/_internal/server/routers/test_users.py

@@ -48,6 +48,7 @@ async def test_admins_see_all_users(self, test_db, session: AsyncSession, client
                 "permissions": {
                     "can_create_projects": True,
                 },
+                "ssh_public_key": None,
             },
             {
                 "id": str(other_user.id),
@@ -59,6 +60,7 @@ async def test_admins_see_all_users(self, test_db, session: AsyncSession, client
                 "permissions": {
                     "can_create_projects": True,
                 },
+                "ssh_public_key": None,
             },
         ]
 
@@ -92,6 +94,7 @@ async def test_non_admins_see_only_themselves(
                 "permissions": {
                     "can_create_projects": True,
                 },
+                "ssh_public_key": None,
             }
         ]
 
@@ -229,7 +232,9 @@ async def test_creates_user(self, test_db, session: AsyncSession, client: AsyncC
                 },
             )
         assert response.status_code == 200
-        assert response.json() == {
+        user_data = response.json()
+        ssh_public_key = user_data["ssh_public_key"]
+        assert user_data == {
             "id": "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e",
             "username": "test",
             "created_at": "2023-01-02T03:04:00+00:00",
@@ -239,6 +244,7 @@ async def test_creates_user(self, test_db, session: AsyncSession, client: AsyncC
             "permissions": {
                 "can_create_projects": True,
             },
+            "ssh_public_key": ssh_public_key,
         }
         res = await session.execute(select(UserModel).where(UserModel.name == "test"))
         assert len(res.scalars().all()) == 1
@@ -264,7 +270,9 @@ async def test_return_400_if_username_taken(
                 },
             )
         assert response.status_code == 200
-        assert response.json() == {
+        user_data = response.json() 
+        ssh_public_key = user_data["ssh_public_key"]
+        assert user_data == {
             "id": "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e",
             "username": "Test",
             "created_at": "2023-01-02T03:04:00+00:00",
@@ -274,6 +282,7 @@ async def test_return_400_if_username_taken(
             "permissions": {
                 "can_create_projects": True,
             },
+            "ssh_public_key": ssh_public_key,
         }
         # Username uniqueness check should be case insensitive
         for username in ["test", "Test", "TesT"]: