users.py

from __future__ import annotations

from collections.abc import AsyncIterator, Sequence
from typing import TYPE_CHECKING, overload

from cognite.client._api_client import APIClient
from cognite.client.constants import DEFAULT_LIMIT_READ
from cognite.client.data_classes.postgres_gateway.users import (
    User,
    UserCreated,
    UserCreatedList,
    UserList,
    UserUpdate,
    UserWrite,
)
from cognite.client.utils._identifier import UsernameSequence
from cognite.client.utils.useful_types import SequenceNotStr

if TYPE_CHECKING:
    from cognite.client import AsyncCogniteClient, ClientConfig


class UsersAPI(APIClient):
    _RESOURCE_PATH = "/postgresgateway"

    def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: AsyncCogniteClient) -> None:
        super().__init__(config, api_version, cognite_client)
        self._CREATE_LIMIT = 1
        self._UPDATE_LIMIT = 1
        self._DELETE_LIMIT = 1
        self._LIST_LIMIT = 100
        self._RETRIEVE_LIMIT = 10

    @overload
    def __call__(self, chunk_size: None = None, limit: int | None = None) -> AsyncIterator[User]: ...

    @overload
    def __call__(self, chunk_size: int, limit: int | None = None) -> AsyncIterator[UserList]: ...

    async def __call__(
        self, chunk_size: int | None = None, limit: int | None = None
    ) -> AsyncIterator[User] | AsyncIterator[UserList]:
        """Iterate over users

        Fetches user as they are iterated over, so you keep a limited number of users in memory.

        Args:
            chunk_size (int | None): Number of users to return in each chunk. Defaults to yielding one user at a time.
            limit (int | None): Maximum number of users to return. Defaults to return all.

        Yields:
            User | UserList: yields User one by one if chunk_size is not specified, else UserList objects.
        """  # noqa: DOC404
        async for item in self._list_generator(
            list_cls=UserList,
            resource_cls=User,
            method="GET",
            chunk_size=chunk_size,
            limit=limit,
        ):
            yield item

    @overload
    async def create(self, user: UserWrite) -> UserCreated: ...

    @overload
    async def create(self, user: Sequence[UserWrite]) -> UserCreatedList: ...

    async def create(self, user: UserWrite | Sequence[UserWrite]) -> UserCreated | UserCreatedList:
        """`Create Users <https://api-docs.cognite.com/20230101-beta/tag/Postgres-Gateway-Users/operation/create_users>`_.

        Create postgres users.

        Args:
            user (UserWrite | Sequence[UserWrite]): The user(s) to create.

        Returns:
            UserCreated | UserCreatedList: The created user(s)

        Examples:

            Create user:

                >>> import os
                >>> from cognite.client import CogniteClient
                >>> from cognite.client.data_classes.postgres_gateway import (
                ...     UserWrite,
                ...     SessionCredentials,
                ... )
                >>> from cognite.client.data_classes import ClientCredentials
                >>> client = CogniteClient()
                >>> # async_client = AsyncCogniteClient()  # another option
                >>> session = client.iam.sessions.create(
                ...     ClientCredentials(os.environ["IDP_CLIENT_ID"], os.environ["IDP_CLIENT_SECRET"]),
                ...     session_type="CLIENT_CREDENTIALS",
                ... )
                >>> user = UserWrite(credentials=SessionCredentials(nonce=session.nonce))
                >>> res = client.postgres_gateway.users.create(user)

        """
        return await self._create_multiple(
            list_cls=UserCreatedList,
            resource_cls=UserCreated,
            items=user,
            input_resource_cls=UserWrite,
        )

    @overload
    async def update(self, items: UserUpdate | UserWrite) -> User: ...

    @overload
    async def update(self, items: Sequence[UserUpdate | UserWrite]) -> UserList: ...

    async def update(self, items: UserUpdate | UserWrite | Sequence[UserUpdate | UserWrite]) -> User | UserList:
        """`Update users <https://api-docs.cognite.com/20230101-beta/tag/Postgres-Gateway-Users/operation/update_users>`_.

        Update postgres users

        Args:
            items (UserUpdate | UserWrite | Sequence[UserUpdate | UserWrite]): The user(s) to update.

        Returns:
            User | UserList: The updated user(s)

        Examples:

            Update user:

                >>> import os
                >>> from cognite.client import CogniteClient
                >>> from cognite.client.data_classes.postgres_gateway import (
                ...     UserUpdate,
                ...     SessionCredentials,
                ... )
                >>> from cognite.client.data_classes import ClientCredentials
                >>> client = CogniteClient()
                >>> # async_client = AsyncCogniteClient()  # another option
                >>> session = client.iam.sessions.create(
                ...     ClientCredentials(os.environ["IDP_CLIENT_ID"], os.environ["IDP_CLIENT_SECRET"]),
                ...     session_type="CLIENT_CREDENTIALS",
                ... )
                >>> update = UserUpdate("myUser").credentials.set(
                ...     SessionCredentials(nonce=session.nonce)
                ... )
                >>> res = client.postgres_gateway.users.update(update)

        """
        return await self._update_multiple(
            items=items,
            list_cls=UserList,
            resource_cls=User,
            update_cls=UserUpdate,
        )

    async def delete(self, username: str | SequenceNotStr[str], ignore_unknown_ids: bool = False) -> None:
        """`Delete postgres user(s) <https://api-docs.cognite.com/20230101-beta/tag/Postgres-Gateway-Users/operation/delete_users>`_.

        Delete postgres users

        Args:
            username (str | SequenceNotStr[str]): Usernames of the users to delete.
            ignore_unknown_ids (bool): Ignore usernames that are not found


        Examples:

            Delete users:

                >>> from cognite.client import CogniteClient, AsyncCogniteClient
                >>> client = CogniteClient()
                >>> # async_client = AsyncCogniteClient()  # another option
                >>> client.postgres_gateway.users.delete(["myUser", "myUser2"])


        """
        extra_body_fields = {"ignore_unknown_ids": ignore_unknown_ids}

        await self._delete_multiple(
            identifiers=UsernameSequence.load(usernames=username),
            wrap_ids=True,
            returns_items=False,
            extra_body_fields=extra_body_fields,
        )

    @overload
    async def retrieve(self, username: str, ignore_unknown_ids: bool = False) -> User: ...

    @overload
    async def retrieve(self, username: SequenceNotStr[str], ignore_unknown_ids: bool = False) -> UserList: ...

    async def retrieve(self, username: str | SequenceNotStr[str], ignore_unknown_ids: bool = False) -> User | UserList:
        """`Retrieve a list of users by their usernames <https://api-docs.cognite.com/20230101-beta/tag/Postgres-Gateway-Users/operation/retreive_users>`_.

        Retrieve a list of postgres users by their usernames, optionally ignoring unknown usernames

        Args:
            username (str | SequenceNotStr[str]): Usernames of the users to retrieve.
            ignore_unknown_ids (bool): Ignore usernames that are not found

        Returns:
            User | UserList: The retrieved user(s).

        Examples:

            Retrieve user:

                    >>> from cognite.client import CogniteClient, AsyncCogniteClient
                    >>> client = CogniteClient()
                    >>> # async_client = AsyncCogniteClient()  # another option
                    >>> res = client.postgres_gateway.users.retrieve("myUser", ignore_unknown_ids=True)

        """
        return await self._retrieve_multiple(
            list_cls=UserList,
            resource_cls=User,
            identifiers=UsernameSequence.load(usernames=username),
            ignore_unknown_ids=ignore_unknown_ids,
        )

    async def list(self, limit: int = DEFAULT_LIMIT_READ) -> UserList:
        """`Fetch scoped users <https://api-docs.cognite.com/20230101-beta/tag/Postgres-Gateway-Users/operation/filter_users>`_.

        List all users in a given project.

        Args:
            limit (int): Limits the number of results to be returned.

        Returns:
            UserList: A list of users

        Examples:

            List users:

                >>> from cognite.client import CogniteClient, AsyncCogniteClient
                >>> client = CogniteClient()
                >>> # async_client = AsyncCogniteClient()  # another option
                >>> user_list = client.postgres_gateway.users.list(limit=5)

            Iterate over users, one-by-one:

                >>> for user in client.postgres_gateway.users():
                ...     user  # do something with the user

            Iterate over chunks of users to reduce memory load:

                >>> for user_list in client.postgres_gateway.users(chunk_size=25):
                ...     user_list  # do something with the users

        """
        return await self._list(
            list_cls=UserList,
            resource_cls=User,
            method="GET",
            limit=limit,
        )