feat(iam): add support for attributes in groups create, groups list (#2231)

Co-authored-by: anders-albert <anders0albert@gmail.com>

Author: stefanmaretic

cognite/client/data_classes/iam.py

@@ -2,11 +2,13 @@
 
 from abc import ABC
 from collections.abc import Iterable
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast
 
 from typing_extensions import Self
 
 from cognite.client.data_classes._base import (
+    CogniteObject,
     CogniteResource,
     CogniteResourceList,
     CogniteResponse,
@@ -24,17 +26,51 @@
 
     from cognite.client import CogniteClient
 
-
 ALL_USER_ACCOUNTS = "allUserAccounts"
 
 
+@dataclass
+class GroupAttributesToken(CogniteObject):
+    """List of applications (represented by their application ID) this group is valid for"""
+
+    app_ids: list[str] = field(default_factory=list)
+
+
+@dataclass
+class GroupAttributes(CogniteObject):
+    """Attributes derived from access token"""
+
+    token: GroupAttributesToken | None = None
+    _unknown_properties: dict[str, Any] = field(default_factory=dict, init=False, repr=False)
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dumps the attributes to a dictionary"""
+        dumped = super().dump(camel_case=camel_case)
+        if self.token is not None:
+            dumped["token"] = self.token.dump(camel_case=camel_case)
+        if self._unknown_properties:
+            dumped.update(self._unknown_properties)
+        return dumped
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        token: GroupAttributesToken | None = None
+        if "token" in resource:
+            token = GroupAttributesToken._load(resource["token"], cognite_client=cognite_client)
+        instance = cls(token=token)
+        existing = {"token"}
+        instance._unknown_properties = {key: value for key, value in resource.items() if key not in existing}
+        return instance
+
+
 class GroupCore(WriteableCogniteResource["GroupWrite"], ABC):
     """No description.
 
     Args:
         name (str): Name of the group.
         source_id (str | None): ID of the group in the source. If this is the same ID as a group in the IdP, a service account in that group will implicitly be a part of this group as well. Can not be used together with 'members'.
         capabilities (list[Capability] | None): List of capabilities (acls) this group should grant its users.
+        attributes (GroupAttributes | None): Attributes of the group, this scopes down access based on the attributes specified.
         metadata (dict[str, str] | None): Custom, immutable application specific metadata. String key -> String value. Limits: Key are at most 32 bytes. Values are at most 512 bytes. Up to 16 key-value pairs. Total size is at most 4096.
         members (Literal['allUserAccounts'] | list[str] | None): Specifies which users are members of the group. Can not be used together with 'source_id'.
     """
@@ -44,6 +80,7 @@ def __init__(
         name: str,
         source_id: str | None = None,
         capabilities: list[Capability] | None = None,
+        attributes: GroupAttributes | None = None,
         metadata: dict[str, str] | None = None,
         members: Literal["allUserAccounts"] | list[str] | None = None,
     ) -> None:
@@ -52,6 +89,7 @@ def __init__(
         self.capabilities = capabilities
         if isinstance(self.capabilities, Capability):
             self.capabilities = [capabilities]
+        self.attributes = attributes
         self.metadata = metadata
         self.members = members
 
@@ -62,6 +100,11 @@ def _load(cls, resource: dict, cognite_client: CogniteClient | None = None, allo
             source_id=resource.get("sourceId"),
             capabilities=[Capability.load(c, allow_unknown=allow_unknown) for c in resource.get("capabilities", [])]
             or None,
+            attributes=(
+                GroupAttributes._load(resource["attributes"], cognite_client=cognite_client)
+                if isinstance(resource.get("attributes"), dict)
+                else None
+            ),
             metadata=resource.get("metadata"),
             members=resource.get("members"),
         )
@@ -70,6 +113,8 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
         dumped = super().dump(camel_case=camel_case)
         if self.capabilities is not None:
             dumped["capabilities"] = [c.dump(camel_case=camel_case) for c in self.capabilities]
+        if self.attributes is not None:
+            dumped["attributes"] = self.attributes.dump(camel_case=camel_case)
         return dumped
 
 
@@ -83,6 +128,7 @@ class Group(GroupCore):
         name (str): Name of the group.
         source_id (str | None): ID of the group in the source. If this is the same ID as a group in the IdP, a service account in that group will implicitly be a part of this group as well. Can not be used together with 'members'.
         capabilities (list[Capability] | None): List of capabilities (acls) this group should grant its users.
+        attributes (GroupAttributes | None): Attributes of the group, this scopes down access based on the attributes specified.
         id (int | None): No description.
         is_deleted (bool | None): No description.
         deleted_time (int | None): No description.
@@ -96,6 +142,7 @@ def __init__(
         name: str,
         source_id: str | None = None,
         capabilities: list[Capability] | None = None,
+        attributes: GroupAttributes | None = None,
         id: int | None = None,
         is_deleted: bool | None = None,
         deleted_time: int | None = None,
@@ -107,6 +154,7 @@ def __init__(
             name=name,
             source_id=source_id,
             capabilities=capabilities,
+            attributes=attributes,
             metadata=metadata,
             members=members,
         )
@@ -125,6 +173,7 @@ def as_write(self) -> GroupWrite:
             name=self.name,
             source_id=self.source_id,
             capabilities=self.capabilities,
+            attributes=self.attributes,
             metadata=self.metadata,
             members=self.members,
         )
@@ -147,6 +196,11 @@ def _load(cls, resource: dict, cognite_client: CogniteClient | None = None, allo
         return cls(
             name=resource["name"],
             source_id=resource.get("sourceId"),
+            attributes=(
+                GroupAttributes._load(resource["attributes"], cognite_client=cognite_client)
+                if isinstance(resource.get("attributes"), dict)
+                else None
+            ),
             capabilities=[Capability.load(c, allow_unknown) for c in resource.get("capabilities", [])] or None,
             id=resource.get("id"),
             is_deleted=resource.get("isDeleted"),
@@ -184,6 +238,7 @@ class GroupWrite(GroupCore):
         name (str): Name of the group.
         source_id (str | None): ID of the group in the source. If this is the same ID as a group in the IdP, a service account in that group will implicitly be a part of this group as well. Can not be used together with 'members'.
         capabilities (list[Capability] | None): List of capabilities (acls) this group should grant its users.
+        attributes (GroupAttributes | None): Attributes of the group, this scopes down access based on the attributes specified.
         metadata (dict[str, str] | None): Custom, immutable application specific metadata. String key -> String value. Limits: Key are at most 32 bytes. Values are at most 512 bytes. Up to 16 key-value pairs. Total size is at most 4096.
         members (Literal['allUserAccounts'] | list[str] | None): Specifies which users are members of the group. Can not be used together with 'source_id'.
     """
@@ -193,12 +248,14 @@ def __init__(
         name: str,
         source_id: str | None = None,
         capabilities: list[Capability] | None = None,
+        attributes: GroupAttributes | None = None,
         metadata: dict[str, str] | None = None,
         members: Literal["allUserAccounts"] | list[str] | None = None,
     ) -> None:
         super().__init__(
             name=name,
             source_id=source_id,
+            attributes=attributes,
             capabilities=capabilities,
             metadata=metadata,
             members=members,
@@ -377,7 +434,10 @@ def __init__(self, subject: str, projects: list[ProjectSpec], capabilities: Proj
 
     @classmethod
     def load(
-        cls, api_response: dict[str, Any], cognite_client: CogniteClient | None = None, allow_unknown: bool = False
+        cls,
+        api_response: dict[str, Any],
+        cognite_client: CogniteClient | None = None,
+        allow_unknown: bool = False,
     ) -> TokenInspection:
         return cls(
             subject=api_response["subject"],

tests/tests_unit/test_api/test_iam.py

@@ -4,7 +4,7 @@
 
 from cognite.client.data_classes import Group, GroupList, SecurityCategory, SecurityCategoryList
 from cognite.client.data_classes.capabilities import AllScope, GroupsAcl, ProjectCapability, ProjectCapabilityList
-from cognite.client.data_classes.iam import ProjectSpec, TokenInspection
+from cognite.client.data_classes.iam import GroupAttributes, ProjectSpec, TokenInspection
 from tests.utils import jsgz_load
 
 
@@ -29,12 +29,50 @@ def mock_groups(rsps, cognite_client):
     yield rsps
 
 
+@pytest.fixture
+def mock_groups_with_attributes(rsps, cognite_client):
+    response_body = {
+        "items": [
+            {
+                "name": "Production Engineers",
+                "sourceId": "b7c9a5a4-99c2-4785-bed3-5e6ad9a78603",
+                "capabilities": [{"groupsAcl": {"actions": ["LIST"], "scope": {"all": {}}}}],
+                "id": 0,
+                "isDeleted": False,
+                "deletedTime": 0,
+                "attributes": {
+                    "token": {
+                        "appIds": ["app1", "app2"],
+                    },
+                    "unknownProperty": "unknownValue",
+                },
+            }
+        ]
+    }
+    url_pattern = re.compile(re.escape(cognite_client.iam._get_base_url_with_base_path()) + "/groups.*")
+    rsps.assert_all_requests_are_fired = False
+    rsps.add(rsps.POST, url_pattern, status=200, json=response_body)
+    rsps.add(rsps.GET, url_pattern, status=200, json=response_body)
+    yield rsps
+
+
 class TestGroups:
     def test_list(self, cognite_client, mock_groups):
         res = cognite_client.iam.groups.list()
         assert isinstance(res, GroupList)
         assert mock_groups.calls[0].response.json()["items"] == res.dump(camel_case=True)
 
+    def test_list_groups_with_attributes(self, cognite_client, mock_groups_with_attributes):
+        res = cognite_client.iam.groups.list()
+        assert isinstance(res, GroupList)
+        assert mock_groups_with_attributes.calls[0].response.json()["items"] == res.dump(camel_case=True)
+        assert mock_groups_with_attributes.calls[0].response.json()["items"][0]["attributes"] == {
+            "token": {
+                "appIds": ["app1", "app2"],
+            },
+            "unknownProperty": "unknownValue",
+        }
+
     def test_create(self, cognite_client, mock_groups):
         my_group = Group(name="My Group", capabilities=[GroupsAcl([GroupsAcl.Action.List], AllScope())])
         res = cognite_client.iam.groups.create(my_group)
@@ -46,6 +84,31 @@ def test_create(self, cognite_client, mock_groups):
         } == jsgz_load(mock_groups.calls[0].request.body)
         assert mock_groups.calls[0].response.json()["items"][0] == res.dump(camel_case=True)
 
+    def test_create_with_attributes(self, cognite_client, mock_groups_with_attributes):
+        # Construct attributes via loader to include unknown properties for pass-through
+        attributes = GroupAttributes.load(
+            {
+                "token": {"appIds": ["app1", "app2"]},
+                "unknownProperty": "unknownValue",
+            }
+        )
+        my_group = Group(
+            name="My Group",
+            capabilities=[GroupsAcl([GroupsAcl.Action.List], AllScope())],
+            attributes=attributes,
+        )
+        res = cognite_client.iam.groups.create(my_group)
+        assert isinstance(res, Group)
+        assert {
+            "items": [
+                {
+                    "name": "My Group",
+                    "capabilities": [{"groupsAcl": {"actions": ["LIST"], "scope": {"all": {}}}}],
+                    "attributes": {"token": {"appIds": ["app1", "app2"]}, "unknownProperty": "unknownValue"},
+                }
+            ]
+        } == jsgz_load(mock_groups_with_attributes.calls[0].request.body)
+
     def test_create_multiple(self, cognite_client, mock_groups):
         res = cognite_client.iam.groups.create([1])
         assert isinstance(res, GroupList)