[Cosmos] make response_hook thread safe for the sync client (#33790)

* only missing queries

* pylint

* upsert method typehints

* Update _cosmos_client_connection.py

* had to include List in return type for batch

* Revert "had to include List in return type for batch"

This reverts commit cacf629.

* read all items is also query - still don't have this logic down

* casting batch

* casting again?

* oops

* return readItems to previous state

* remove casting in client_connection Batch

* update __queryFeed to be thread safe

* forgot to save headers

* Update _cosmos_client_connection.py

* typehint for non-used method was missing

* Updated User operations

* make response_hook internal in cosmos_client

* make response_hook internal in database

* make response_hook internal container

* small fixes

* Update database.py

* Update _cosmos_client_connection.py

* fixes from comments

* wrong naming on these samples pls let me squeeze it in <3

* revert changes on cosmos_client methods

* Update cosmos_client.py

---------

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

Author: simorenoh

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

@@ -153,7 +153,6 @@ def read(  # pylint:disable=docstring-missing-param
         :rtype: dict[str, Any]
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         if populate_query_metrics:
             warnings.warn(
                 "the populate_query_metrics flag does not apply to this method and will be removed in the future",
@@ -167,8 +166,6 @@ def read(  # pylint:disable=docstring-missing-param
         self._properties = self.client_connection.ReadContainer(
             collection_link, options=request_options, **kwargs
         )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, self._properties)
         return self._properties
 
     @distributed_trace
@@ -211,7 +208,6 @@ def read_item(  # pylint:disable=docstring-missing-param
         """
         doc_link = self._get_document_link(item)
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
 
         if partition_key is not None:
             request_options["partitionKey"] = self._set_partition_key(partition_key)
@@ -229,10 +225,7 @@ def read_item(  # pylint:disable=docstring-missing-param
             validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
             request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
 
-        result = self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs)
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
+        return self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs)
 
     @distributed_trace
     def read_all_items(  # pylint:disable=docstring-missing-param
@@ -275,8 +268,7 @@ def read_all_items(  # pylint:disable=docstring-missing-param
             response_hook.clear()
 
         items = self.client_connection.ReadItems(
-            collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs
-        )
+            collection_link=self.container_link, feed_options=feed_options, response_hook=response_hook, **kwargs)
         if response_hook:
             response_hook(self.client_connection.last_response_headers, items)
         return items
@@ -494,7 +486,6 @@ def replace_item(  # pylint:disable=docstring-missing-param
         """
         item_link = self._get_document_link(item)
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         request_options["disableAutomaticIdGeneration"] = True
         if populate_query_metrics is not None:
             warnings.warn(
@@ -507,12 +498,9 @@ def replace_item(  # pylint:disable=docstring-missing-param
         if post_trigger_include is not None:
             request_options["postTriggerInclude"] = post_trigger_include
 
-        result = self.client_connection.ReplaceItem(
+        return self.client_connection.ReplaceItem(
             document_link=item_link, new_document=body, options=request_options, **kwargs
         )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
 
     @distributed_trace
     def upsert_item(  # pylint:disable=docstring-missing-param
@@ -546,7 +534,6 @@ def upsert_item(  # pylint:disable=docstring-missing-param
         :rtype: Dict[str, Any]
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         request_options["disableAutomaticIdGeneration"] = True
         if populate_query_metrics is not None:
             warnings.warn(
@@ -559,15 +546,12 @@ def upsert_item(  # pylint:disable=docstring-missing-param
         if post_trigger_include is not None:
             request_options["postTriggerInclude"] = post_trigger_include
 
-        result = self.client_connection.UpsertItem(
+        return self.client_connection.UpsertItem(
             database_or_container_link=self.container_link,
             document=body,
             options=request_options,
             **kwargs
         )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
 
     @distributed_trace
     def create_item(  # pylint:disable=docstring-missing-param
@@ -606,7 +590,6 @@ def create_item(  # pylint:disable=docstring-missing-param
         :rtype: Dict[str, Any]
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
 
         request_options["disableAutomaticIdGeneration"] = not kwargs.pop('enable_automatic_id_generation', False)
         if populate_query_metrics:
@@ -622,12 +605,8 @@ def create_item(  # pylint:disable=docstring-missing-param
         if indexing_directive is not None:
             request_options["indexingDirective"] = indexing_directive
 
-        result = self.client_connection.CreateItem(
-            database_or_container_link=self.container_link, document=body, options=request_options, **kwargs
-        )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
+        return self.client_connection.CreateItem(
+            database_or_container_link=self.container_link, document=body, options=request_options, **kwargs)
 
     @distributed_trace
     def patch_item(
@@ -662,19 +641,15 @@ def patch_item(
         :rtype: dict[str, Any]
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         request_options["disableAutomaticIdGeneration"] = True
         request_options["partitionKey"] = self._set_partition_key(partition_key)
         filter_predicate = kwargs.pop("filter_predicate", None)
         if filter_predicate is not None:
             request_options["filterPredicate"] = filter_predicate
 
         item_link = self._get_document_link(item)
-        result = self.client_connection.PatchItem(
+        return self.client_connection.PatchItem(
             document_link=item_link, operations=patch_operations, options=request_options, **kwargs)
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
 
     @distributed_trace
     def execute_item_batch(
@@ -703,14 +678,10 @@ def execute_item_batch(
         """
         request_options = build_options(kwargs)
         request_options["partitionKey"] = self._set_partition_key(partition_key)
-        response_hook = kwargs.pop('response_hook', None)
         request_options["disableAutomaticIdGeneration"] = True
 
-        result = self.client_connection.Batch(
+        return self.client_connection.Batch(
             collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs)
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
 
     @distributed_trace
     def delete_item(  # pylint:disable=docstring-missing-param
@@ -743,7 +714,6 @@ def delete_item(  # pylint:disable=docstring-missing-param
         :rtype: None
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         if partition_key is not None:
             request_options["partitionKey"] = self._set_partition_key(partition_key)
         if populate_query_metrics is not None:
@@ -759,8 +729,6 @@ def delete_item(  # pylint:disable=docstring-missing-param
 
         document_link = self._get_document_link(item)
         self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs)
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, None)
 
     @distributed_trace
     def read_offer(self, **kwargs: Any) -> Offer:
@@ -827,7 +795,6 @@ def replace_throughput(
             or the throughput properties could not be updated.
         :rtype: ~azure.cosmos.ThroughputProperties
         """
-        response_hook = kwargs.pop('response_hook', None)
         properties = self._get_properties()
         link = properties["_self"]
         query_spec = {
@@ -844,8 +811,6 @@ def replace_throughput(
         data = self.client_connection.ReplaceOffer(
             offer_link=throughput_properties[0]["_self"], offer=throughput_properties[0], **kwargs)
 
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, data)
         return ThroughputProperties(offer_throughput=data["content"]["offerThroughput"], properties=data)
 
     @distributed_trace
@@ -936,16 +901,12 @@ def get_conflict(
         :rtype: Dict[str, Any]
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         if partition_key is not None:
             request_options["partitionKey"] = self._set_partition_key(partition_key)
 
-        result = self.client_connection.ReadConflict(
+        return self.client_connection.ReadConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, result)
-        return result
 
     @distributed_trace
     def delete_conflict(
@@ -968,15 +929,12 @@ def delete_conflict(
         :rtype: None
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         if partition_key is not None:
             request_options["partitionKey"] = self._set_partition_key(partition_key)
 
         self.client_connection.DeleteConflict(
             conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
         )
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, None)
 
     @distributed_trace
     def delete_all_items_by_partition_key(
@@ -1002,11 +960,8 @@ def delete_all_items_by_partition_key(
         :rtype: None
         """
         request_options = build_options(kwargs)
-        response_hook = kwargs.pop('response_hook', None)
         # regardless if partition key is valid we set it as invalid partition keys are set to a default empty value
         request_options["partitionKey"] = self._set_partition_key(partition_key)
 
-        self.client_connection.DeleteAllItemsByPartitionKey(collection_link=self.container_link,
-                                                            options=request_options, **kwargs)
-        if response_hook:
-            response_hook(self.client_connection.last_response_headers, None)
+        self.client_connection.DeleteAllItemsByPartitionKey(
+            collection_link=self.container_link, options=request_options, **kwargs)