Adding pagination support

DataGateway.Service.Tests/QueryTests.cs

@@ -20,5 +20,52 @@ public async Task TestSimpleQuery()
             // Validate results
             Assert.IsFalse(response.ToString().Contains("Error"));
         }
+
+        /// <summary>
+        /// This test runs a query to list all the items in a container. Then, gets all the items by
+        /// running a paginated query that gets n items per page. We then make sure the number of documents match
+        /// </summary>
+        [TestMethod]
+        public async Task TestPaginatedQuery()
+        {
+            // Add query resolver
+            _metadataStoreProvider.StoreQueryResolver(TestHelper.SimplePaginatedQueryResolver());
+            _metadataStoreProvider.StoreQueryResolver(TestHelper.SimpleListQueryResolver());
+
+            // Run query
+            int actualElements = 0;
+            _controller.ControllerContext.HttpContext = GetHttpContextWithBody(TestHelper.SimpleListQuery);
+            using (JsonDocument fullQueryResponse = await _controller.PostAsync())
+            {
+                actualElements = fullQueryResponse.RootElement.GetProperty("data").GetProperty("queryAll").GetArrayLength();
+            }
+            // Run paginated query
+            int totalElementsFromPaginatedQuery = 0;
+            string continuationToken = "null";
+            const int pagesize = 15;
+
+            do
+            {
+                if (continuationToken != "null")
+                {
+                    // We need to append an escape quote to continuation token because of the way we are using string.format
+                    // for generating the graphql paginated query stringformat for this test.
+                    continuationToken = "\\\"" + continuationToken + "\\\"";
+                }
+
+                string paginatedQuery = string.Format(TestHelper.SimplePaginatedQueryFormat, arg0: pagesize, arg1: continuationToken);
+                _controller.ControllerContext.HttpContext = GetHttpContextWithBody(paginatedQuery);
+                using JsonDocument paginatedQueryResponse = await _controller.PostAsync();
+                JsonElement page = paginatedQueryResponse.RootElement
+                    .GetProperty("data")
+                    .GetProperty("paginatedQuery");
+                JsonElement continuation = page.GetProperty("endCursor");
+                continuationToken = continuation.ToString();
+                totalElementsFromPaginatedQuery += page.GetProperty("nodes").GetArrayLength();
+            } while (!string.IsNullOrEmpty(continuationToken));
+
+            // Validate results
+            Assert.AreEqual(actualElements, totalElementsFromPaginatedQuery);
+        }
     }
 }

DataGateway.Service.Tests/TestBase.cs

@@ -31,7 +31,8 @@ private void Init()
             string uid = Guid.NewGuid().ToString();
             dynamic sourceItem = TestHelper.GetItem(uid);
 
-            _clientProvider.Client.GetContainer(TestHelper.DB_NAME, TestHelper.COL_NAME).CreateItemAsync(sourceItem, new PartitionKey(uid));
+            _clientProvider.Client.GetContainer(TestHelper.DB_NAME, TestHelper.COL_NAME)
+                .CreateItemAsync(sourceItem, new PartitionKey(uid)).Wait();
             _metadataStoreProvider = new MetadataStoreProviderForTest();
             _queryEngine = new CosmosQueryEngine(_clientProvider, _metadataStoreProvider);
             _mutationEngine = new CosmosMutationEngine(_clientProvider, _metadataStoreProvider);

DataGateway.Service.Tests/TestHelper.cs

@@ -15,30 +15,53 @@ class TestHelper
         public static readonly string QUERY_NAME = "myQuery";
         public static readonly string MUTATION_NAME = "addPost";
         public static string GraphQLTestSchema = @"
-                                        type Query {
-                                          myQuery: MyPojo
-                                        }
-
-                                        type MyPojo {
-                                            myProp : String
-                                            id : String
-                                        }
-
-                                        type Mutation {
-                                            addPost(
-                                                myProp: String!
-                                                id : String!
-                                            ): MyPojo
-                                        }";
+            type Query {
+                myQuery: MyPojo
+                queryAll:  [MyPojo]
+                paginatedQuery(first: Int, after: String): MyPojoConnection
+            }
 
-        public static string SampleQuery = "{\"query\": \"{myQuery { myProp    }}\" } ";
+            type MyPojoConnection {
+                 nodes: [MyPojo]
+                 endCursor: String
+                 hasNextPage: Boolean
+            }
+
+            type MyPojo {
+                myProp : String
+                id : String
+            }
 
+            type Mutation {
+                addPost(
+                    myProp: String!
+                    id : String!
+                ): MyPojo
+            }";
+
+        public static string SampleQuery = "{\"query\": \"{myQuery { myProp    }}\" } ";
+        public static string SimpleListQuery = "{\"query\": \"{queryAll { myProp    }}\" } ";
         public static string SampleMutation = "{\"query\": \"mutation addPost {addPost(myProp : \\\"myValueBM \\\"id : \\\"myIdBM \\\") { myProp}}\"}";
+        public static string SimplePaginatedQueryFormat = "{{\"query\": \"{{paginatedQuery (first: {0}, after: {1}){{" +
+            " nodes{{ id  myProp }}    endCursor    hasNextPage  }} }}\" }}";
 
+        // Resolvers
         public static GraphQLQueryResolver SampleQueryResolver()
         {
             string raw =
-                "{\r\n    \"id\" : \"myQuery\",\r\n    \"databaseName\": \"" + DB_NAME + "\",\r\n    \"containerName\": \"" + COL_NAME + "\",\r\n    \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
+                "{\n    \"id\" : \"myQuery\",\n    \"databaseName\": \"" +
+                DB_NAME + "\",\n    \"containerName\": \"" + COL_NAME +
+                "\",\n    \"parametrizedQuery\": \"SELECT * FROM r\"\n}";
+
+            return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
+        }
+
+        public static GraphQLQueryResolver SimpleListQueryResolver()
+        {
+            string raw =
+                "{\r\n    \"id\" : \"queryAll\",\r\n    \"databaseName\": \"" +
+                DB_NAME + "\",\r\n    \"containerName\": \"" + COL_NAME +
+                "\",\r\n    \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
 
             return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
         }
@@ -50,6 +73,17 @@ public static MutationResolver SampleMutationResolver()
             return JsonConvert.DeserializeObject<MutationResolver>(raw);
         }
 
+        public static GraphQLQueryResolver SimplePaginatedQueryResolver()
+        {
+            string raw =
+                "{\r\n    \"id\" : \"paginatedQuery\",\r\n    \"databaseName\": \"" +
+                DB_NAME + "\",\r\n    \"containerName\": \"" + COL_NAME +
+                "\" ,\n \"isPaginated\" : true " +
+                ",\r\n    \"parametrizedQuery\": \"SELECT * FROM r\"\r\n}";
+
+            return JsonConvert.DeserializeObject<GraphQLQueryResolver>(raw);
+        }
+
         private static Lazy<IOptions<DataGatewayConfig>> _dataGatewayConfig = new(() => TestHelper.LoadConfig());
 
         private static IOptions<DataGatewayConfig> LoadConfig()

DataGateway.Service/Models/GraphQLQueryResolver.cs

@@ -10,6 +10,7 @@ public class GraphQLQueryResolver
         public string ParametrizedQuery { get; set; }
         public QuerySpec QuerySpec { get; set; }
         public bool IsList { get; set; }
+        public bool IsPaginated { get; set; }
     }
 
     public class QuerySpec

DataGateway.Service/Resolvers/CosmosQueryEngine.cs

@@ -31,7 +31,7 @@ public CosmosQueryEngine(CosmosClientProvider clientProvider, IMetadataStoreProv
         /// Executes the given IMiddlewareContext of the GraphQL query and
         /// expecting a single Json back.
         /// </summary>
-        public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters)
+        public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isPaginatedQuery)
         {
             // TODO: fixme we have multiple rounds of serialization/deserialization JsomDocument/JObject
             // TODO: add support for nesting
@@ -41,6 +41,10 @@ public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictio
             string graphQLQueryName = context.Selection.Field.Name.Value;
             GraphQLQueryResolver resolver = this._metadataStoreProvider.GetQueryResolver(graphQLQueryName);
             Container container = this._clientProvider.Client.GetDatabase(resolver.DatabaseName).GetContainer(resolver.ContainerName);
+
+            QueryRequestOptions queryRequestOptions = new();
+            string requestContinuation = null;
+
             QueryDefinition querySpec = new(resolver.ParametrizedQuery);
 
             if (parameters != null)
@@ -51,7 +55,42 @@ public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictio
                 }
             }
 
-            FeedResponse<JObject> firstPage = await container.GetItemQueryIterator<JObject>(querySpec).ReadNextAsync();
+            if (parameters.TryGetValue("first", out object maxSize))
+            {
+                queryRequestOptions.MaxItemCount = Convert.ToInt32(maxSize);
+            }
+
+            if (parameters.TryGetValue("after", out object after))
+            {
+                requestContinuation = Base64Decode(after as string);
+            }
+
+            FeedResponse<JObject> firstPage = await container.GetItemQueryIterator<JObject>(querySpec, requestContinuation, queryRequestOptions).ReadNextAsync();
+
+            if (isPaginatedQuery)
+            {
+                JArray jarray = new();
+                IEnumerator<JObject> enumerator = firstPage.GetEnumerator();
+                while (enumerator.MoveNext())
+                {
+                    JObject item = enumerator.Current;
+                    jarray.Add(item);
+                }
+
+                string responseContinuation = firstPage.ContinuationToken;
+                if (string.IsNullOrEmpty(responseContinuation))
+                {
+                    responseContinuation = null;
+                }
+
+                JObject res = new(
+                   new JProperty("endCursor", Base64Encode(responseContinuation)),
+                   new JProperty("hasNextPage", responseContinuation != null),
+                   new JProperty("nodes", jarray));
+
+                // This extra deserialize/serialization will be removed after moving to Newtonsoft from System.Text.Json
+                return JsonDocument.Parse(res.ToString());
+            }
 
             JObject firstItem = null;
 
@@ -111,5 +150,28 @@ public Task<JsonDocument> ExecuteAsync(FindRequestContext queryStructure)
         {
             throw new NotImplementedException();
         }
+
+        private static string Base64Encode(string plainText)
+        {
+            if (plainText == default)
+            {
+                return null;
+            }
+
+            byte[] plainTextBytes = System.Text.Encoding.UTF8.GetBytes(plainText);
+            return System.Convert.ToBase64String(plainTextBytes);
+        }
+
+        private static string Base64Decode(string base64EncodedData)
+        {
+            if (base64EncodedData == default)
+            {
+                return null;
+            }
+
+            byte[] base64EncodedBytes = System.Convert.FromBase64String(base64EncodedData);
+            return System.Text.Encoding.UTF8.GetString(base64EncodedBytes);
+        }
+
     }
 }

DataGateway.Service/Resolvers/IQueryEngine.cs

@@ -14,7 +14,7 @@ public interface IQueryEngine
         /// Executes the given IMiddlewareContext of the GraphQL query and
         /// expecting a single Json back.
         /// </summary>
-        public Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters);
+        public Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isPaginationQuery);
 
         /// <summary>
         /// Executes the given IMiddlewareContext of the GraphQL and expecting a

DataGateway.Service/Resolvers/SqlQueryEngine.cs

@@ -50,7 +50,7 @@ public static async Task<string> GetJsonStringFromDbReader(DbDataReader dbDataRe
         /// Executes the given IMiddlewareContext of the GraphQL query and
         /// expecting a single Json back.
         /// </summary>
-        public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters)
+        public async Task<JsonDocument> ExecuteAsync(IMiddlewareContext context, IDictionary<string, object> parameters, bool isContinuationQuery)
         {
             SqlQueryStructure structure = new(context, parameters, _metadataStoreProvider, _queryBuilder);
             return await ExecuteAsync(structure);

DataGateway.Service/Services/GraphQLService.cs

@@ -33,7 +33,7 @@ public void ParseAsync(String data)
             ISchema schema = SchemaBuilder.New()
                .AddDocumentFromString(data)
                .AddAuthorizeDirectiveType()
-               .Use((services, next) => new ResolverMiddleware(next, _queryEngine, _mutationEngine))
+               .Use((services, next) => new ResolverMiddleware(next, _queryEngine, _mutationEngine, _metadataStoreProvider))
                .Create();
 
             // Below is pretty much an inlined version of

DataGateway.Service/Services/ResolverMiddleware.cs

@@ -1,6 +1,8 @@
+using System;
 using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading.Tasks;
+using Azure.DataGateway.Service.Models;
 using Azure.DataGateway.Service.Resolvers;
 using HotChocolate.Language;
 using HotChocolate.Resolvers;
@@ -17,12 +19,17 @@ public class ResolverMiddleware
         private readonly FieldDelegate _next;
         private readonly IQueryEngine _queryEngine;
         private readonly IMutationEngine _mutationEngine;
+        private readonly IMetadataStoreProvider _metadataStoreProvider;
 
-        public ResolverMiddleware(FieldDelegate next, IQueryEngine queryEngine, IMutationEngine mutationEngine)
+        public ResolverMiddleware(FieldDelegate next,
+            IQueryEngine queryEngine,
+            IMutationEngine mutationEngine,
+            IMetadataStoreProvider metadataStoreProvider)
         {
             _next = next;
             _queryEngine = queryEngine;
             _mutationEngine = mutationEngine;
+            _metadataStoreProvider = metadataStoreProvider;
         }
 
         public ResolverMiddleware(FieldDelegate next)
@@ -42,14 +49,15 @@ public async Task InvokeAsync(IMiddlewareContext context)
             else if (context.Selection.Field.Coordinate.TypeName.Value == "Query")
             {
                 IDictionary<string, object> parameters = GetParametersFromContext(context);
+                bool isPaginatedQuery = IsPaginatedQuery(context.Selection.Field.Name.Value);
 
                 if (context.Selection.Type.IsListType())
                 {
                     context.Result = await _queryEngine.ExecuteListAsync(context, parameters);
                 }
                 else
                 {
-                    context.Result = await _queryEngine.ExecuteAsync(context, parameters);
+                    context.Result = await _queryEngine.ExecuteAsync(context, parameters, isPaginatedQuery);
                 }
             }
             else if (context.Selection.Field.Type.IsLeafType())
@@ -89,6 +97,23 @@ public async Task InvokeAsync(IMiddlewareContext context)
             await _next(context);
         }
 
+        /// <summary>
+        /// Identifies if a query is paginated or not by checking the IsPaginated param on the respective resolver.
+        /// </summary>
+        /// <param name="queryName the name of the query"></param>
+        /// <returns></returns>
+        private bool IsPaginatedQuery(string queryName)
+        {
+            GraphQLQueryResolver resolver = _metadataStoreProvider.GetQueryResolver(queryName);
+            if (resolver == null)
+            {
+                string message = string.Format("There is no resolver for the query: {0}", queryName);
+                throw new InvalidOperationException(message);
+            }
+
+            return resolver.IsPaginated;
+        }
+
         private static bool TryGetPropertyFromParent(IMiddlewareContext context, out JsonElement jsonElement)
         {
             JsonDocument result = context.Parent<JsonDocument>();

DataGateway.Service/cosmos-config.json

@@ -24,6 +24,13 @@
       "databaseName": "graphqldb",
       "containerName": "planet",
       "parametrizedQuery": "select * FROM c"
+    },
+    {
+      "id": "planets",
+      "isPaginated":  true,
+      "databaseName": "graphqldb",
+      "containerName": "planet",
+      "parametrizedQuery": "select * FROM c"
     }
   ],
   "MutationResolvers": [

DataGateway.Service/schema.gql

@@ -3,8 +3,15 @@ type Query {
     characterById (id : ID!): Character
     planetById (id: ID! = 1): Planet
     planetList: [Planet]
+    planets(first: Int, after: String): PlanetConnection
     }
 
+type PlanetConnection {
+        nodes: [Planet]
+        endCursor: String
+        hasNextPage: Boolean
+}
+
 type Character {
         id : ID,
         name : String,