[BUGFIX] [ENHANCEMENT] create users and workspaces with predefined ID (#5786)

# Description
<!-- Please include a summary of the changes and the related issue.
Please also include relevant motivation and context. List any
dependencies that are required for this change. -->

This PR allows users and workspaces to be created with predefined IDs.
Passing a predefined ID when creating workspaces and users simplifies
the migration process since users don't need to review responses and map
them to new users, avoiding potential issues in the process.


**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- Bug fix (non-breaking change which fixes an issue)
- Improvement (change adding some improvement to an existing
functionality)
- Documentation update

**How Has This Been Tested**
<!-- Please add some reference about how your feature has been tested.
-->

**Checklist**
<!-- Please go over the list and make sure you've taken everything into
account -->

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

Author: frascuchon

argilla-server/CHANGELOG.md

@@ -16,6 +16,11 @@ These are the section headers that we use:
 
 ## [Unreleased]()
 
+### Added
+
+- Added support to create users with predefined ids. ([#5786](https://github.com/argilla-io/argilla/pull/5786))
+- Added support to create workspaces with predefined ids. ([#5786](https://github.com/argilla-io/argilla/pull/5786))
+
 ## [2.6.0](https://github.com/argilla-io/argilla/compare/v2.5.0...v2.6.0)
 
 ### Added

argilla-server/src/argilla_server/api/schemas/v1/users.py

@@ -59,6 +59,7 @@ class User(BaseModel):
 
 
 class UserCreate(BaseModel):
+    id: Optional[UUID] = None
     first_name: UserFirstName
     last_name: Optional[UserLastName] = None
     username: UserUsername

argilla-server/src/argilla_server/api/schemas/v1/workspaces.py

@@ -13,7 +13,8 @@
 #  limitations under the License.
 
 from datetime import datetime
-from typing import List
+from optparse import Option
+from typing import List, Optional
 from uuid import UUID
 
 from pydantic import BaseModel, Field, ConfigDict
@@ -29,6 +30,7 @@ class Workspace(BaseModel):
 
 
 class WorkspaceCreate(BaseModel):
+    id: Optional[UUID] = None
     name: str = Field(min_length=1)
 
 

argilla-server/src/argilla_server/contexts/accounts.py

@@ -66,7 +66,15 @@ async def create_workspace(db: AsyncSession, workspace_attrs: dict) -> Workspace
     if await Workspace.get_by(db, name=workspace_attrs["name"]) is not None:
         raise NotUniqueError(f"Workspace name `{workspace_attrs['name']}` is not unique")
 
-    return await Workspace.create(db, name=workspace_attrs["name"])
+    if workspace_id := workspace_attrs.get("id"):
+        if await Workspace.get(db, id=workspace_id) is not None:
+            raise NotUniqueError(f"Workspace with id `{workspace_id}` is not unique")
+
+    return await Workspace.create(
+        db,
+        id=workspace_attrs.get("id"),
+        name=workspace_attrs["name"],
+    )
 
 
 async def delete_workspace(db: AsyncSession, workspace: Workspace):
@@ -108,8 +116,13 @@ async def create_user(db: AsyncSession, user_attrs: dict, workspaces: Union[List
     if await get_user_by_username(db, user_attrs["username"]) is not None:
         raise NotUniqueError(f"User username `{user_attrs['username']}` is not unique")
 
+    if user_id := user_attrs.get("id"):
+        if await User.get(db, id=user_id) is not None:
+            raise NotUniqueError(f"User with id `{user_id}` is not unique")
+
     user = await User.create(
         db,
+        id=user_attrs.get("id"),
         first_name=user_attrs["first_name"],
         last_name=user_attrs["last_name"],
         username=user_attrs["username"],

argilla-server/tests/unit/api/handlers/v1/users/test_create_user.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+from uuid import uuid4
 
 import pytest
 from argilla_server.constants import API_KEY_HEADER_NAME
@@ -146,6 +146,86 @@ async def test_create_user_with_non_default_role(
         assert response.json()["role"] == UserRole.owner
         assert user.role == UserRole.owner
 
+    async def test_create_user_with_predefined_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        user_id = uuid4()
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={
+                "id": str(user_id),
+                "first_name": "First name",
+                "last_name": "Last name",
+                "username": "username",
+                "password": "12345678",
+            },
+        )
+
+        assert response.status_code == 201
+
+        user = (await db.execute(select(User).filter_by(username="username"))).scalar_one()
+        assert user.id == user_id
+
+    async def test_create_user_with_none_user_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={
+                "id": None,
+                "first_name": "First name",
+                "last_name": "Last name",
+                "username": "username",
+                "password": "12345678",
+            },
+        )
+
+        assert response.status_code == 201
+
+        user = (await db.execute(select(User).filter_by(username="username"))).scalar_one()
+        assert user.id is not None
+
+    async def test_create_user_with_wrong_user_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={
+                "id": "wrong_id",
+                "first_name": "First name",
+                "last_name": "Last name",
+                "username": "username",
+                "password": "12345678",
+            },
+        )
+
+        assert response.status_code == 422
+        assert (await db.execute(select(func.count(User.id)))).scalar() == 1
+
+    async def test_create_user_with_existing_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        user_id = uuid4()
+        await UserFactory.create(id=user_id)
+
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={
+                "id": str(user_id),
+                "first_name": "First name",
+                "last_name": "Last name",
+                "username": "username",
+                "password": "12345678",
+            },
+        )
+
+        assert response.status_code == 409
+        assert (await db.execute(select(func.count(User.id)))).scalar() == 2
+
     async def test_create_user_without_authentication(self, db: AsyncSession, async_client: AsyncClient):
         response = await async_client.post(
             self.url(),

argilla-server/tests/unit/api/handlers/v1/workspaces/test_create_workspace.py

@@ -11,6 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from uuid import uuid4
 
 import pytest
 from argilla_server.constants import API_KEY_HEADER_NAME
@@ -47,6 +48,77 @@ async def test_create_workspace(self, db: AsyncSession, async_client: AsyncClien
             "updated_at": workspace.updated_at.isoformat(),
         }
 
+    async def test_create_workspace_with_predefined_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        workspace_id = uuid4()
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={"id": str(workspace_id), "name": "workspace"},
+        )
+
+        assert response.status_code == 201
+
+        assert (await db.execute(select(func.count(Workspace.id)))).scalar() == 1
+        workspace = (await db.execute(select(Workspace).filter_by(name="workspace"))).scalar_one()
+
+        assert response.json() == {
+            "id": str(workspace_id),
+            "name": "workspace",
+            "inserted_at": workspace.inserted_at.isoformat(),
+            "updated_at": workspace.updated_at.isoformat(),
+        }
+
+    async def test_create_workspace_with_none_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={"id": None, "name": "workspace"},
+        )
+
+        assert response.status_code == 201
+
+        assert (await db.execute(select(func.count(Workspace.id)))).scalar() == 1
+        workspace = (await db.execute(select(Workspace).filter_by(name="workspace"))).scalar_one()
+
+        assert response.json() == {
+            "id": str(workspace.id),
+            "name": "workspace",
+            "inserted_at": workspace.inserted_at.isoformat(),
+            "updated_at": workspace.updated_at.isoformat(),
+        }
+
+    async def test_create_workspace_with_wrong_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={"id": "wrong_id", "name": "workspace"},
+        )
+
+        assert response.status_code == 422
+
+        assert (await db.execute(select(func.count(Workspace.id)))).scalar() == 0
+
+    async def test_create_workspace_with_existing_id(
+        self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict
+    ):
+        workspace_id = uuid4()
+        await WorkspaceFactory.create(id=workspace_id)
+
+        response = await async_client.post(
+            self.url(),
+            headers=owner_auth_header,
+            json={"id": str(workspace_id), "name": "workspace"},
+        )
+
+        assert response.status_code == 409
+        assert (await db.execute(select(func.count(Workspace.id)))).scalar() == 1
+
     async def test_create_workspace_without_authentication(self, db: AsyncSession, async_client: AsyncClient):
         response = await async_client.post(
             self.url(),

argilla/CHANGELOG.md

@@ -19,6 +19,9 @@ These are the section headers that we use:
 ### Added
 
 - Return similarity score when searching by similarity. ([#5778](https://github.com/argilla-io/argilla/pull/5778))
+- Added support to create users with predefined ids. ([#5786](https://github.com/argilla-io/argilla/pull/5786))
+- Added support to create workspaces with predefined ids. ([#5786](https://github.com/argilla-io/argilla/pull/5786))
+
 
 ## [2.6.0](https://github.com/argilla-io/argilla/compare/v2.5.0...v2.6.0)
 

argilla/docs/how_to_guides/migrate_from_legacy_datasets.md

@@ -69,26 +69,30 @@ Next, recreate the users and workspaces on the Argilla V2 server:
 ```python
 for workspace in workspaces_v1:
     rg.Workspace(
-        name=workspace.name
+        id=workspace.id,
+        name=workspace.name,
     ).create()
 ```
 
 ```python
 for user in users_v1:
-    user = rg.User(
+    user_v2 = rg.User(
+        id=user.id,
         username=user.username,
         first_name=user.first_name,
         last_name=user.last_name,
         role=user.role,
         password="<your_chosen_password>" # (1)
     ).create()
+
     if user.role == "owner":
        continue
 
-    for workspace_name in user.workspaces:
-        if workspace_name != user.name:
-            workspace = client.workspaces(name=workspace_name)
-            user.add_to_workspace(workspace)
+    for workspace in user.workspaces:
+        workspace_v2 = client.workspaces(name=workspace.name)
+        if workspace_v2 is None:
+            continue
+        user.add_to_workspace(workspace_v2)
 ```
 
 1. You need to chose a new password for the user, to do this programmatically you can use the `uuid` package to generate a random password. Take care to keep track of the passwords you chose, since you will not be able to retrieve them later.

argilla/src/argilla/_api/_workspaces.py

@@ -35,7 +35,7 @@ class WorkspacesAPI(ResourceAPI[WorkspaceModel]):
     @api_error_handler
     def create(self, workspace: WorkspaceModel) -> WorkspaceModel:
         # TODO: Unify API endpoint
-        response = self.http_client.post(url="/api/v1/workspaces", json={"name": workspace.name})
+        response = self.http_client.post(url="/api/v1/workspaces", json=workspace.model_dump())
         response.raise_for_status()
         response_json = response.json()
         workspace = self._model_from_json(json_workspace=response_json)

argilla/src/argilla/users/_resource.py

@@ -45,8 +45,8 @@ def __init__(
         last_name: Optional[str] = None,
         role: Optional[str] = None,
         password: Optional[str] = None,
-        client: Optional["Argilla"] = None,
         id: Optional[UUID] = None,
+        client: Optional["Argilla"] = None,
         _model: Optional[UserModel] = None,
     ) -> None:
         """Initializes a User object with a client and a username
@@ -57,6 +57,7 @@ def __init__(
             last_name (str): The last name of the user
             role (str): The role of the user, either 'annotator', admin, or 'owner'
             password (str): The password of the user
+            id (UUID): The ID of the user. If provided before a .create, the will be created with this ID
             client (Argilla): The client used to interact with Argilla
 
         Returns: