Implement max_concurrency optional and fix database method overloads

- Make max_concurrency Optional[int] in read_items (container.py)
- Update default to None to use Python's ThreadPoolExecutor default
- Update _read_items_helper.py to accept Optional[int] for max_concurrency
- Update _cosmos_client_connection.py to default to None
- Fix create_database overloads (sync): remove populate_query_metrics, make offer_throughput keyword-only, add response_hook
- Fix create_database_if_not_exists overloads (sync): same changes
- Fix create_database overloads (async): add response_hook parameter
- Fix create_database_if_not_exists overloads (async): add response_hook parameter
- Update implementations to handle response_hook and remove positional offer_throughput/populate_query_metrics

Co-authored-by: bambriz <[email protected]>

Author: Copilot

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -1081,7 +1081,7 @@ def read_items(
             raise ValueError("Could not find partition key definition for collection.")
 
         # Extract and remove max_concurrency from kwargs
-        max_concurrency = kwargs.pop('max_concurrency', 10)
+        max_concurrency = kwargs.pop('max_concurrency', None)
 
         helper = ReadItemsHelperSync(
             client=self,

sdk/cosmos/azure-cosmos/azure/cosmos/_read_items_helper.py

@@ -48,7 +48,7 @@ def __init__(
             partition_key_definition: Dict[str, Any],
             *,
             executor: Optional[ThreadPoolExecutor] = None,
-            max_concurrency: int = 10,
+            max_concurrency: Optional[int] = None,
             **kwargs: Any
     ):
         self.client = client
@@ -81,6 +81,7 @@ def read_items(self) -> CosmosList:
         if self.executor is not None:
             return self._execute_with_executor(self.executor, query_chunks)
 
+        # Create a new executor; if max_concurrency is None, use ThreadPoolExecutor's default
         with ThreadPoolExecutor(max_workers=self.max_concurrency) as executor:
             return self._execute_with_executor(executor, query_chunks)
 

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py

@@ -263,6 +263,7 @@ async def create_database(
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
+        response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
         return_properties: Literal[False] = False,
         **kwargs: Any
@@ -302,6 +303,7 @@ async def create_database(
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
+        response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
         return_properties: Literal[True],
         **kwargs: Any
@@ -392,11 +394,14 @@ async def create_database( # pylint:disable=docstring-should-be-keyword
 
         offer_throughput = kwargs.pop("offer_throughput", None)
         return_properties = kwargs.pop("return_properties", False)
+        response_hook = kwargs.pop("response_hook", None)
 
         request_options = _build_options(kwargs)
         _set_throughput_options(offer=offer_throughput, request_options=request_options)
 
         result = await self.client_connection.CreateDatabase(database={"id": id}, options=request_options, **kwargs)
+        if response_hook:
+            response_hook(self.client_connection.last_response_headers)
         if not return_properties:
             return DatabaseProxy(self.client_connection, id=result["id"], properties=result)
         return  DatabaseProxy(self.client_connection, id=result["id"], properties=result), result
@@ -408,6 +413,7 @@ async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
+        response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
         return_properties: Literal[False] = False,
         **kwargs: Any
@@ -443,6 +449,7 @@ async def create_database_if_not_exists(  # pylint: disable=redefined-builtin
         *,
         offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
+        response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
         return_properties: Literal[True],
         **kwargs: Any
@@ -527,10 +534,13 @@ async def create_database_if_not_exists( # pylint:disable=docstring-should-be-ke
 
         offer_throughput = kwargs.pop("offer_throughput", None)
         return_properties = kwargs.pop("return_properties", False)
+        response_hook = kwargs.pop("response_hook", None)
 
         try:
             database_proxy = self.get_database_client(id)
             result = await database_proxy.read(**kwargs)
+            if response_hook:
+                response_hook(self.client_connection.last_response_headers)
             if not return_properties:
                 return database_proxy
             return database_proxy, result
@@ -539,6 +549,7 @@ async def create_database_if_not_exists( # pylint:disable=docstring-should-be-ke
                 id,
                 offer_throughput=offer_throughput,
                 return_properties=return_properties,
+                response_hook=response_hook,
                 **kwargs
             )
 

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -293,7 +293,7 @@ def read_items(
             items: Sequence[Tuple[str, _PartitionKeyType]],
             *,
             executor: Optional[ThreadPoolExecutor] = None,
-            max_concurrency: int = 10,
+            max_concurrency: Optional[int] = None,
             consistency_level: Optional[str] = None,
             session_token: Optional[str] = None,
             initial_headers: Optional[Dict[str, str]] = None,
@@ -312,7 +312,8 @@ def read_items(
         :keyword executor: Optional ThreadPoolExecutor for handling concurrent operations.
                       If not provided, a new executor will be created as needed.
         :keyword int max_concurrency: The maximum number of concurrent operations for the
-                      items request. This value is ignored if an executor is provided. Defaults to 10.
+                      items request. This value is ignored if an executor is provided. If not specified,
+                      defaults to Python's ThreadPoolExecutor default of min(32, (os.cpu_count() or 1) + 4).
         :keyword str consistency_level: The consistency level to use for the request.
         :keyword str session_token: Token for use with Session consistency.
         :keyword dict[str, str] initial_headers: Initial headers to be sent as part of the request.

sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py

@@ -264,13 +264,13 @@ def from_connection_string(
             **kwargs
         )
 
+    @overload
     @overload
     def create_database(  # pylint:disable=docstring-missing-param
         self,
         id: str,
-        populate_query_metrics: Optional[bool] = None,
-        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         *,
+        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
         response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
@@ -307,9 +307,8 @@ def create_database(  # pylint:disable=docstring-missing-param
     def create_database(  # pylint:disable=docstring-missing-param
         self,
         id: str,
-        populate_query_metrics: Optional[bool] = None,
-        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         *,
+        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
         response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
@@ -394,10 +393,10 @@ def create_database(  # pylint:disable=docstring-missing-param, docstring-should
                 UserWarning)
 
         id = args[0] if args else kwargs.pop("id")
-        populate_query_metrics = args[1] if len(args) > 1 else kwargs.pop("populate_query_metrics", None)
-        offer_throughput = args[2] if len(args) > 2 else kwargs.pop("offer_throughput", None)
-        if len(args) > 3:
-            raise TypeError(f"Unexpected positional arguments: {args[3:]}")
+        if len(args) > 1:
+            raise TypeError(f"Unexpected positional arguments: {args[1:]}")
+        populate_query_metrics = kwargs.pop("populate_query_metrics", None)
+        offer_throughput = kwargs.pop("offer_throughput", None)
         return_properties = kwargs.pop("return_properties", False)
         response_hook = kwargs.pop("response_hook", None)
 
@@ -421,10 +420,10 @@ def create_database(  # pylint:disable=docstring-missing-param, docstring-should
     def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
         self,
         id: str,
-        populate_query_metrics: Optional[bool] = None,
-        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         *,
+        offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
         initial_headers: Optional[Dict[str, str]] = None,
+        response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
         throughput_bucket: Optional[int] = None,
         return_properties: Literal[False] = False,
         **kwargs: Any
@@ -455,10 +454,10 @@ def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
     def create_database_if_not_exists(  # pylint:disable=docstring-missing-param
             self,
             id: str,
-            populate_query_metrics: Optional[bool] = None,
-            offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
             *,
+            offer_throughput: Optional[Union[int, ThroughputProperties]] = None,
             initial_headers: Optional[Dict[str, str]] = None,
+            response_hook: Optional[Callable[[Mapping[str, Any]], None]] = None,
             throughput_bucket: Optional[int] = None,
             return_properties: Literal[True],
             **kwargs: Any
@@ -534,26 +533,29 @@ def create_database_if_not_exists(  # pylint:disable=docstring-missing-param, do
                 UserWarning)
 
         id = args[0] if args else kwargs.pop("id")
-        populate_query_metrics = args[1] if len(args) > 1 else kwargs.pop("populate_query_metrics", None)
-        offer_throughput = args[2] if len(args) > 2 else kwargs.pop("offer_throughput", None)
-        if len(args) > 3:
-            raise TypeError(f"Unexpected positional arguments: {args[3:]}")
+        if len(args) > 1:
+            raise TypeError(f"Unexpected positional arguments: {args[1:]}")
+        populate_query_metrics = kwargs.pop("populate_query_metrics", None)
+        offer_throughput = kwargs.pop("offer_throughput", None)
         return_properties = kwargs.pop("return_properties", False)
+        response_hook = kwargs.pop("response_hook", None)
         try:
             database_proxy = self.get_database_client(id)
             result = database_proxy.read(
                 populate_query_metrics=populate_query_metrics,
                 **kwargs
             )
+            if response_hook:
+                response_hook(self.client_connection.last_response_headers)
             if not return_properties:
                 return database_proxy
             return database_proxy, result
         except CosmosResourceNotFoundError:
             return self.create_database(
                 id,
-                populate_query_metrics=populate_query_metrics,
                 offer_throughput=offer_throughput,
                 return_properties=return_properties,
+                response_hook=response_hook,
                 **kwargs
             )