chore: [CDF-25309] Introduce service principals data classes (#2213)

doctrino · web-flow · commit 9c276116f98a · 2025-09-08T10:06:16.000Z
diff --git a/cognite/client/data_classes/principals.py b/cognite/client/data_classes/principals.py
@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+from typing_extensions import Self
+
+from cognite.client.data_classes._base import (
+    CogniteObject,
+    CogniteResource,
+    CogniteResourceList,
+)
+
+if TYPE_CHECKING:
+    from cognite.client import CogniteClient
+
+
+class Principal(CogniteResource, ABC):
+    _type: ClassVar[str]
+
+    def __init__(self, id: str) -> None:
+        self.id = id
+
+    @classmethod
+    @abstractmethod
+    def _load_principal(cls, resource: dict[str, Any]) -> Self:
+        """Load a principal from a resource dictionary."""
+        raise NotImplementedError("This method should be implemented in subclasses.")
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        type_ = resource.get("type")
+        if type_ is None and hasattr(cls, "_type"):
+            type_ = cls._type
+        elif type_ is None:
+            raise KeyError("type")
+        principal_cls = _PRINCIPAL_CLS_BY_TYPE.get(type_.lower())
+        if principal_cls is None:
+            return UnknownPrincipal._load_principal(resource)  # type: ignore[return-value]
+        return cast(
+            Self,
+            principal_cls._load_principal(
+                resource,
+            ),
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the principal to a dictionary."""
+        output = super().dump(camel_case=camel_case)
+        output["type"] = self._type.upper()
+        return output
+
+
+class UserPrincipal(Principal):
+    """Represents a user principal in Cognite Data Fusion (CDF).
+
+    Arguments:
+        id (str): The ID of an organization user
+        name (str): Human-readable name of the principal
+        picture_url (str): URL to a picture of the principal
+        email (str | None): User email. Do not use this to uniquely identify a user, as it can be changed
+            and is not guaranteed to be unique. Use the id field instead.
+        given_name (str | None): The given name of the user
+        middle_name (str | None): The middle name of the user
+        family_name (str | None): The family name of the user
+
+    """
+
+    _type = "user"
+
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        picture_url: str,
+        email: str | None = None,
+        given_name: str | None = None,
+        middle_name: str | None = None,
+        family_name: str | None = None,
+    ) -> None:
+        super().__init__(id)
+        self.name = name
+        self.picture_url = picture_url
+        self.email = email
+        self.given_name = given_name
+        self.middle_name = middle_name
+        self.family_name = family_name
+
+    @classmethod
+    def _load_principal(cls, resource: dict[str, Any]) -> Self:
+        return cls(
+            id=resource["id"],
+            name=resource["name"],
+            picture_url=resource["pictureUrl"],
+            email=resource.get("email"),
+            given_name=resource.get("givenName"),
+            middle_name=resource.get("middleName"),
+            family_name=resource.get("familyName"),
+        )
+
+
+@dataclass
+class ServiceAccountCreator(CogniteObject):
+    """The creator of a service account.
+
+    Arguments:
+        org_id (str): The ID of an organization.
+        user_id (str): The ID of an organization user
+
+    """
+
+    org_id: str
+    user_id: str
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> Self:
+        return cls(org_id=resource["orgId"], user_id=resource["userId"])
+
+
+class ServicePrincipal(Principal):
+    """Represents a service account principal in Cognite Data Fusion (CDF).
+
+    Arguments:
+        id (str): Unique identifier of a service account
+        name (str): Human-readable name of the service account
+        created_by (ServiceAccountCreator): The creator of the service account
+        created_time (int): When the principal was created. It is given as the number of milliseconds
+            since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
+        last_updated_time (int): When the principal was last updated.  It is given as the number of milliseconds
+            since 00:00:00 Thursday, 1 January 1970, Coordinated Universal Time (UTC), minus leap seconds.
+        picture_url (str): URL to a picture of the principal.
+        external_id (str | None): The external ID provided by the client. Must be unique for the resource type.
+        description (str | None): A description of the service account.
+
+    """
+
+    _type = "service_account"
+
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        created_by: ServiceAccountCreator,
+        created_time: int,
+        last_updated_time: int,
+        picture_url: str,
+        external_id: str | None = None,
+        description: str | None = None,
+    ) -> None:
+        super().__init__(id)
+        self.name = name
+        self.created_by = created_by
+        self.created_time = created_time
+        self.last_updated_time = last_updated_time
+        self.picture_url = picture_url
+        self.external_id = external_id
+        self.description = description
+
+    @classmethod
+    def _load_principal(cls, resource: dict[str, Any]) -> Self:
+        return cls(
+            id=resource["id"],
+            name=resource["name"],
+            created_by=ServiceAccountCreator._load(resource["createdBy"]),
+            created_time=resource["createdTime"],
+            last_updated_time=resource["lastUpdatedTime"],
+            picture_url=resource["pictureUrl"],
+            external_id=resource.get("externalId"),
+            description=resource.get("description"),
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        output = super().dump(camel_case=camel_case)
+        output["createdBy" if camel_case else "created_by"] = self.created_by.dump(camel_case=camel_case)
+        return output
+
+
+class UnknownPrincipal(Principal):
+    """Represents an unknown principal in Cognite Data Fusion (CDF).
+
+    This class is used when the principal type is not recognized or not defined in the SDK.
+    Typically, this can happen when a new type of principal is introduced in CDF that is not yet supported by the SDK.
+
+    Arguments:
+        id (str): Unique identifier of the principal.
+        type (str): The type of the principal, which is not recognized by the SDK.
+        data (dict[str, Any]): Additional data associated with the principal, excluding the 'id' and 'type' fields.
+
+    """
+
+    _type = "unknown"
+
+    def __init__(self, id: str, type: str, data: dict[str, Any]) -> None:
+        super().__init__(id)
+        self.type = type
+        self.__data = data
+
+    @classmethod
+    def _load_principal(cls, resource: dict[str, Any]) -> Self:
+        return cls(
+            id=resource["id"],
+            type=resource["type"],
+            data={k: v for k, v in resource.items() if k not in {"id", "type"}},
+        )
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        return {"id": self.id, "type": self.type, **self.__data}
+
+
+class PrincipalList(CogniteResourceList[Principal]):
+    _RESOURCE = Principal
+
+    def as_ids(self) -> list[str]:
+        """Returns a list of principal IDs."""
+        return [principal.id for principal in self]
+
+
+# Build the mapping AFTER all classes are defined
+_PRINCIPAL_CLS_BY_TYPE: dict[str, type[Principal]] = {
+    subclass._type: subclass  # type: ignore[type-abstract]
+    for subclass in Principal.__subclasses__()
+    if hasattr(subclass, "_type") and not getattr(subclass, "__abstractmethods__", None)
+}
diff --git a/tests/tests_unit/test_data_classes/test_principals.py b/tests/tests_unit/test_data_classes/test_principals.py
@@ -0,0 +1,62 @@
+import pytest
+
+from cognite.client.data_classes.principals import Principal, ServicePrincipal, UnknownPrincipal, UserPrincipal
+
+
+class TestPrincipals:
+    @pytest.mark.parametrize(
+        "data, expected_cls",
+        [
+            pytest.param(
+                {
+                    "id": "123",
+                    "name": "test_principal",
+                    "type": "USER",
+                    "givenName": "Test",
+                    "familyName": "Principal",
+                    "middleName": "Middle",
+                    "pictureUrl": "http://example.com/picture.jpg",
+                },
+                UserPrincipal,
+                id="user_principal",
+            ),
+            pytest.param(
+                {
+                    "id": "456",
+                    "name": "another_principal",
+                    "type": "SERVICE_ACCOUNT",
+                    "createdBy": {
+                        "orgId": "org_123",
+                        "userId": "user_456",
+                    },
+                    "createdTime": 1622547800,
+                    "lastUpdatedTime": 1622547800,
+                    "pictureUrl": "http://example.com/service_principal.jpg",
+                    "externalId": "service_principal_123",
+                    "description": "A service principal for testing purposes",
+                },
+                ServicePrincipal,
+                id="service principal",
+            ),
+            pytest.param(
+                {
+                    "id": "some_id",
+                    "type": "NEW_TYPE",
+                    "name": "new_name",
+                    "someOtherField": {"nestedField": "value"},
+                },
+                UnknownPrincipal,
+                id="unknown_principal",
+            ),
+        ],
+    )
+    def test_dump_load(self, data: dict[str, object], expected_cls: type[Principal]) -> None:
+        loaded = Principal._load(data)
+
+        assert isinstance(loaded, expected_cls), (
+            f"Loaded principal is not of type {expected_cls.__name__}: {type(loaded).__name__}"
+        )
+
+        dumped = loaded.dump(camel_case=True)
+
+        assert dumped == data, f"Dumped data does not match original: {dumped} != {data}"
diff --git a/tests/tests_unit/test_meta.py b/tests/tests_unit/test_meta.py
@@ -12,6 +12,7 @@
     InternalIdTransformerMixin,
 )
 from cognite.client.data_classes.datapoints import DatapointsArrayList, DatapointsList
+from cognite.client.data_classes.principals import PrincipalList
 from tests.utils import all_concrete_subclasses, all_subclasses
 
 ALL_FILEPATHS = Path("cognite/client/").rglob("*.py")
@@ -87,7 +88,15 @@ def test_all_base_api_paths_have_retry_or_specifically_no_set(
     assert not (has_retry and no_retry_needed)
 
 
-@pytest.mark.parametrize("lst_cls", all_concrete_subclasses(CogniteResourceList))
+@pytest.mark.parametrize(
+    "lst_cls",
+    [
+        list_cls
+        # Principal list .as_ids() returns a list of strings and not integers,
+        # so we skip the check for it.
+        for list_cls in all_concrete_subclasses(CogniteResourceList, exclude={PrincipalList})
+    ],
+)
 def test_ensure_identifier_mixins(lst_cls):
     # TODO: Data Modeling uses "as_ids()" even though existing classes use the same for "integer internal ids"
     if "data_modeling" in str(lst_cls):
diff --git a/tests/utils.py b/tests/utils.py
@@ -97,10 +97,10 @@ def all_subclasses(base: T_Type, exclude: set[type] | None = None) -> list[T_Typ
     )
 
 
-def all_concrete_subclasses(base: T_Type) -> list[T_Type]:
+def all_concrete_subclasses(base: T_Type, exclude: set[type] | None = None) -> list[T_Type]:
     return [
         sub
-        for sub in all_subclasses(base)
+        for sub in all_subclasses(base, exclude=exclude)
         if all(base is not abc.ABC for base in sub.__bases__)
         and not inspect.isabstract(sub)
         # The FakeCogniteResourceGenerator does not support descriptors, so we exclude the Typed classes
