anaplan-sdk 0.5.0a1__py3-none-any.whl → 0.5.0a3__py3-none-any.whl

anaplan_sdk/_async_clients/_bulk.py

@@ -1,13 +1,13 @@
 import logging
-from asyncio import gather, sleep
+from asyncio import gather
 from copy import copy
-from typing import AsyncIterator, Iterator
+from typing import AsyncIterator, Iterator, Literal
 
 import httpx
 from typing_extensions import Self
 
 from anaplan_sdk._auth import _create_auth
-from anaplan_sdk._base import _AsyncBaseClient, action_url
+from anaplan_sdk._services import _AsyncHttpService, action_url, sort_params
 from anaplan_sdk.exceptions import AnaplanActionError, InvalidIdentifierException
 from anaplan_sdk.models import (
     Action,
@@ -30,7 +30,7 @@ from ._transactional import _AsyncTransactionalClient
 logger = logging.getLogger("anaplan_sdk")
 
 
-class AsyncClient(_AsyncBaseClient):
+class AsyncClient:
     """
     Asynchronous Anaplan Client. For guides and examples
     refer to https://vinzenzklass.github.io/anaplan-sdk.
@@ -49,9 +49,13 @@ class AsyncClient(_AsyncBaseClient):
         auth: httpx.Auth | None = None,
         timeout: float | httpx.Timeout = 30,
         retry_count: int = 2,
+        backoff: float = 1.0,
+        backoff_factor: float = 2.0,
+        page_size: int = 5_000,
         status_poll_delay: int = 1,
         upload_chunk_size: int = 25_000_000,
         allow_file_creation: bool = False,
+        **httpx_kwargs,
     ) -> None:
         """
         Asynchronous Anaplan Client. For guides and examples
@@ -84,6 +88,14 @@ class AsyncClient(_AsyncBaseClient):
         :param retry_count: The number of times to retry an HTTP request if it fails. Set this to 0
                to never retry. Defaults to 2, meaning each HTTP Operation will be tried a total
                number of 2 times.
+        :param backoff: The initial backoff time in seconds for the retry mechanism. This is the
+               time to wait before the first retry.
+        :param backoff_factor: The factor by which the backoff time is multiplied after each retry.
+               For example, if the initial backoff is 1 second and the factor is 2, the second
+               retry will wait 2 seconds, the third retry will wait 4 seconds, and so on.
+        :param page_size: The number of items to return per page when paginating through results.
+               Defaults to 5000. This is the maximum number of items that can be returned per
+               request. If you pass a value greater than 5000, it will be capped to 5000.
         :param status_poll_delay: The delay between polling the status of a task.
         :param upload_chunk_size: The size of the chunks to upload. This is the maximum size of
                each chunk. Defaults to 25MB.
@@ -92,39 +104,41 @@ class AsyncClient(_AsyncBaseClient):
                altogether. A file that is created this way will not be referenced by any action in
                anaplan until manually assigned so there is typically no value in dynamically
                creating new files and uploading content to them.
-        """
-        _client = httpx.AsyncClient(
-            auth=(
-                auth
-                or _create_auth(
-                    token=token,
-                    user_email=user_email,
-                    password=password,
-                    certificate=certificate,
-                    private_key=private_key,
-                    private_key_password=private_key_password,
-                )
-            ),
-            timeout=timeout,
+        :param httpx_kwargs: Additional keyword arguments to pass to the `httpx.AsyncClient`.
+               This can be used to set additional options such as proxies, headers, etc. See
+               https://www.python-httpx.org/api/#asyncclient for the full list of arguments.
+        """
+        _auth = auth or _create_auth(
+            token=token,
+            user_email=user_email,
+            password=password,
+            certificate=certificate,
+            private_key=private_key,
+            private_key_password=private_key_password,
+        )
+        _client = httpx.AsyncClient(auth=_auth, timeout=timeout, **httpx_kwargs)
+        self._http = _AsyncHttpService(
+            _client,
+            retry_count=retry_count,
+            backoff=backoff,
+            backoff_factor=backoff_factor,
+            page_size=page_size,
+            poll_delay=status_poll_delay,
         )
         self._workspace_id = workspace_id
         self._model_id = model_id
-        self._retry_count = retry_count
         self._url = f"https://api.anaplan.com/2/0/workspaces/{workspace_id}/models/{model_id}"
         self._transactional_client = (
-            _AsyncTransactionalClient(_client, model_id, retry_count) if model_id else None
-        )
-        self._alm_client = (
-            _AsyncAlmClient(_client, model_id, self._retry_count, status_poll_delay)
-            if model_id
-            else None
+            _AsyncTransactionalClient(self._http, model_id) if model_id else None
         )
-        self._audit = _AsyncAuditClient(_client, self._retry_count)
-        self._cloud_works = _AsyncCloudWorksClient(_client, self._retry_count)
-        self.status_poll_delay = status_poll_delay
+        self._alm_client = _AsyncAlmClient(self._http, model_id) if model_id else None
+        self._audit = _AsyncAuditClient(self._http)
+        self._cloud_works = _AsyncCloudWorksClient(self._http)
         self.upload_chunk_size = upload_chunk_size
         self.allow_file_creation = allow_file_creation
-        super().__init__(retry_count, _client)
+        logger.debug(
+            f"Initialized AsyncClient with workspace_id={workspace_id}, model_id={model_id}"
+        )
 
     @classmethod
     def from_existing(
@@ -152,12 +166,8 @@ class AsyncClient(_AsyncBaseClient):
             f"with workspace_id={new_ws_id}, model_id={new_model_id}."
         )
         client._url = f"https://api.anaplan.com/2/0/workspaces/{new_ws_id}/models/{new_model_id}"
-        client._transactional_client = _AsyncTransactionalClient(
-            existing._client, new_model_id, existing._retry_count
-        )
-        client._alm_client = _AsyncAlmClient(
-            existing._client, new_model_id, existing._retry_count, existing.status_poll_delay
-        )
+        client._transactional_client = _AsyncTransactionalClient(existing._http, new_model_id)
+        client._alm_client = _AsyncAlmClient(existing._http, new_model_id)
         return client
 
     @property
@@ -177,14 +187,14 @@ class AsyncClient(_AsyncBaseClient):
         return self._cloud_works
 
     @property
-    def transactional(self) -> _AsyncTransactionalClient:
+    def tr(self) -> _AsyncTransactionalClient:
         """
         The Transactional Client provides access to the Anaplan Transactional API. This is useful
         for more advanced use cases where you need to interact with the Anaplan Model in a more
         granular way.
 
         If you instantiated the client without the field `model_id`, this will raise a
-        :py:class:`ValueError`, since none of the endpoints can be invoked without the model Id.
+        `ValueError`, since none of the endpoints can be invoked without the model Id.
         :return: The Transactional Client.
         """
         if not self._transactional_client:
@@ -215,40 +225,54 @@ class AsyncClient(_AsyncBaseClient):
             )
         return self._alm_client
 
-    async def get_workspaces(self, search_pattern: str | None = None) -> list[Workspace]:
+    async def get_workspaces(
+        self,
+        search_pattern: str | None = None,
+        sort_by: Literal["size_allowance", "name"] = "name",
+        descending: bool = False,
+    ) -> list[Workspace]:
         """
         Lists all the Workspaces the authenticated user has access to.
         :param search_pattern: Optionally filter for specific workspaces. When provided,
                case-insensitive matches workspaces with names containing this string.
                You can use the wildcards `%` for 0-n characters, and `_` for exactly 1 character.
                When None (default), returns all users.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Workspaces.
         """
-        params = {"tenantDetails": "true"}
+        params = {"tenantDetails": "true"} | sort_params(sort_by, descending)
         if search_pattern:
             params["s"] = search_pattern
         return [
             Workspace.model_validate(e)
-            for e in await self._get_paginated(
+            for e in await self._http.get_paginated(
                 "https://api.anaplan.com/2/0/workspaces", "workspaces", params=params
             )
         ]
 
-    async def get_models(self, search_pattern: str | None = None) -> list[Model]:
+    async def get_models(
+        self,
+        search_pattern: str | None = None,
+        sort_by: Literal["active_state", "name"] = "name",
+        descending: bool = False,
+    ) -> list[Model]:
         """
         Lists all the Models the authenticated user has access to.
         :param search_pattern: Optionally filter for specific models. When provided,
                case-insensitive matches model names containing this string.
                You can use the wildcards `%` for 0-n characters, and `_` for exactly 1 character.
                When None (default), returns all models.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Models.
         """
-        params = {"modelDetails": "true"}
+        params = {"modelDetails": "true"} | sort_params(sort_by, descending)
         if search_pattern:
             params["s"] = search_pattern
         return [
             Model.model_validate(e)
-            for e in await self._get_paginated(
+            for e in await self._http.get_paginated(
                 "https://api.anaplan.com/2/0/models", "models", params=params
             )
         ]
@@ -261,62 +285,83 @@ class AsyncClient(_AsyncBaseClient):
         :return:
         """
         logger.info(f"Deleting Models: {', '.join(model_ids)}.")
-        res = await self._post(
+        res = await self._http.post(
             f"https://api.anaplan.com/2/0/workspaces/{self._workspace_id}/bulkDeleteModels",
             json={"modelIdsToDelete": model_ids},
         )
         return ModelDeletionResult.model_validate(res)
 
-    async def get_files(self) -> list[File]:
+    async def get_files(
+        self, sort_by: Literal["id", "name"] = "id", descending: bool = False
+    ) -> list[File]:
         """
         Lists all the Files in the Model.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Files.
         """
-        return [
-            File.model_validate(e) for e in await self._get_paginated(f"{self._url}/files", "files")
-        ]
+        res = await self._http.get_paginated(
+            f"{self._url}/files", "files", params=sort_params(sort_by, descending)
+        )
+        return [File.model_validate(e) for e in res]
 
-    async def get_actions(self) -> list[Action]:
+    async def get_actions(
+        self, sort_by: Literal["id", "name"] = "id", descending: bool = False
+    ) -> list[Action]:
         """
         Lists all the Actions in the Model. This will only return the Actions listed under
         `Other Actions` in Anaplan. For Imports, exports, and processes, see their respective
         methods instead.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Actions.
         """
-        return [
-            Action.model_validate(e)
-            for e in await self._get_paginated(f"{self._url}/actions", "actions")
-        ]
+        res = await self._http.get_paginated(
+            f"{self._url}/actions", "actions", params=sort_params(sort_by, descending)
+        )
+        return [Action.model_validate(e) for e in res]
 
-    async def get_processes(self) -> list[Process]:
+    async def get_processes(
+        self, sort_by: Literal["id", "name"] = "id", descending: bool = False
+    ) -> list[Process]:
         """
         Lists all the Processes in the Model.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Processes.
         """
-        return [
-            Process.model_validate(e)
-            for e in await self._get_paginated(f"{self._url}/processes", "processes")
-        ]
+        res = await self._http.get_paginated(
+            f"{self._url}/processes", "processes", params=sort_params(sort_by, descending)
+        )
+        return [Process.model_validate(e) for e in res]
 
-    async def get_imports(self) -> list[Import]:
+    async def get_imports(
+        self, sort_by: Literal["id", "name"] = "id", descending: bool = False
+    ) -> list[Import]:
         """
         Lists all the Imports in the Model.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Imports.
         """
-        return [
-            Import.model_validate(e)
-            for e in await self._get_paginated(f"{self._url}/imports", "imports")
-        ]
+        res = await self._http.get_paginated(
+            f"{self._url}/imports", "imports", params=sort_params(sort_by, descending)
+        )
+        return [Import.model_validate(e) for e in res]
 
-    async def get_exports(self) -> list[Export]:
+    async def get_exports(
+        self, sort_by: Literal["id", "name"] = "id", descending: bool = False
+    ) -> list[Export]:
         """
         Lists all the Exports in the Model.
+        :param sort_by: The field to sort the results by.
+        :param descending: If True, the results will be sorted in descending order.
         :return: The List of Exports.
         """
-        return [
-            Export.model_validate(e)
-            for e in await self._get_paginated(f"{self._url}/exports", "exports")
-        ]
+        res = await self._http.get_paginated(
+            f"{self._url}/exports", "exports", params=sort_params(sort_by, descending)
+        )
+        return [Export.model_validate(e) for e in res]
 
     async def run_action(self, action_id: int, wait_for_completion: bool = True) -> TaskStatus:
         """
@@ -330,18 +375,17 @@ class AsyncClient(_AsyncBaseClient):
                until the task is complete. If False, it will spawn the task and return immediately.
         """
         body = {"localeName": "en_US"}
-        res = await self._post(f"{self._url}/{action_url(action_id)}/{action_id}/tasks", json=body)
+        res = await self._http.post(
+            f"{self._url}/{action_url(action_id)}/{action_id}/tasks", json=body
+        )
         task_id = res["task"]["taskId"]
         logger.info(f"Invoked Action '{action_id}', spawned Task: '{task_id}'.")
 
         if not wait_for_completion:
             return TaskStatus.model_validate(await self.get_task_status(action_id, task_id))
-
-        while (status := await self.get_task_status(action_id, task_id)).task_state != "COMPLETE":
-            await sleep(self.status_poll_delay)
-
+        status = await self._http.poll_task(self.get_task_status, action_id, task_id)
         if status.task_state == "COMPLETE" and not status.result.successful:
-            logger.error(f"Task '{task_id}' completed with errors: {status.result.error_message}")
+            logger.error(f"Task '{task_id}' completed with errors.")
             raise AnaplanActionError(f"Task '{task_id}' completed with errors.")
 
         logger.info(f"Task '{task_id}' of '{action_id}' completed successfully.")
@@ -356,11 +400,11 @@ class AsyncClient(_AsyncBaseClient):
         chunk_count = await self._file_pre_check(file_id)
         logger.info(f"File {file_id} has {chunk_count} chunks.")
         if chunk_count <= 1:
-            return await self._get_binary(f"{self._url}/files/{file_id}")
+            return await self._http.get_binary(f"{self._url}/files/{file_id}")
         return b"".join(
             await gather(
                 *(
-                    self._get_binary(f"{self._url}/files/{file_id}/chunks/{i}")
+                    self._http.get_binary(f"{self._url}/files/{file_id}/chunks/{i}")
                     for i in range(chunk_count)
                 )
             )
@@ -380,13 +424,13 @@ class AsyncClient(_AsyncBaseClient):
         chunk_count = await self._file_pre_check(file_id)
         logger.info(f"File {file_id} has {chunk_count} chunks.")
         if chunk_count <= 1:
-            yield await self._get_binary(f"{self._url}/files/{file_id}")
+            yield await self._http.get_binary(f"{self._url}/files/{file_id}")
             return
 
         for batch_start in range(0, chunk_count, batch_size):
             batch_chunks = await gather(
                 *(
-                    self._get_binary(f"{self._url}/files/{file_id}/chunks/{i}")
+                    self._http.get_binary(f"{self._url}/files/{file_id}/chunks/{i}")
                     for i in range(batch_start, min(batch_start + batch_size, chunk_count))
                 )
             )
@@ -462,21 +506,26 @@ class AsyncClient(_AsyncBaseClient):
             logger.info(
                 f"Completed final upload stream batch of size {len(tasks)} for file {file_id}."
             )
-        await self._post(f"{self._url}/files/{file_id}/complete", json={"id": file_id})
+        await self._http.post(f"{self._url}/files/{file_id}/complete", json={"id": file_id})
         logger.info(f"Completed upload stream for '{file_id}'.")
 
-    async def upload_and_import(self, file_id: int, content: str | bytes, action_id: int) -> None:
+    async def upload_and_import(
+        self, file_id: int, content: str | bytes, action_id: int, wait_for_completion: bool = True
+    ) -> TaskStatus:
         """
         Convenience wrapper around `upload_file()` and `run_action()` to upload content to a file
         and run an import action in one call.
         :param file_id: The identifier of the file to upload to.
         :param content: The content to upload. **This Content will be compressed before uploading.
-                        If you are passing the Input as bytes, pass it uncompressed to avoid
-                        redundant work.**
+               If you are passing the Input as bytes, pass it uncompressed to avoid redundant
+               work.**
         :param action_id: The identifier of the action to run after uploading the content.
+        :param wait_for_completion: If True, the method will poll the import task status and not
+               return until the task is complete. If False, it will spawn the import task and
+               return immediately.
         """
         await self.upload_file(file_id, content)
-        await self.run_action(action_id)
+        return await self.run_action(action_id, wait_for_completion)
 
     async def export_and_download(self, action_id: int) -> bytes:
         """
@@ -496,7 +545,7 @@ class AsyncClient(_AsyncBaseClient):
         """
         return [
             TaskSummary.model_validate(e)
-            for e in await self._get_paginated(
+            for e in await self._http.get_paginated(
                 f"{self._url}/{action_url(action_id)}/{action_id}/tasks", "tasks"
             )
         ]
@@ -510,7 +559,9 @@ class AsyncClient(_AsyncBaseClient):
         """
         return TaskStatus.model_validate(
             (
-                await self._get(f"{self._url}/{action_url(action_id)}/{action_id}/tasks/{task_id}")
+                await self._http.get(
+                    f"{self._url}/{action_url(action_id)}/{action_id}/tasks/{task_id}"
+                )
             ).get("task")
         )
 
@@ -521,7 +572,7 @@ class AsyncClient(_AsyncBaseClient):
         :param task_id: The Task identifier, sometimes also referred to as the Correlation Id.
         :return: The content of the solution logs.
         """
-        return await self._get_binary(
+        return await self._http.get_binary(
             f"{self._url}/optimizeActions/{action_id}/tasks/{task_id}/solutionLogs"
         )
 
@@ -532,7 +583,7 @@ class AsyncClient(_AsyncBaseClient):
         return file.chunk_count
 
     async def _upload_chunk(self, file_id: int, index: int, chunk: str | bytes) -> None:
-        await self._put_binary_gzip(f"{self._url}/files/{file_id}/chunks/{index}", chunk)
+        await self._http.put_binary_gzip(f"{self._url}/files/{file_id}/chunks/{index}", chunk)
         logger.debug(f"Chunk {index} loaded to file '{file_id}'.")
 
     async def _set_chunk_count(self, file_id: int, num_chunks: int) -> None:
@@ -542,7 +593,9 @@ class AsyncClient(_AsyncBaseClient):
                 "to avoid this error, set `allow_file_creation=True` on the calling instance. "
                 "Make sure you have understood the implications of this before doing so. "
             )
-        response = await self._post(f"{self._url}/files/{file_id}", json={"chunkCount": num_chunks})
+        response = await self._http.post(
+            f"{self._url}/files/{file_id}", json={"chunkCount": num_chunks}
+        )
         optionally_new_file = int(response.get("file").get("id"))
         if optionally_new_file != file_id:
             if self.allow_file_creation: