From 5d67f5c68a4ce664e7b08e37ae61da184aba3398 Mon Sep 17 00:00:00 2001 From: "Justin \"J.R.\" Hill" Date: Wed, 1 Nov 2023 14:48:05 -0700 Subject: [PATCH 1/2] feat(python-sdk): run sync batch_check with ThreadPoolExecutor for parallelism --- openfga_sdk/sync/client/client.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/openfga_sdk/sync/client/client.py b/openfga_sdk/sync/client/client.py index c6241ba3..5912a466 100644 --- a/openfga_sdk/sync/client/client.py +++ b/openfga_sdk/sync/client/client.py @@ -40,9 +40,9 @@ from openfga_sdk.models.write_request import WriteRequest from openfga_sdk.validation import is_well_formed_ulid_string -import time import uuid from typing import List +from concurrent.futures import ThreadPoolExecutor CLIENT_METHOD_HEADER = "X-OpenFGA-Client-Method" CLIENT_BULK_REQUEST_ID_HEADER = "X-OpenFGA-Client-Bulk-Request-Id" @@ -543,12 +543,16 @@ def batch_check(self, body: List[ClientCheckRequest], options: dict[str, str] = max_parallel_requests = 10 if options is not None and "max_parallel_requests" in options: max_parallel_requests = options["max_parallel_requests"] - # Break the batch into chunks - request_batches = _chuck_array(body, max_parallel_requests) + batch_check_response = [] - for request_batch in request_batches: - response = [self._single_batch_check(i, options) for i in request_batch] - batch_check_response.extend(response) + + def single_batch_check(request): + return self._single_batch_check(request, options) + + with ThreadPoolExecutor(max_workers=max_parallel_requests) as executor: + for response in executor.map(single_batch_check, body): + batch_check_response.append(response) + return batch_check_response def expand(self, body: ClientExpandRequest, options: dict[str, str] = None): # noqa: E501 From 095c3ddb6be885c3f77dc67e1cf52a5504385ac7 Mon Sep 17 00:00:00 2001 From: "Justin \"J.R.\" Hill" Date: Wed, 1 Nov 2023 17:11:48 -0700 Subject: [PATCH 2/2] feat(client): use asyncio semaphore for streaming batch check --- openfga_sdk/client/client.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/openfga_sdk/client/client.py b/openfga_sdk/client/client.py index 930b439b..75db4ec3 100644 --- a/openfga_sdk/client/client.py +++ b/openfga_sdk/client/client.py @@ -511,17 +511,20 @@ async def check(self, body: ClientCheckRequest, options: dict[str, str] = None): ) return api_response - async def _single_batch_check(self, body: ClientCheckRequest, options: dict[str, str] = None): # noqa: E501 + async def _single_batch_check(self, body: ClientCheckRequest, semaphore: asyncio.Semaphore, options: dict[str, str] = None): # noqa: E501 """ Run a single batch request and return body in a SingleBatchCheckResponse :param body - ClientCheckRequest defining check request :param authorization_model_id(options) - Overrides the authorization model id in the configuration """ + await semaphore.acquire() try: api_response = await self.check(body, options) return BatchCheckResponse(allowed=api_response.allowed, request=body, response=api_response, error=None) except Exception as err: return BatchCheckResponse(allowed=False, request=body, response=None, error=err) + finally: + semaphore.release() async def batch_check(self, body: List[ClientCheckRequest], options: dict[str, str] = None): # noqa: E501 """ @@ -543,13 +546,11 @@ async def batch_check(self, body: List[ClientCheckRequest], options: dict[str, s max_parallel_requests = 10 if options is not None and "max_parallel_requests" in options: max_parallel_requests = options["max_parallel_requests"] - # Break the batch into chunks - request_batches = _chuck_array(body, max_parallel_requests) - batch_check_response = [] - for request_batch in request_batches: - request = [self._single_batch_check(i, options) for i in request_batch] - response = await asyncio.gather(*request) - batch_check_response.extend(response) + + sem = asyncio.Semaphore(max_parallel_requests) + batch_check_coros = [self._single_batch_check(request, sem, options) for request in body] + batch_check_response = await asyncio.gather(*batch_check_coros) + return batch_check_response async def expand(self, body: ClientExpandRequest, options: dict[str, str] = None): # noqa: E501