# The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. import unittest import uuid import pytest from azure.core.paging import ItemPaged import azure.cosmos._base as base import azure.cosmos.cosmos_client as cosmos_client import test_config from azure.cosmos import _query_iterable as query_iterable, DatabaseProxy, ContainerProxy from azure.cosmos.exceptions import CosmosHttpResponseError from azure.cosmos.partition_key import PartitionKey @pytest.mark.cosmosQuery class TestCrossPartitionTopOrderBy(unittest.TestCase): """Orderby Tests. """ document_definitions = None created_container: ContainerProxy = None client: cosmos_client.CosmosClient = None created_db: DatabaseProxy = None host = test_config.TestConfig.host masterKey = test_config.TestConfig.masterKey connectionPolicy = test_config.TestConfig.connectionPolicy TEST_DATABASE_ID = test_config.TestConfig.TEST_DATABASE_ID TEST_CONTAINER_ID = test_config.TestConfig.TEST_MULTI_PARTITION_CONTAINER_ID @classmethod def setUpClass(cls): if (cls.masterKey == '[YOUR_KEY_HERE]' or cls.host == '[YOUR_ENDPOINT_HERE]'): raise Exception( "You must specify your Azure Cosmos account values for " "'masterKey' and 'host' at the top of this class to run the " "tests.") cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey) cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) cls.created_container = cls.created_db.create_container( id='orderby_tests collection ' + str(uuid.uuid4()), indexing_policy={ 'includedPaths': [ { 'path': '/', 'indexes': [ { 'kind': 'Range', 'dataType': 'Number' }, { 'kind': 'Range', 'dataType': 'String' } ] } ] }, partition_key=PartitionKey(path='/id'), offer_throughput=30000) cls.collection_link = cls.GetDocumentCollectionLink(cls.created_db, cls.created_container) # create a document using the document definition cls.document_definitions = [] for i in range(20): d = {'id': str(i), 'name': 'sample document', 'spam': 'eggs' + str(i), 'cnt': i, 'key': 'value', 'spam2': 'eggs' + str(i) if (i == 3) else i, 'boolVar': (i % 2 == 0), 'number': 1.1 * i } cls.created_container.create_item(d) cls.document_definitions.append(d) @classmethod def tearDownClass(cls): try: cls.created_db.delete_container(cls.created_container.id) except CosmosHttpResponseError: pass def test_orderby_query(self): # test a simple order by query # an order by query query = { 'query': 'SELECT * FROM root r order by r.spam', } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_query_as_string(self): # test a simple order by query as string # an order by query query = 'SELECT * FROM root r order by r.spam' def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_asc_query(self): # test an order by query with explicit ascending ordering # an ascending order by query (ascending explicitly mentioned in the query) query = { 'query': 'SELECT * FROM root r order by r.spam ASC', } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_desc_query(self): # test an order by query with explicit descending ordering # a descending order by query query = { 'query': 'SELECT * FROM root r order by r.spam DESC', } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key, reverse=True)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_top_query(self): # test an order by query combined with top top_count = 9 # sanity check self.assertLess(top_count, len(self.document_definitions)) # an order by query with top, total existing docs more than requested top count query = { 'query': 'SELECT top %d * FROM root r order by r.spam' % top_count # nosec } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)[:top_count]] self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_top_query_less_results_than_top_counts(self): # test an order by query combined with top. where top is greater than the total number of docs top_count = 30 # sanity check self.assertGreater(top_count, len(self.document_definitions)) # an order by query with top, total existing docs less than requested top count query = { 'query': 'SELECT top %d * FROM root r order by r.spam' % top_count # nosec } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] self.execute_query_and_validate_results(query, expected_ordered_ids) def test_top_query(self): # test a simple top query without order by. # The rewrittenQuery in the query execution info responded by backend will be empty partition_key_ranges = list(self.client.client_connection._ReadPartitionKeyRanges(self.collection_link)) docs_by_partition_key_range_id = self.find_docs_by_partition_key_range_id() # find the first two non-empty target partition key ranges cnt = 0 first_two_ranges_results = [] for r in partition_key_ranges: if cnt >= 2: break p_id = r['id'] if len(docs_by_partition_key_range_id[p_id]) > 0: first_two_ranges_results.extend(docs_by_partition_key_range_id[p_id]) cnt += 1 # sanity checks self.assertEqual(cnt, 2) self.assertLess(2, len(partition_key_ranges)) # sanity check self.assertLess(len(first_two_ranges_results), len(self.document_definitions)) self.assertGreater(len(first_two_ranges_results), 1) expected_ordered_ids = [d['id'] for d in first_two_ranges_results] # a top query, the results will be sorted based on the target partition key range query = { 'query': 'SELECT top %d * FROM root r' % len(expected_ordered_ids) # nosec } self.execute_query_and_validate_results(query, expected_ordered_ids) def test_top_query_as_string(self): # test a simple top query without order by. # The rewrittenQuery in the query execution info responded by backend will be empty partition_key_ranges = list(self.client.client_connection._ReadPartitionKeyRanges(self.collection_link)) docs_by_partition_key_range_id = self.find_docs_by_partition_key_range_id() # find the first two non-empty target partition key ranges cnt = 0 first_two_ranges_results = [] for r in partition_key_ranges: if cnt >= 2: break p_id = r['id'] if len(docs_by_partition_key_range_id[p_id]) > 0: first_two_ranges_results.extend(docs_by_partition_key_range_id[p_id]) cnt += 1 # sanity checks self.assertEqual(cnt, 2) self.assertLess(2, len(partition_key_ranges)) # sanity check self.assertLess(len(first_two_ranges_results), len(self.document_definitions)) self.assertGreater(len(first_two_ranges_results), 1) expected_ordered_ids = [d['id'] for d in first_two_ranges_results] # a top query, the results will be sorted based on the target partition key range query = 'SELECT top %d * FROM root r' % len(expected_ordered_ids) # nosec self.execute_query_and_validate_results(query, expected_ordered_ids) def test_parametrized_top_query(self): # test a simple parameterized query without order by. # The rewrittenQuery in the query execution info responded by backend will be empty partition_key_ranges = list(self.client.client_connection._ReadPartitionKeyRanges(self.collection_link)) docs_by_partition_key_range_id = self.find_docs_by_partition_key_range_id() # find the first two non-empty target partition key ranges cnt = 0 first_two_ranges_results = [] for r in partition_key_ranges: if cnt >= 2: break p_id = r['id'] if len(docs_by_partition_key_range_id[p_id]) > 0: first_two_ranges_results.extend(docs_by_partition_key_range_id[p_id]) cnt += 1 # sanity checks self.assertEqual(cnt, 2) self.assertLess(2, len(partition_key_ranges)) # sanity check self.assertLess(len(first_two_ranges_results), len(self.document_definitions)) self.assertGreater(len(first_two_ranges_results), 1) expected_ordered_ids = [d['id'] for d in first_two_ranges_results] # a top query, the results will be sorted based on the target partition key range query = { 'query': 'SELECT top @n * FROM root r', "parameters": [ {"name": "@n", "value": len(expected_ordered_ids)} ] } self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_query_with_parametrized_top(self): # test an order by query combined with parametrized top top_count = 9 # sanity check self.assertLess(top_count, len(self.document_definitions)) def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)[:top_count]] # a parametrized top order by query query = { 'query': 'SELECT top @n * FROM root r order by r.spam', "parameters": [ {"name": "@n", "value": top_count} ] } self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_query_with_parametrized_predicate(self): # test an order by query combined with parametrized predicate # an order by query with parametrized predicate query = { 'query': 'SELECT * FROM root r where r.cnt > @cnt order by r.spam', "parameters": [ {"name": "@cnt", "value": 5} ] } def get_order_by_key(r): return r['spam'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key) if r['cnt'] > 5] self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_query_noncomparable_orderby_item(self): # test orderby with different order by item type # an order by query query = { 'query': 'SELECT * FROM root r order by r.spam2 DESC', } def get_order_by_key(r): return r['id'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order try: self.execute_query_and_validate_results(query, expected_ordered_ids) self.fail('non comparable order by items did not result in failure.') except ValueError as e: self.assertTrue( e.args[0] == "Expected String, but got Number." or e.message == "Expected Number, but got String.") def test_orderby_integer_query(self): # an order by integer query query = { 'query': 'SELECT * FROM root r order by r.cnt', } def get_order_by_key(r): return r['cnt'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_floating_point_number_query(self): # an orderby by floating point number query query = { 'query': 'SELECT * FROM root r order by r.number', } def get_order_by_key(r): return r['number'] expected_ordered_ids = [r['id'] for r in sorted(self.document_definitions, key=get_order_by_key)] # validates the results size and order self.execute_query_and_validate_results(query, expected_ordered_ids) def test_orderby_boolean_query(self): # an orderby by floating point number query query = { 'query': 'SELECT * FROM root r order by r.boolVar', } result_iterable = self.created_container.query_items( query=query, enable_cross_partition_query=True, max_item_count=2 ) results = list(result_iterable) # validates the results size and order self.assertEqual(len(results), len(self.document_definitions)) # false values before true values index = 0 while index < len(results): if results[index]['boolVar']: break self.assertTrue(int(results[index]['id']) % 2 == 1) index = index + 1 while index < len(results): self.assertTrue(results[index]['boolVar']) self.assertTrue(int(results[index]['id']) % 2 == 0) index = index + 1 def find_docs_by_partition_key_range_id(self): query = { 'query': 'SELECT * FROM root r' } partition_key_range = list(self.client.client_connection._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.client_connection.QueryFeed(path, collection_id, query, options, r['id']) docResultsIterable = ItemPaged( self.client.client_connection, query, options, fetch_function=fetch_fn, collection_link=self.collection_link, page_iterator_class=query_iterable.QueryIterable ) 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 execute_query_and_validate_results(self, query, expected_ordered_ids): # executes the query and validates the results against the expected results page_size = 2 result_iterable = self.created_container.query_items( query=query, enable_cross_partition_query=True, max_item_count=page_size ) self.assertTrue(isinstance(result_iterable, ItemPaged)) self.assertEqual(result_iterable._page_iterator_class, query_iterable.QueryIterable) ###################################### # test next() behavior ###################################### it = result_iterable.__iter__() def invokeNext(): return next(it) # validate that invocations of next() produces the same results as expected_ordered_ids for i in range(len(expected_ordered_ids)): item = invokeNext() self.assertEqual(item['id'], expected_ordered_ids[i]) ###################################### # test by_page() behavior ###################################### results = {} cnt = 0 page_iter = result_iterable.by_page() for page in page_iter: fetched_res = list(page) fetched_size = len(fetched_res) for item in fetched_res: self.assertEqual(item['id'], expected_ordered_ids[cnt]) results[cnt] = item cnt = cnt + 1 if (cnt < len(expected_ordered_ids)): self.assertEqual(fetched_size, page_size, "page size") else: if cnt == len(expected_ordered_ids): 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), len(expected_ordered_ids)) # no more results will be returned with self.assertRaises(StopIteration): next(page_iter) @classmethod def GetDatabaseLink(cls, database, is_name_based=True): if is_name_based: return 'dbs/' + database.id else: return database['_self'] @classmethod def GetDocumentCollectionLink(cls, database, document_collection, is_name_based=True): if is_name_based: return cls.GetDatabaseLink(database) + '/colls/' + document_collection.id else: return document_collection['_self'] @classmethod def GetDocumentLink(cls, database, document_collection, document, is_name_based=True): if is_name_based: return cls.GetDocumentCollectionLink(database, document_collection) + '/docs/' + document['id'] else: return document['_self'] if __name__ == "__main__": unittest.main()