Allow getAll with partition keys for entities not having sort key (#254)

* Allow getAll with partition keys for entities not having sort key * fixed key in async dynamo db
agorapulse · Jul 11, 2024 · 5d5dae1 · 5d5dae1
1 parent 4844b91
commit 5d5dae1
Show file tree

Hide file tree

Showing 12 changed files with 73 additions and 0 deletions.
diff --git a/docs/guide/src/docs/asciidoc/dynamodb.adoc b/docs/guide/src/docs/asciidoc/dynamodb.adoc
@@ -498,6 +498,8 @@ The following table summarizes the supported method signatures:
 
     `List<DynamoDBEntity> getAll(@HashKey String parentId, String... rangeKeys);`
 
+    `List<DynamoDBEntityNoRange> getAll(@HashKey List<String> parentIds);`
+
 | Loads a single entity or a list of entities from the table. Range key is required for tables which defines the range key
 
 

diff --git a/...b/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/AsyncDynamoDbService.java b/...b/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/AsyncDynamoDbService.java
@@ -145,6 +145,8 @@ default <R> Publisher<R> updateAll(Publisher<T> items, Function<UpdateBuilder<T,
 
     Publisher<T> getAll(Object partitionKey, Publisher<?> sortKeys);
 
+    Publisher<T> getAll(Publisher<?> partitionKeys);
+
     Publisher<T> get(Key key);
 
     Publisher<Long> count(DetachedQuery<T> query);

diff --git a/...ava/com/agorapulse/micronaut/amazon/awssdk/dynamodb/AsyncDynamoDbServiceIntroduction.java b/...ava/com/agorapulse/micronaut/amazon/awssdk/dynamodb/AsyncDynamoDbServiceIntroduction.java
@@ -284,6 +284,9 @@ private <T> Publisher<T> handleGet(AsyncDynamoDbService<T> service, MethodInvoca
         Object partitionValue = partitionAndSort.getPartitionValue(params);
 
         if (!partitionAndSort.hasSortKey()) {
+            if (partitionAndSort.isPartitionKeyPublisherOrIterable()) {
+                return service.getAll(partitionAndSort.getPartitionAttributeValues(conversionService, params));
+            }
             return service.get(partitionValue, null);
         }
 

diff --git a/...ain/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultAsyncDynamoDbService.java b/...ain/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultAsyncDynamoDbService.java
@@ -230,6 +230,11 @@ public Publisher<T> getAll(Object partitionKey, Publisher<?> sortKeys) {
         return doWithKeys(partitionKey, sortKeys, this::getAll);
     }
 
+    @Override
+    public Publisher<T> getAll(Publisher<?> partitionKeys) {
+        return doWithKeys(partitionKeys, this::getAllByAttributeValue);
+    }
+
     @Override
     public Publisher<T> get(Key key) {
         return Mono.fromFuture(table.getItem(key)).map(this::postLoad);
@@ -289,6 +294,19 @@ private DetachedQuery<T> simplePartitionAndSort(Object partitionKey, Object sort
         });
     }
 
+    private Publisher<T> getAllByAttributeValue(Publisher<AttributeValue> partitionKeys) {
+        TableSchema<T> tableSchema = table.tableSchema();
+        Map<AttributeValue, Integer> order = new ConcurrentHashMap<>();
+        AtomicInteger counter = new AtomicInteger();
+        Comparator<T> comparator = Comparator.comparingInt(i -> order.getOrDefault(tableSchema.attributeValue(i, tableSchema.tableMetadata().primaryPartitionKey()), 0));
+
+        return Flux.from(partitionKeys).buffer(BATCH_SIZE).map(batchRangeKeys -> enhancedClient.batchGetItem(b -> b.readBatches(batchRangeKeys.stream().map(k -> {
+                order.put(k, counter.getAndIncrement());
+                return ReadBatch.builder(tableSchema.itemType().rawClass()).mappedTableResource(table).addGetItem(Key.builder().partitionValue(k).build()).build();
+        }).toList()))).flatMap(r -> Flux.from(r.resultsForTable(table)).map(this::postLoad)).sort(comparator);
+
+    }
+
     private Publisher<T> getAll(AttributeValue hashKey, Publisher<AttributeValue> rangeKeys) {
         TableSchema<T> tableSchema = table.tableSchema();
         Map<AttributeValue, Integer> order = new ConcurrentHashMap<>();
@@ -377,4 +395,9 @@ private <R> Publisher<R> doWithKeys(Object partitionKey, Publisher<?> sortKeys,
         Optional<String> sortKeyName = table.tableSchema().tableMetadata().primarySortKey();
         return function.apply(partitionKeyValue, Flux.from(sortKeys).map(key -> attributeConversionHelper.convert(table, sortKeyName.get(), key)));
     }
+
+    private <R> Publisher<R> doWithKeys(Publisher<?> partitionKeys, Function<Publisher<AttributeValue>, Publisher<R>> function) {
+        String hashKeyName = table.tableSchema().tableMetadata().primaryPartitionKey();
+        return function.apply(Flux.from(partitionKeys).map(key -> attributeConversionHelper.convert(table, hashKeyName, key)));
+    }
 }
diff --git a/...src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultDynamoDbService.java b/...src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultDynamoDbService.java
@@ -224,6 +224,11 @@ public Publisher<T> getAll(Object partitionKey, Publisher<?> sortKeys) {
         return doWithKeys(partitionKey, sortKeys, this::getAll);
     }
 
+    @Override
+    public Publisher<T> getAll(Publisher<?> partitionKeys) {
+        return doWithKeys(partitionKeys, this::getAllByAttributeValue);
+    }
+
     @Override
     public T get(Key key) {
         T item = table.getItem(key);
@@ -286,6 +291,22 @@ private DetachedQuery<T> simplePartitionAndSort(Object partitionKey, Object sort
         });
     }
 
+    private Publisher<T> getAllByAttributeValue(Publisher<AttributeValue> partitionKeys) {
+        TableSchema<T> tableSchema = table.tableSchema();
+        Map<AttributeValue, Integer> order = new ConcurrentHashMap<>();
+        AtomicInteger counter = new AtomicInteger();
+
+        return Flux.from(partitionKeys).buffer(BATCH_SIZE).map(batchRangeKeys -> enhancedClient.batchGetItem(b -> b.readBatches(batchRangeKeys.stream().map(k -> {
+                order.put(k, counter.getAndIncrement());
+                return ReadBatch.builder(tableSchema.itemType().rawClass()).mappedTableResource(table).addGetItem(Key.builder().partitionValue(k).build()).build();
+            })
+            .collect(Collectors.toList())))).flatMap(r -> {
+            Comparator<T> comparator = Comparator.comparingInt(i -> order.getOrDefault(tableSchema.attributeValue(i, tableSchema.tableMetadata().primaryPartitionKey()), 0));
+            List<T> it = r.resultsForTable(table).stream().sorted(comparator).collect(Collectors.toList());
+            return Flux.fromIterable(it).map(this::postLoad);
+        });
+    }
+
     private Publisher<T> getAll(AttributeValue hashKey, Publisher<AttributeValue> rangeKeys) {
         TableSchema<T> tableSchema = table.tableSchema();
         Map<AttributeValue, Integer> order = new ConcurrentHashMap<>();
@@ -377,4 +398,9 @@ private <R> Publisher<R> doWithKeys(Object partitionKey, Publisher<?> sortKeys,
         Optional<String> sortKeyName = table.tableSchema().tableMetadata().primarySortKey();
         return function.apply(partitionKeyValue, Flux.from(sortKeys).map(key -> attributeConversionHelper.convert(table, sortKeyName.get(), key)));
     }
+
+    private <R> Publisher<R> doWithKeys(Publisher<?> partitionKeys, Function<Publisher<AttributeValue>, Publisher<R>> function) {
+        String hashKeyName = table.tableSchema().tableMetadata().primaryPartitionKey();
+        return function.apply(Flux.from(partitionKeys).map(key -> attributeConversionHelper.convert(table, hashKeyName, key)));
+    }
 }
diff --git a/...namodb/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DynamoDbService.java b/...namodb/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DynamoDbService.java
@@ -145,6 +145,8 @@ default <R> int updateAll(Publisher<T> items, Function<UpdateBuilder<T, T>, Upda
 
     Publisher<T> getAll(Object partitionKey, Publisher<?> sortKeys);
 
+    Publisher<T> getAll(Publisher<?> partitionKeys);
+
     T get(Key key);
 
     int count(DetachedQuery<T> query);

diff --git a/...java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/SyncDynamoDbServiceIntroduction.java b/...java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/SyncDynamoDbServiceIntroduction.java
@@ -232,6 +232,9 @@ private <T> Object handleGet(DynamoDbService<T> service, MethodInvocationContext
         Object partitionValue = partitionAndSort.getPartitionValue(params);
 
         if (!partitionAndSort.hasSortKey()) {
+            if (partitionAndSort.isPartitionKeyPublisherOrIterable()) {
+                return publisherOrIterable(service.getAll(partitionAndSort.getPartitionAttributeValues(conversionService, params)), context.getReturnType().getType());
+            }
             return service.get(partitionValue, null);
         }
 

diff --git a/...odb/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/annotation/HashKey.java b/...odb/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/annotation/HashKey.java
@@ -29,6 +29,7 @@
  *
  * @deprecated use @{@link PartitionKey} instead
  */
+@Deprecated
 @Inherited
 @Documented
 @Retention(RetentionPolicy.RUNTIME)

diff --git a/...db/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/annotation/RangeKey.java b/...db/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/annotation/RangeKey.java
@@ -29,6 +29,7 @@
  *
  * @deprecated use @{@link SortKey} instead
  */
+@Deprecated
 @Inherited
 @Documented
 @Retention(RetentionPolicy.RUNTIME)

diff --git a/...db/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/util/QueryArguments.java b/...db/src/main/java/com/agorapulse/micronaut/amazon/awssdk/dynamodb/util/QueryArguments.java
@@ -157,6 +157,10 @@ public Publisher<?> getSortAttributeValues(ConversionService conversionService,
         return sortKey == null ? Flux.empty() : toPublisher(conversionService, Object.class, sortKey.getFirstArgument(), params);
     }
 
+    public Publisher<?> getPartitionAttributeValues(ConversionService conversionService, Map<String, MutableArgumentValue<?>> params) {
+        return partitionKey == null ? Flux.empty() : toPublisher(conversionService, Object.class, partitionKey, params);
+    }
+
 
     public <T> Consumer<QueryBuilder<T>> generateQuery(MethodInvocationContext<Object, Object> context, ConversionService conversionService) {
         return q -> {
@@ -217,6 +221,10 @@ public boolean isSortKeyPublisherOrIterable() {
         return sortKey.getFirstArgument().getType().isArray() || Iterable.class.isAssignableFrom(sortKey.getFirstArgument().getType()) || Publisher.class.isAssignableFrom(sortKey.getFirstArgument().getType());
     }
 
+    public boolean isPartitionKeyPublisherOrIterable() {
+        return partitionKey.getType().isArray() || Iterable.class.isAssignableFrom(partitionKey.getType()) || Publisher.class.isAssignableFrom(partitionKey.getType());
+    }
+
     public boolean isCustomized() {
         return index != null || consistent || descending || !filters.isEmpty() || sortKey != null && sortKey.getOperator() != Filter.Operator.EQ || lastEvaluatedKey != null || limit != null || page != null;
     }

diff --git a/.../com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultDynamoDBServiceNoRangeSpec.groovy b/.../com/agorapulse/micronaut/amazon/awssdk/dynamodb/DefaultDynamoDBServiceNoRangeSpec.groovy
@@ -60,6 +60,7 @@ class DefaultDynamoDBServiceNoRangeSpec extends Specification {
             service.count('5') == 1
             service.get('6')
             service.count('6') == 1
+            service.getAll(['1', '2', '3', '4', '5', '6']).parentId == ['1', '2', '3', '4', '5', '6']
 
         when:
             service.increment('1')

diff --git a/.../groovy/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DynamoDBEntityNoRangeService.java b/.../groovy/com/agorapulse/micronaut/amazon/awssdk/dynamodb/DynamoDBEntityNoRangeService.java
@@ -35,6 +35,7 @@
 public interface DynamoDBEntityNoRangeService {
 
     DynamoDBEntityNoRange get(@PartitionKey String parentId);
+    List<DynamoDBEntityNoRange> getAll(@PartitionKey List<String> parentIds);
     DynamoDBEntityNoRange save(DynamoDBEntityNoRange entity);
     Flowable<DynamoDBEntityNoRange> saveAll(DynamoDBEntityNoRange... entities);