feat(records): stream-type-aware concurrency limits

Model the Records API's hierarchical rate-limit budgets for mutable and
immutable streams. Write semaphore (20) is shared across both stream
types. Query, retrieve, and aggregate each have separate mutable and
immutable semaphores matching the API's documented limits.

Retrieve and aggregate have dedicated budgets checked before the shared
query budget (both must pass). The semaphore helpers on RecordsAPI are
wired up for all three tiers, ready for use by endpoint implementations.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: andersfylling

cognite/client/_api/data_modeling/records.py

@@ -14,6 +14,8 @@
     from cognite.client import AsyncCogniteClient
     from cognite.client.config import ClientConfig
 
+StreamType = Literal["mutable", "immutable"]
+
 
 class RecordsAPI(APIClient):
     def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: AsyncCogniteClient) -> None:
@@ -22,16 +24,47 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
             api_maturity="General Availability", sdk_maturity="alpha", feature_name="Records"
         )
 
-    _OPERATION_TO_RATE_LIMIT: ClassVar[dict[str, RecordsConcurrencyOperation]] = {
-        "write": RecordsConcurrencyOperation.WRITE,
-        "delete": RecordsConcurrencyOperation.WRITE,
-    }
-
     def _get_semaphore(self, operation: Literal["write", "delete"]) -> asyncio.BoundedSemaphore:
         from cognite.client import global_config
 
         return global_config.concurrency_settings.records._semaphore_factory(
-            self._OPERATION_TO_RATE_LIMIT[operation], project=self._cognite_client.config.project
+            RecordsConcurrencyOperation.WRITE, project=self._cognite_client.config.project
+        )
+
+    def _get_query_semaphore(self, stream_type: StreamType) -> asyncio.BoundedSemaphore:
+        from cognite.client import global_config
+
+        op = (
+            RecordsConcurrencyOperation.QUERY_MUTABLE
+            if stream_type == "mutable"
+            else RecordsConcurrencyOperation.QUERY_IMMUTABLE
+        )
+        return global_config.concurrency_settings.records._semaphore_factory(
+            op, project=self._cognite_client.config.project
+        )
+
+    def _get_retrieve_semaphore(self, stream_type: StreamType) -> asyncio.BoundedSemaphore:
+        from cognite.client import global_config
+
+        op = (
+            RecordsConcurrencyOperation.RETRIEVE_MUTABLE
+            if stream_type == "mutable"
+            else RecordsConcurrencyOperation.RETRIEVE_IMMUTABLE
+        )
+        return global_config.concurrency_settings.records._semaphore_factory(
+            op, project=self._cognite_client.config.project
+        )
+
+    def _get_aggregate_semaphore(self, stream_type: StreamType) -> asyncio.BoundedSemaphore:
+        from cognite.client import global_config
+
+        op = (
+            RecordsConcurrencyOperation.AGGREGATE_MUTABLE
+            if stream_type == "mutable"
+            else RecordsConcurrencyOperation.AGGREGATE_IMMUTABLE
+        )
+        return global_config.concurrency_settings.records._semaphore_factory(
+            op, project=self._cognite_client.config.project
         )
 
     def _records_url(self, stream_id: str, suffix: str = "") -> str:

cognite/client/_sync_api/data_modeling/records.py

@@ -224,24 +224,114 @@ def __repr__(self) -> str:
 
 class RecordsConcurrencyOperation(Enum):
     WRITE = "write"
+    QUERY_MUTABLE = "query_mutable"
+    QUERY_IMMUTABLE = "query_immutable"
+    RETRIEVE_MUTABLE = "retrieve_mutable"
+    RETRIEVE_IMMUTABLE = "retrieve_immutable"
+    AGGREGATE_MUTABLE = "aggregate_mutable"
+    AGGREGATE_IMMUTABLE = "aggregate_immutable"
 
 
 class RecordsGlobalConcurrencyConfig(ConcurrencyConfig):
     """
-    Global concurrency settings for the Records API. Named "global" to distinguish from
-    future per-endpoint rate limits that may be added later.
+    Global concurrency settings for the Records API.
+
+    The Records API has separate rate-limit budgets for reads and writes, and read budgets
+    differ between mutable and immutable streams. Read budgets are hierarchical: the
+    retrieve and aggregate endpoints each have a dedicated budget that is checked *before*
+    the shared query budget (both must pass).
+
+    - **write**: Shared across ingest, upsert, and delete (same limit for both stream types).
+    - **query_mutable / query_immutable**: Shared read budget consumed by all query endpoints
+      (list/filter, sync, retrieve, aggregate).
+    - **retrieve_mutable / retrieve_immutable**: Dedicated budget for retrieve, checked
+      *in addition to* the shared query budget.
+    - **aggregate_mutable / aggregate_immutable**: Dedicated budget for aggregate, checked
+      *in addition to* the shared query budget.
 
     Args:
         concurrency_settings (ConcurrencySettings): Reference to the parent settings object.
-        write (int): Maximum concurrent write requests (ingest, delete).
+        write (int): Maximum concurrent write requests (ingest, upsert, delete).
+        query_mutable (int): Maximum concurrent query requests against mutable streams.
+        query_immutable (int): Maximum concurrent query requests against immutable streams.
+        retrieve_mutable (int): Dedicated retrieve concurrency for mutable streams.
+        retrieve_immutable (int): Dedicated retrieve concurrency for immutable streams.
+        aggregate_mutable (int): Dedicated aggregate concurrency for mutable streams.
+        aggregate_immutable (int): Dedicated aggregate concurrency for immutable streams.
     """
 
     def __init__(
         self,
         concurrency_settings: ConcurrencySettings,
         write: int,
+        query_mutable: int,
+        query_immutable: int,
+        retrieve_mutable: int,
+        retrieve_immutable: int,
+        aggregate_mutable: int,
+        aggregate_immutable: int,
     ) -> None:
         super().__init__(concurrency_settings, "records", read=0, write=write, delete=0)
+        self._query_mutable = query_mutable
+        self._query_immutable = query_immutable
+        self._retrieve_mutable = retrieve_mutable
+        self._retrieve_immutable = retrieve_immutable
+        self._aggregate_mutable = aggregate_mutable
+        self._aggregate_immutable = aggregate_immutable
+
+    @property
+    def query_mutable(self) -> int:
+        return self._query_mutable
+
+    @query_mutable.setter
+    def query_mutable(self, value: int) -> None:
+        self._check_frozen("query_mutable")
+        self._query_mutable = value
+
+    @property
+    def query_immutable(self) -> int:
+        return self._query_immutable
+
+    @query_immutable.setter
+    def query_immutable(self, value: int) -> None:
+        self._check_frozen("query_immutable")
+        self._query_immutable = value
+
+    @property
+    def retrieve_mutable(self) -> int:
+        return self._retrieve_mutable
+
+    @retrieve_mutable.setter
+    def retrieve_mutable(self, value: int) -> None:
+        self._check_frozen("retrieve_mutable")
+        self._retrieve_mutable = value
+
+    @property
+    def retrieve_immutable(self) -> int:
+        return self._retrieve_immutable
+
+    @retrieve_immutable.setter
+    def retrieve_immutable(self, value: int) -> None:
+        self._check_frozen("retrieve_immutable")
+        self._retrieve_immutable = value
+
+    @property
+    def aggregate_mutable(self) -> int:
+        return self._aggregate_mutable
+
+    @aggregate_mutable.setter
+    def aggregate_mutable(self, value: int) -> None:
+        self._check_frozen("aggregate_mutable")
+        self._aggregate_mutable = value
+
+    @property
+    def aggregate_immutable(self) -> int:
+        return self._aggregate_immutable
+
+    @aggregate_immutable.setter
+    def aggregate_immutable(self, value: int) -> None:
+        self._check_frozen("aggregate_immutable")
+        self._aggregate_immutable = value
 
     def _semaphore_factory(self, operation: RecordsConcurrencyOperation, project: str) -> asyncio.BoundedSemaphore:
         key = (operation.value, project, asyncio.get_running_loop())
@@ -254,13 +344,31 @@ def _semaphore_factory(self, operation: RecordsConcurrencyOperation, project: st
         match operation:
             case RecordsConcurrencyOperation.WRITE:
                 sem = asyncio.BoundedSemaphore(self._write)
+            case RecordsConcurrencyOperation.QUERY_MUTABLE:
+                sem = asyncio.BoundedSemaphore(self._query_mutable)
+            case RecordsConcurrencyOperation.QUERY_IMMUTABLE:
+                sem = asyncio.BoundedSemaphore(self._query_immutable)
+            case RecordsConcurrencyOperation.RETRIEVE_MUTABLE:
+                sem = asyncio.BoundedSemaphore(self._retrieve_mutable)
+            case RecordsConcurrencyOperation.RETRIEVE_IMMUTABLE:
+                sem = asyncio.BoundedSemaphore(self._retrieve_immutable)
+            case RecordsConcurrencyOperation.AGGREGATE_MUTABLE:
+                sem = asyncio.BoundedSemaphore(self._aggregate_mutable)
+            case RecordsConcurrencyOperation.AGGREGATE_IMMUTABLE:
+                sem = asyncio.BoundedSemaphore(self._aggregate_immutable)
             case _:
                 assert_never(operation)
         self._semaphore_cache[key] = sem
         return sem
 
     def __repr__(self) -> str:
-        return f"Concurrency[records](write={self._write})"
+        return (
+            f"Concurrency[records]("
+            f"write={self._write}, "
+            f"query_mutable={self._query_mutable}, query_immutable={self._query_immutable}, "
+            f"retrieve_mutable={self._retrieve_mutable}, retrieve_immutable={self._retrieve_immutable}, "
+            f"aggregate_mutable={self._aggregate_mutable}, aggregate_immutable={self._aggregate_immutable})"
+        )
 
 
 class FileConcurrencyConfig(ConcurrencyConfig):
@@ -425,7 +533,16 @@ def __init__(self) -> None:
             write_schema=1,
         )
         self._files = FileConcurrencyConfig(self, read=4, write=2, upload=5, download=5, delete=2, open_files=15)
-        self._records = RecordsGlobalConcurrencyConfig(self, write=20)
+        self._records = RecordsGlobalConcurrencyConfig(
+            self,
+            write=20,
+            query_mutable=30,
+            query_immutable=10,
+            retrieve_mutable=20,
+            retrieve_immutable=10,
+            aggregate_mutable=10,
+            aggregate_immutable=5,
+        )
 
     @functools.cached_property
     def _all_concurrency_configs(self) -> list[ConcurrencyConfig]:

cognite/client/utils/_concurrency.py

@@ -14,6 +14,7 @@
     ConcurrencyConfig,
     ConcurrencySettings,
     EventLoopThreadExecutor,
+    RecordsConcurrencyOperation,
     _get_event_loop_executor,
     execute_async_tasks,
 )
@@ -136,6 +137,95 @@ async def test_invalid_operation_hits_assert_never(self) -> None:
             self.cs.general._semaphore_factory("totally_invalid", "proj-a")  # type: ignore[arg-type]
 
 
+class TestRecordsConcurrencyConfig:
+    def test_defaults(self) -> None:
+        cs = ConcurrencySettings()
+        assert cs.records.write == 20
+        assert cs.records.query_mutable == 30
+        assert cs.records.query_immutable == 10
+        assert cs.records.retrieve_mutable == 20
+        assert cs.records.retrieve_immutable == 10
+        assert cs.records.aggregate_mutable == 10
+        assert cs.records.aggregate_immutable == 5
+
+    def test_setters_work_before_freeze(self) -> None:
+        cs = ConcurrencySettings()
+        cs.records.write = 10
+        cs.records.query_mutable = 15
+        cs.records.query_immutable = 5
+        cs.records.retrieve_mutable = 12
+        cs.records.retrieve_immutable = 6
+        cs.records.aggregate_mutable = 8
+        cs.records.aggregate_immutable = 3
+        assert cs.records.write == 10
+        assert cs.records.query_mutable == 15
+        assert cs.records.query_immutable == 5
+        assert cs.records.retrieve_mutable == 12
+        assert cs.records.retrieve_immutable == 6
+        assert cs.records.aggregate_mutable == 8
+        assert cs.records.aggregate_immutable == 3
+
+    @pytest.mark.parametrize(
+        "attr",
+        ["write", "query_mutable", "query_immutable", "retrieve_mutable", "retrieve_immutable",
+         "aggregate_mutable", "aggregate_immutable"],
+    )
+    def test_setter_raises_after_freeze(self, attr: str) -> None:
+        cs = ConcurrencySettings()
+        cs._freeze()
+        with pytest.raises(RuntimeError, match="Cannot modify"):
+            setattr(cs.records, attr, 1)
+
+    def test_repr(self) -> None:
+        cs = ConcurrencySettings()
+        r = repr(cs.records)
+        assert "write=20" in r
+        assert "query_mutable=30" in r
+        assert "query_immutable=10" in r
+        assert "retrieve_mutable=20" in r
+        assert "retrieve_immutable=10" in r
+        assert "aggregate_mutable=10" in r
+        assert "aggregate_immutable=5" in r
+
+
+@pytest.mark.usefixtures("fresh_unfrozen_global_concurrency")
+class TestRecordsSemaphoreFactory:
+    cs: ClassVar[ConcurrencySettings] = global_config.concurrency_settings
+
+    @pytest.mark.parametrize(
+        "operation, expected_value",
+        [
+            (RecordsConcurrencyOperation.WRITE, 20),
+            (RecordsConcurrencyOperation.QUERY_MUTABLE, 30),
+            (RecordsConcurrencyOperation.QUERY_IMMUTABLE, 10),
+            (RecordsConcurrencyOperation.RETRIEVE_MUTABLE, 20),
+            (RecordsConcurrencyOperation.RETRIEVE_IMMUTABLE, 10),
+            (RecordsConcurrencyOperation.AGGREGATE_MUTABLE, 10),
+            (RecordsConcurrencyOperation.AGGREGATE_IMMUTABLE, 5),
+        ],
+    )
+    async def test_semaphore_values(self, operation: RecordsConcurrencyOperation, expected_value: int) -> None:
+        sem = self.cs.records._semaphore_factory(operation, "proj-a")
+        assert sem._value == expected_value
+
+    async def test_all_operations_produce_distinct_semaphores(self) -> None:
+        sems = {
+            op: self.cs.records._semaphore_factory(op, "proj-a")
+            for op in RecordsConcurrencyOperation
+        }
+        assert len(set(id(s) for s in sems.values())) == len(RecordsConcurrencyOperation)
+
+    async def test_cache_hit(self) -> None:
+        sem1 = self.cs.records._semaphore_factory(RecordsConcurrencyOperation.QUERY_MUTABLE, "proj-a")
+        sem2 = self.cs.records._semaphore_factory(RecordsConcurrencyOperation.QUERY_MUTABLE, "proj-a")
+        assert sem1 is sem2
+
+    async def test_different_project_different_semaphore(self) -> None:
+        sem_a = self.cs.records._semaphore_factory(RecordsConcurrencyOperation.QUERY_MUTABLE, "proj-a")
+        sem_b = self.cs.records._semaphore_factory(RecordsConcurrencyOperation.QUERY_MUTABLE, "proj-b")
+        assert sem_a is not sem_b
+
+
 async def i_dont_like_5(i: int) -> int:
     if i < 5:
         return i