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 @@ ...@@ -21,7 +21,7 @@
<groupId>org.opengroup.osdu</groupId> <groupId>org.opengroup.osdu</groupId>
<artifactId>core-lib-azure</artifactId> <artifactId>core-lib-azure</artifactId>
<packaging>jar</packaging> <packaging>jar</packaging>
<version>0.0.29</version> <version>0.0.30</version>
<name>core-lib-azure</name> <name>core-lib-azure</name>
<properties> <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; ...@@ -26,6 +26,7 @@ import com.azure.cosmos.NotFoundException;
import com.azure.cosmos.SqlQuerySpec; import com.azure.cosmos.SqlQuerySpec;
import com.azure.cosmos.internal.AsyncDocumentClient; import com.azure.cosmos.internal.AsyncDocumentClient;
import com.azure.cosmos.internal.Document; import com.azure.cosmos.internal.Document;
import org.opengroup.osdu.core.common.model.http.AppException; import org.opengroup.osdu.core.common.model.http.AppException;
import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy; import org.springframework.context.annotation.Lazy;
...@@ -186,11 +187,34 @@ public class CosmosStore { ...@@ -186,11 +187,34 @@ public class CosmosStore {
final SqlQuerySpec query, final SqlQuerySpec query,
final FeedOptions options, final FeedOptions options,
final Class<T> clazz) { 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); CosmosContainer cosmosContainer = getCosmosContainer(dataPartitionId, cosmosDBName, collection);
Iterator<FeedResponse<CosmosItemProperties>> paginatedResponse = cosmosContainer.queryItems(query, options); Iterator<FeedResponse<CosmosItemProperties>> paginatedResponse = cosmosContainer.queryItems(query, options);
while (paginatedResponse.hasNext()) { while (paginatedResponse.hasNext()) {
for (CosmosItemProperties properties : paginatedResponse.next().getResults()) { FeedResponse<CosmosItemProperties> next = paginatedResponse.next();
continuationToken = next.getContinuationToken();
for (CosmosItemProperties properties : next.getResults()) {
try { try {
results.add(properties.getObject(clazz)); results.add(properties.getObject(clazz));
} catch (IOException e) { } catch (IOException e) {
...@@ -200,8 +224,10 @@ public class CosmosStore { ...@@ -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 dataPartitionId Data partition id to fetch appropriate cosmos client for each partition
* @param cosmosDBName Database to be used * @param cosmosDBName Database to be used
......
Supports Markdown
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