Skip to content

Commit 40196e9

Browse files
mbhaskarmbhaskar
authored
Adding pagination support (#59)
* Pagination cleanup and refactoring * Fixing formatting * Adding tests for pagination * Resolving merge conflicts Pr comments * Cleanup * Fixing flaky test cleanup * Paginated query resolver now expects isPaginated to be set to true Refactoring tests cosmos continuation token is now base64encoded. * fixing test * reformat and cleanup * Handling IDisposable objects properly Cleanup
1 parent 9ee03c0 commit 40196e9

File tree

11 files changed

+209
-25
lines changed

11 files changed

+209
-25
lines changed

‎DataGateway.Service.Tests/QueryTests.cs‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,5 +20,52 @@ public async Task TestSimpleQuery()
2020
// Validate results
2121
Assert.IsFalse(response.ToString().Contains("Error"));
2222
}
23+
24+
/// <summary>
25+
/// This test runs a query to list all the items in a container. Then, gets all the items by
26+
/// running a paginated query that gets n items per page. We then make sure the number of documents match
27+
/// </summary>
28+
[TestMethod]
29+
public async Task TestPaginatedQuery()
30+
{
31+
// Add query resolver
32+
_metadataStoreProvider.StoreQueryResolver(TestHelper.SimplePaginatedQueryResolver());
33+
_metadataStoreProvider.StoreQueryResolver(TestHelper.SimpleListQueryResolver());
34+
35+
// Run query
36+
int actualElements = 0;
37+
_controller.ControllerContext.HttpContext = GetHttpContextWithBody(TestHelper.SimpleListQuery);
38+
using (JsonDocument fullQueryResponse = await _controller.PostAsync())
39+
{
40+
actualElements = fullQueryResponse.RootElement.GetProperty("data").GetProperty("queryAll").GetArrayLength();
41+
}
42+
// Run paginated query
43+
int totalElementsFromPaginatedQuery = 0;
44+
string continuationToken = "null";
45+
const int pagesize = 15;
46+
47+
do
48+
{
49+
if (continuationToken != "null")
50+
{
51+
// We need to append an escape quote to continuation token because of the way we are using string.format
52+
// for generating the graphql paginated query stringformat for this test.
53+
continuationToken = "\\\"" + continuationToken + "\\\"";
54+
}
55+
56+
string paginatedQuery = string.Format(TestHelper.SimplePaginatedQueryFormat, arg0: pagesize, arg1: continuationToken);
57+
_controller.ControllerContext.HttpContext = GetHttpContextWithBody(paginatedQuery);
58+
using JsonDocument paginatedQueryResponse = await _controller.PostAsync();
59+
JsonElement page = paginatedQueryResponse.RootElement
60+
.GetProperty("data")
61+
.GetProperty("paginatedQuery");
62+
JsonElement continuation = page.GetProperty("endCursor");
63+
continuationToken = continuation.ToString();
64+
totalElementsFromPaginatedQuery += page.GetProperty("nodes").GetArrayLength();
65+
} while (!string.IsNullOrEmpty(continuationToken));
66+
67+
// Validate results
68+
Assert.AreEqual(actualElements, totalElementsFromPaginatedQuery);
69+
}
2370
}
2471
}

‎DataGateway.Service.Tests/TestBase.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ private void Init()
3131
string uid = Guid.NewGuid().ToString();
3232
dynamic sourceItem = TestHelper.GetItem(uid);
3333

34-
_clientProvider.Client.GetContainer(TestHelper.DB_NAME, TestHelper.COL_NAME).CreateItemAsync(sourceItem, new PartitionKey(uid));
34+
_clientProvider.Client.GetContainer(TestHelper.DB_NAME, TestHelper.COL_NAME)
35+
.CreateItemAsync(sourceItem, new PartitionKey(uid)).Wait();
3536
_metadataStoreProvider = new MetadataStoreProviderForTest();
3637
_queryEngine = new CosmosQueryEngine(_clientProvider, _metadataStoreProvider);
3738
_mutationEngine = new CosmosMutationEngine(_clientProvider, _metadataStoreProvider);

‎DataGateway.Service.Tests/TestHelper.cs‎

Lines changed: 51 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,30 +15,53 @@ class TestHelper
1515
public static readonly string QUERY_NAME = "myQuery";
1616
public static readonly string MUTATION_NAME = "addPost";
1717
public static string GraphQLTestSchema = @"
18-
type Query {
19-
myQuery: MyPojo
20-
}
21-
22-
type MyPojo {
23-
myProp : String
24-
id : String
25-
}
26-
27-
type Mutation {
28-
addPost(
29-
myProp: String!
30-
id : String!
31-
): MyPojo
32-
}";
18+
type Query {
19+
myQuery: MyPojo
20+
queryAll: [MyPojo]
21+
paginatedQuery(first: Int, after: String): MyPojoConnection
22+
}
3323
34-
public static string SampleQuery = "{\"query\": \"{myQuery { myProp }}\" } ";
24+
type MyPojoConnection {
25+
nodes: [MyPojo]
26+
endCursor: String
27+
hasNextPage: Boolean
28+
}
29+
30+
type MyPojo {
31+
myProp : String
32+
id : String
33+
}
3534
35+
type Mutation {
36+
addPost(
37+
myProp: String!
38+
id : String!
39+
): MyPojo
40+
}";
41+
42+
public static string SampleQuery = "{\"query\": \"{myQuery { myProp }}\" } ";
43+
public static string SimpleListQuery = "{\"query\": \"{queryAll { myProp }}\" } ";
3644
public static string SampleMutation = "{\"query\": \"mutation addPost {addPost(myProp : \\\"myValueBM \\\"id : \\\"myIdBM \\\") { myProp}}\"}";
45+
public static string SimplePaginatedQueryFormat = "{{\"query\": \"{{paginatedQuery (first: {0}, after: {1}){{" +
46+
" nodes{{ id myProp }} endCursor hasNextPage }} }}\" }}";
3747

48+
// Resolvers
3849
public static GraphQLQueryResolver SampleQueryResolver()
3950
{
4051
string raw =
41-
"{\r\n \"id\" : \"myQuery\",\r\n \"databaseName\": \"" + DB_NAME + "\",\r\n \"containerName\": \"" + COL_NAME + "\",\r\n \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
52+
"{\n \"id\" : \"myQuery\",\n \"databaseName\": \"" +
53+
DB_NAME + "\",\n \"containerName\": \"" + COL_NAME +
54+
"\",\n \"parametrizedQuery\": \"SELECT * FROM r\"\n}";
55+
56+
return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
57+
}
58+
59+
public static GraphQLQueryResolver SimpleListQueryResolver()
60+
{
61+
string raw =
62+
"{\r\n \"id\" : \"queryAll\",\r\n \"databaseName\": \"" +
63+
DB_NAME + "\",\r\n \"containerName\": \"" + COL_NAME +
64+
"\",\r\n \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
4265

4366
return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
4467
}
@@ -50,6 +73,17 @@ public static MutationResolver SampleMutationResolver()
5073
return JsonConvert.DeserializeObject<MutationResolver>(raw);
5174
}
5275

76+
public static GraphQLQueryResolver SimplePaginatedQueryResolver()
77+
{
78+
string raw =
79+
"{\r\n \"id\" : \"paginatedQuery\",\r\n \"databaseName\": \"" +
80+
DB_NAME + "\",\r\n \"containerName\": \"" + COL_NAME +
81+
"\" ,\n \"isPaginated\" : true " +
82+
",\r\n \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
83+
84+
return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
85+
}
86+
5387
private static Lazy<IOptions<DataGatewayConfig>> _dataGatewayConfig = new(() => TestHelper.LoadConfig());
5488

5589
private static IOptions<DataGatewayConfig> LoadConfig()

‎DataGateway.Service/Models/GraphQLQueryResolver.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ public class GraphQLQueryResolver
1010
public string ParametrizedQuery { get; set; }
1111
public QuerySpec QuerySpec { get; set; }
1212
public bool IsList { get; set; }
13+
public bool IsPaginated { get; set; }
1314
}
1415

1516
public class QuerySpec

‎DataGateway.Service/Resolvers/CosmosQueryEngine.cs‎

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ public CosmosQueryEngine(CosmosClientProvider clientProvider, IMetadataStoreProv
3131
/// Executes the given IMiddlewareContext of the GraphQL query and
3232
/// expecting a single Json back.
3333
/// </summary>
34-
public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters)
34+
public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isPaginatedQuery)
3535
{
3636
// TODO: fixme we have multiple rounds of serialization/deserialization JsomDocument/JObject
3737
// TODO: add support for nesting
@@ -41,6 +41,10 @@ public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictio
4141
string graphQLQueryName = context.Selection.Field.Name.Value;
4242
GraphQLQueryResolver resolver = this._metadataStoreProvider.GetQueryResolver(graphQLQueryName);
4343
Container container = this._clientProvider.Client.GetDatabase(resolver.DatabaseName).GetContainer(resolver.ContainerName);
44+
45+
QueryRequestOptions queryRequestOptions = new();
46+
string requestContinuation = null;
47+
4448
QueryDefinition querySpec = new(resolver.ParametrizedQuery);
4549

4650
if (parameters != null)
@@ -51,7 +55,42 @@ public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictio
5155
}
5256
}
5357

54-
FeedResponse<JObject> firstPage = await container.GetItemQueryIterator<JObject>(querySpec).ReadNextAsync();
58+
if (parameters.TryGetValue("first", out object maxSize))
59+
{
60+
queryRequestOptions.MaxItemCount = Convert.ToInt32(maxSize);
61+
}
62+
63+
if (parameters.TryGetValue("after", out object after))
64+
{
65+
requestContinuation = Base64Decode(after as string);
66+
}
67+
68+
FeedResponse<JObject> firstPage = await container.GetItemQueryIterator<JObject>(querySpec, requestContinuation, queryRequestOptions).ReadNextAsync();
69+
70+
if (isPaginatedQuery)
71+
{
72+
JArray jarray = new();
73+
IEnumerator<JObject> enumerator = firstPage.GetEnumerator();
74+
while (enumerator.MoveNext())
75+
{
76+
JObject item = enumerator.Current;
77+
jarray.Add(item);
78+
}
79+
80+
string responseContinuation = firstPage.ContinuationToken;
81+
if (string.IsNullOrEmpty(responseContinuation))
82+
{
83+
responseContinuation = null;
84+
}
85+
86+
JObject res = new(
87+
new JProperty("endCursor", Base64Encode(responseContinuation)),
88+
new JProperty("hasNextPage", responseContinuation != null),
89+
new JProperty("nodes", jarray));
90+
91+
// This extra deserialize/serialization will be removed after moving to Newtonsoft from System.Text.Json
92+
return JsonDocument.Parse(res.ToString());
93+
}
5594

5695
JObject firstItem = null;
5796

@@ -111,5 +150,28 @@ public Task<JsonDocument> ExecuteAsync(FindRequestContext queryStructure)
111150
{
112151
throw new NotImplementedException();
113152
}
153+
154+
private static string Base64Encode(string plainText)
155+
{
156+
if (plainText == default)
157+
{
158+
return null;
159+
}
160+
161+
byte[] plainTextBytes = System.Text.Encoding.UTF8.GetBytes(plainText);
162+
return System.Convert.ToBase64String(plainTextBytes);
163+
}
164+
165+
private static string Base64Decode(string base64EncodedData)
166+
{
167+
if (base64EncodedData == default)
168+
{
169+
return null;
170+
}
171+
172+
byte[] base64EncodedBytes = System.Convert.FromBase64String(base64EncodedData);
173+
return System.Text.Encoding.UTF8.GetString(base64EncodedBytes);
174+
}
175+
114176
}
115177
}

‎DataGateway.Service/Resolvers/IQueryEngine.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public interface IQueryEngine
1414
/// Executes the given IMiddlewareContext of the GraphQL query and
1515
/// expecting a single Json back.
1616
/// </summary>
17-
public Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters);
17+
public Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isPaginationQuery);
1818

1919
/// <summary>
2020
/// Executes the given IMiddlewareContext of the GraphQL and expecting a

‎DataGateway.Service/Resolvers/SqlQueryEngine.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public static async Task<string> GetJsonStringFromDbReader(DbDataReader dbDataRe
5050
/// Executes the given IMiddlewareContext of the GraphQL query and
5151
/// expecting a single Json back.
5252
/// </summary>
53-
public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters)
53+
public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isContinuationQuery)
5454
{
5555
SqlQueryStructure structure = new(context, parameters, _metadataStoreProvider, _queryBuilder);
5656
return await ExecuteAsync(structure);

‎DataGateway.Service/Services/GraphQLService.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ public void ParseAsync(String data)
3333
ISchema schema = SchemaBuilder.New()
3434
.AddDocumentFromString(data)
3535
.AddAuthorizeDirectiveType()
36-
.Use((services, next) => new ResolverMiddleware(next, _queryEngine, _mutationEngine))
36+
.Use((services, next) => new ResolverMiddleware(next, _queryEngine, _mutationEngine, _metadataStoreProvider))
3737
.Create();
3838

3939
// Below is pretty much an inlined version of

‎DataGateway.Service/Services/ResolverMiddleware.cs‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1+
using System;
12
using System.Collections.Generic;
23
using System.Text.Json;
34
using System.Threading.Tasks;
5+
using Azure.DataGateway.Service.Models;
46
using Azure.DataGateway.Service.Resolvers;
57
using HotChocolate.Language;
68
using HotChocolate.Resolvers;
@@ -17,12 +19,17 @@ public class ResolverMiddleware
1719
private readonly FieldDelegate _next;
1820
private readonly IQueryEngine _queryEngine;
1921
private readonly IMutationEngine _mutationEngine;
22+
private readonly IMetadataStoreProvider _metadataStoreProvider;
2023

21-
public ResolverMiddleware(FieldDelegate next, IQueryEngine queryEngine, IMutationEngine mutationEngine)
24+
public ResolverMiddleware(FieldDelegate next,
25+
IQueryEngine queryEngine,
26+
IMutationEngine mutationEngine,
27+
IMetadataStoreProvider metadataStoreProvider)
2228
{
2329
_next = next;
2430
_queryEngine = queryEngine;
2531
_mutationEngine = mutationEngine;
32+
_metadataStoreProvider = metadataStoreProvider;
2633
}
2734

2835
public ResolverMiddleware(FieldDelegate next)
@@ -42,14 +49,15 @@ public async Task InvokeAsync(IMiddlewareContext context)
4249
else if (context.Selection.Field.Coordinate.TypeName.Value == "Query")
4350
{
4451
IDictionary<string, object> parameters = GetParametersFromContext(context);
52+
bool isPaginatedQuery = IsPaginatedQuery(context.Selection.Field.Name.Value);
4553

4654
if (context.Selection.Type.IsListType())
4755
{
4856
context.Result = await _queryEngine.ExecuteListAsync(context, parameters);
4957
}
5058
else
5159
{
52-
context.Result = await _queryEngine.ExecuteAsync(context, parameters);
60+
context.Result = await _queryEngine.ExecuteAsync(context, parameters, isPaginatedQuery);
5361
}
5462
}
5563
else if (context.Selection.Field.Type.IsLeafType())
@@ -89,6 +97,23 @@ public async Task InvokeAsync(IMiddlewareContext context)
8997
await _next(context);
9098
}
9199

100+
/// <summary>
101+
/// Identifies if a query is paginated or not by checking the IsPaginated param on the respective resolver.
102+
/// </summary>
103+
/// <param name="queryName the name of the query"></param>
104+
/// <returns></returns>
105+
private bool IsPaginatedQuery(string queryName)
106+
{
107+
GraphQLQueryResolver resolver = _metadataStoreProvider.GetQueryResolver(queryName);
108+
if (resolver == null)
109+
{
110+
string message = string.Format("There is no resolver for the query: {0}", queryName);
111+
throw new InvalidOperationException(message);
112+
}
113+
114+
return resolver.IsPaginated;
115+
}
116+
92117
private static bool TryGetPropertyFromParent(IMiddlewareContext context, out JsonElement jsonElement)
93118
{
94119
JsonDocument result = context.Parent<JsonDocument>();

‎DataGateway.Service/cosmos-config.json‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,13 @@
2424
"databaseName": "graphqldb",
2525
"containerName": "planet",
2626
"parametrizedQuery": "select * FROM c"
27+
},
28+
{
29+
"id": "planets",
30+
"isPaginated": true,
31+
"databaseName": "graphqldb",
32+
"containerName": "planet",
33+
"parametrizedQuery": "select * FROM c"
2734
}
2835
],
2936
"MutationResolvers": [

0 commit comments

Comments
 (0)