[Cosmos] fix response_hook not being passed through for complex queries (#40696)

* pass response_hook through aggregate query pipelines

* Update CHANGELOG.md

* add tests, remove comments, pylint

* leftover print statements

* add response hook to global statistics queries

* fix tests

Author: simorenoh

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -12,6 +12,7 @@
 * Fixed bug where change feed requests would not respect the partition key filter. See [PR 40677](https://github.com/Azure/azure-sdk-for-python/pull/40677).
 * Fixed how the environment variables in the sdk are parsed. See [PR 40303](https://github.com/Azure/azure-sdk-for-python/pull/40303).
 * Fixed health check to check the first write region when it is not specified in the preferred regions. See [PR 40588](https://github.com/Azure/azure-sdk-for-python/pull/40588).
+* Fixed `response_hook` not getting called for aggregate queries. See [PR 40696](https://github.com/Azure/azure-sdk-for-python/pull/40696).
 * Fixed bug where writes were being retried for 5xx status codes for patch and replace. See [PR 40672](https://github.com/Azure/azure-sdk-for-python/pull/40672).
 
 #### Other Changes

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

@@ -1112,7 +1112,8 @@ def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], CaseInse
             options,
             fetch_function=fetch_fn,
             collection_link=database_or_container_link,
-            page_iterator_class=query_iterable.QueryIterable
+            page_iterator_class=query_iterable.QueryIterable,
+            response_hook=response_hook
         )
 
     def QueryItemsChangeFeed(

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/document_producer.py

@@ -40,7 +40,8 @@ class _DocumentProducer(object):
     result of each.
     """
 
-    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp, options):
+    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp, options,
+                 response_hook):
         """
         Constructor
         """
@@ -60,7 +61,8 @@ def __init__(self, partition_key_target_range, client, collection_link, query, d
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
 
         async def fetch_fn(options):
-            return await self._client.QueryFeed(path, collection_id, query, options, partition_key_target_range["id"])
+            return await self._client.QueryFeed(path, collection_id, query, options, partition_key_target_range["id"],
+                                                response_hook=response_hook)
 
         self._ex_context = _DefaultQueryExecutionContext(client, self._options, fetch_fn)
 

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/execution_dispatcher.py

@@ -47,7 +47,7 @@ class _ProxyQueryExecutionContext(_QueryExecutionContextBase):  # pylint: disabl
     to _MultiExecutionContextAggregator
     """
 
-    def __init__(self, client, resource_link, query, options, fetch_function):
+    def __init__(self, client, resource_link, query, options, fetch_function, response_hook):
         """
         Constructor
         """
@@ -57,6 +57,7 @@ def __init__(self, client, resource_link, query, options, fetch_function):
         self._resource_link = resource_link
         self._query = query
         self._fetch_function = fetch_function
+        self._response_hook = response_hook
 
     async def __anext__(self):
         """Returns the next query result.
@@ -127,7 +128,8 @@ async def _create_pipelined_execution_context(self, query_execution_info):
                                                                                         self._resource_link,
                                                                                         self._query,
                                                                                         self._options,
-                                                                                        query_execution_info)
+                                                                                        query_execution_info,
+                                                                                        self._response_hook)
             await execution_context_aggregator._configure_partition_ranges()
         elif query_execution_info.has_hybrid_search_query_info():
             hybrid_search_query_info = query_execution_info._query_execution_info['hybridSearchQueryInfo']
@@ -137,14 +139,16 @@ async def _create_pipelined_execution_context(self, query_execution_info):
                                                                         self._resource_link,
                                                                         self._options,
                                                                         query_execution_info,
-                                                                        hybrid_search_query_info)
+                                                                        hybrid_search_query_info,
+                                                                        self._response_hook)
             await execution_context_aggregator._run_hybrid_search()
         else:
             execution_context_aggregator = multi_execution_aggregator._MultiExecutionContextAggregator(self._client,
                                                                                                    self._resource_link,
                                                                                                    self._query,
                                                                                                    self._options,
-                                                                                                   query_execution_info)
+                                                                                                   query_execution_info,
+                                                                                                   self._response_hook)
             await execution_context_aggregator._configure_partition_ranges()
         return _PipelineExecutionContext(self._client, self._options, execution_context_aggregator,
                                          query_execution_info)

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py

@@ -45,7 +45,7 @@ class _HybridSearchContextAggregator(_QueryExecutionContextBase):
     """
 
     def __init__(self, client, resource_link, options, partitioned_query_execution_info,
-                 hybrid_search_query_info):
+                 hybrid_search_query_info, response_hook):
         super(_HybridSearchContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -57,6 +57,7 @@ def __init__(self, client, resource_link, options, partitioned_query_execution_i
         self._final_results = []
         self._aggregated_global_statistics = None
         self._document_producer_comparator = None
+        self._response_hook = response_hook
 
     async def _run_hybrid_search(self):
         # Check if we need to run global statistics queries, and if so do for every partition in the container
@@ -75,6 +76,7 @@ async def _run_hybrid_search(self):
                         global_statistics_query,
                         self._document_producer_comparator,
                         self._options,
+                        self._response_hook
                     )
                 )
 
@@ -117,6 +119,7 @@ async def _run_hybrid_search(self):
                         rewritten_query['rewrittenQuery'],
                         self._document_producer_comparator,
                         self._options,
+                        self._response_hook
                     )
                 )
         # verify all document producers have items/ no splits
@@ -222,6 +225,7 @@ async def _repair_document_producer(self, query, target_all_ranges=False):
                     query,
                     self._document_producer_comparator,
                     self._options,
+                    self._response_hook
                 )
             )
 

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/multi_execution_aggregator.py

@@ -62,7 +62,7 @@ def peek(self):
         def size(self):
             return len(self._heap)
 
-    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info):
+    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info, response_hook):
         super(_MultiExecutionContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -72,6 +72,7 @@ def __init__(self, client, resource_link, query, options, partitioned_query_ex_i
         self._query = query
         self._partitioned_query_ex_info = partitioned_query_ex_info
         self._sort_orders = partitioned_query_ex_info.get_order_by()
+        self._response_hook = response_hook
 
         if self._sort_orders:
             self._document_producer_comparator = document_producer._OrderByDocumentProducerComparator(self._sort_orders)
@@ -154,6 +155,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
             query,
             self._document_producer_comparator,
             self._options,
+            self._response_hook
         )
 
     async def _get_target_partition_key_range(self):

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/non_streaming_order_by_aggregator.py

@@ -22,7 +22,7 @@ class _NonStreamingOrderByContextAggregator(_QueryExecutionContextBase):
     by the user.
     """
 
-    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info):
+    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info, response_hook):
         super(_NonStreamingOrderByContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -35,6 +35,7 @@ def __init__(self, client, resource_link, query, options, partitioned_query_ex_i
         self._orderByPQ = _MultiExecutionContextAggregator.PriorityQueue()
         self._doc_producers = []
         self._document_producer_comparator = document_producer._NonStreamingOrderByComparator(self._sort_orders)
+        self._response_hook = response_hook
 
 
     async def __anext__(self):
@@ -99,6 +100,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
             query,
             self._document_producer_comparator,
             self._options,
+            self._response_hook
         )
 
     async def _get_target_partition_key_range(self):

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/document_producer.py

@@ -39,7 +39,8 @@ class _DocumentProducer(object):
     result of each.
     """
 
-    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp, options):
+    def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp, options,
+                 response_hook):
         """
         Constructor
         """
@@ -59,7 +60,8 @@ def __init__(self, partition_key_target_range, client, collection_link, query, d
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
 
         def fetch_fn(options):
-            return self._client.QueryFeed(path, collection_id, query, options, partition_key_target_range["id"])
+            return self._client.QueryFeed(path, collection_id, query, options, partition_key_target_range["id"],
+                                          response_hook=response_hook)
 
         self._ex_context = _DefaultQueryExecutionContext(client, self._options, fetch_fn)
 

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/execution_dispatcher.py

@@ -77,7 +77,7 @@ class _ProxyQueryExecutionContext(_QueryExecutionContextBase):  # pylint: disabl
     to _MultiExecutionContextAggregator
     """
 
-    def __init__(self, client, resource_link, query, options, fetch_function):
+    def __init__(self, client, resource_link, query, options, fetch_function, response_hook):
         """
         Constructor
         """
@@ -87,6 +87,7 @@ def __init__(self, client, resource_link, query, options, fetch_function):
         self._resource_link = resource_link
         self._query = query
         self._fetch_function = fetch_function
+        self._response_hook = response_hook
 
     def __next__(self):
         """Returns the next query result.
@@ -160,7 +161,8 @@ def _create_pipelined_execution_context(self, query_execution_info):
                                                                                         self._resource_link,
                                                                                         self._query,
                                                                                         self._options,
-                                                                                        query_execution_info)
+                                                                                        query_execution_info,
+                                                                                        self._response_hook)
         elif query_execution_info.has_hybrid_search_query_info():
             hybrid_search_query_info = query_execution_info._query_execution_info['hybridSearchQueryInfo']
             _verify_valid_hybrid_search_query(hybrid_search_query_info)
@@ -169,15 +171,17 @@ def _create_pipelined_execution_context(self, query_execution_info):
                                                                         self._resource_link,
                                                                         self._options,
                                                                         query_execution_info,
-                                                                        hybrid_search_query_info)
+                                                                        hybrid_search_query_info,
+                                                                        self._response_hook)
             execution_context_aggregator._run_hybrid_search()
         else:
             execution_context_aggregator = \
                 multi_execution_aggregator._MultiExecutionContextAggregator(self._client,
                                                                             self._resource_link,
                                                                             self._query,
                                                                             self._options,
-                                                                            query_execution_info)
+                                                                            query_execution_info,
+                                                                            self._response_hook)
         return _PipelineExecutionContext(self._client, self._options, execution_context_aggregator,
                                          query_execution_info)
 

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py

@@ -150,7 +150,7 @@ class _HybridSearchContextAggregator(_QueryExecutionContextBase):
     """
 
     def __init__(self, client, resource_link, options,
-                 partitioned_query_execution_info, hybrid_search_query_info):
+                 partitioned_query_execution_info, hybrid_search_query_info, response_hook):
         super(_HybridSearchContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -162,6 +162,7 @@ def __init__(self, client, resource_link, options,
         self._final_results = []
         self._aggregated_global_statistics = None
         self._document_producer_comparator = None
+        self._response_hook = response_hook
 
     def _run_hybrid_search(self):
         # Check if we need to run global statistics queries, and if so do for every partition in the container
@@ -180,6 +181,7 @@ def _run_hybrid_search(self):
                         global_statistics_query,
                         self._document_producer_comparator,
                         self._options,
+                        self._response_hook
                     )
                 )
 
@@ -221,6 +223,7 @@ def _run_hybrid_search(self):
                         rewritten_query['rewrittenQuery'],
                         self._document_producer_comparator,
                         self._options,
+                        self._response_hook
                     )
                 )
         # verify all document producers have items/ no splits
@@ -347,6 +350,7 @@ def _repair_document_producer(self, query, target_all_ranges=False):
                     query,
                     self._document_producer_comparator,
                     self._options,
+                    self._response_hook
                 )
             )
 

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py

@@ -63,7 +63,7 @@ def peek(self):
         def size(self):
             return len(self._heap)
 
-    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info):
+    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info, response_hook):
         super(_MultiExecutionContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -73,6 +73,7 @@ def __init__(self, client, resource_link, query, options, partitioned_query_ex_i
         self._query = query
         self._partitioned_query_ex_info = partitioned_query_ex_info
         self._sort_orders = partitioned_query_ex_info.get_order_by()
+        self._response_hook = response_hook
 
         if self._sort_orders:
             self._document_producer_comparator = document_producer._OrderByDocumentProducerComparator(self._sort_orders)
@@ -186,6 +187,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
             query,
             self._document_producer_comparator,
             self._options,
+            self._response_hook
         )
 
     def _get_target_partition_key_range(self):

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/non_streaming_order_by_aggregator.py

@@ -22,7 +22,7 @@ class _NonStreamingOrderByContextAggregator(_QueryExecutionContextBase):
     by the user.
     """
 
-    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info):
+    def __init__(self, client, resource_link, query, options, partitioned_query_ex_info, response_hook):
         super(_NonStreamingOrderByContextAggregator, self).__init__(client, options)
 
         # use the routing provider in the client
@@ -33,6 +33,7 @@ def __init__(self, client, resource_link, query, options, partitioned_query_ex_i
         self._partitioned_query_ex_info = partitioned_query_ex_info
         self._sort_orders = partitioned_query_ex_info.get_order_by()
         self._orderByPQ = _MultiExecutionContextAggregator.PriorityQueue()
+        self._response_hook = response_hook
 
         # will be a list of (partition_min, partition_max) tuples
         targetPartitionRanges = self._get_target_partition_key_range()
@@ -142,6 +143,7 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
             query,
             self._document_producer_comparator,
             self._options,
+            self._response_hook
         )
 
     def _get_target_partition_key_range(self):

sdk/cosmos/azure-cosmos/azure/cosmos/_query_iterable.py

@@ -43,6 +43,7 @@ def __init__(
         database_link=None,
         partition_key=None,
         continuation_token=None,
+        response_hook=None,
     ):
         """Instantiates a QueryIterable for non-client side partitioning queries.
 
@@ -73,8 +74,7 @@ def __init__(
         self._database_link = database_link
         self._partition_key = partition_key
         self._ex_context = execution_dispatcher._ProxyQueryExecutionContext(
-            self._client, self._collection_link, self._query, self._options, self._fetch_function
-        )
+            self._client, self._collection_link, self._query, self._options, self._fetch_function, response_hook)
         super(QueryIterable, self).__init__(self._fetch_next, self._unpack, continuation_token=continuation_token)
 
     def _unpack(self, block):

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

@@ -2278,7 +2278,8 @@ async def fetch_fn(options: Mapping[str, Any]) -> Tuple[List[Dict[str, Any]], Ca
             options,
             fetch_function=fetch_fn,
             collection_link=database_or_container_link,
-            page_iterator_class=query_iterable.QueryIterable
+            page_iterator_class=query_iterable.QueryIterable,
+            response_hook=response_hook
         )
 
     def QueryItemsChangeFeed(