diff --git a/README.md b/README.md index 78f6dc8d..928151a5 100644 --- a/README.md +++ b/README.md @@ -875,7 +875,9 @@ Similar to [check](#check), but instead of checking a single user-object relatio options = { # You can rely on the model id set in the configuration or override it for this specific request - "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1" + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1", + # Optionally collapse checks that resolve through pure relation aliases + "optimize_relation_aliases": True, } checks = [ClientBatchCheckItem( user="user:81684243-9356-4421-8fbf-a4f8d36aa31b", @@ -965,6 +967,13 @@ response = await fga_client.batch_check(ClientBatchCheckRequest(checks=checks), # ] ``` +When `optimize_relation_aliases` is enabled, the SDK reads and caches the pinned +authorization model, then combines checks only when their user, object, context, +and contextual tuples are identical and their relations are pure aliases of the +same relation. The response still contains an entry for every original +`correlation_id` and request. This optimization is disabled by default and +requires an authorization model ID. + ##### Client Batch Check @@ -1569,4 +1578,3 @@ See [CONTRIBUTING](./CONTRIBUTING.md) for details. This project is licensed under the Apache-2.0 license. See the [LICENSE](https://github.com/openfga/python-sdk/blob/main/LICENSE) file for more info. The code in this repo was auto generated by [OpenAPI Generator](https://github.com/OpenAPITools/openapi-generator) from a template based on the [python legacy template](https://github.com/OpenAPITools/openapi-generator/tree/master/modules/openapi-generator/src/main/resources/python-legacy), licensed under the [Apache License 2.0](https://github.com/OpenAPITools/openapi-generator/blob/master/LICENSE). - diff --git a/openfga_sdk/client/client.py b/openfga_sdk/client/client.py index b48a3ec8..66dce488 100644 --- a/openfga_sdk/client/client.py +++ b/openfga_sdk/client/client.py @@ -2,7 +2,7 @@ import uuid from collections.abc import AsyncIterator -from typing import Any +from typing import Any, cast from openfga_sdk.api.open_fga_api import OpenFgaApi from openfga_sdk.api_client import ApiClient @@ -37,6 +37,10 @@ construct_write_single_response, ) from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import ( + build_relation_aliases, + group_batch_checks, +) from openfga_sdk.constants import ( CLIENT_BULK_REQUEST_ID_HEADER, CLIENT_MAX_BATCH_SIZE, @@ -172,6 +176,9 @@ def __init__(self, configuration: ClientConfiguration): self._client_configuration = configuration self._api_client = ApiClient(configuration) self._api = OpenFgaApi(self._api_client) + self._relation_alias_cache: dict[ + tuple[str, str], asyncio.Task[dict[str, dict[str, str]]] + ] = {} # Set default headers from configuration if configuration.headers: @@ -185,6 +192,14 @@ async def __aexit__(self, exc_type, exc_value, traceback): await self.close() async def close(self): + """Cancel cached model loads and close the API client.""" + tasks = list(self._relation_alias_cache.values()) + self._relation_alias_cache.clear() + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) await self._api.close() def _get_authorization_model_id( @@ -247,6 +262,53 @@ def get_authorization_model_id(self): """ return self._client_configuration.authorization_model_id + async def _get_relation_aliases( + self, + options: dict[str, int | str | dict[str, int | str]] | None, + ) -> dict[str, dict[str, str]]: + """Return cached relation aliases for the configured model.""" + authorization_model_id = self._get_authorization_model_id(options) + if authorization_model_id is None: + raise FgaValidationException( + "authorization_model_id is required when optimizing BatchCheck" + ) + + store_id = self.get_store_id() + if store_id is None or store_id == "": + raise FgaValidationException("store_id is required but not configured") + + cache_key = (store_id, authorization_model_id) + task = self._relation_alias_cache.get(cache_key) + if task is None: + task = asyncio.create_task(self._load_relation_aliases(options)) + self._relation_alias_cache[cache_key] = task + + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + if task.cancelled() and self._relation_alias_cache.get(cache_key) is task: + self._relation_alias_cache.pop(cache_key, None) + raise + except Exception: + if self._relation_alias_cache.get(cache_key) is task: + self._relation_alias_cache.pop(cache_key, None) + raise + + async def _load_relation_aliases( + self, + options: dict[str, int | str | dict[str, int | str]] | None, + ) -> dict[str, dict[str, str]]: + """Read the configured model and build relation alias mappings.""" + model_options = { + key: options[key] + for key in ("authorization_model_id", "headers", "retry_params") + if options is not None and key in options + } + response = await self.read_authorization_model(model_options) + if response.authorization_model is None: + raise FgaValidationException("authorization model was not returned") + return build_relation_aliases(response.authorization_model) + ################# # Stores ################# @@ -784,6 +846,7 @@ async def batch_check( :param retryParams(options) - Override the retry parameters for this request :param retryParams.maxRetry(options) - Override the max number of retries on each API request :param retryParams.minWaitInMs(options) - Override the minimum wait before a retry is initiated + :param optimize_relation_aliases(options) - Collapse equivalent pure relation aliases. Defaults to false """ options = set_heading_if_not_set( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) @@ -809,42 +872,73 @@ async def batch_check( elif isinstance(options["max_batch_size"], int): max_batch_size = options["max_batch_size"] - id_to_check: dict[str, ClientBatchCheckItem] = {} - - def track_and_transform(checks): - transformed = [] - for check in checks: - if check.correlation_id is None: - check.correlation_id = str(uuid.uuid4()) - - if check.correlation_id in id_to_check: - raise FgaValidationException( - f"Duplicate correlation_id ({check.correlation_id}) provided" - ) - - id_to_check[check.correlation_id] = check + optimize_relation_aliases = options.get("optimize_relation_aliases") is True + if ( + optimize_relation_aliases + and self._get_authorization_model_id(options) is None + ): + raise FgaValidationException( + "authorization_model_id is required when optimizing BatchCheck" + ) - transformed.append(construct_batch_item(check)) - return transformed + id_to_check: dict[str, ClientBatchCheckItem] = {} + for check in body.checks: + if check.correlation_id is None: + check.correlation_id = str(uuid.uuid4()) + + correlation_id = cast(str, check.correlation_id) + if correlation_id in id_to_check: + raise FgaValidationException( + f"Duplicate correlation_id ({correlation_id}) provided" + ) + id_to_check[correlation_id] = check + + aliases_by_type = ( + await self._get_relation_aliases(options) + if optimize_relation_aliases + else {} + ) + groups = group_batch_checks(body.checks, aliases_by_type) + submitted_to_original_ids: dict[str, tuple[str, ...]] = {} + submitted_checks: list[ClientBatchCheckItem] = [] + + for group in groups: + representative = body.checks[group.indexes[0]] + submitted_check = representative + if len(group.indexes) > 1: + submitted_check = ClientBatchCheckItem( + user=representative.user, + relation=group.relation, + object=representative.object, + correlation_id=representative.correlation_id, + contextual_tuples=representative.contextual_tuples, + context=representative.context, + ) + + submitted_id = cast(str, submitted_check.correlation_id) + submitted_to_original_ids[submitted_id] = tuple( + cast(str, body.checks[index].correlation_id) for index in group.indexes + ) + submitted_checks.append(submitted_check) checks = [ - track_and_transform( - body.checks[i * max_batch_size : (i + 1) * max_batch_size] - ) - for i in range((len(body.checks) + max_batch_size - 1) // max_batch_size) + [construct_batch_item(check) for check in chunk] + for chunk in _chuck_array(submitted_checks, max_batch_size) ] result = [] sem = asyncio.Semaphore(max_parallel_requests) def map_response(id, result): - check = id_to_check[id] - return ClientBatchCheckSingleResponse( - allowed=result.allowed, - request=check, - correlation_id=id, - error=result.error, - ) + return [ + ClientBatchCheckSingleResponse( + allowed=result.allowed, + request=id_to_check[original_id], + correlation_id=original_id, + error=result.error, + ) + for original_id in submitted_to_original_ids[id] + ] async def coro(checks): res = await self._single_batch_check( @@ -857,9 +951,8 @@ async def coro(checks): options, ) - result.extend( - [map_response(c_id, c_result) for c_id, c_result in res.result.items()] - ) + for correlation_id, check_result in res.result.items(): + result.extend(map_response(correlation_id, check_result)) batch_check_coros = [coro(request) for request in checks] await asyncio.gather(*batch_check_coros) diff --git a/openfga_sdk/client/relation_optimizer.py b/openfga_sdk/client/relation_optimizer.py new file mode 100644 index 00000000..e4d9b5e8 --- /dev/null +++ b/openfga_sdk/client/relation_optimizer.py @@ -0,0 +1,124 @@ +from dataclasses import dataclass + +from openfga_sdk.client.models.batch_check_item import ClientBatchCheckItem +from openfga_sdk.models.authorization_model import AuthorizationModel +from openfga_sdk.models.userset import Userset + + +@dataclass(frozen=True) +class RelationCheckGroup: + """BatchCheck items that can share one relation evaluation.""" + + relation: str + indexes: tuple[int, ...] + + +def build_relation_aliases( + authorization_model: AuthorizationModel, +) -> dict[str, dict[str, str]]: + """Build canonical targets for pure computed-userset relation aliases.""" + aliases_by_type: dict[str, dict[str, str]] = {} + + for type_definition in authorization_model.type_definitions or []: + relations = type_definition.relations or {} + direct_aliases = { + relation: target + for relation, rewrite in relations.items() + if (target := _pure_computed_userset_target(rewrite)) is not None + and target in relations + } + + canonical_aliases: dict[str, str] = {} + for relation in direct_aliases: + target = _resolve_alias(relation, direct_aliases) + if target is not None and target != relation: + canonical_aliases[relation] = target + + aliases_by_type[type_definition.type] = canonical_aliases + + return aliases_by_type + + +def group_batch_checks( + checks: list[ClientBatchCheckItem], aliases_by_type: dict[str, dict[str, str]] +) -> list[RelationCheckGroup]: + """Group checks that differ only by pure aliases of the same relation.""" + candidates: list[tuple[str, ClientBatchCheckItem, list[int]]] = [] + + for index, check in enumerate(checks): + object_type, separator, _ = check.object.partition(":") + aliases = aliases_by_type.get(object_type, {}) if separator else {} + target = aliases.get(check.relation, check.relation) + + for candidate_target, representative, indexes in candidates: + if candidate_target == target and _same_check_inputs(check, representative): + indexes.append(index) + break + else: + candidates.append((target, check, [index])) + + groups: list[RelationCheckGroup] = [] + for target, _, indexes in candidates: + if len(indexes) > 1 and any( + checks[index].relation != target for index in indexes + ): + groups.append(RelationCheckGroup(relation=target, indexes=tuple(indexes))) + else: + groups.extend( + RelationCheckGroup( + relation=checks[index].relation, + indexes=(index,), + ) + for index in indexes + ) + + groups.sort(key=lambda group: group.indexes[0]) + return groups + + +def _same_check_inputs(left: ClientBatchCheckItem, right: ClientBatchCheckItem) -> bool: + """Return whether two checks differ only in relation and correlation ID.""" + return ( + left.user == right.user + and left.object == right.object + and left.context == right.context + and left.contextual_tuples == right.contextual_tuples + ) + + +def _pure_computed_userset_target(rewrite: Userset) -> str | None: + """Return the target when a rewrite is only a same-object computed userset.""" + computed_userset = rewrite.computed_userset + if computed_userset is None or not computed_userset.relation: + return None + + if any( + value is not None + for value in ( + rewrite.this, + rewrite.tuple_to_userset, + rewrite.union, + rewrite.intersection, + rewrite.difference, + ) + ): + return None + + if computed_userset.object not in (None, ""): + return None + + return computed_userset.relation + + +def _resolve_alias(relation: str, direct_aliases: dict[str, str]) -> str | None: + """Resolve an alias chain, returning none when it contains a cycle.""" + visited = set() + current = relation + + while current in direct_aliases: + if current in visited: + return None + visited.add(current) + current = direct_aliases[current] + + return current diff --git a/openfga_sdk/sync/client/client.py b/openfga_sdk/sync/client/client.py index c736daec..566d438b 100644 --- a/openfga_sdk/sync/client/client.py +++ b/openfga_sdk/sync/client/client.py @@ -1,8 +1,9 @@ import uuid from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor -from typing import Any +from concurrent.futures import Future, ThreadPoolExecutor +from threading import Lock +from typing import Any, cast from openfga_sdk.client.configuration import ClientConfiguration from openfga_sdk.client.models.assertion import ClientAssertion @@ -35,6 +36,10 @@ construct_write_single_response, ) from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import ( + build_relation_aliases, + group_batch_checks, +) from openfga_sdk.constants import ( CLIENT_BULK_REQUEST_ID_HEADER, CLIENT_MAX_BATCH_SIZE, @@ -172,6 +177,10 @@ def __init__(self, configuration: ClientConfiguration) -> None: self._client_configuration = configuration self._api_client = ApiClient(configuration) self._api = OpenFgaApi(self._api_client) + self._relation_alias_cache: dict[ + tuple[str, str], Future[dict[str, dict[str, str]]] + ] = {} + self._relation_alias_cache_lock = Lock() # Set default headers from configuration if configuration.headers: @@ -247,6 +256,49 @@ def get_authorization_model_id(self): """ return self._client_configuration.authorization_model_id + def _get_relation_aliases( + self, + options: dict[str, int | str | dict[str, int | str]] | None, + ) -> dict[str, dict[str, str]]: + """Return cached relation aliases for the configured model.""" + authorization_model_id = self._get_authorization_model_id(options) + if authorization_model_id is None: + raise FgaValidationException( + "authorization_model_id is required when optimizing BatchCheck" + ) + + store_id = self.get_store_id() + if store_id is None or store_id == "": + raise FgaValidationException("store_id is required but not configured") + + cache_key = (store_id, authorization_model_id) + with self._relation_alias_cache_lock: + future = self._relation_alias_cache.get(cache_key) + should_load = future is None + if future is None: + future = Future() + self._relation_alias_cache[cache_key] = future + + if should_load: + try: + model_options = { + key: options[key] + for key in ("authorization_model_id", "headers", "retry_params") + if options is not None and key in options + } + response = self.read_authorization_model(model_options) + if response.authorization_model is None: + raise FgaValidationException("authorization model was not returned") + future.set_result(build_relation_aliases(response.authorization_model)) + except BaseException as error: + future.set_exception(error) + with self._relation_alias_cache_lock: + if self._relation_alias_cache.get(cache_key) is future: + self._relation_alias_cache.pop(cache_key, None) + raise + + return future.result() + ################# # Stores ################# @@ -780,6 +832,7 @@ def batch_check( :param retryParams(options) - Override the retry parameters for this request :param retryParams.maxRetry(options) - Override the max number of retries on each API request :param retryParams.minWaitInMs(options) - Override the minimum wait before a retry is initiated + :param optimize_relation_aliases(options) - Collapse equivalent pure relation aliases. Defaults to false """ options = set_heading_if_not_set( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) @@ -805,39 +858,68 @@ def batch_check( elif isinstance(options["max_batch_size"], int): max_batch_size = options["max_batch_size"] - id_to_check: dict[str, ClientBatchCheckItem] = {} - - def track_and_transform(checks): - transformed = [] - for check in checks: - if check.correlation_id is None: - check.correlation_id = str(uuid.uuid4()) + optimize_relation_aliases = options.get("optimize_relation_aliases") is True + if ( + optimize_relation_aliases + and self._get_authorization_model_id(options) is None + ): + raise FgaValidationException( + "authorization_model_id is required when optimizing BatchCheck" + ) - if check.correlation_id in id_to_check: - raise FgaValidationException( - f"Duplicate correlation_id ({check.correlation_id}) provided" - ) + id_to_check: dict[str, ClientBatchCheckItem] = {} + for check in body.checks: + if check.correlation_id is None: + check.correlation_id = str(uuid.uuid4()) + + correlation_id = cast(str, check.correlation_id) + if correlation_id in id_to_check: + raise FgaValidationException( + f"Duplicate correlation_id ({correlation_id}) provided" + ) + id_to_check[correlation_id] = check - id_to_check[check.correlation_id] = check + aliases_by_type = ( + self._get_relation_aliases(options) if optimize_relation_aliases else {} + ) + groups = group_batch_checks(body.checks, aliases_by_type) + submitted_to_original_ids: dict[str, tuple[str, ...]] = {} + submitted_checks: list[ClientBatchCheckItem] = [] + + for group in groups: + representative = body.checks[group.indexes[0]] + submitted_check = representative + if len(group.indexes) > 1: + submitted_check = ClientBatchCheckItem( + user=representative.user, + relation=group.relation, + object=representative.object, + correlation_id=representative.correlation_id, + contextual_tuples=representative.contextual_tuples, + context=representative.context, + ) - transformed.append(construct_batch_item(check)) - return transformed + submitted_id = cast(str, submitted_check.correlation_id) + submitted_to_original_ids[submitted_id] = tuple( + cast(str, body.checks[index].correlation_id) for index in group.indexes + ) + submitted_checks.append(submitted_check) checks = [ - track_and_transform( - body.checks[i * max_batch_size : (i + 1) * max_batch_size] - ) - for i in range((len(body.checks) + max_batch_size - 1) // max_batch_size) + [construct_batch_item(check) for check in chunk] + for chunk in _chuck_array(submitted_checks, max_batch_size) ] def map_response(id, result): - check = id_to_check[id] - return ClientBatchCheckSingleResponse( - allowed=result.allowed, - request=check, - correlation_id=id, - error=result.error, - ) + return [ + ClientBatchCheckSingleResponse( + allowed=result.allowed, + request=id_to_check[original_id], + correlation_id=original_id, + error=result.error, + ) + for original_id in submitted_to_original_ids[id] + ] def single_batch_check(checks): res = self._single_batch_check( @@ -855,12 +937,8 @@ def single_batch_check(checks): with ThreadPoolExecutor(max_workers=max_parallel_requests) as executor: for response in executor.map(single_batch_check, checks): - result.extend( - [ - map_response(c_id, c_result) - for c_id, c_result in response.result.items() - ] - ) + for correlation_id, check_result in response.result.items(): + result.extend(map_response(correlation_id, check_result)) return ClientBatchCheckResponse(result) diff --git a/test/client/client_test.py b/test/client/client_test.py index 9850b5d1..bf80bb1f 100644 --- a/test/client/client_test.py +++ b/test/client/client_test.py @@ -1,3 +1,4 @@ +import asyncio import copy import json import uuid @@ -2486,6 +2487,424 @@ def mock_v4(val: str): ) await api_client.close() + @patch.object(rest.RESTClientObject, "request") + async def test_batch_check_optimizes_relation_aliases(self, mock_request): + """BatchCheck can collapse pure aliases using a cached model.""" + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + model_response = json.dumps( + { + "authorization_model": { + "id": authorization_model_id, + "schema_version": "1.1", + "type_definitions": [ + { + "type": "document", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_edit"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_edit"} + }, + "can_edit": {"this": {}}, + }, + } + ], + } + } + ) + + def mock_optimized_requests(method, url, **kwargs): + if method == "GET": + return mock_response(model_response, 200) + + checks = kwargs["body"]["checks"] + self.assertEqual(len(checks), 1) + self.assertEqual(checks[0]["tuple_key"]["relation"], "can_edit") + self.assertEqual(checks[0]["correlation_id"], "child") + return mock_response( + json.dumps({"result": {"child": {"allowed": True}}}), 200 + ) + + mock_request.side_effect = mock_optimized_requests + configuration = self.configuration + configuration.store_id = store_id + checks = [ + ClientBatchCheckItem( + user="user:anne", + relation="can_add_child", + object="document:roadmap", + correlation_id="child", + context={"view_count": 1}, + ), + ClientBatchCheckItem( + user="user:anne", + relation="can_add_records", + object="document:roadmap", + correlation_id="records", + context={"view_count": 1}, + ), + ] + options = { + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + } + + async with OpenFgaClient(configuration) as api_client: + allowed_response = await api_client.batch_check( + ClientBatchCheckRequest(checks=checks), options + ) + + self.assertEqual( + [item.correlation_id for item in allowed_response.result], + ["child", "records"], + ) + self.assertEqual( + [item.request for item in allowed_response.result], + checks, + ) + self.assertTrue(all(item.allowed for item in allowed_response.result)) + model_requests = [ + call + for call in mock_request.call_args_list + if "/authorization-models/" in call.args[1] + ] + batch_requests = [ + call + for call in mock_request.call_args_list + if call.args[1].endswith("/batch-check") + ] + self.assertEqual(len(model_requests), 1) + self.assertEqual(len(batch_requests), 1) + + @patch.object(rest.RESTClientObject, "request") + async def test_batch_check_optimization_preserves_group_boundaries( + self, mock_request + ): + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + model_response = json.dumps( + { + "authorization_model": { + "id": authorization_model_id, + "schema_version": "1.1", + "type_definitions": [ + { + "type": "document", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_edit"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_edit"} + }, + "can_edit": {"this": {}}, + }, + }, + { + "type": "folder", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_manage"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_manage"} + }, + "can_manage": {"this": {}}, + }, + }, + ], + } + } + ) + expected_submitted_checks = [ + { + "tuple_key": { + "user": "user:anne", + "relation": "can_edit", + "object": "document:roadmap", + }, + "correlation_id": "document-roadmap-anne-child", + }, + { + "tuple_key": { + "user": "user:bob", + "relation": "can_edit", + "object": "document:roadmap", + }, + "correlation_id": "document-roadmap-bob-child", + }, + { + "tuple_key": { + "user": "user:anne", + "relation": "can_edit", + "object": "document:budget", + }, + "correlation_id": "document-budget-anne-child", + }, + { + "tuple_key": { + "user": "user:anne", + "relation": "can_manage", + "object": "folder:roadmap", + }, + "correlation_id": "folder-roadmap-anne-child", + }, + ] + + def mock_optimized_requests(method, url, **kwargs): + if method == "GET": + return mock_response(model_response, 200) + + self.assertEqual(kwargs["body"]["checks"], expected_submitted_checks) + return mock_response( + json.dumps( + { + "result": { + "document-roadmap-anne-child": {"allowed": True}, + "document-roadmap-bob-child": { + "error": { + "input_error": "validation_error", + "message": "bob check failed", + } + }, + "document-budget-anne-child": {"allowed": False}, + "folder-roadmap-anne-child": {"allowed": True}, + } + } + ), + 200, + ) + + mock_request.side_effect = mock_optimized_requests + checks = [ + ClientBatchCheckItem( + user=user, + relation=relation, + object=object_name, + correlation_id=correlation_id, + ) + for user, relation, object_name, correlation_id in ( + ( + "user:anne", + "can_add_child", + "document:roadmap", + "document-roadmap-anne-child", + ), + ( + "user:anne", + "can_add_records", + "document:roadmap", + "document-roadmap-anne-records", + ), + ( + "user:bob", + "can_add_child", + "document:roadmap", + "document-roadmap-bob-child", + ), + ( + "user:bob", + "can_add_records", + "document:roadmap", + "document-roadmap-bob-records", + ), + ( + "user:anne", + "can_add_child", + "document:budget", + "document-budget-anne-child", + ), + ( + "user:anne", + "can_add_records", + "document:budget", + "document-budget-anne-records", + ), + ( + "user:anne", + "can_add_child", + "folder:roadmap", + "folder-roadmap-anne-child", + ), + ( + "user:anne", + "can_add_records", + "folder:roadmap", + "folder-roadmap-anne-records", + ), + ) + ] + configuration = self.configuration + configuration.store_id = store_id + + async with OpenFgaClient(configuration) as api_client: + response = await api_client.batch_check( + ClientBatchCheckRequest(checks=checks), + options={ + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + }, + ) + + results_by_id = {item.correlation_id: item for item in response.result} + self.assertEqual(set(results_by_id), {check.correlation_id for check in checks}) + expected_allowed = (True, True, False, False, False, False, True, True) + for check, allowed in zip(checks, expected_allowed, strict=True): + result = results_by_id[check.correlation_id] + self.assertIs(result.request, check) + self.assertEqual(result.request.user, check.user) + self.assertEqual(result.request.object, check.object) + self.assertEqual(result.request.relation, check.relation) + self.assertEqual(result.allowed, allowed) + + self.assertEqual( + results_by_id["document-roadmap-bob-child"].error.message, + "bob check failed", + ) + self.assertEqual( + results_by_id["document-roadmap-bob-records"].error.message, + "bob check failed", + ) + batch_requests = [ + call + for call in mock_request.call_args_list + if call.args[1].endswith("/batch-check") + ] + self.assertEqual(len(batch_requests), 1) + + @patch.object(rest.RESTClientObject, "request") + async def test_batch_check_optimization_requires_model_id(self, mock_request): + configuration = self.configuration + configuration.store_id = store_id + body = ClientBatchCheckRequest( + checks=[ + ClientBatchCheckItem( + user="user:anne", + relation="can_view", + object="document:roadmap", + ) + ] + ) + + async with OpenFgaClient(configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing BatchCheck", + ): + await api_client.batch_check( + body, + options={"optimize_relation_aliases": True}, + ) + + mock_request.assert_not_called() + + async def test_relation_alias_cache_requires_store_id(self): + async with OpenFgaClient(self.configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing BatchCheck", + ): + await api_client._get_relation_aliases(None) + + with self.assertRaisesRegex( + FgaValidationException, + "store_id is required but not configured", + ): + await api_client._get_relation_aliases( + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"} + ) + + async def test_relation_alias_cache_evicts_model_load_errors(self): + configuration = self.configuration + configuration.store_id = store_id + options = { + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1", + "continuation_token": "ignored", + "headers": {"x-test": "value"}, + "optimize_relation_aliases": True, + "page_size": 10, + } + expected_model_options = { + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1", + "headers": {"x-test": "value"}, + } + + async with OpenFgaClient(configuration) as api_client: + with patch.object( + api_client, + "read_authorization_model", + return_value=ReadAuthorizationModelResponse(), + ) as mock_read_model: + for _ in range(2): + with self.assertRaisesRegex( + FgaValidationException, + "authorization model was not returned", + ): + await api_client._get_relation_aliases(options) + + self.assertEqual(mock_read_model.await_count, 2) + self.assertEqual( + [request.args[0] for request in mock_read_model.await_args_list], + [expected_model_options, expected_model_options], + ) + self.assertEqual(api_client._relation_alias_cache, {}) + + async def test_relation_alias_cache_evicts_cancelled_loads(self): + configuration = self.configuration + configuration.store_id = store_id + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + options = {"authorization_model_id": authorization_model_id} + + async with OpenFgaClient(configuration) as api_client: + task = asyncio.create_task(asyncio.sleep(10)) + task.cancel() + cache_key = (store_id, authorization_model_id) + api_client._relation_alias_cache[cache_key] = task + + with self.assertRaises(asyncio.CancelledError): + await api_client._get_relation_aliases(options) + + self.assertNotIn(cache_key, api_client._relation_alias_cache) + + async def test_close_cancels_shared_relation_alias_loads(self): + configuration = self.configuration + configuration.store_id = store_id + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + options = {"authorization_model_id": authorization_model_id} + load_started = asyncio.Event() + load_cancelled = asyncio.Event() + wait_forever = asyncio.Event() + + async def load_model(_options): + load_started.set() + try: + await wait_forever.wait() + except asyncio.CancelledError: + load_cancelled.set() + raise + + api_client = OpenFgaClient(configuration) + with patch.object( + api_client, + "read_authorization_model", + side_effect=load_model, + ): + caller = asyncio.create_task(api_client._get_relation_aliases(options)) + await asyncio.wait_for(load_started.wait(), timeout=1) + caller.cancel() + with self.assertRaises(asyncio.CancelledError): + await caller + + cache_key = (store_id, authorization_model_id) + cached_load = api_client._relation_alias_cache[cache_key] + self.assertFalse(cached_load.done()) + + await api_client.close() + + self.assertTrue(cached_load.cancelled()) + self.assertTrue(load_cancelled.is_set()) + self.assertEqual(api_client._relation_alias_cache, {}) + async def test_batch_check_errors_dupe_cor_id(self): """Test case for duplicate correlation_id being provided to batch_check""" diff --git a/test/client/relation_optimizer_test.py b/test/client/relation_optimizer_test.py new file mode 100644 index 00000000..c022a77f --- /dev/null +++ b/test/client/relation_optimizer_test.py @@ -0,0 +1,148 @@ +from openfga_sdk.client.models.batch_check_item import ClientBatchCheckItem +from openfga_sdk.client.models.tuple import ClientTuple +from openfga_sdk.client.relation_optimizer import ( + build_relation_aliases, + group_batch_checks, +) +from openfga_sdk.models.authorization_model import AuthorizationModel +from openfga_sdk.models.object_relation import ObjectRelation +from openfga_sdk.models.type_definition import TypeDefinition +from openfga_sdk.models.userset import Userset +from openfga_sdk.models.usersets import Usersets + + +def test_build_relation_aliases_only_includes_pure_aliases(): + model = AuthorizationModel( + id="01GXSA8YR785C4FYS3C0RTG7B1", + schema_version="1.1", + type_definitions=[ + TypeDefinition( + type="document", + relations={ + "can_comment": Userset( + computed_userset=ObjectRelation(relation="commenter") + ), + "can_reply": Userset( + computed_userset=ObjectRelation(relation="can_comment") + ), + "commenter": Userset(this={}), + "viewer": Userset(this={}), + "can_view": Userset( + union=Usersets( + child=[ + Userset( + computed_userset=ObjectRelation(relation="viewer") + ) + ] + ) + ), + "mixed_rewrite": Userset( + computed_userset=ObjectRelation(relation="viewer"), + union=Usersets(child=[Userset(this={})]), + ), + "other_object": Userset( + computed_userset=ObjectRelation( + object="document:other", relation="viewer" + ) + ), + "missing_target": Userset( + computed_userset=ObjectRelation(relation="missing") + ), + "cycle_a": Userset( + computed_userset=ObjectRelation(relation="cycle_b") + ), + "cycle_b": Userset( + computed_userset=ObjectRelation(relation="cycle_a") + ), + }, + ) + ], + ) + + assert build_relation_aliases(model) == { + "document": { + "can_comment": "commenter", + "can_reply": "commenter", + } + } + + +def test_group_batch_checks_requires_all_non_relation_inputs_to_match(): + contextual_tuple = ClientTuple( + user="user:bob", + relation="viewer", + object="document:roadmap", + ) + matching = { + "user": "user:anne", + "object": "document:roadmap", + "context": {"view_count": 1}, + "contextual_tuples": [contextual_tuple], + } + checks = [ + ClientBatchCheckItem(relation="can_comment", **matching), + ClientBatchCheckItem( + relation="can_reply", + **{ + **matching, + "contextual_tuples": [ + ClientTuple( + user="user:bob", + relation="viewer", + object="document:roadmap", + ) + ], + }, + ), + ClientBatchCheckItem(relation="commenter", **matching), + ClientBatchCheckItem( + user="user:bob", + relation="can_comment", + object="document:roadmap", + context={"view_count": 1}, + contextual_tuples=[contextual_tuple], + ), + ClientBatchCheckItem( + user="user:anne", + relation="can_reply", + object="document:roadmap", + context={"view_count": 2}, + contextual_tuples=[contextual_tuple], + ), + ClientBatchCheckItem( + user="user:anne", + relation="can_comment", + object="document:roadmap", + context={"view_count": 1}, + contextual_tuples=[], + ), + ClientBatchCheckItem( + user="user:anne", relation="viewer", object="document:roadmap" + ), + ClientBatchCheckItem( + user="user:anne", relation="viewer", object="document:roadmap" + ), + ClientBatchCheckItem( + user="user:anne", relation="can_comment", object="document" + ), + ] + + groups = group_batch_checks( + checks, + { + "document": { + "can_comment": "commenter", + "can_reply": "commenter", + } + }, + ) + + assert [(group.relation, group.indexes) for group in groups] == [ + ("commenter", (0, 1, 2)), + ("can_comment", (3,)), + ("can_reply", (4,)), + ("can_comment", (5,)), + ("viewer", (6,)), + ("viewer", (7,)), + ("can_comment", (8,)), + ] diff --git a/test/models_contract_test.py b/test/models_contract_test.py new file mode 100644 index 00000000..7b4cdf16 --- /dev/null +++ b/test/models_contract_test.py @@ -0,0 +1,95 @@ +import importlib +import inspect +import pkgutil + +import pytest + +import openfga_sdk.models + +from openfga_sdk.configuration import Configuration + + +def _generated_model_classes(): + classes = [] + for module_info in pkgutil.iter_modules(openfga_sdk.models.__path__): + if module_info.name.startswith("_"): + continue + + module = importlib.import_module(f"openfga_sdk.models.{module_info.name}") + classes.extend( + model_class + for _, model_class in inspect.getmembers(module, inspect.isclass) + if model_class.__module__ == module.__name__ + and hasattr(model_class, "openapi_types") + ) + + return sorted(classes, key=lambda model_class: model_class.__name__) + + +GENERATED_MODEL_CLASSES = _generated_model_classes() + + +def _sample_value(openapi_type): + if openapi_type.startswith("list["): + return ["sample"] + if openapi_type.startswith("dict["): + return {"key": "sample"} + if openapi_type == "bool": + return True + if openapi_type in ("int", "float"): + return 1 + return "sample" + + +def _different_value(value): + if isinstance(value, list): + return [*value, "different"] + if isinstance(value, dict): + return {**value, "different": "value"} + if isinstance(value, bool): + return not value + if isinstance(value, (int, float)): + return value + 1 + return "different" + + +@pytest.mark.parametrize( + "model_class", + GENERATED_MODEL_CLASSES, + ids=lambda model_class: model_class.__name__, +) +def test_generated_model_common_contract(model_class): + configuration = Configuration(api_url="http://api.fga.example") + configuration.client_side_validation = False + instance = model_class(local_vars_configuration=configuration) + equal_instance = model_class(local_vars_configuration=configuration) + expected = {} + + for attribute, openapi_type in model_class.openapi_types.items(): + value = _sample_value(openapi_type) + setattr(instance, attribute, value) + setattr(equal_instance, attribute, value) + expected[attribute] = value + + assert instance.to_dict() == expected + assert instance.to_dict(serialize=True) == { + model_class.attribute_map.get(attribute, attribute): value + for attribute, value in expected.items() + } + assert instance.to_str() == repr(instance) + assert instance == equal_instance + assert not instance != equal_instance + assert instance != object() + assert not instance == object() + + first_attribute = next(iter(model_class.openapi_types), None) + if first_attribute is None: + return + + setattr( + equal_instance, + first_attribute, + _different_value(getattr(equal_instance, first_attribute)), + ) + assert instance != equal_instance + assert not instance == equal_instance diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index ec2686af..1a98d184 100644 --- a/test/sync/client/client_test.py +++ b/test/sync/client/client_test.py @@ -2,7 +2,9 @@ import json import uuid +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from threading import Event from unittest import IsolatedAsyncioTestCase, TestCase from unittest.mock import ANY, patch @@ -2489,6 +2491,403 @@ def mock_v4(val: str): ) api_client.close() + @patch.object(rest.RESTClientObject, "request") + def test_batch_check_optimizes_relation_aliases(self, mock_request): + """BatchCheck can collapse pure aliases using a cached model.""" + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + model_response = json.dumps( + { + "authorization_model": { + "id": authorization_model_id, + "schema_version": "1.1", + "type_definitions": [ + { + "type": "document", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_edit"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_edit"} + }, + "can_edit": {"this": {}}, + }, + } + ], + } + } + ) + + def mock_optimized_requests(method, url, **kwargs): + if method == "GET": + return mock_response(model_response, 200) + + checks = kwargs["body"]["checks"] + self.assertEqual(len(checks), 1) + self.assertEqual(checks[0]["tuple_key"]["relation"], "can_edit") + self.assertEqual(checks[0]["correlation_id"], "child") + return mock_response( + json.dumps({"result": {"child": {"allowed": True}}}), 200 + ) + + mock_request.side_effect = mock_optimized_requests + configuration = self.configuration + configuration.store_id = store_id + checks = [ + ClientBatchCheckItem( + user="user:anne", + relation="can_add_child", + object="document:roadmap", + correlation_id="child", + context={"view_count": 1}, + ), + ClientBatchCheckItem( + user="user:anne", + relation="can_add_records", + object="document:roadmap", + correlation_id="records", + context={"view_count": 1}, + ), + ] + options = { + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + } + + with OpenFgaClient(configuration) as api_client: + allowed_response = api_client.batch_check( + ClientBatchCheckRequest(checks=checks), options + ) + + self.assertEqual( + [item.correlation_id for item in allowed_response.result], + ["child", "records"], + ) + self.assertEqual( + [item.request for item in allowed_response.result], + checks, + ) + self.assertTrue(all(item.allowed for item in allowed_response.result)) + model_requests = [ + call + for call in mock_request.call_args_list + if "/authorization-models/" in call.args[1] + ] + batch_requests = [ + call + for call in mock_request.call_args_list + if call.args[1].endswith("/batch-check") + ] + self.assertEqual(len(model_requests), 1) + self.assertEqual(len(batch_requests), 1) + + @patch.object(rest.RESTClientObject, "request") + def test_batch_check_optimization_preserves_group_boundaries(self, mock_request): + authorization_model_id = "01GXSA8YR785C4FYS3C0RTG7B1" + model_response = json.dumps( + { + "authorization_model": { + "id": authorization_model_id, + "schema_version": "1.1", + "type_definitions": [ + { + "type": "document", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_edit"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_edit"} + }, + "can_edit": {"this": {}}, + }, + }, + { + "type": "folder", + "relations": { + "can_add_child": { + "computedUserset": {"relation": "can_manage"} + }, + "can_add_records": { + "computedUserset": {"relation": "can_manage"} + }, + "can_manage": {"this": {}}, + }, + }, + ], + } + } + ) + expected_submitted_checks = [ + { + "tuple_key": { + "user": "user:anne", + "relation": "can_edit", + "object": "document:roadmap", + }, + "correlation_id": "document-roadmap-anne-child", + }, + { + "tuple_key": { + "user": "user:bob", + "relation": "can_edit", + "object": "document:roadmap", + }, + "correlation_id": "document-roadmap-bob-child", + }, + { + "tuple_key": { + "user": "user:anne", + "relation": "can_edit", + "object": "document:budget", + }, + "correlation_id": "document-budget-anne-child", + }, + { + "tuple_key": { + "user": "user:anne", + "relation": "can_manage", + "object": "folder:roadmap", + }, + "correlation_id": "folder-roadmap-anne-child", + }, + ] + + def mock_optimized_requests(method, url, **kwargs): + if method == "GET": + return mock_response(model_response, 200) + + self.assertEqual(kwargs["body"]["checks"], expected_submitted_checks) + return mock_response( + json.dumps( + { + "result": { + "document-roadmap-anne-child": {"allowed": True}, + "document-roadmap-bob-child": { + "error": { + "input_error": "validation_error", + "message": "bob check failed", + } + }, + "document-budget-anne-child": {"allowed": False}, + "folder-roadmap-anne-child": {"allowed": True}, + } + } + ), + 200, + ) + + mock_request.side_effect = mock_optimized_requests + checks = [ + ClientBatchCheckItem( + user=user, + relation=relation, + object=object_name, + correlation_id=correlation_id, + ) + for user, relation, object_name, correlation_id in ( + ( + "user:anne", + "can_add_child", + "document:roadmap", + "document-roadmap-anne-child", + ), + ( + "user:anne", + "can_add_records", + "document:roadmap", + "document-roadmap-anne-records", + ), + ( + "user:bob", + "can_add_child", + "document:roadmap", + "document-roadmap-bob-child", + ), + ( + "user:bob", + "can_add_records", + "document:roadmap", + "document-roadmap-bob-records", + ), + ( + "user:anne", + "can_add_child", + "document:budget", + "document-budget-anne-child", + ), + ( + "user:anne", + "can_add_records", + "document:budget", + "document-budget-anne-records", + ), + ( + "user:anne", + "can_add_child", + "folder:roadmap", + "folder-roadmap-anne-child", + ), + ( + "user:anne", + "can_add_records", + "folder:roadmap", + "folder-roadmap-anne-records", + ), + ) + ] + configuration = self.configuration + configuration.store_id = store_id + + with OpenFgaClient(configuration) as api_client: + response = api_client.batch_check( + ClientBatchCheckRequest(checks=checks), + options={ + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + }, + ) + + results_by_id = {item.correlation_id: item for item in response.result} + self.assertEqual(set(results_by_id), {check.correlation_id for check in checks}) + expected_allowed = (True, True, False, False, False, False, True, True) + for check, allowed in zip(checks, expected_allowed, strict=True): + result = results_by_id[check.correlation_id] + self.assertIs(result.request, check) + self.assertEqual(result.request.user, check.user) + self.assertEqual(result.request.object, check.object) + self.assertEqual(result.request.relation, check.relation) + self.assertEqual(result.allowed, allowed) + + self.assertEqual( + results_by_id["document-roadmap-bob-child"].error.message, + "bob check failed", + ) + self.assertEqual( + results_by_id["document-roadmap-bob-records"].error.message, + "bob check failed", + ) + batch_requests = [ + call + for call in mock_request.call_args_list + if call.args[1].endswith("/batch-check") + ] + self.assertEqual(len(batch_requests), 1) + + @patch.object(rest.RESTClientObject, "request") + def test_batch_check_optimization_requires_model_id(self, mock_request): + configuration = self.configuration + configuration.store_id = store_id + body = ClientBatchCheckRequest( + checks=[ + ClientBatchCheckItem( + user="user:anne", + relation="can_view", + object="document:roadmap", + ) + ] + ) + + with OpenFgaClient(configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing BatchCheck", + ): + api_client.batch_check( + body, + options={"optimize_relation_aliases": True}, + ) + + mock_request.assert_not_called() + + def test_relation_alias_cache_requires_store_id(self): + with OpenFgaClient(self.configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing BatchCheck", + ): + api_client._get_relation_aliases(None) + + with self.assertRaisesRegex( + FgaValidationException, + "store_id is required but not configured", + ): + api_client._get_relation_aliases( + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"} + ) + + def test_relation_alias_cache_evicts_model_load_errors(self): + configuration = self.configuration + configuration.store_id = store_id + options = { + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1", + "continuation_token": "ignored", + "headers": {"x-test": "value"}, + "optimize_relation_aliases": True, + "page_size": 10, + } + expected_model_options = { + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1", + "headers": {"x-test": "value"}, + } + + with OpenFgaClient(configuration) as api_client: + with patch.object( + api_client, + "read_authorization_model", + return_value=ReadAuthorizationModelResponse(), + ) as mock_read_model: + for _ in range(2): + with self.assertRaisesRegex( + FgaValidationException, + "authorization model was not returned", + ): + api_client._get_relation_aliases(options) + + self.assertEqual(mock_read_model.call_count, 2) + self.assertEqual( + [request.args[0] for request in mock_read_model.call_args_list], + [expected_model_options, expected_model_options], + ) + self.assertEqual(api_client._relation_alias_cache, {}) + + def test_relation_alias_cache_shares_concurrent_model_loads(self): + configuration = self.configuration + configuration.store_id = store_id + options = {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"} + load_started = Event() + release_load = Event() + model_response = ReadAuthorizationModelResponse( + AuthorizationModel( + id="01GXSA8YR785C4FYS3C0RTG7B1", + schema_version="1.1", + type_definitions=[], + ) + ) + + def load_model(_options): + load_started.set() + if not release_load.wait(timeout=2): + raise AssertionError("timed out waiting to release model load") + return model_response + + with OpenFgaClient(configuration) as api_client: + with patch.object( + api_client, + "read_authorization_model", + side_effect=load_model, + ) as mock_read_model: + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(api_client._get_relation_aliases, options) + self.assertTrue(load_started.wait(timeout=1)) + second = executor.submit(api_client._get_relation_aliases, options) + release_load.set() + first_result = first.result(timeout=2) + second_result = second.result(timeout=2) + + self.assertIs(first_result, second_result) + self.assertEqual(mock_read_model.call_count, 1) + def test_batch_check_errors_dupe_cor_id(self): """Test case for duplicate correlation_id being provided to batch_check"""