feat: optimize aliased ListRelations checks

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).
-

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,

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