feat: optimize aliased BatchCheck requests

README.md

@@ -875,7 +875,9 @@ Similar to [check](#check), but instead of checking a single user-object relatio
 
 options = {
     # You can rely on the model id set in the configuration or override it for this specific request
-    "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1"
+    "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1",
+    # Optionally collapse checks that resolve through pure relation aliases
+    "optimize_relation_aliases": True,
 }
 checks = [ClientBatchCheckItem(
     user="user:81684243-9356-4421-8fbf-a4f8d36aa31b",
@@ -965,6 +967,13 @@ response = await fga_client.batch_check(ClientBatchCheckRequest(checks=checks),
 # ]
 ```
 
+When `optimize_relation_aliases` is enabled, the SDK reads and caches the pinned
+authorization model, then combines checks only when their user, object, context,
+and contextual tuples are identical and their relations are pure aliases of the
+same relation. The response still contains an entry for every original
+`correlation_id` and request. This optimization is disabled by default and
+requires an authorization model ID.
+
 ##### Client Batch Check
 
 
@@ -1569,4 +1578,3 @@ See [CONTRIBUTING](./CONTRIBUTING.md) for details.
 This project is licensed under the Apache-2.0 license. See the [LICENSE](https://github.com/openfga/python-sdk/blob/main/LICENSE) file for more info.
 
 The code in this repo was auto generated by [OpenAPI Generator](https://github.com/OpenAPITools/openapi-generator) from a template based on the [python legacy template](https://github.com/OpenAPITools/openapi-generator/tree/master/modules/openapi-generator/src/main/resources/python-legacy), licensed under the [Apache License 2.0](https://github.com/OpenAPITools/openapi-generator/blob/master/LICENSE).
-

openfga_sdk/client/client.py

@@ -2,7 +2,7 @@
 import uuid
 
 from collections.abc import AsyncIterator
-from typing import Any
+from typing import Any, cast
 
 from openfga_sdk.api.open_fga_api import OpenFgaApi
 from openfga_sdk.api_client import ApiClient
@@ -37,6 +37,10 @@
     construct_write_single_response,
 )
 from openfga_sdk.client.models.write_transaction_opts import WriteTransactionOpts
+from openfga_sdk.client.relation_optimizer import (
+    build_relation_aliases,
+    group_batch_checks,
+)
 from openfga_sdk.constants import (
     CLIENT_BULK_REQUEST_ID_HEADER,
     CLIENT_MAX_BATCH_SIZE,
@@ -172,6 +176,9 @@ def __init__(self, configuration: ClientConfiguration):
         self._client_configuration = configuration
         self._api_client = ApiClient(configuration)
         self._api = OpenFgaApi(self._api_client)
+        self._relation_alias_cache: dict[
+            tuple[str, str], asyncio.Task[dict[str, dict[str, str]]]
+        ] = {}
 
         # Set default headers from configuration
         if configuration.headers:
@@ -185,6 +192,14 @@ async def __aexit__(self, exc_type, exc_value, traceback):
         await self.close()
 
     async def close(self):
+        """Cancel cached model loads and close the API client."""
+        tasks = list(self._relation_alias_cache.values())
+        self._relation_alias_cache.clear()
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
         await self._api.close()
 
     def _get_authorization_model_id(
@@ -247,6 +262,53 @@ def get_authorization_model_id(self):
         """
         return self._client_configuration.authorization_model_id
 
+    async def _get_relation_aliases(
+        self,
+        options: dict[str, int | str | dict[str, int | str]] | None,
+    ) -> dict[str, dict[str, str]]:
+        """Return cached relation aliases for the configured model."""
+        authorization_model_id = self._get_authorization_model_id(options)
+        if authorization_model_id is None:
+            raise FgaValidationException(
+                "authorization_model_id is required when optimizing BatchCheck"
+            )
+
+        store_id = self.get_store_id()
+        if store_id is None or store_id == "":
+            raise FgaValidationException("store_id is required but not configured")
+
+        cache_key = (store_id, authorization_model_id)
+        task = self._relation_alias_cache.get(cache_key)
+        if task is None:
+            task = asyncio.create_task(self._load_relation_aliases(options))
+            self._relation_alias_cache[cache_key] = task
+
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            if task.cancelled() and self._relation_alias_cache.get(cache_key) is task:
+                self._relation_alias_cache.pop(cache_key, None)
+            raise
+        except Exception:
+            if self._relation_alias_cache.get(cache_key) is task:
+                self._relation_alias_cache.pop(cache_key, None)
+            raise
+
+    async def _load_relation_aliases(
+        self,
+        options: dict[str, int | str | dict[str, int | str]] | None,
+    ) -> dict[str, dict[str, str]]:
+        """Read the configured model and build relation alias mappings."""
+        model_options = {
+            key: options[key]
+            for key in ("authorization_model_id", "headers", "retry_params")
+            if options is not None and key in options
+        }
+        response = await self.read_authorization_model(model_options)
+        if response.authorization_model is None:
+            raise FgaValidationException("authorization model was not returned")
+        return build_relation_aliases(response.authorization_model)
+
     #################
     # Stores
     #################
@@ -784,6 +846,7 @@ async def batch_check(
         :param retryParams(options) - Override the retry parameters for this request
         :param retryParams.maxRetry(options) - Override the max number of retries on each API request
         :param retryParams.minWaitInMs(options) - Override the minimum wait before a retry is initiated
+        :param optimize_relation_aliases(options) - Collapse equivalent pure relation aliases. Defaults to false
         """
         options = set_heading_if_not_set(
             options, CLIENT_BULK_REQUEST_ID_HEADER, str(uuid.uuid4())
@@ -809,42 +872,73 @@ async def batch_check(
             elif isinstance(options["max_batch_size"], int):
                 max_batch_size = options["max_batch_size"]
 
-        id_to_check: dict[str, ClientBatchCheckItem] = {}
-
-        def track_and_transform(checks):
-            transformed = []
-            for check in checks:
-                if check.correlation_id is None:
-                    check.correlation_id = str(uuid.uuid4())
-
-                if check.correlation_id in id_to_check:
-                    raise FgaValidationException(
-                        f"Duplicate correlation_id ({check.correlation_id}) provided"
-                    )
-
-                id_to_check[check.correlation_id] = check
+        optimize_relation_aliases = options.get("optimize_relation_aliases") is True
+        if (
+            optimize_relation_aliases
+            and self._get_authorization_model_id(options) is None
+        ):
+            raise FgaValidationException(
+                "authorization_model_id is required when optimizing BatchCheck"
+            )
 
-                transformed.append(construct_batch_item(check))
-            return transformed
+        id_to_check: dict[str, ClientBatchCheckItem] = {}
+        for check in body.checks:
+            if check.correlation_id is None:
+                check.correlation_id = str(uuid.uuid4())
+
+            correlation_id = cast(str, check.correlation_id)
+            if correlation_id in id_to_check:
+                raise FgaValidationException(
+                    f"Duplicate correlation_id ({correlation_id}) provided"
+                )
+            id_to_check[correlation_id] = check
+
+        aliases_by_type = (
+            await self._get_relation_aliases(options)
+            if optimize_relation_aliases
+            else {}
+        )
+        groups = group_batch_checks(body.checks, aliases_by_type)
+        submitted_to_original_ids: dict[str, tuple[str, ...]] = {}
+        submitted_checks: list[ClientBatchCheckItem] = []
+
+        for group in groups:
+            representative = body.checks[group.indexes[0]]
+            submitted_check = representative
+            if len(group.indexes) > 1:
+                submitted_check = ClientBatchCheckItem(
+                    user=representative.user,
+                    relation=group.relation,
+                    object=representative.object,
+                    correlation_id=representative.correlation_id,
+                    contextual_tuples=representative.contextual_tuples,
+                    context=representative.context,
+                )
+
+            submitted_id = cast(str, submitted_check.correlation_id)
+            submitted_to_original_ids[submitted_id] = tuple(
+                cast(str, body.checks[index].correlation_id) for index in group.indexes
+            )
+            submitted_checks.append(submitted_check)
 
         checks = [
-            track_and_transform(
-                body.checks[i * max_batch_size : (i + 1) * max_batch_size]
-            )
-            for i in range((len(body.checks) + max_batch_size - 1) // max_batch_size)
+            [construct_batch_item(check) for check in chunk]
+            for chunk in _chuck_array(submitted_checks, max_batch_size)
         ]
 
         result = []
         sem = asyncio.Semaphore(max_parallel_requests)
 
         def map_response(id, result):
-            check = id_to_check[id]
-            return ClientBatchCheckSingleResponse(
-                allowed=result.allowed,
-                request=check,
-                correlation_id=id,
-                error=result.error,
-            )
+            return [
+                ClientBatchCheckSingleResponse(
+                    allowed=result.allowed,
+                    request=id_to_check[original_id],
+                    correlation_id=original_id,
+                    error=result.error,
+                )
+                for original_id in submitted_to_original_ids[id]
+            ]
 
         async def coro(checks):
             res = await self._single_batch_check(
@@ -857,9 +951,8 @@ async def coro(checks):
                 options,
             )
 
-            result.extend(
-                [map_response(c_id, c_result) for c_id, c_result in res.result.items()]
-            )
+            for correlation_id, check_result in res.result.items():
+                result.extend(map_response(correlation_id, check_result))
 
         batch_check_coros = [coro(request) for request in checks]
         await asyncio.gather(*batch_check_coros)

openfga_sdk/client/relation_optimizer.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+
+from openfga_sdk.client.models.batch_check_item import ClientBatchCheckItem
+from openfga_sdk.models.authorization_model import AuthorizationModel
+from openfga_sdk.models.userset import Userset
+
+
+@dataclass(frozen=True)
+class RelationCheckGroup:
+    """BatchCheck items that can share one relation evaluation."""
+
+    relation: str
+    indexes: tuple[int, ...]
+
+
+def build_relation_aliases(
+    authorization_model: AuthorizationModel,
+) -> dict[str, dict[str, str]]:
+    """Build canonical targets for pure computed-userset relation aliases."""
+    aliases_by_type: dict[str, dict[str, str]] = {}
+
+    for type_definition in authorization_model.type_definitions or []:
+        relations = type_definition.relations or {}
+        direct_aliases = {
+            relation: target
+            for relation, rewrite in relations.items()
+            if (target := _pure_computed_userset_target(rewrite)) is not None
+            and target in relations
+        }
+
+        canonical_aliases: dict[str, str] = {}
+        for relation in direct_aliases:
+            target = _resolve_alias(relation, direct_aliases)
+            if target is not None and target != relation:
+                canonical_aliases[relation] = target
+
+        aliases_by_type[type_definition.type] = canonical_aliases
+
+    return aliases_by_type
+
+
+def group_batch_checks(
+    checks: list[ClientBatchCheckItem], aliases_by_type: dict[str, dict[str, str]]
+) -> list[RelationCheckGroup]:
+    """Group checks that differ only by pure aliases of the same relation."""
+    candidates: list[tuple[str, ClientBatchCheckItem, list[int]]] = []
+
+    for index, check in enumerate(checks):
+        object_type, separator, _ = check.object.partition(":")
+        aliases = aliases_by_type.get(object_type, {}) if separator else {}
+        target = aliases.get(check.relation, check.relation)
+
+        for candidate_target, representative, indexes in candidates:
+            if candidate_target == target and _same_check_inputs(check, representative):
+                indexes.append(index)
+                break
+        else:
+            candidates.append((target, check, [index]))
+
+    groups: list[RelationCheckGroup] = []
+    for target, _, indexes in candidates:
+        if len(indexes) > 1 and any(
+            checks[index].relation != target for index in indexes
+        ):
+            groups.append(RelationCheckGroup(relation=target, indexes=tuple(indexes)))
+        else:
+            groups.extend(
+                RelationCheckGroup(
+                    relation=checks[index].relation,
+                    indexes=(index,),
+                )
+                for index in indexes
+            )
+
+    groups.sort(key=lambda group: group.indexes[0])
+    return groups
+
+
+def _same_check_inputs(left: ClientBatchCheckItem, right: ClientBatchCheckItem) -> bool:
+    """Return whether two checks differ only in relation and correlation ID."""
+    return (
+        left.user == right.user
+        and left.object == right.object
+        and left.context == right.context
+        and left.contextual_tuples == right.contextual_tuples
+    )
+
+
+def _pure_computed_userset_target(rewrite: Userset) -> str | None:
+    """Return the target when a rewrite is only a same-object computed userset."""
+    computed_userset = rewrite.computed_userset
+    if computed_userset is None or not computed_userset.relation:
+        return None
+
+    if any(
+        value is not None
+        for value in (
+            rewrite.this,
+            rewrite.tuple_to_userset,
+            rewrite.union,
+            rewrite.intersection,
+            rewrite.difference,
+        )
+    ):
+        return None
+
+    if computed_userset.object not in (None, ""):
+        return None
+
+    return computed_userset.relation
+
+
+def _resolve_alias(relation: str, direct_aliases: dict[str, str]) -> str | None:
+    """Resolve an alias chain, returning none when it contains a cycle."""
+    visited = set()
+    current = relation
+
+    while current in direct_aliases:
+        if current in visited:
+            return None
+        visited.add(current)
+        current = direct_aliases[current]
+
+    return current