feat: merge equivalent checks in BatchCheck (beta)

Makefile

@@ -1,4 +1,7 @@
-.PHONY: sync test lint fmt fix check doc
+.PHONY: sync test lint fmt fix check doc update-equivalent-checks-cases
+
+# openfga/language commit to copy the shared equivalent-checks cases from
+LANGUAGE_REF ?= 5afcf707e44aa21d8ae17e47938b3f0164a13cc7
 
 # Install/update dependencies
 sync:
@@ -26,6 +29,11 @@ fix:
 	uv run ruff check --fix .
 	uv run ruff format .
 
+# Copy the shared equivalent-checks cases from openfga/language
+# Usage: make update-equivalent-checks-cases LANGUAGE_REF=<commit>
+update-equivalent-checks-cases:
+	curl -fsSL https://raw.githubusercontent.com/openfga/language/$(LANGUAGE_REF)/tests/data/equivalent-checks-cases.json -o test/data/equivalent-checks-cases.json
+
 # Run checks (lint + test)
 check: lint fmt test
 
@@ -39,4 +47,5 @@ doc:
 	@echo "  fmt     - Format code"
 	@echo "  fix     - Fix fixable linting and formatting issues"
 	@echo "  check   - Run both linting and tests"
+	@echo "  update-equivalent-checks-cases - Copy the shared equivalent-checks cases from openfga/language (LANGUAGE_REF=<commit>)"
 	@echo "  doc     - Show this help message"

README.md

@@ -965,6 +965,31 @@ response = await fga_client.batch_check(ClientBatchCheckRequest(checks=checks),
 # ]
 ```
 
+###### Merge equivalent checks (beta)
+
+Set `merge_equivalent_checks` to `True` to send checks that are certain to get the same answer only once. The SDK copies each answer back, so you still get one result per check you passed in.
+
+It merges:
+- **Aliases:** with `define can_view: viewer`, checking both `viewer` and `can_view` for the same user and object sends one check. Chains like `define can_share: can_view` work too.
+- **Exact repeats:** the same user, relation, object, context and contextual tuples.
+
+```python
+options = {
+    "authorization_model_id": "01GXSA8YR785C4FYS3C0RTG7B1",  # required, here or in the configuration
+    "merge_equivalent_checks": True,
+}
+response = await fga_client.batch_check(ClientBatchCheckRequest(checks=checks), options)
+# response.checks_sent: how many checks this call actually sent
+```
+
+- **Off by default.** With it off, nothing changes. The value must be a bool, or the call fails with a validation error.
+- **Needs an authorization model ID.** Without one, the call fails before anything is sent.
+- **Reads each model once per client.** The SDK reads a model the first time it needs it and caches it. Create the client once and reuse it.
+- **Falls back safely.** If the model can't be read (for example, no read access), the SDK logs a warning and sends the checks unmerged.
+- **Errors are shared.** If a merged check returns an error, every check it stands for gets that error.
+- **Correlation IDs** must match `^[\w\d-]{1,36}$` (ASCII). Missing ones are generated.
+- **Telemetry:** see the `fga-client.batch_check.*` and `fga-client.model_cache.lookup` metrics in [OpenTelemetry](docs/opentelemetry.md).
+
 ##### Client Batch Check
 
 

docs/opentelemetry.md

@@ -21,17 +21,22 @@ If you configure the OpenTelemetry SDK, these metrics will be exported and sent
 
 ### Supported Metrics
 
-| Metric Name                      | Type      | Enabled by Default | Description                                                                       |
-| -------------------------------- | --------- | ------------------ | --------------------------------------------------------------------------------- |
-| `fga-client.request.duration`    | Histogram | Yes                | Total request time for FGA requests, in milliseconds                              |
-| `fga-client.query.duration`      | Histogram | Yes                | Time taken by the FGA server to process and evaluate the request, in milliseconds |
-| `fga-client.credentials.request` | Counter   | Yes                | Total number of new token requests initiated using the Client Credentials flow    |
-| `fga-client.request`             | Counter   | No                 | Total number of requests made to the FGA server                                   |
+| Metric Name                               | Type      | Enabled by Default | Description                                                                                   |
+| ----------------------------------------- | --------- | ------------------ | --------------------------------------------------------------------------------------------- |
+| `fga-client.request.duration`             | Histogram | Yes                | Total request time for FGA requests, in milliseconds                                          |
+| `fga-client.query.duration`               | Histogram | Yes                | Time taken by the FGA server to process and evaluate the request, in milliseconds             |
+| `fga-client.credentials.request`          | Counter   | Yes                | Total number of new token requests initiated using the Client Credentials flow                |
+| `fga-client.request`                      | Counter   | No                 | Total number of requests made to the FGA server                                               |
+| `fga-client.batch_check.checks_requested` | Counter   | Yes                | Checks passed to `batch_check` with `merge_equivalent_checks` on                              |
+| `fga-client.batch_check.checks_sent`      | Counter   | Yes                | Checks actually sent by `batch_check` with `merge_equivalent_checks` on                       |
+| `fga-client.batch_check.merge_fallback`   | Counter   | Yes                | `batch_check` calls sent unmerged because the authorization model could not be read           |
+| `fga-client.model_cache.lookup`           | Counter   | Yes                | Authorization model cache lookups by `merge_equivalent_checks`, with a `hit` or `miss` result |
 
 ### Supported Attributes
 
 | Attribute Name                        | Type   | Enabled by Default | Description                                                                       |
 | ------------------------------------- | ------ | ------------------ | --------------------------------------------------------------------------------- |
+| `fga-client.model_cache.result`       | string | Yes                | Result of an authorization model cache lookup (`hit`/`miss`)                      |
 | `fga-client.request.batch_check_size` | int    | No                 | The total size of the `check` list in a `BatchCheck` call                         |
 | `fga-client.request.client_id`        | string | Yes                | Client ID associated with the request, if any                                     |
 | `fga-client.request.method`           | string | Yes                | FGA method/action that was performed (e.g., Check, ListObjects) in TitleCase      |

openfga_sdk/client/client.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import uuid
 
 from collections.abc import AsyncIterator
@@ -7,6 +8,14 @@
 from openfga_sdk.api.open_fga_api import OpenFgaApi
 from openfga_sdk.api_client import ApiClient
 from openfga_sdk.client.configuration import ClientConfiguration
+from openfga_sdk.client.equivalent_checks import (
+    expand,
+    merge_enabled,
+    resolve,
+    unmerged,
+    validate_correlation_ids,
+)
+from openfga_sdk.client.model_cache import AsyncModelCache
 from openfga_sdk.client.models.assertion import ClientAssertion
 from openfga_sdk.client.models.batch_check_item import (
     ClientBatchCheckItem,
@@ -71,9 +80,13 @@
     WriteAuthorizationModelRequest,
 )
 from openfga_sdk.models.write_request import WriteRequest
+from openfga_sdk.telemetry.attributes import TelemetryAttributes
 from openfga_sdk.validation import is_well_formed_ulid_string
 
 
+logger = logging.getLogger(__name__)
+
+
 def _chuck_array(array, max_size):
     """
     Helper function to chuck array into arrays of max_size
@@ -127,6 +140,16 @@ def options_to_kwargs(
     return kwargs
 
 
+def _model_read_kwargs(
+    options: dict[str, int | str | dict[str, int | str]] | None = None,
+) -> dict[str, int | str | dict[str, int | str]]:
+    """
+    Return the kwargs a BatchCheck call passes on to its authorization model read
+    """
+    kwargs = options_to_kwargs(options)
+    return {key: kwargs[key] for key in ("_headers", "_retry_params") if key in kwargs}
+
+
 def options_to_transaction_info(
     options: dict[str, int | str | dict[str, int | str]] | None = None,
 ):
@@ -172,6 +195,7 @@ def __init__(self, configuration: ClientConfiguration):
         self._client_configuration = configuration
         self._api_client = ApiClient(configuration)
         self._api = OpenFgaApi(self._api_client)
+        self._model_cache = AsyncModelCache()
 
         # Set default headers from configuration
         if configuration.headers:
@@ -185,6 +209,7 @@ async def __aexit__(self, exc_type, exc_value, traceback):
         await self.close()
 
     async def close(self):
+        await self._model_cache.close()
         await self._api.close()
 
     def _get_authorization_model_id(
@@ -780,6 +805,7 @@ async def batch_check(
         :param authorization_model_id(options) - Overrides the authorization model id in the configuration
         :param max_parallel_requests(options) - Max number of requests to issue in parallel. Defaults to 10
         :param max_batch_size(options) - Max number of checks to include in a request. Defaults to 50
+        :param merge_equivalent_checks(options) - Beta. Send checks that are certain to get the same answer only once. Requires an authorization model id
         :param header(options) - Custom headers to send alongside the request
         :param retryParams(options) - Override the retry parameters for this request
         :param retryParams.maxRetry(options) - Override the max number of retries on each API request
@@ -809,6 +835,11 @@ async def batch_check(
             elif isinstance(options["max_batch_size"], int):
                 max_batch_size = options["max_batch_size"]
 
+        if merge_enabled(options):
+            return await self._merged_batch_check(
+                body, options, max_parallel_requests, max_batch_size
+            )
+
         id_to_check: dict[str, ClientBatchCheckItem] = {}
 
         def track_and_transform(checks):
@@ -866,6 +897,116 @@ async def coro(checks):
 
         return ClientBatchCheckResponse(result)
 
+    async def _merged_batch_check(
+        self,
+        body: ClientBatchCheckRequest,
+        options: dict[str, int | str | dict[str, int | str]],
+        max_parallel_requests: int,
+        max_batch_size: int,
+    ) -> ClientBatchCheckResponse:
+        """
+        Run a BatchCheck with merge_equivalent_checks on
+        """
+        # Reading the model and sending the checks must use the same model.
+        authorization_model_id = self._get_authorization_model_id(options)
+        if authorization_model_id is None:
+            raise FgaValidationException(
+                "authorization_model_id is required when merge_equivalent_checks is on"
+            )
+        store_id = self._api_client._get_store_id()
+
+        checks = body.checks
+        for check in checks:
+            if check.correlation_id is None:
+                check.correlation_id = str(uuid.uuid4())
+        validate_correlation_ids(checks)
+
+        if not checks:
+            return ClientBatchCheckResponse([], checks_sent=0)
+
+        telemetry = self._api_client._telemetry.metrics
+        telemetry_configuration = self._api_client.configuration.telemetry
+        attributes = {
+            TelemetryAttributes.fga_client_request_store_id: store_id,
+            TelemetryAttributes.fga_client_request_model_id: authorization_model_id,
+        }
+
+        async def load_model():
+            try:
+                response = await self._api.read_authorization_model(
+                    authorization_model_id, **_model_read_kwargs(options)
+                )
+                if response.authorization_model is None:
+                    raise ValueError("The authorization model response is empty")
+                return response.authorization_model
+            except Exception as err:
+                logger.warning(
+                    "merge_equivalent_checks: couldn't read authorization model %s in store %s (status: %s). "
+                    "Sending checks unmerged.",
+                    authorization_model_id,
+                    store_id,
+                    getattr(err, "status", None),
+                )
+                raise
+
+        lookup = await self._model_cache.get(
+            (store_id, authorization_model_id), load_model
+        )
+        telemetry.modelCacheLookup(
+            attributes={
+                **attributes,
+                TelemetryAttributes.fga_client_model_cache_result: (
+                    "hit" if lookup.hit else "miss"
+                ),
+            },
+            configuration=telemetry_configuration,
+        )
+
+        if lookup.model is not None:
+            groups = resolve(lookup.model, checks)
+        else:
+            groups = unmerged(checks)
+            telemetry.batchCheckMergeFallback(
+                attributes={
+                    **attributes,
+                    TelemetryAttributes.http_response_status_code: getattr(
+                        lookup.error, "status", None
+                    ),
+                },
+                configuration=telemetry_configuration,
+            )
+
+        results = {}
+        sem = asyncio.Semaphore(max_parallel_requests)
+        consistency = self._get_consistency(options)
+
+        async def send(chunk):
+            response = await self._single_batch_check(
+                BatchCheckRequest(
+                    checks=[construct_batch_item(group.check) for group in chunk],
+                    authorization_model_id=authorization_model_id,
+                    consistency=consistency,
+                ),
+                sem,
+                options,
+            )
+            results.update(response.result or {})
+
+        await asyncio.gather(
+            *[send(chunk) for chunk in _chuck_array(groups, max_batch_size)]
+        )
+
+        telemetry.batchCheckChecksRequested(
+            len(checks), attributes=attributes, configuration=telemetry_configuration
+        )
+        telemetry.batchCheckChecksSent(
+            len(groups), attributes=attributes, configuration=telemetry_configuration
+        )
+
+        return ClientBatchCheckResponse(
+            expand(checks, groups, results), checks_sent=len(groups)
+        )
+
     async def expand(
         self,
         body: ClientExpandRequest,