diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java index 20828ad632f1a..26a1a849b9c7e 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java @@ -4,8 +4,12 @@ package com.azure.spring.data.cosmos.core; import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosBridgeInternal; import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlQuerySpec; import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.IntegrationTestCollectionManager; import com.azure.spring.data.cosmos.common.PageTestUtils; @@ -13,6 +17,7 @@ import com.azure.spring.data.cosmos.common.TestUtils; import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; +import com.azure.spring.data.cosmos.core.generator.FindQuerySpecGenerator; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; import com.azure.spring.data.cosmos.core.query.CosmosPageRequest; import com.azure.spring.data.cosmos.core.query.CosmosQuery; @@ -38,6 +43,7 @@ import java.util.Collections; import java.util.List; import java.util.UUID; +import java.util.concurrent.ConcurrentMap; import static com.azure.spring.data.cosmos.common.TestConstants.ADDRESSES; import static com.azure.spring.data.cosmos.common.TestConstants.FIRST_NAME; @@ -63,11 +69,12 @@ public class CosmosTemplatePartitionIT { HOBBIES, ADDRESSES); private static final PartitionPerson TEST_PERSON_2 = new PartitionPerson(ID_2, NEW_FIRST_NAME, - TEST_PERSON.getZipCode(), HOBBIES, ADDRESSES); + NEW_ZIP_CODE, HOBBIES, ADDRESSES); @ClassRule public static final IntegrationTestCollectionManager collectionManager = new IntegrationTestCollectionManager(); + private static CosmosFactory cosmosFactory; private static CosmosTemplate cosmosTemplate; private static String containerName; private static CosmosEntityInformation personInfo; @@ -82,8 +89,10 @@ public class CosmosTemplatePartitionIT { @Before public void setUp() throws ClassNotFoundException { if (cosmosTemplate == null) { + // Enable Query plan caching for testing + System.setProperty("COSMOS.QUERYPLAN_CACHING_ENABLED", "true"); CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - final CosmosFactory cosmosFactory = new CosmosFactory(client, TestConstants.DB_NAME); + cosmosFactory = new CosmosFactory(client, TestConstants.DB_NAME); final CosmosMappingContext mappingContext = new CosmosMappingContext(); personInfo = new CosmosEntityInformation<>(PartitionPerson.class); @@ -121,6 +130,41 @@ public void testFindWithPartition() { assertEquals(TEST_PERSON, result.get(0)); } + @Test + public void testFindWithPartitionWithQueryPlanCachingEnabled() { + Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, PROPERTY_ZIP_CODE, + Collections.singletonList(ZIP_CODE), Part.IgnoreCaseType.NEVER); + CosmosQuery query = new CosmosQuery(criteria); + SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query); + List result = TestUtils.toList(cosmosTemplate.find(query, PartitionPerson.class, + PartitionPerson.class.getSimpleName())); + + assertThat(result.size()).isEqualTo(1); + assertEquals(TEST_PERSON, result.get(0)); + + CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); + AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient); + ConcurrentMap initialCache = asyncDocumentClient.getQueryPlanCache(); + assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); + int initialSize = initialCache.size(); + + cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(TEST_PERSON_2.getZipCode())); + + criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, PROPERTY_ZIP_CODE, + Collections.singletonList(NEW_ZIP_CODE), Part.IgnoreCaseType.NEVER); + query = new CosmosQuery(criteria); + // Fire the same query but with different partition key value to make sure query plan caching is enabled + sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query); + result = TestUtils.toList(cosmosTemplate.find(query, PartitionPerson.class, + PartitionPerson.class.getSimpleName())); + + ConcurrentMap postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); + assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); + assertThat(postQueryCallCache.size()).isEqualTo(initialSize); + assertThat(result.size()).isEqualTo(1); + assertEquals(TEST_PERSON_2, result.get(0)); + } + @Test public void testFindIgnoreCaseWithPartition() { Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, PROPERTY_ZIP_CODE, @@ -196,8 +240,6 @@ public void testDeleteByIdPartition() { final List inserted = TestUtils.toList(cosmosTemplate.findAll(PartitionPerson.class)); assertThat(inserted.size()).isEqualTo(2); - assertThat(inserted.get(0).getZipCode()).isEqualTo(TEST_PERSON.getZipCode()); - assertThat(inserted.get(1).getZipCode()).isEqualTo(TEST_PERSON.getZipCode()); cosmosTemplate.deleteById(PartitionPerson.class.getSimpleName(), TEST_PERSON.getId(), new PartitionKey(TEST_PERSON.getZipCode())); diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java index 6ea2f33d8cf6a..3c980e9318ac2 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java @@ -3,13 +3,18 @@ package com.azure.spring.data.cosmos.core; import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosBridgeInternal; import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlQuerySpec; import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.ReactiveIntegrationTestCollectionManager; import com.azure.spring.data.cosmos.common.TestConstants; import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; +import com.azure.spring.data.cosmos.core.generator.FindQuerySpecGenerator; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; import com.azure.spring.data.cosmos.core.query.CosmosQuery; import com.azure.spring.data.cosmos.core.query.Criteria; @@ -35,7 +40,9 @@ import java.util.Collections; import java.util.UUID; +import java.util.concurrent.ConcurrentMap; +import static org.assertj.core.api.Assertions.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; @@ -49,11 +56,12 @@ public class ReactiveCosmosTemplatePartitionIT { private static final PartitionPerson TEST_PERSON_2 = new PartitionPerson(TestConstants.ID_2, TestConstants.NEW_FIRST_NAME, - TEST_PERSON.getZipCode(), TestConstants.HOBBIES, TestConstants.ADDRESSES); + TestConstants.NEW_ZIP_CODE, TestConstants.HOBBIES, TestConstants.ADDRESSES); @ClassRule public static final ReactiveIntegrationTestCollectionManager collectionManager = new ReactiveIntegrationTestCollectionManager(); + private static CosmosFactory cosmosFactory; private static ReactiveCosmosTemplate cosmosTemplate; private static String containerName; private static CosmosEntityInformation personInfo; @@ -68,8 +76,10 @@ public class ReactiveCosmosTemplatePartitionIT { @Before public void setUp() throws ClassNotFoundException { if (cosmosTemplate == null) { + // Enable Query plan caching for testing + System.setProperty("COSMOS.QUERYPLAN_CACHING_ENABLED", "true"); CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - final CosmosFactory dbFactory = new CosmosFactory(client, TestConstants.DB_NAME); + cosmosFactory = new CosmosFactory(client, TestConstants.DB_NAME); final CosmosMappingContext mappingContext = new CosmosMappingContext(); personInfo = @@ -80,7 +90,7 @@ public void setUp() throws ClassNotFoundException { final MappingCosmosConverter dbConverter = new MappingCosmosConverter(mappingContext, null); - cosmosTemplate = new ReactiveCosmosTemplate(dbFactory, cosmosConfig, dbConverter); + cosmosTemplate = new ReactiveCosmosTemplate(cosmosFactory, cosmosConfig, dbConverter); } collectionManager.ensureContainersCreatedAndEmpty(cosmosTemplate, PartitionPerson.class); cosmosTemplate.insert(TEST_PERSON).block(); @@ -100,6 +110,46 @@ public void testFindWithPartition() { }).verifyComplete(); } + @Test + public void testFindWithPartitionWithQueryPlanCachingEnabled() { + Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, TestConstants.PROPERTY_ZIP_CODE, + Collections.singletonList(TestConstants.ZIP_CODE), Part.IgnoreCaseType.NEVER); + CosmosQuery query = new CosmosQuery(criteria); + SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query); + Flux partitionPersonFlux = cosmosTemplate.find(query, + PartitionPerson.class, + PartitionPerson.class.getSimpleName()); + StepVerifier.create(partitionPersonFlux).consumeNextWith(actual -> { + Assert.assertThat(actual.getFirstName(), is(equalTo(TEST_PERSON.getFirstName()))); + Assert.assertThat(actual.getZipCode(), is(equalTo(TEST_PERSON.getZipCode()))); + }).verifyComplete(); + + CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); + AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient); + ConcurrentMap initialCache = asyncDocumentClient.getQueryPlanCache(); + assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); + int initialSize = initialCache.size(); + + cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(TEST_PERSON_2.getZipCode())).block(); + + // Fire the same query with different partition key value to make sure query plan caching is enabled + criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, TestConstants.PROPERTY_ZIP_CODE, + Collections.singletonList(TestConstants.NEW_ZIP_CODE), Part.IgnoreCaseType.NEVER); + query = new CosmosQuery(criteria); + sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query); + partitionPersonFlux = cosmosTemplate.find(query, + PartitionPerson.class, + PartitionPerson.class.getSimpleName()); + StepVerifier.create(partitionPersonFlux).consumeNextWith(actual -> { + Assert.assertThat(actual.getFirstName(), is(equalTo(TEST_PERSON_2.getFirstName()))); + Assert.assertThat(actual.getZipCode(), is(equalTo(TEST_PERSON_2.getZipCode()))); + }).verifyComplete(); + + ConcurrentMap postQueryCallCache = asyncDocumentClient.getQueryPlanCache(); + assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue(); + assertThat(postQueryCallCache.size()).isEqualTo(initialSize); + } + @Test public void testFindIgnoreCaseWithPartition() { final Criteria criteria = Criteria.getInstance(CriteriaType.IS_EQUAL, TestConstants.PROPERTY_ZIP_CODE, diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/generator/AbstractQueryGenerator.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/generator/AbstractQueryGenerator.java index be2666ed7a1f1..3e432bd71296b 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/generator/AbstractQueryGenerator.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/generator/AbstractQueryGenerator.java @@ -18,7 +18,7 @@ import java.util.ArrayList; import java.util.Collection; import java.util.List; -import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter.toCosmosDbValue; @@ -31,9 +31,9 @@ public abstract class AbstractQueryGenerator { protected AbstractQueryGenerator() { } - private String generateQueryParameter(@NonNull String subject) { + private String generateQueryParameter(@NonNull String subject, int counter) { // user.name, user['name'] or user["first name"] are not valid sql parameter identifiers. - return subject.replaceAll("[^a-zA-Z\\d]", "_") + UUID.randomUUID().toString().replaceAll("-", "_"); + return subject.replaceAll("[^a-zA-Z\\d]", "_") + counter; } private String generateUnaryQuery(@NonNull Criteria criteria) { @@ -48,14 +48,14 @@ private String generateUnaryQuery(@NonNull Criteria criteria) { } } - private String generateBinaryQuery(@NonNull Criteria criteria, @NonNull List> parameters) { + private String generateBinaryQuery(@NonNull Criteria criteria, @NonNull List> parameters, int counter) { Assert.isTrue(criteria.getSubjectValues().size() == 1, "Binary criteria should have only one subject value"); Assert.isTrue(CriteriaType.isBinary(criteria.getType()), "Criteria type should be binary operation"); final String subject = criteria.getSubject(); final Object subjectValue = toCosmosDbValue(criteria.getSubjectValues().get(0)); - final String parameter = generateQueryParameter(subject); + final String parameter = generateQueryParameter(subject, counter); final Part.IgnoreCaseType ignoreCase = criteria.getIgnoreCase(); final String sqlKeyword = criteria.getType().getSqlKeyword(); parameters.add(Pair.of(parameter, subjectValue)); @@ -103,14 +103,14 @@ private String getFunctionCondition(final Part.IgnoreCaseType ignoreCase, final } } - private String generateBetween(@NonNull Criteria criteria, @NonNull List> parameters) { + private String generateBetween(@NonNull Criteria criteria, @NonNull List> parameters, int counter) { final String subject = criteria.getSubject(); final Object value1 = toCosmosDbValue(criteria.getSubjectValues().get(0)); final Object value2 = toCosmosDbValue(criteria.getSubjectValues().get(1)); final String subject1 = subject + "start"; final String subject2 = subject + "end"; - final String parameter1 = generateQueryParameter(subject1); - final String parameter2 = generateQueryParameter(subject2); + final String parameter1 = generateQueryParameter(subject1, counter); + final String parameter2 = generateQueryParameter(subject2, counter); final String keyword = criteria.getType().getSqlKeyword(); parameters.add(Pair.of(parameter1, value1)); @@ -152,7 +152,7 @@ private String generateInQuery(@NonNull Criteria criteria, @NonNull List> parameters) { + private String generateQueryBody(@NonNull Criteria criteria, @NonNull List> parameters, @NonNull final AtomicInteger counter) { final CriteriaType type = criteria.getType(); switch (type) { @@ -162,7 +162,7 @@ private String generateQueryBody(@NonNull Criteria criteria, @NonNull List>> generateQueryBody(@NonNull CosmosQuery query) { + private Pair>> generateQueryBody(@NonNull CosmosQuery query, @NonNull final AtomicInteger counter) { final List> parameters = new ArrayList<>(); - String queryString = this.generateQueryBody(query.getCriteria(), parameters); + String queryString = this.generateQueryBody(query.getCriteria(), parameters, counter); if (StringUtils.hasText(queryString)) { queryString = String.join(" ", "WHERE", queryString); @@ -248,7 +248,8 @@ private String generateQueryTail(@NonNull CosmosQuery query) { protected SqlQuerySpec generateCosmosQuery(@NonNull CosmosQuery query, @NonNull String queryHead) { - final Pair>> queryBody = generateQueryBody(query); + final AtomicInteger counter = new AtomicInteger(); + final Pair>> queryBody = generateQueryBody(query, counter); String queryString = String.join(" ", queryHead, queryBody.getFirst(), generateQueryTail(query)); final List> parameters = queryBody.getSecond();