chore: make OrgAPI into a base client (#2359)

Author: haakonvt

cognite/client/_api/iam.py

@@ -6,7 +6,7 @@
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast, overload
 
-from cognite.client._api.organization import PrincipalsAPI
+from cognite.client._api.org_apis.principals import PrincipalsAPI
 from cognite.client._api.user_profiles import UserProfilesAPI
 from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ

cognite/client/_api/organization.py …gnite/client/_api/org_apis/principals.pycognite/client/_api/organization.py renamed to cognite/client/_api/org_apis/principals.py cognite/client/_api/organization.py renamed to cognite/client/_api/org_apis/principals.py

@@ -1,50 +1,18 @@
 from __future__ import annotations
 
 import warnings
-from abc import ABC
 from collections.abc import Sequence
-from functools import cached_property
 from typing import overload
 from urllib.parse import urljoin
 
-from cognite.client._api_client import APIClient
 from cognite.client._constants import DEFAULT_LIMIT_READ
+from cognite.client._org_client import OrgAPIClient
 from cognite.client.data_classes.principals import Principal, PrincipalList
-from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils._identifier import PrincipalIdentifierSequence
 from cognite.client.utils.useful_types import SequenceNotStr
 
 
-class OrgAPI(APIClient, ABC):
-    _auth_url = "https://auth.cognite.com"
-
-    def _get_base_url_with_base_path(self) -> str:
-        """Get base URL with base path including organization and api version if applicable"""
-        base_path = ""
-        if self._api_version:
-            base_path = f"/api/{self._api_version}/orgs/{self._organization}"
-        # The OrganizationAPi uses the auth_url as the base for these endpoints instead of the
-        # base_url like the rest of the SDK.
-        return urljoin(self._auth_url, base_path)
-
-    @cached_property
-    def _organization(self) -> str:
-        headers = self._configure_headers(
-            "application/json",
-            additional_headers=self._config.headers.copy(),
-            api_subversion=self._api_subversion,
-        )
-        # This is an internal endpoint, not part of the public API
-        full_url = urljoin(self._config.base_url, f"/api/v1/projects/{self._config.project}")
-        response = self._http_client_with_retry.request(method="GET", url=full_url, headers=headers)
-        if response.status_code != 200:
-            raise CogniteAPIError(
-                "Could not look-up organization", response.status_code, response.headers.get("x-request-id")
-            )
-        return response.json()["organization"]
-
-
-class PrincipalsAPI(OrgAPI):
+class PrincipalsAPI(OrgAPIClient):
     _RESOURCE_PATH = "/principals"
 
     def me(self) -> Principal:
@@ -61,9 +29,10 @@ def me(self) -> Principal:
                 >>> res = client.iam.principals.me()
         """
         # the /me endpoint is not using the /orgs/{org} base path, so we have to construct the URL manually
-        path = "/principals/me"
+        path = f"{self._RESOURCE_PATH}/me"
         if self._api_version:
             path = f"/api/{self._api_version}{path}"
+
         full_url = urljoin(self._auth_url, path)
         headers = self._configure_headers(
             "application/json",

cognite/client/_org_client.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from abc import ABC
+from functools import cached_property
+from urllib.parse import urljoin
+
+from cognite.client._api_client import APIClient
+from cognite.client.exceptions import CogniteAPIError
+
+
+class OrgAPIClient(APIClient, ABC):
+    _auth_url = "https://auth.cognite.com"
+
+    def _get_base_url_with_base_path(self) -> str:
+        """Get base URL with base path including organization and api version if applicable"""
+        base_path = ""
+        if self._api_version:
+            base_path = f"/api/{self._api_version}/orgs/{self._organization}"
+        # The OrganizationAPi uses the auth_url as the base for these endpoints instead of the
+        # base_url like the rest of the SDK.
+        return urljoin(self._auth_url, base_path)
+
+    @cached_property
+    def _organization(self) -> str:
+        headers = self._configure_headers(
+            "application/json",
+            additional_headers=self._config.headers.copy(),
+            api_subversion=self._api_subversion,
+        )
+        # This is an internal endpoint, not part of the public API
+        full_url = urljoin(self._config.base_url, f"/api/v1/projects/{self._config.project}")
+        response = self._http_client_with_retry.request(method="GET", url=full_url, headers=headers)
+        if response.status_code != 200:
+            raise CogniteAPIError(
+                "Could not look-up organization", response.status_code, response.headers.get("x-request-id")
+            )
+        return response.json()["organization"]

cognite/client/testing.py

@@ -42,7 +42,7 @@
 from cognite.client._api.hosted_extractors.sources import SourcesAPI
 from cognite.client._api.iam import IAMAPI, GroupsAPI, SecurityCategoriesAPI, SessionsAPI, TokenAPI
 from cognite.client._api.labels import LabelsAPI
-from cognite.client._api.organization import PrincipalsAPI
+from cognite.client._api.org_apis.principals import PrincipalsAPI
 from cognite.client._api.postgres_gateway import PostgresGatewaysAPI
 from cognite.client._api.postgres_gateway.tables import TablesAPI as PostgresTablesAPI
 from cognite.client._api.postgres_gateway.users import UsersAPI as PostgresUsersAPI

docs/source/identity_and_access_management.rst

@@ -15,15 +15,15 @@ Principals
 
 Get the current caller's principal
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. automethod:: cognite.client._api.organization.PrincipalsAPI.me
+.. automethod:: cognite.client._api.org_apis.principals.PrincipalsAPI.me
 
 Retrieve a principal
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. automethod:: cognite.client._api.organization.PrincipalsAPI.retrieve
+.. automethod:: cognite.client._api.org_apis.principals.PrincipalsAPI.retrieve
 
 List principals
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. automethod:: cognite.client._api.organization.PrincipalsAPI.list
+.. automethod:: cognite.client._api.org_apis.principals.PrincipalsAPI.list
 
 Tokens
 ^^^^^^

tests/tests_integration/test_api/test_organization.py

@@ -1,17 +1,17 @@
 import pytest
 
 from cognite.client import CogniteClient
-from cognite.client._api.organization import OrgAPI
+from cognite.client._org_client import OrgAPIClient
 
 
 @pytest.fixture(scope="session")
-def org_api(cognite_client_cog_idp: CogniteClient) -> OrgAPI:
+def org_api(cognite_client_cog_idp: CogniteClient) -> OrgAPIClient:
     client = cognite_client_cog_idp
-    return OrgAPI(client.config, client._API_VERSION, client)
+    return OrgAPIClient(client.config, client._API_VERSION, client)
 
 
 class TestOrganizationAPI:
-    def test_get_organization(self, org_api: OrgAPI) -> None:
+    def test_get_organization(self, org_api: OrgAPIClient) -> None:
         organization = org_api._organization
         assert isinstance(organization, str)
         assert organization != ""
@@ -20,6 +20,6 @@ def test_get_organization(self, org_api: OrgAPI) -> None:
         else:
             pytest.skip("Organization not supported")
 
-    def test_get_request_org_endpoint(self, org_api: OrgAPI) -> None:
+    def test_get_request_org_endpoint(self, org_api: OrgAPIClient) -> None:
         response = org_api._get("/principals")
         assert response.status_code == 200

tests/tests_unit/test_testing.py

@@ -1,8 +1,8 @@
 import pytest
 
 from cognite.client import ClientConfig, CogniteClient
-from cognite.client._api.organization import OrgAPI
 from cognite.client._api_client import APIClient
+from cognite.client._org_client import OrgAPIClient
 from cognite.client.credentials import Token
 from cognite.client.testing import CogniteClientMock, monkeypatch_cognite_client
 from tests.utils import all_mock_children, all_subclasses, get_api_class_by_attribute
@@ -11,8 +11,8 @@
 def test_ensure_all_apis_are_available_on_cognite_mock():
     mocked_apis = all_mock_children(CogniteClientMock())
     available = {v.__class__ for v in mocked_apis.values()}
-    # OrgAPI is a base API and should not mocked directly.
-    expected = set(all_subclasses(APIClient, exclude={OrgAPI}))
+    # OrgAPIClient is a base API and should not mocked directly.
+    expected = set(all_subclasses(APIClient, exclude={OrgAPIClient}))
     # Any new APIs that have not been added to CogniteClientMock?
     assert not expected.difference(available), f"Missing APIs: {expected.difference(available)}"
     # Any removed APIs that are still available on CogniteClientMock?