def __init__(self, partition_key_target_range, client, collection_link, query, document_producer_comp): ''' Constructor ''' # TODO: is that fine we build the options dict and we don't inherit it? self._options = {} self._partition_key_target_range = partition_key_target_range self._doc_producer_comp = document_producer_comp self._client = client self._buffer = deque() self._is_finished = False self._has_started = False self._cur_item = None # initiate execution context path = base.GetPathFromLink(collection_link, 'docs') collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) def fetch_fn(options): return self._client.QueryFeed(path, collection_id, query, options, partition_key_target_range['id']) self._ex_context = _DefaultQueryExecutionContext( client, self._options, fetch_fn)
def fetch_next_block(self): """Fetches the next block of query results. :Returns: list of fetched items. """ # Fetch next block of results by executing the query against the current document collection fetched_items = self.fetch_items() # If there are multiple document collections to query for(in case of partitioning), keep looping through each one of them, # creating separate feed queries for each collection and fetching the items while not fetched_items: if self._collection_links and self._current_collection_index < self._collection_links_length: path = base.GetPathFromLink( self._collection_links[self._current_collection_index], 'docs') collection_id = base.GetResourceIdOrFullNameFromLink( self._collection_links[self._current_collection_index]) self._continuation = None self._has_started = False def fetch_fn(options): return self._client.QueryFeed(path, collection_id, self._query, options) self._fetch_function = fetch_fn fetched_items = self.fetch_items() self._current_collection_index += 1 else: break return fetched_items
def PartitioningQueryIterable(cls, client, options, retry_policy, database_link, query, partition_key): """ :Parameters: - `client`: DocumentClient, instance of document client - `options`: dict, the request options for the request. - `retry_policy`: documents.RetryPolicy, retry policy to be used - `database_link`: str, database self link or ID based link - `query`: str or dict - `partition_key`: str, partition key for the query Example of `fetch_function`: >>> def result_fn(result): >>> return result['Databases'] """ # This will call the base constructor(__init__ method above) self = cls(options, retry_policy, None) self._query = query self._client = client partition_resolver = client.GetPartitionResolver(database_link) if (partition_resolver is None): raise ValueError(client.PartitionResolverErrorMessage) else: self._collection_links = partition_resolver.ResolveForRead( partition_key) self._collection_links_length = len(self._collection_links) if self._collection_links is None: raise ValueError("_collection_links is None.") if self._collection_links_length <= 0: raise ValueError("_collection_links_length is not greater than 0.") # Creating the QueryFeed for the first collection path = base.GetPathFromLink( self._collection_links[self._current_collection_index], 'docs') collection_id = base.GetResourceIdOrFullNameFromLink( self._collection_links[self._current_collection_index]) self._current_collection_index += 1 def fetch_fn(options): return client.QueryFeed(path, collection_id, query, options) self._fetch_function = fetch_fn return self
def __init__(self, client, options, database_link, query, partition_key): """ Constructor :param DocumentClient client: :param dict options: The request options for the request. :param str database_link: database self link or ID based link :param (str or dict) query: Partition_key (str): partition key for the query """ super(self.__class__, self).__init__(client, options) self._current_collection_index = 0 self._collection_links = [] self._collection_links_length = 0 self._query = query self._client = client partition_resolver = client.GetPartitionResolver(database_link) if (partition_resolver is None): raise ValueError(client.PartitionResolverErrorMessage) else: self._collection_links = partition_resolver.ResolveForRead( partition_key) self._collection_links_length = len(self._collection_links) if self._collection_links is None: raise ValueError("_collection_links is None.") if self._collection_links_length <= 0: raise ValueError("_collection_links_length is not greater than 0.") # Creating the QueryFeed for the first collection path = base.GetPathFromLink( self._collection_links[self._current_collection_index], 'docs') collection_id = base.GetResourceIdOrFullNameFromLink( self._collection_links[self._current_collection_index]) self._current_collection_index += 1 def fetch_fn(options): return client.QueryFeed(path, collection_id, query, options) self._fetch_function = fetch_fn
def find_docs_by_partition_key_range_id(self): query = { 'query': 'SELECT * FROM root r' } partition_key_range = list(self.client._ReadPartitionKeyRanges(self.collection_link)) docs_by_partition_key_range_id = {} for r in partition_key_range: options = {} path = base.GetPathFromLink(self.collection_link, 'docs') collection_id = base.GetResourceIdOrFullNameFromLink(self.collection_link) def fetch_fn(options): return self.client.QueryFeed(path, collection_id, query, options, r['id']) docResultsIterable = query_iterable.QueryIterable(self.client, query, options, fetch_fn, self.collection_link) docs = list(docResultsIterable) self.assertFalse(r['id'] in docs_by_partition_key_range_id) docs_by_partition_key_range_id[r['id']] = docs return docs_by_partition_key_range_id
def _test_default_execution_context(self, options, query, expected_number_of_results): page_size = options['maxItemCount'] collection_link = self.GetDocumentCollectionLink(self.created_db, self.created_collection) path = base.GetPathFromLink(collection_link, 'docs') collection_id = base.GetResourceIdOrFullNameFromLink(collection_link) def fetch_fn(options): return self.client.QueryFeed(path, collection_id, query, options) ###################################### # test next() behavior ###################################### ex = base_execution_context._DefaultQueryExecutionContext(self.client, options, fetch_fn) it = ex.__iter__() def invokeNext(): return next(it) results = {} # validate that invocations of next() produces the same results as expected for i in xrange(expected_number_of_results): item = invokeNext() results[item['id']] = item self.assertEqual(len(results), expected_number_of_results) # after the result set is exhausted, invoking next must raise a StopIteration exception self.assertRaises(StopIteration, invokeNext) ###################################### # test fetch_next_block() behavior ###################################### ex = base_execution_context._DefaultQueryExecutionContext(self.client, options, fetch_fn) results = {} cnt = 0 while True: fetched_res = ex.fetch_next_block() fetched_size = len(fetched_res) for item in fetched_res: results[item['id']] = item cnt += fetched_size if (cnt < expected_number_of_results): # backend may not necessarily return exactly page_size of results self.assertEqual(fetched_size, page_size, "page size") else: if cnt == expected_number_of_results: self.assertTrue(fetched_size <= page_size, "last page size") break else: #cnt > expected_number_of_results self.fail("more results than expected") # validate the number of collected results self.assertEqual(len(results), expected_number_of_results) # no more results will be returned self.assertEqual(ex.fetch_next_block(), [])