From cb6a76cb239b053d5634df37da44211b4945e809 Mon Sep 17 00:00:00 2001 From: Kushagra Thapar Date: Fri, 16 Aug 2019 02:05:46 -0700 Subject: [PATCH] Removed dependency from reactive cosmos template because it was initiating two cosmos clients (#412) --- .../cosmosdb/core/DocumentDbTemplate.java | 74 ++++++++++++++++--- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java index 38e8960..37ecd50 100644 --- a/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java +++ b/src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java @@ -18,6 +18,7 @@ import com.microsoft.azure.documentdb.DocumentCollection; import com.microsoft.azure.documentdb.PartitionKey; import com.microsoft.azure.spring.data.cosmosdb.CosmosDbFactory; import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter; +import com.microsoft.azure.spring.data.cosmosdb.core.generator.CountQueryGenerator; import com.microsoft.azure.spring.data.cosmosdb.core.generator.FindQuerySpecGenerator; import com.microsoft.azure.spring.data.cosmosdb.core.query.Criteria; import com.microsoft.azure.spring.data.cosmosdb.core.query.CriteriaType; @@ -47,8 +48,9 @@ import java.util.stream.Collectors; @Slf4j public class DocumentDbTemplate implements DocumentDbOperations, ApplicationContextAware { + private static final String COUNT_VALUE_KEY = "_aggregate"; + private final MappingDocumentDbConverter mappingDocumentDbConverter; - private final ReactiveCosmosTemplate reactiveCosmosTemplate; private final String databaseName; private final CosmosClient cosmosClient; @@ -61,7 +63,6 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont this.mappingDocumentDbConverter = mappingDocumentDbConverter; - this.reactiveCosmosTemplate = new ReactiveCosmosTemplate(cosmosDbFactory, mappingDocumentDbConverter, dbName); this.databaseName = dbName; this.cosmosClient = cosmosDbFactory.getCosmosClient(); } @@ -201,7 +202,12 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont @Override public void deleteCollection(@NonNull String collectionName) { Assert.hasText(collectionName, "collectionName should have text."); - reactiveCosmosTemplate.deleteContainer(collectionName); + try { + cosmosClient.getDatabase(this.databaseName).getContainer(collectionName).delete().block(); + } catch (Exception e) { + throw new DocumentDBAccessException("failed to delete collection: " + collectionName, + e); + } } public String getCollectionName(Class domainClass) { @@ -212,8 +218,13 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont @Override public DocumentCollection createCollectionIfNotExists(@NonNull DocumentDbEntityInformation information) { - final CosmosContainerResponse response = reactiveCosmosTemplate - .createCollectionIfNotExists(information) + final CosmosContainerResponse response = cosmosClient + .createDatabaseIfNotExists(this.databaseName) + .flatMap(cosmosDatabaseResponse -> cosmosDatabaseResponse + .database() + .createContainerIfNotExists(information.getCollectionName(), + "/" + information.getPartitionKeyFieldName()) + .map(cosmosContainerResponse -> cosmosContainerResponse)) .block(); if (response == null) { throw new DocumentDBAccessException("Failed to create collection"); @@ -232,7 +243,15 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont pk = com.azure.data.cosmos.PartitionKey.None; } try { - reactiveCosmosTemplate.deleteById(collectionName, id, pk).block(); + final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + options.partitionKey(pk); + cosmosClient.getDatabase(this.databaseName) + .getContainer(collectionName) + .getItem(id.toString(), partitionKey) + .delete(options) + .onErrorResume(Mono::error) + .then() + .block(); } catch (Exception e) { throw new DocumentDBAccessException("deleteById exception", e); } @@ -255,7 +274,10 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont Assert.hasText(collectionName, "collection should not be null, empty or only whitespaces"); try { - return reactiveCosmosTemplate.find(query, domainClass, collectionName).collectList().block(); + return findDocuments(query, domainClass, collectionName) + .stream() + .map(cosmosItemProperties -> cosmosItemProperties.toObject(domainClass)) + .collect(Collectors.toList()); } catch (Exception e) { throw new DocumentDBAccessException("Failed to execute find operation from " + collectionName, e); } @@ -355,7 +377,8 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont public long count(String collectionName) { Assert.hasText(collectionName, "collectionName should not be empty"); - final Long count = reactiveCosmosTemplate.count(collectionName).block(); + final DocumentQuery query = new DocumentQuery(Criteria.getInstance(CriteriaType.ALL)); + final Long count = getCountValue(query, true, collectionName); if (count == null) { throw new DocumentDBAccessException("Failed to get count for collectionName: " + collectionName); } @@ -367,8 +390,9 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont Assert.notNull(domainClass, "domainClass should not be null"); Assert.hasText(collectionName, "collectionName should not be empty"); - final boolean isCrossPartitionQuery = query.isCrossPartitionQuery(getPartitionKeyNames(domainClass)); - final Long count = reactiveCosmosTemplate.count(query, isCrossPartitionQuery, collectionName).block(); + final boolean isCrossPartitionQuery = + query.isCrossPartitionQuery(getPartitionKeyNames(domainClass)); + final Long count = getCountValue(query, isCrossPartitionQuery, collectionName); if (count == null) { throw new DocumentDBAccessException("Failed to get count for collectionName: " + collectionName); } @@ -380,6 +404,30 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont return this.mappingDocumentDbConverter; } + private Long getCountValue(DocumentQuery query, boolean isCrossPartitionQuery, String containerName) { + final SqlQuerySpec querySpec = new CountQueryGenerator().generateCosmos(query); + final FeedOptions options = new FeedOptions(); + + options.enableCrossPartitionQuery(isCrossPartitionQuery); + + return executeQuery(querySpec, containerName, options) + .onErrorResume(this::databaseAccessExceptionHandler) + .next() + .map(r -> r.results().get(0).getLong(COUNT_VALUE_KEY)) + .block(); + } + + private Mono databaseAccessExceptionHandler(Throwable e) { + throw new DocumentDBAccessException("failed to access cosmosdb database", e); + } + + private Flux> executeQuery(SqlQuerySpec sqlQuerySpec, String collectionName, + FeedOptions options) { + return cosmosClient.getDatabase(this.databaseName) + .getContainer(collectionName) + .queryItems(sqlQuerySpec, options); + } + @SuppressWarnings("unchecked") private List getPartitionKeyNames(Class domainClass) { final DocumentDbEntityInformation entityInfo = new DocumentDbEntityInformation(domainClass); @@ -405,10 +453,12 @@ public class DocumentDbTemplate implements DocumentDbOperations, ApplicationCont } } - private List findDocuments(@NonNull DocumentQuery query, @NonNull Class domainClass, + private List findDocuments(@NonNull DocumentQuery query, + @NonNull Class domainClass, @NonNull String containerName) { final SqlQuerySpec sqlQuerySpec = new FindQuerySpecGenerator().generateCosmos(query); - final boolean isCrossPartitionQuery = query.isCrossPartitionQuery(getPartitionKeyNames(domainClass)); + final boolean isCrossPartitionQuery = + query.isCrossPartitionQuery(getPartitionKeyNames(domainClass)); final FeedOptions feedOptions = new FeedOptions(); feedOptions.enableCrossPartitionQuery(isCrossPartitionQuery); return cosmosClient