diff --git a/README.md b/README.md index 78f6dc8d..c37573ff 100644 --- a/README.md +++ b/README.md @@ -1168,7 +1168,10 @@ List the relations a user has on an object. 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 relations that are pure aliases of the same relation. + # This requires an OpenFGA server with BatchCheck support. + "optimize_relation_aliases": True, } body = ClientListRelationsRequest( user="user:81684243-9356-4421-8fbf-a4f8d36aa31b", @@ -1191,6 +1194,12 @@ response = await fga_client.list_relations(body, options) # response.relations = ["can_view", "can_edit"] ``` +When `optimize_relation_aliases` is enabled, the SDK reads and caches the +specified immutable authorization model. If multiple requested relations are +pure aliases of the same relation, it evaluates the shared relation once with +the BatchCheck API and returns the result under the original relation names. +The option is disabled by default and requires an `authorization_model_id`. + #### List Users List the users who have a certain relation to a particular type. @@ -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..11e0bf55 100644 --- a/openfga_sdk/client/client.py +++ b/openfga_sdk/client/client.py @@ -37,6 +37,12 @@ construct_write_single_response, ) from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import ( + RelationCheckGroup, + build_relation_aliases, + group_relations, + is_concrete_user, +) from openfga_sdk.constants import ( CLIENT_BULK_REQUEST_ID_HEADER, CLIENT_MAX_BATCH_SIZE, @@ -172,6 +178,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 +194,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 +264,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 ListRelations" + ) + + 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 ################# @@ -983,12 +1047,30 @@ async def list_relations( :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 consistency(options) - The type of consistency preferred for the request + :param optimize_relation_aliases(options) - Collapse pure relation aliases before evaluation. Defaults to false """ options = set_heading_if_not_set(options, CLIENT_METHOD_HEADER, "ListRelations") options = set_heading_if_not_set( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) ) + if options.get("optimize_relation_aliases") is True: + if self._get_authorization_model_id(options) is None: + raise FgaValidationException( + "authorization_model_id is required when optimizing ListRelations" + ) + if is_concrete_user(body.user): + object_type, separator, _ = body.object.partition(":") + if separator and object_type: + aliases_by_type = await self._get_relation_aliases(options) + groups = group_relations( + body.relations, aliases_by_type.get(object_type, {}) + ) + if any(len(group.indexes) > 1 for group in groups): + return await self._list_relations_with_groups( + body, options, groups + ) + request_body = [ construct_check_request( user=body.user, @@ -1012,6 +1094,68 @@ async def list_relations( result_list = list(result_iterator) return [i.request.relation for i in result_list] + async def _list_relations_with_groups( + self, + body: ClientListRelationsRequest, + options: dict[str, int | str | dict[str, int | str]], + groups: list[RelationCheckGroup], + ) -> list[str]: + """Evaluate grouped checks and preserve requested relation names.""" + checks = [ + ClientBatchCheckItem( + user=body.user, + relation=group.relation, + object=body.object, + contextual_tuples=body.contextual_tuples, + context=body.context, + ) + for group in groups + ] + batch_response = await self.batch_check( + ClientBatchCheckRequest(checks=checks), options + ) + responses_by_relation = { + response.request.relation: response for response in batch_response.result + } + allowed = [False] * len(body.relations) + + for group in groups: + response = responses_by_relation.get(group.relation) + if response is None or response.error is not None: + fallback_checks = [ + construct_check_request( + user=body.user, + relation=body.relations[index], + object=body.object, + contextual_tuples=body.contextual_tuples, + context=body.context, + ) + for index in group.indexes + ] + fallback_responses = await self.client_batch_check( + fallback_checks, options + ) + first_error = next( + ( + fallback.error + for fallback in fallback_responses + if fallback.error is not None + ), + None, + ) + if first_error is not None: + raise first_error + for index, fallback in zip(group.indexes, fallback_responses): + allowed[index] = fallback.allowed + continue + + for index in group.indexes: + allowed[index] = response.allowed + + return [ + relation for index, relation in enumerate(body.relations) if allowed[index] + ] + async def list_users( self, body: ClientListUsersRequest, diff --git a/openfga_sdk/client/relation_optimizer.py b/openfga_sdk/client/relation_optimizer.py new file mode 100644 index 00000000..a1e347e0 --- /dev/null +++ b/openfga_sdk/client/relation_optimizer.py @@ -0,0 +1,102 @@ +from dataclasses import dataclass + +from openfga_sdk.models.authorization_model import AuthorizationModel +from openfga_sdk.models.userset import Userset + + +@dataclass(frozen=True) +class RelationCheckGroup: + """Relations that can share a single authorization check.""" + + relation: str + indexes: tuple[int, ...] + + +def is_concrete_user(user: str) -> bool: + """Return whether a user string represents one concrete object.""" + return user != "*" and not user.endswith(":*") and "#" not in user + + +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_relations( + relations: list[str], aliases: dict[str, str] +) -> list[RelationCheckGroup]: + """Group requested relations by their canonical evaluation target.""" + indexes_by_target: dict[str, list[int]] = {} + for index, relation in enumerate(relations): + target = aliases.get(relation, relation) + indexes_by_target.setdefault(target, []).append(index) + + groups = [] + for target, indexes in indexes_by_target.items(): + submitted_relation = target if len(indexes) > 1 else relations[indexes[0]] + groups.append( + RelationCheckGroup( + relation=submitted_relation, + indexes=tuple(indexes), + ) + ) + return groups + + +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..4b04c2dd 100644 --- a/openfga_sdk/sync/client/client.py +++ b/openfga_sdk/sync/client/client.py @@ -1,7 +1,8 @@ import uuid from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor +from threading import Lock from typing import Any from openfga_sdk.client.configuration import ClientConfiguration @@ -35,6 +36,12 @@ construct_write_single_response, ) from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import ( + RelationCheckGroup, + build_relation_aliases, + group_relations, + is_concrete_user, +) from openfga_sdk.constants import ( CLIENT_BULK_REQUEST_ID_HEADER, CLIENT_MAX_BATCH_SIZE, @@ -172,6 +179,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 +258,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 ListRelations" + ) + + 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 ################# @@ -981,12 +1035,28 @@ def list_relations( :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 consistency(options) - The type of consistency preferred for the request + :param optimize_relation_aliases(options) - Collapse pure relation aliases before evaluation. Defaults to false """ options = set_heading_if_not_set(options, CLIENT_METHOD_HEADER, "ListRelations") options = set_heading_if_not_set( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) ) + if options.get("optimize_relation_aliases") is True: + if self._get_authorization_model_id(options) is None: + raise FgaValidationException( + "authorization_model_id is required when optimizing ListRelations" + ) + if is_concrete_user(body.user): + object_type, separator, _ = body.object.partition(":") + if separator and object_type: + aliases_by_type = self._get_relation_aliases(options) + groups = group_relations( + body.relations, aliases_by_type.get(object_type, {}) + ) + if any(len(group.indexes) > 1 for group in groups): + return self._list_relations_with_groups(body, options, groups) + request_body = [ construct_check_request( user=body.user, @@ -1010,6 +1080,66 @@ def list_relations( result_list = list(result_iterator) return [i.request.relation for i in result_list] + def _list_relations_with_groups( + self, + body: ClientListRelationsRequest, + options: dict[str, int | str | dict[str, int | str]], + groups: list[RelationCheckGroup], + ) -> list[str]: + """Evaluate grouped checks and preserve requested relation names.""" + checks = [ + ClientBatchCheckItem( + user=body.user, + relation=group.relation, + object=body.object, + contextual_tuples=body.contextual_tuples, + context=body.context, + ) + for group in groups + ] + batch_response = self.batch_check( + ClientBatchCheckRequest(checks=checks), options + ) + responses_by_relation = { + response.request.relation: response for response in batch_response.result + } + allowed = [False] * len(body.relations) + + for group in groups: + response = responses_by_relation.get(group.relation) + if response is None or response.error is not None: + fallback_checks = [ + construct_check_request( + user=body.user, + relation=body.relations[index], + object=body.object, + contextual_tuples=body.contextual_tuples, + context=body.context, + ) + for index in group.indexes + ] + fallback_responses = self.client_batch_check(fallback_checks, options) + first_error = next( + ( + fallback.error + for fallback in fallback_responses + if fallback.error is not None + ), + None, + ) + if first_error is not None: + raise first_error + for index, fallback in zip(group.indexes, fallback_responses): + allowed[index] = fallback.allowed + continue + + for index in group.indexes: + allowed[index] = response.allowed + + return [ + relation for index, relation in enumerate(body.relations) if allowed[index] + ] + def list_users( self, body: ClientListUsersRequest, diff --git a/test/client/client_test.py b/test/client/client_test.py index 9850b5d1..a8b2e7fe 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 @@ -15,7 +16,14 @@ from openfga_sdk.client.models.assertion import ClientAssertion from openfga_sdk.client.models.batch_check_item import ClientBatchCheckItem from openfga_sdk.client.models.batch_check_request import ClientBatchCheckRequest +from openfga_sdk.client.models.batch_check_response import ClientBatchCheckResponse +from openfga_sdk.client.models.batch_check_single_response import ( + ClientBatchCheckSingleResponse, +) from openfga_sdk.client.models.check_request import ClientCheckRequest +from openfga_sdk.client.models.client_batch_check_response import ( + ClientBatchCheckClientResponse, +) from openfga_sdk.client.models.expand_request import ClientExpandRequest from openfga_sdk.client.models.list_objects_request import ClientListObjectsRequest from openfga_sdk.client.models.list_relations_request import ClientListRelationsRequest @@ -26,6 +34,7 @@ from openfga_sdk.client.models.write_request import ClientWriteRequest from openfga_sdk.client.models.write_single_response import ClientWriteSingleResponse from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import RelationCheckGroup from openfga_sdk.configuration import RetryParams from openfga_sdk.exceptions import ( FgaValidationException, @@ -35,6 +44,7 @@ ) from openfga_sdk.models.assertion import Assertion from openfga_sdk.models.authorization_model import AuthorizationModel +from openfga_sdk.models.check_error import CheckError from openfga_sdk.models.check_response import CheckResponse from openfga_sdk.models.consistency_preference import ConsistencyPreference from openfga_sdk.models.create_store_request import CreateStoreRequest @@ -2879,6 +2889,397 @@ def mock_check_requests(*args, **kwargs): ) await api_client.close() + @patch.object(rest.RESTClientObject, "request") + async def test_list_relations_optimizes_relation_aliases(self, mock_request): + """ListRelations 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") + correlation_id = checks[0]["correlation_id"] + return mock_response( + json.dumps({"result": {correlation_id: {"allowed": True}}}), + 200, + ) + + mock_request.side_effect = mock_optimized_requests + configuration = self.configuration + configuration.store_id = store_id + request = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ) + + async with OpenFgaClient(configuration) as api_client: + for _ in range(2): + response = await api_client.list_relations( + request, + options={ + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + }, + ) + self.assertEqual(response, ["can_add_child", "can_add_records"]) + + 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), 2) + + @patch.object(rest.RESTClientObject, "request") + async def test_list_relations_optimization_requires_model_id(self, mock_request): + configuration = self.configuration + configuration.store_id = store_id + requests = [ + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + ClientListRelationsRequest( + user="user:*", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document", + ), + ] + + async with OpenFgaClient(configuration) as api_client: + for request in requests: + with self.subTest(user=request.user, object=request.object): + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing ListRelations", + ): + await api_client.list_relations( + request, + 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 ListRelations", + ): + 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_optimized_list_relations_rechecks_missing_batch_results(self): + configuration = self.configuration + configuration.store_id = store_id + body = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ) + fallback_responses = [ + ClientBatchCheckClientResponse( + allowed=allowed, + request=ClientCheckRequest( + user=body.user, + relation=body.relations[index], + object=body.object, + ), + ) + for index, allowed in enumerate((True, False)) + ] + + async with OpenFgaClient(configuration) as api_client: + with ( + patch.object( + api_client, + "batch_check", + return_value=ClientBatchCheckResponse(result=[]), + ), + patch.object( + api_client, + "client_batch_check", + return_value=fallback_responses, + ), + ): + response = await api_client._list_relations_with_groups( + body, + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"}, + [RelationCheckGroup(relation="can_edit", indexes=(0, 1))], + ) + + self.assertEqual(response, ["can_add_child"]) + + async def test_optimized_list_relations_raises_fallback_errors(self): + configuration = self.configuration + configuration.store_id = store_id + body = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ) + optimized_error = CheckError(message="optimized check failed") + fallback_error = ValueError("fallback check failed") + batch_response = ClientBatchCheckResponse( + result=[ + ClientBatchCheckSingleResponse( + allowed=False, + request=ClientTuple( + user=body.user, + relation="can_edit", + object=body.object, + ), + correlation_id="optimized", + error=optimized_error, + ) + ] + ) + fallback_response = ClientBatchCheckClientResponse( + allowed=False, + request=ClientCheckRequest( + user=body.user, + relation=body.relations[0], + object=body.object, + ), + error=fallback_error, + ) + + async with OpenFgaClient(configuration) as api_client: + with ( + patch.object( + api_client, + "batch_check", + return_value=batch_response, + ), + patch.object( + api_client, + "client_batch_check", + return_value=[fallback_response], + ), + ): + with self.assertRaisesRegex(ValueError, "fallback check failed"): + await api_client._list_relations_with_groups( + body, + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"}, + [RelationCheckGroup(relation="can_edit", indexes=(0, 1))], + ) + + @patch.object(rest.RESTClientObject, "request") + async def test_list_relations_rechecks_aliases_after_batch_error( + 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": {}}, + }, + } + ], + } + } + ) + + def mock_error_then_checks(method, url, **kwargs): + if method == "GET": + return mock_response(model_response, 200) + if url.endswith("/batch-check"): + correlation_id = kwargs["body"]["checks"][0]["correlation_id"] + return mock_response( + json.dumps( + { + "result": { + correlation_id: { + "error": { + "input_error": "validation_error", + "message": "check failed", + } + } + } + } + ), + 200, + ) + + relation = kwargs["body"]["tuple_key"]["relation"] + return mock_response( + json.dumps({"allowed": relation == "can_add_child"}), 200 + ) + + mock_request.side_effect = mock_error_then_checks + configuration = self.configuration + configuration.store_id = store_id + + async with OpenFgaClient(configuration) as api_client: + response = await api_client.list_relations( + ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ), + options={ + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + }, + ) + + self.assertEqual(response, ["can_add_child"]) + check_requests = [ + call + for call in mock_request.call_args_list + if call.args[1].endswith("/check") + ] + self.assertEqual(len(check_requests), 2) + @patch.object(rest.RESTClientObject, "request") async def test_list_relations_unauthorized(self, mock_request): """Test case for list relations with 401 response""" diff --git a/test/client/relation_optimizer_test.py b/test/client/relation_optimizer_test.py new file mode 100644 index 00000000..302b5245 --- /dev/null +++ b/test/client/relation_optimizer_test.py @@ -0,0 +1,88 @@ +from openfga_sdk.client.relation_optimizer import ( + build_relation_aliases, + group_relations, + is_concrete_user, +) +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_relations_collapses_aliases_and_preserves_singletons(): + groups = group_relations( + ["can_comment", "can_reply", "viewer", "commenter"], + { + "can_comment": "commenter", + "can_reply": "commenter", + }, + ) + + assert [(group.relation, group.indexes) for group in groups] == [ + ("commenter", (0, 1, 3)), + ("viewer", (2,)), + ] + + +def test_is_concrete_user_rejects_usersets_and_wildcards(): + assert is_concrete_user("user:anne") + assert not is_concrete_user("team:eng#member") + assert not is_concrete_user("user:*") + assert not is_concrete_user("*") diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index ec2686af..7c96c74b 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 @@ -13,7 +15,14 @@ from openfga_sdk.client.models.assertion import ClientAssertion from openfga_sdk.client.models.batch_check_item import ClientBatchCheckItem from openfga_sdk.client.models.batch_check_request import ClientBatchCheckRequest +from openfga_sdk.client.models.batch_check_response import ClientBatchCheckResponse +from openfga_sdk.client.models.batch_check_single_response import ( + ClientBatchCheckSingleResponse, +) from openfga_sdk.client.models.check_request import ClientCheckRequest +from openfga_sdk.client.models.client_batch_check_response import ( + ClientBatchCheckClientResponse, +) from openfga_sdk.client.models.expand_request import ClientExpandRequest from openfga_sdk.client.models.list_objects_request import ClientListObjectsRequest from openfga_sdk.client.models.list_relations_request import ClientListRelationsRequest @@ -24,6 +33,7 @@ from openfga_sdk.client.models.write_request import ClientWriteRequest from openfga_sdk.client.models.write_single_response import ClientWriteSingleResponse from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts +from openfga_sdk.client.relation_optimizer import RelationCheckGroup from openfga_sdk.configuration import RetryParams from openfga_sdk.exceptions import ( FgaValidationException, @@ -33,6 +43,7 @@ ) from openfga_sdk.models.assertion import Assertion from openfga_sdk.models.authorization_model import AuthorizationModel +from openfga_sdk.models.check_error import CheckError from openfga_sdk.models.check_response import CheckResponse from openfga_sdk.models.consistency_preference import ConsistencyPreference from openfga_sdk.models.create_store_request import CreateStoreRequest @@ -2882,6 +2893,323 @@ def mock_check_requests(*args, **kwargs): ) api_client.close() + @patch.object(rest.RESTClientObject, "request") + def test_list_relations_optimizes_relation_aliases(self, mock_request): + """ListRelations 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") + correlation_id = checks[0]["correlation_id"] + return mock_response( + json.dumps({"result": {correlation_id: {"allowed": True}}}), + 200, + ) + + mock_request.side_effect = mock_optimized_requests + configuration = self.configuration + configuration.store_id = store_id + request = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ) + + with OpenFgaClient(configuration) as api_client: + for _ in range(2): + response = api_client.list_relations( + request, + options={ + "authorization_model_id": authorization_model_id, + "optimize_relation_aliases": True, + }, + ) + self.assertEqual(response, ["can_add_child", "can_add_records"]) + + 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), 2) + + @patch.object(rest.RESTClientObject, "request") + def test_list_relations_optimization_requires_model_id(self, mock_request): + configuration = self.configuration + configuration.store_id = store_id + requests = [ + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + ClientListRelationsRequest( + user="user:*", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document", + ), + ] + + with OpenFgaClient(configuration) as api_client: + for request in requests: + with self.subTest(user=request.user, object=request.object): + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing ListRelations", + ): + api_client.list_relations( + request, + 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 ListRelations", + ): + 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_optimized_list_relations_rechecks_errors_and_missing_results(self): + configuration = self.configuration + configuration.store_id = store_id + body = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records", "can_view"], + object="document:roadmap", + ) + batch_response = ClientBatchCheckResponse( + result=[ + ClientBatchCheckSingleResponse( + allowed=False, + request=ClientTuple( + user=body.user, + relation="can_edit", + object=body.object, + ), + correlation_id="optimized", + error=CheckError(message="optimized check failed"), + ) + ] + ) + fallback_responses = [ + ClientBatchCheckClientResponse( + allowed=allowed, + request=ClientCheckRequest( + user=body.user, + relation=body.relations[index], + object=body.object, + ), + ) + for index, allowed in enumerate((True, False)) + ] + missing_response = ClientBatchCheckClientResponse( + allowed=False, + request=ClientCheckRequest( + user=body.user, + relation=body.relations[2], + object=body.object, + ), + ) + + with OpenFgaClient(configuration) as api_client: + with ( + patch.object( + api_client, + "batch_check", + return_value=batch_response, + ), + patch.object( + api_client, + "client_batch_check", + side_effect=[fallback_responses, [missing_response]], + ) as mock_fallback, + ): + response = api_client._list_relations_with_groups( + body, + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"}, + [ + RelationCheckGroup(relation="can_edit", indexes=(0, 1)), + RelationCheckGroup(relation="can_view", indexes=(2,)), + ], + ) + + self.assertEqual(response, ["can_add_child"]) + self.assertEqual(mock_fallback.call_count, 2) + + def test_optimized_list_relations_raises_fallback_errors(self): + configuration = self.configuration + configuration.store_id = store_id + body = ClientListRelationsRequest( + user="user:anne", + relations=["can_add_child", "can_add_records"], + object="document:roadmap", + ) + batch_response = ClientBatchCheckResponse( + result=[ + ClientBatchCheckSingleResponse( + allowed=False, + request=ClientTuple( + user=body.user, + relation="can_edit", + object=body.object, + ), + correlation_id="optimized", + error=CheckError(message="optimized check failed"), + ) + ] + ) + fallback_response = ClientBatchCheckClientResponse( + allowed=False, + request=ClientCheckRequest( + user=body.user, + relation=body.relations[0], + object=body.object, + ), + error=ValueError("fallback check failed"), + ) + + with OpenFgaClient(configuration) as api_client: + with ( + patch.object( + api_client, + "batch_check", + return_value=batch_response, + ), + patch.object( + api_client, + "client_batch_check", + return_value=[fallback_response], + ), + ): + with self.assertRaisesRegex(ValueError, "fallback check failed"): + api_client._list_relations_with_groups( + body, + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"}, + [RelationCheckGroup(relation="can_edit", indexes=(0, 1))], + ) + @patch.object(rest.RESTClientObject, "request") def test_list_relations_unauthorized(self, mock_request): """Test case for list relations with 401 response"""