Merge branch 'master' into CDF-28136-fix-workflow-integration-test-teardown

Author: vincent-cognite

cognite/client/_api/data_modeling/records.py

@@ -2,11 +2,10 @@
 
 import asyncio
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from cognite.client._api_client import APIClient
 from cognite.client.data_classes.data_modeling.records import RecordId, RecordIdSequence, RecordWrite
-from cognite.client.utils._concurrency import RecordsConcurrencyOperation
 from cognite.client.utils._experimental import FeaturePreviewWarning
 from cognite.client.utils._url import interpolate_and_url_encode
 
@@ -22,16 +21,11 @@ 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:
+    def _get_semaphore(self, operation: Any) -> 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
+            operation, project=self._cognite_client.config.project
         )
 
     def _records_url(self, stream_id: str, suffix: str = "") -> str:

cognite/client/_sync_api/data_modeling/records.py

@@ -7,7 +7,6 @@
 from abc import ABC, abstractmethod
 from collections import UserList
 from collections.abc import Callable, Coroutine
-from enum import Enum
 from typing import (
     Any,
     Literal,
@@ -222,45 +221,47 @@ def __repr__(self) -> str:
         )
 
 
-class RecordsConcurrencyOperation(Enum):
-    WRITE = "write"
-
-
 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.
 
     Args:
         concurrency_settings (ConcurrencySettings): Reference to the parent settings object.
+        read (int): Maximum concurrent read requests (list, retrieve, sync).
         write (int): Maximum concurrent write requests (ingest, delete).
     """
 
     def __init__(
         self,
         concurrency_settings: ConcurrencySettings,
+        read: int,
         write: int,
     ) -> None:
-        super().__init__(concurrency_settings, "records", read=0, write=write, delete=0)
+        super().__init__(concurrency_settings, "records", read=read, write=write, delete=0)
 
-    def _semaphore_factory(self, operation: RecordsConcurrencyOperation, project: str) -> asyncio.BoundedSemaphore:
-        key = (operation.value, project, asyncio.get_running_loop())
+    def _semaphore_factory(
+        self, operation: Literal["read", "write", "delete"], project: str
+    ) -> asyncio.BoundedSemaphore:
+        key = (operation, project, asyncio.get_running_loop())
         if key in self._semaphore_cache:
             return self._semaphore_cache[key]
 
         from cognite.client import global_config
 
         global_config.concurrency_settings._freeze()
         match operation:
-            case RecordsConcurrencyOperation.WRITE:
+            case "read":
+                sem = asyncio.BoundedSemaphore(self._read)
+            case "write" | "delete":
                 sem = asyncio.BoundedSemaphore(self._write)
             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](read={self._read}, write={self._write})"
 
 
 class FileConcurrencyConfig(ConcurrencyConfig):
@@ -425,7 +426,7 @@ 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, read=4, write=20)
 
     @functools.cached_property
     def _all_concurrency_configs(self) -> list[ConcurrencyConfig]: