From 487c9b30d42d6f65f6d719c808efddeb556d5e60 Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Mon, 17 Aug 2026 07:01:26 +0000 Subject: [PATCH 1/5] feat: optimize aliased ListRelations checks --- README.md | 12 +- openfga_sdk/client/client.py | 125 ++++++++++++++++ openfga_sdk/client/relation_optimizer.py | 100 +++++++++++++ openfga_sdk/sync/client/client.py | 124 +++++++++++++++- test/client/client_test.py | 174 +++++++++++++++++++++++ test/client/relation_optimizer_test.py | 84 +++++++++++ test/sync/client/client_test.py | 95 +++++++++++++ 7 files changed, 711 insertions(+), 3 deletions(-) create mode 100644 openfga_sdk/client/relation_optimizer.py create mode 100644 test/client/relation_optimizer_test.py diff --git a/README.md b/README.md index 78f6dc8d..4a9ca093 100644 --- a/README.md +++ b/README.md @@ -691,7 +691,10 @@ By default, write runs in a transaction mode where any invalid operation (deleti 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 = ClientWriteRequest( writes=[ @@ -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..2b9df29a 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: @@ -247,6 +256,46 @@ 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]]: + 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]]: + response = await self.read_authorization_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 +1032,25 @@ 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 and 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 +1074,69 @@ 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]: + 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: + continue + if 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..aed38d00 --- /dev/null +++ b/openfga_sdk/client/relation_optimizer.py @@ -0,0 +1,100 @@ +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: + 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: + 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..1ecf82d3 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,43 @@ 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]]: + 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: + response = self.read_authorization_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 +1029,25 @@ 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 and 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 +1071,67 @@ 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]: + 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: + continue + if 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..02390954 100644 --- a/test/client/client_test.py +++ b/test/client/client_test.py @@ -2879,6 +2879,180 @@ 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 + + async with OpenFgaClient(configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing ListRelations", + ): + await api_client.list_relations( + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + options={"optimize_relation_aliases": True}, + ) + + mock_request.assert_not_called() + + @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..d39f84e9 --- /dev/null +++ b/test/client/relation_optimizer_test.py @@ -0,0 +1,84 @@ +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") + ) + ] + ) + ), + "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..81ce2a45 100644 --- a/test/sync/client/client_test.py +++ b/test/sync/client/client_test.py @@ -2882,6 +2882,101 @@ 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 + + with OpenFgaClient(configuration) as api_client: + with self.assertRaisesRegex( + FgaValidationException, + "authorization_model_id is required when optimizing ListRelations", + ): + api_client.list_relations( + ClientListRelationsRequest( + user="user:anne", + relations=["can_view", "viewer"], + object="document:roadmap", + ), + options={"optimize_relation_aliases": True}, + ) + + mock_request.assert_not_called() + @patch.object(rest.RESTClientObject, "request") def test_list_relations_unauthorized(self, mock_request): """Test case for list relations with 401 response""" From 25d32015b42f6dbd4e05cf680aa0de8397c1028c Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Mon, 17 Aug 2026 07:13:42 +0000 Subject: [PATCH 2/5] test: cover ListRelations optimization edge cases --- test/client/client_test.py | 135 ++++++++++++++++++++++ test/client/relation_optimizer_test.py | 4 + test/sync/client/client_test.py | 150 +++++++++++++++++++++++++ 3 files changed, 289 insertions(+) diff --git a/test/client/client_test.py b/test/client/client_test.py index 02390954..281b1052 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 @@ -2974,6 +2984,131 @@ async def test_list_relations_optimization_requires_model_id(self, mock_request) 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, + "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"} + + 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(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_optimized_list_relations_ignores_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", + ) + + async with OpenFgaClient(configuration) as api_client: + with patch.object( + api_client, + "batch_check", + return_value=ClientBatchCheckResponse(result=[]), + ): + response = await api_client._list_relations_with_groups( + body, + {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"}, + [RelationCheckGroup(relation="can_edit", indexes=(0, 1))], + ) + + self.assertEqual(response, []) + + 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 diff --git a/test/client/relation_optimizer_test.py b/test/client/relation_optimizer_test.py index d39f84e9..302b5245 100644 --- a/test/client/relation_optimizer_test.py +++ b/test/client/relation_optimizer_test.py @@ -35,6 +35,10 @@ def test_build_relation_aliases_only_includes_pure_aliases(): ] ) ), + "mixed_rewrite": Userset( + computed_userset=ObjectRelation(relation="viewer"), + union=Usersets(child=[Userset(this={})]), + ), "other_object": Userset( computed_userset=ObjectRelation( object="document:other", relation="viewer" diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index 81ce2a45..b86b67df 100644 --- a/test/sync/client/client_test.py +++ b/test/sync/client/client_test.py @@ -13,7 +13,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 +31,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 +41,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 @@ -2977,6 +2986,147 @@ def test_list_relations_optimization_requires_model_id(self, mock_request): 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, + "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"} + + 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(api_client._relation_alias_cache, {}) + + 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)) + ] + + 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_responses, + ), + ): + 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"]) + + 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""" From 0de3ad57a3d396e7ffdce039d46b8adc75f34439 Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Mon, 17 Aug 2026 07:23:13 +0000 Subject: [PATCH 3/5] fix: address ListRelations review feedback --- README.md | 10 ++-- openfga_sdk/client/client.py | 14 +++-- openfga_sdk/client/relation_optimizer.py | 2 + openfga_sdk/sync/client/client.py | 13 +++-- test/client/client_test.py | 46 +++++++++++++--- test/sync/client/client_test.py | 68 ++++++++++++++++++++++-- 6 files changed, 130 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 4a9ca093..c37573ff 100644 --- a/README.md +++ b/README.md @@ -691,10 +691,7 @@ By default, write runs in a transaction mode where any invalid operation (deleti options = { # You can rely on the model id set in the configuration or override it for this specific request - "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, + "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1" } body = ClientWriteRequest( writes=[ @@ -1171,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", diff --git a/openfga_sdk/client/client.py b/openfga_sdk/client/client.py index 2b9df29a..dcb3b4ac 100644 --- a/openfga_sdk/client/client.py +++ b/openfga_sdk/client/client.py @@ -260,6 +260,7 @@ 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( @@ -291,7 +292,13 @@ async def _load_relation_aliases( self, options: dict[str, int | str | dict[str, int | str]] | None, ) -> dict[str, dict[str, str]]: - response = await self.read_authorization_model(options) + """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) @@ -1080,6 +1087,7 @@ async def _list_relations_with_groups( 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, @@ -1100,9 +1108,7 @@ async def _list_relations_with_groups( for group in groups: response = responses_by_relation.get(group.relation) - if response is None: - continue - if response.error is not None: + if response is None or response.error is not None: fallback_checks = [ construct_check_request( user=body.user, diff --git a/openfga_sdk/client/relation_optimizer.py b/openfga_sdk/client/relation_optimizer.py index aed38d00..a1e347e0 100644 --- a/openfga_sdk/client/relation_optimizer.py +++ b/openfga_sdk/client/relation_optimizer.py @@ -65,6 +65,7 @@ def group_relations( 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 @@ -88,6 +89,7 @@ def _pure_computed_userset_target(rewrite: Userset) -> str | None: 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 diff --git a/openfga_sdk/sync/client/client.py b/openfga_sdk/sync/client/client.py index 1ecf82d3..770e14c9 100644 --- a/openfga_sdk/sync/client/client.py +++ b/openfga_sdk/sync/client/client.py @@ -262,6 +262,7 @@ 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( @@ -282,7 +283,12 @@ def _get_relation_aliases( if should_load: try: - response = self.read_authorization_model(options) + 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)) @@ -1077,6 +1083,7 @@ def _list_relations_with_groups( 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, @@ -1097,9 +1104,7 @@ def _list_relations_with_groups( for group in groups: response = responses_by_relation.get(group.relation) - if response is None: - continue - if response.error is not None: + if response is None or response.error is not None: fallback_checks = [ construct_check_request( user=body.user, diff --git a/test/client/client_test.py b/test/client/client_test.py index 281b1052..7de91bf7 100644 --- a/test/client/client_test.py +++ b/test/client/client_test.py @@ -2997,7 +2997,17 @@ async def test_relation_alias_cache_requires_store_id(self): async def test_relation_alias_cache_evicts_model_load_errors(self): configuration = self.configuration configuration.store_id = store_id - options = {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"} + 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( @@ -3013,6 +3023,10 @@ async def test_relation_alias_cache_evicts_model_load_errors(self): 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): @@ -3032,7 +3046,7 @@ async def test_relation_alias_cache_evicts_cancelled_loads(self): self.assertNotIn(cache_key, api_client._relation_alias_cache) - async def test_optimized_list_relations_ignores_missing_batch_results(self): + async def test_optimized_list_relations_rechecks_missing_batch_results(self): configuration = self.configuration configuration.store_id = store_id body = ClientListRelationsRequest( @@ -3040,12 +3054,30 @@ async def test_optimized_list_relations_ignores_missing_batch_results(self): 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=[]), + 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, @@ -3053,7 +3085,7 @@ async def test_optimized_list_relations_ignores_missing_batch_results(self): [RelationCheckGroup(relation="can_edit", indexes=(0, 1))], ) - self.assertEqual(response, []) + self.assertEqual(response, ["can_add_child"]) async def test_optimized_list_relations_raises_fallback_errors(self): configuration = self.configuration diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index b86b67df..b111d606 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 @@ -2999,7 +3001,17 @@ def test_relation_alias_cache_requires_store_id(self): def test_relation_alias_cache_evicts_model_load_errors(self): configuration = self.configuration configuration.store_id = store_id - options = {"authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"} + 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( @@ -3015,8 +3027,49 @@ def test_relation_alias_cache_evicts_model_load_errors(self): 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 @@ -3050,6 +3103,14 @@ def test_optimized_list_relations_rechecks_errors_and_missing_results(self): ) 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 ( @@ -3061,8 +3122,8 @@ def test_optimized_list_relations_rechecks_errors_and_missing_results(self): patch.object( api_client, "client_batch_check", - return_value=fallback_responses, - ), + side_effect=[fallback_responses, [missing_response]], + ) as mock_fallback, ): response = api_client._list_relations_with_groups( body, @@ -3074,6 +3135,7 @@ def test_optimized_list_relations_rechecks_errors_and_missing_results(self): ) 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 From 34b6d902e08293caf3e2e851bef3a73f6764f55f Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Mon, 17 Aug 2026 07:26:10 +0000 Subject: [PATCH 4/5] fix: validate optimization and clean up model loads --- openfga_sdk/client/client.py | 33 +++++++++---- openfga_sdk/sync/client/client.py | 23 +++++---- test/client/client_test.py | 78 ++++++++++++++++++++++++++----- test/sync/client/client_test.py | 39 +++++++++++----- 4 files changed, 129 insertions(+), 44 deletions(-) diff --git a/openfga_sdk/client/client.py b/openfga_sdk/client/client.py index dcb3b4ac..11e0bf55 100644 --- a/openfga_sdk/client/client.py +++ b/openfga_sdk/client/client.py @@ -194,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( @@ -1046,17 +1054,22 @@ async def list_relations( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) ) - if options.get("optimize_relation_aliases") is True and 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 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 any(len(group.indexes) > 1 for group in groups): - return await self._list_relations_with_groups(body, options, groups) + 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( diff --git a/openfga_sdk/sync/client/client.py b/openfga_sdk/sync/client/client.py index 770e14c9..4b04c2dd 100644 --- a/openfga_sdk/sync/client/client.py +++ b/openfga_sdk/sync/client/client.py @@ -1042,17 +1042,20 @@ def list_relations( options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4()) ) - if options.get("optimize_relation_aliases") is True and 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 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 any(len(group.indexes) > 1 for group in groups): - return self._list_relations_with_groups(body, options, groups) + 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( diff --git a/test/client/client_test.py b/test/client/client_test.py index 7de91bf7..b017be27 100644 --- a/test/client/client_test.py +++ b/test/client/client_test.py @@ -2967,20 +2967,35 @@ def mock_optimized_requests(method, url, **kwargs): 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: - with self.assertRaisesRegex( - FgaValidationException, - "authorization_model_id is required when optimizing ListRelations", - ): - await api_client.list_relations( - ClientListRelationsRequest( - user="user:anne", - relations=["can_view", "viewer"], - object="document:roadmap", - ), - options={"optimize_relation_aliases": True}, - ) + 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() @@ -3046,6 +3061,45 @@ async def test_relation_alias_cache_evicts_cancelled_loads(self): 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 diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index b111d606..d741fff4 100644 --- a/test/sync/client/client_test.py +++ b/test/sync/client/client_test.py @@ -2971,20 +2971,35 @@ def mock_optimized_requests(method, url, **kwargs): 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: - with self.assertRaisesRegex( - FgaValidationException, - "authorization_model_id is required when optimizing ListRelations", - ): - api_client.list_relations( - ClientListRelationsRequest( - user="user:anne", - relations=["can_view", "viewer"], - object="document:roadmap", - ), - options={"optimize_relation_aliases": True}, - ) + 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() From 8be9c7d82f29720ec7811e304cd7a0f6ae61336b Mon Sep 17 00:00:00 2001 From: Siddhant Khare Date: Mon, 17 Aug 2026 07:28:24 +0000 Subject: [PATCH 5/5] test: cover relation alias cache validation --- test/client/client_test.py | 6 ++++++ test/sync/client/client_test.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/test/client/client_test.py b/test/client/client_test.py index b017be27..a8b2e7fe 100644 --- a/test/client/client_test.py +++ b/test/client/client_test.py @@ -3001,6 +3001,12 @@ async def test_list_relations_optimization_requires_model_id(self, mock_request) 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", diff --git a/test/sync/client/client_test.py b/test/sync/client/client_test.py index d741fff4..7c96c74b 100644 --- a/test/sync/client/client_test.py +++ b/test/sync/client/client_test.py @@ -3005,6 +3005,12 @@ def test_list_relations_optimization_requires_model_id(self, mock_request): 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",