Commit 16b6bdbe authored by Aliaksei Darafeyeu's avatar Aliaksei Darafeyeu
Browse files

adds queryItemsWithContinuationToken to CosmosStore

parent 8683e419
Pipeline #9586 failed with stage
in 1 minute and 59 seconds
......@@ -21,7 +21,7 @@
<groupId>org.opengroup.osdu</groupId>
<artifactId>core-lib-azure</artifactId>
<packaging>jar</packaging>
<version>0.0.29</version>
<version>0.0.30</version>
<name>core-lib-azure</name>
<properties>
......
package org.opengroup.osdu.azure;
import java.util.List;
import lombok.AllArgsConstructor;
import lombok.Data;
/**
* CosmosPagingResult class.
* @param <T> Type to return
*/
@Data
@AllArgsConstructor
public class CosmosPagingResult<T> {
private List<T> items;
private String continuationToken;
}
......@@ -26,6 +26,7 @@ import com.azure.cosmos.NotFoundException;
import com.azure.cosmos.SqlQuerySpec;
import com.azure.cosmos.internal.AsyncDocumentClient;
import com.azure.cosmos.internal.Document;
import org.opengroup.osdu.core.common.model.http.AppException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
......@@ -186,11 +187,34 @@ public class CosmosStore {
final SqlQuerySpec query,
final FeedOptions options,
final Class<T> clazz) {
ArrayList<T> results = new ArrayList<>();
return queryItemsWithContinuationToken(dataPartitionId, cosmosDBName, collection, query, options, clazz).getItems();
}
/**
* @param dataPartitionId Data partition id to fetch appropriate cosmos client for each partition
* @param cosmosDBName Database to be used
* @param collection Collection to be used
* @param query {@link SqlQuerySpec} to execute
* @param options Query options
* @param clazz Class type of response
* @param <T> Type of response
* @return Pair of continuationToken and items found in container
*/
public <T> CosmosPagingResult<T> queryItemsWithContinuationToken(
final String dataPartitionId,
final String cosmosDBName,
final String collection,
final SqlQuerySpec query,
final FeedOptions options,
final Class<T> clazz) {
List<T> results = new ArrayList<>();
String continuationToken = null;
CosmosContainer cosmosContainer = getCosmosContainer(dataPartitionId, cosmosDBName, collection);
Iterator<FeedResponse<CosmosItemProperties>> paginatedResponse = cosmosContainer.queryItems(query, options);
while (paginatedResponse.hasNext()) {
for (CosmosItemProperties properties : paginatedResponse.next().getResults()) {
FeedResponse<CosmosItemProperties> next = paginatedResponse.next();
continuationToken = next.getContinuationToken();
for (CosmosItemProperties properties : next.getResults()) {
try {
results.add(properties.getObject(clazz));
} catch (IOException e) {
......@@ -200,8 +224,10 @@ public class CosmosStore {
}
}
}
return results;
return new CosmosPagingResult<>(results, continuationToken);
}
/**
* @param dataPartitionId Data partition id to fetch appropriate cosmos client for each partition
* @param cosmosDBName Database to be used
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment