[PLT-3393] Add support of groups to invite_user()

libs/labelbox/src/labelbox/__init__.py

@@ -55,7 +55,7 @@
 from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
 from labelbox.schema.tool_building.prompt_issue_tool import PromptIssueTool
 from labelbox.schema.tool_building.relationship_tool import RelationshipTool
-from labelbox.schema.role import Role, ProjectRole
+from labelbox.schema.role import Role, ProjectRole, UserGroupRole
 from labelbox.schema.invite import Invite, InviteLimit
 from labelbox.schema.data_row_metadata import (
     DataRowMetadataOntology,

libs/labelbox/src/labelbox/orm/model.py

@@ -399,6 +399,7 @@ class Entity(metaclass=EntityMeta):
         CatalogSlice: Type[labelbox.CatalogSlice]
         ModelSlice: Type[labelbox.ModelSlice]
         TaskQueue: Type[labelbox.TaskQueue]
+        UserGroupRole: Type[labelbox.UserGroupRole]
 
     @classmethod
     def _attributes_of_type(cls, attr_type):

libs/labelbox/src/labelbox/schema/organization.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Set, Optional, Union
 
 from lbox.exceptions import LabelboxError
 
@@ -22,6 +22,7 @@
         ProjectRole,
         Role,
         User,
+        UserGroupRole,
     )
 
 
@@ -65,6 +66,7 @@ def invite_user(
         email: str,
         role: "Role",
         project_roles: Optional[List["ProjectRole"]] = None,
+        user_group_roles: Optional[List["UserGroupRole"]] = None,
     ) -> "Invite":
         """
         Invite a new member to the org. This will send the user an email invite
@@ -88,6 +90,40 @@ def invite_user(
                 f"Project roles cannot be set for a user with organization level permissions. Found role name `{role.name}`, expected `NONE`"
             )
 
+        if user_group_roles and role.name != "NONE":
+            raise ValueError(
+                f"User Group roles cannot be set for a user with organization level permissions. Found role name `{role.name}`, expected `NONE`"
+            )
+
+        if user_group_roles:
+            # The backend can 500 if the same groupId appears more than once.
+            # We dedupe exact duplicates (same groupId+roleId), but reject
+            # conflicting assignments (same groupId with different roleId).
+
+            deduped_user_group_roles: Dict[str, "UserGroupRole"] = {}
+            conflicting_user_group_ids: Set[str] = set()
+
+            for user_group_role in user_group_roles:
+                user_group_id = user_group_role.user_group.id
+                role_id = user_group_role.role.uid
+
+                existing = deduped_user_group_roles.get(user_group_id)
+                if existing is None:
+                    deduped_user_group_roles[user_group_id] = user_group_role
+                else:
+                    if existing.role.uid != role_id:
+                        conflicting_user_group_ids.add(user_group_id)
+
+            if conflicting_user_group_ids:
+                conflicts_str = ", ".join(sorted(conflicting_user_group_ids))
+                raise ValueError(
+                    "user_group_roles contains conflicting role assignments for "
+                    "the same UserGroup. Each UserGroup may only appear once. "
+                    f"Conflicting user_group.id values: {conflicts_str}"
+                )
+
+            user_group_roles = list(deduped_user_group_roles.values())
+
         data_param = "data"
         query_str = """mutation createInvitesPyApi($%s: [CreateInviteInput!]){
                     createInvites(data: $%s){  invite { id createdAt organizationRoleName inviteeEmail inviter { %s } }}}""" % (
@@ -104,6 +140,19 @@ def invite_user(
             for project_role in project_roles or []
         ]
 
+        user_group_ids = [
+            user_group_role.user_group.id
+            for user_group_role in user_group_roles or []
+        ]
+
+        user_group_with_role_ids = [
+            {
+                "groupId": user_group_role.user_group.id,
+                "roleId": user_group_role.role.uid,
+            }
+            for user_group_role in user_group_roles or []
+        ]
+
         res = self.client.execute(
             query_str,
             {
@@ -114,6 +163,8 @@ def invite_user(
                         "organizationId": self.uid,
                         "organizationRoleId": role.uid,
                         "projects": projects,
+                        "userGroupIds": user_group_ids,
+                        "userGroupWithRoleIds": user_group_with_role_ids,
                     }
                 ]
             },

libs/labelbox/src/labelbox/schema/role.py

@@ -6,6 +6,7 @@
 
 if TYPE_CHECKING:
     from labelbox import Client, Project
+    from labelbox.schema.user_group import UserGroup
 
 _ROLES: Optional[Dict[str, "Role"]] = None
 
@@ -45,3 +46,9 @@ class UserRole(Role): ...
 class ProjectRole:
     project: "Project"
     role: Role
+
+
+@dataclass
+class UserGroupRole:
+    user_group: "UserGroup"
+    role: Role

libs/labelbox/tests/unit/schema/test_organization_invite_user.py

@@ -0,0 +1,151 @@
+import pytest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from labelbox.schema.role import UserGroupRole
+from labelbox.schema.organization import Organization
+
+
+def test_invite_user_duplicate_user_group_roles_same_role_is_deduped():
+    client = MagicMock()
+    client.get_user.return_value = SimpleNamespace(uid="inviter-id")
+    client.execute.return_value = {
+        "createInvites": [
+            {
+                "invite": {
+                    "id": "invite-id",
+                    "createdAt": "2020-01-01T00:00:00.000Z",
+                    "organizationRoleName": "NONE",
+                    "inviteeEmail": "someone@example.com",
+                    "inviter": {"id": "inviter-id"},
+                }
+            }
+        ]
+    }
+
+    organization = Organization(
+        client,
+        {
+            "id": "org-id",
+            "name": "Test Org",
+            "createdAt": "2020-01-01T00:00:00.000Z",
+            "updatedAt": "2020-01-01T00:00:00.000Z",
+        },
+    )
+
+    org_role_none = SimpleNamespace(uid="org-role-none-id", name="NONE")
+    reviewer_role = SimpleNamespace(uid="reviewer-role-id", name="REVIEWER")
+    user_group = SimpleNamespace(id="user-group-id")
+
+    user_group_roles = [
+        UserGroupRole(user_group=user_group, role=reviewer_role),
+        UserGroupRole(user_group=user_group, role=reviewer_role),
+    ]
+
+    organization.invite_user(
+        email="someone@example.com",
+        role=org_role_none,
+        user_group_roles=user_group_roles,
+    )
+
+    # ensure we only send one entry per group
+    args, kwargs = client.execute.call_args
+    assert kwargs == {}
+    payload = args[1]["data"][0]
+    assert payload["userGroupIds"] == ["user-group-id"]
+    assert payload["userGroupWithRoleIds"] == [
+        {"groupId": "user-group-id", "roleId": "reviewer-role-id"}
+    ]
+
+
+def test_invite_user_duplicate_user_group_roles_conflicting_roles_raises_value_error():
+    client = MagicMock()
+    client.get_user.return_value = SimpleNamespace(uid="inviter-id")
+
+    organization = Organization(
+        client,
+        {
+            "id": "org-id",
+            "name": "Test Org",
+            "createdAt": "2020-01-01T00:00:00.000Z",
+            "updatedAt": "2020-01-01T00:00:00.000Z",
+        },
+    )
+
+    org_role_none = SimpleNamespace(uid="org-role-none-id", name="NONE")
+    reviewer_role = SimpleNamespace(uid="reviewer-role-id", name="REVIEWER")
+    team_manager_role = SimpleNamespace(
+        uid="team-manager-role-id", name="TEAM_MANAGER"
+    )
+    user_group = SimpleNamespace(id="user-group-id")
+
+    user_group_roles = [
+        UserGroupRole(user_group=user_group, role=reviewer_role),
+        UserGroupRole(user_group=user_group, role=team_manager_role),
+    ]
+
+    with pytest.raises(ValueError, match="conflicting role assignments"):
+        organization.invite_user(
+            email="someone@example.com",
+            role=org_role_none,
+            user_group_roles=user_group_roles,
+        )
+
+    client.execute.assert_not_called()
+
+
+def test_invite_user_user_group_roles_payload_contains_all_groups():
+    client = MagicMock()
+    client.get_user.return_value = SimpleNamespace(uid="inviter-id")
+    client.execute.return_value = {
+        "createInvites": [
+            {
+                "invite": {
+                    "id": "invite-id",
+                    "createdAt": "2020-01-01T00:00:00.000Z",
+                    "organizationRoleName": "NONE",
+                    "inviteeEmail": "someone@example.com",
+                    "inviter": {"id": "inviter-id"},
+                }
+            }
+        ]
+    }
+
+    organization = Organization(
+        client,
+        {
+            "id": "org-id",
+            "name": "Test Org",
+            "createdAt": "2020-01-01T00:00:00.000Z",
+            "updatedAt": "2020-01-01T00:00:00.000Z",
+        },
+    )
+
+    org_role_none = SimpleNamespace(uid="org-role-none-id", name="NONE")
+    reviewer_role = SimpleNamespace(uid="reviewer-role-id", name="REVIEWER")
+    team_manager_role = SimpleNamespace(
+        uid="team-manager-role-id", name="TEAM_MANAGER"
+    )
+
+    ug1 = SimpleNamespace(id="user-group-1")
+    ug2 = SimpleNamespace(id="user-group-2")
+
+    user_group_roles = [
+        UserGroupRole(user_group=ug1, role=reviewer_role),
+        UserGroupRole(user_group=ug2, role=team_manager_role),
+    ]
+
+    organization.invite_user(
+        email="someone@example.com",
+        role=org_role_none,
+        user_group_roles=user_group_roles,
+    )
+
+    args, kwargs = client.execute.call_args
+    assert kwargs == {}
+    payload = args[1]["data"][0]
+    assert payload["userGroupIds"] == ["user-group-1", "user-group-2"]
+    assert payload["userGroupWithRoleIds"] == [
+        {"groupId": "user-group-1", "roleId": "reviewer-role-id"},
+        {"groupId": "user-group-2", "roleId": "team-manager-role-id"},
+    ]