Add DataLoader functionality (#1566)

ballerina-tests/Dependencies.toml

@@ -76,16 +76,19 @@ dependencies = [
 	{org = "ballerina", name = "jballerina.java"},
 	{org = "ballerina", name = "jwt"},
 	{org = "ballerina", name = "lang.array"},
+	{org = "ballerina", name = "lang.regexp"},
 	{org = "ballerina", name = "lang.value"},
 	{org = "ballerina", name = "log"},
 	{org = "ballerina", name = "mime"},
 	{org = "ballerina", name = "oauth2"},
 	{org = "ballerina", name = "task"},
 	{org = "ballerina", name = "time"},
+	{org = "ballerina", name = "uuid"},
 	{org = "ballerina", name = "websocket"}
 ]
 modules = [
 	{org = "ballerina", packageName = "graphql", moduleName = "graphql"},
+	{org = "ballerina", packageName = "graphql", moduleName = "graphql.dataloader"},
 	{org = "ballerina", packageName = "graphql", moduleName = "graphql.parser"},
 	{org = "ballerina", packageName = "graphql", moduleName = "graphql.subgraph"}
 ]
@@ -367,7 +370,6 @@ modules = [
 org = "ballerina"
 name = "uuid"
 version = "1.6.0"
-scope = "testOnly"
 dependencies = [
 	{org = "ballerina", name = "crypto"},
 	{org = "ballerina", name = "jballerina.java"},

ballerina-tests/tests/41_dataloader.bal

@@ -0,0 +1,150 @@
+// Copyright (c) 2023 WSO2 LLC. (http://www.wso2.com) All Rights Reserved.
+//
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import ballerina/graphql;
+import ballerina/websocket;
+import ballerina/test;
+
+@test:Config {
+    groups: ["dataloader", "query"],
+    after: resetDispatchCounters
+}
+isolated function testDataLoaderWithQuery() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader");
+    string document = check getGraphqlDocumentFromFile("dataloader_with_query");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("dataloader_with_query");
+    assertJsonValuesWithOrder(response, expectedPayload);
+    assertDispatchCountForAuthorLoader(1);
+    assertDispatchCountForBookLoader(1);
+}
+
+@test:Config {
+    groups: ["dataloader", "query"],
+    after: resetDispatchCounters
+}
+isolated function testDataLoaderWithDifferentAliasForSameField() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader");
+    string document = check getGraphqlDocumentFromFile("dataloader_with_different_alias_for_same_field");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("dataloader_with_different_alias_for_same_field");
+    assertJsonValuesWithOrder(response, expectedPayload);
+    assertDispatchCountForAuthorLoader(1);
+    assertDispatchCountForBookLoader(1);
+}
+
+@test:Config {
+    groups: ["dataloader", "subscription"],
+    after: resetDispatchCounters
+}
+isolated function testDataLoaderWithSubscription() returns error? {
+    string document = check getGraphqlDocumentFromFile("dataloader_with_subscription");
+    websocket:ClientConfiguration config = {subProtocols: [GRAPHQL_TRANSPORT_WS]};
+    websocket:Client wsClient = check new ("ws://localhost:9090/dataloader", config);
+    check initiateGraphqlWsConnection(wsClient);
+    check sendSubscriptionMessage(wsClient, document, "1");
+    json[] authorSequence = check getJsonContentFromFile("data_loader_with_subscription").ensureType();
+    foreach int i in 0 ..< 5 {
+        json expectedMsgPayload = {data: {authors: authorSequence[i]}};
+        check validateNextMessage(wsClient, expectedMsgPayload, id = "1");
+    }
+    assertDispatchCountForBookLoader(5);
+}
+
+@test:Config {
+    groups: ["dataloader", "mutation"],
+    dependsOn: [testDataLoaderWithQuery, testDataLoaderWithSubscription],
+    after: resetDispatchCounters
+}
+isolated function testDataLoaderWithMutation() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader");
+    string document = check getGraphqlDocumentFromFile("dataloader_with_mutation");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("dataloader_with_mutation");
+    assertJsonValuesWithOrder(response, expectedPayload);
+    assertDispatchCountForUpdateAuthorLoader(1);
+    assertDispatchCountForBookLoader(1);
+}
+
+@test:Config {
+    groups: ["dataloader", "interceptor"],
+    after: resetDispatchCounters
+}
+isolated function testDataLoaderWithInterceptors() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader_with_interceptor");
+    string document = check getGraphqlDocumentFromFile("dataloader_with_interceptor");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("dataloader_with_interceptor");
+    assertJsonValuesWithOrder(response, expectedPayload);
+    assertDispatchCountForAuthorLoader(1);
+    assertDispatchCountForBookLoader(1);
+}
+
+@test:Config {
+    groups: ["dataloader", "dispatch-error"],
+    after: resetDispatchCounters
+}
+isolated function testBatchFunctionReturningErrors() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader");
+    string document = check getGraphqlDocumentFromFile("batch_function_returing_errors");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("batch_function_returing_errors");
+    assertJsonValuesWithOrder(response, expectedPayload);
+    assertDispatchCountForAuthorLoader(1);
+    assertDispatchCountForBookLoader(0);
+}
+
+@test:Config {
+    groups: ["dataloader", "dispatch-error"],
+    after: resetDispatchCounters
+}
+isolated function testBatchFunctionReturingNonMatchingNumberOfResults() returns error? {
+    graphql:Client graphqlClient = check new ("localhost:9090/dataloader_with_faulty_batch_function");
+    string document = check getGraphqlDocumentFromFile("batch_function_returning_non_matcing_number_of_results");
+    json response = check graphqlClient->execute(document);
+    json expectedPayload = check getJsonContentFromFile("batch_function_returning_non_matcing_number_of_results");
+    assertJsonValuesWithOrder(response, expectedPayload);
+}
+
+isolated function resetDispatchCounters() {
+    lock {
+        dispatchCountOfAuthorLoader = 0;
+    }
+    lock {
+        dispatchCountOfBookLoader = 0;
+    }
+    lock {
+        dispatchCountOfUpdateAuthorLoader = 0;
+    }
+}
+
+isolated function assertDispatchCountForBookLoader(int expectedCount) {
+    lock {
+        test:assertEquals(dispatchCountOfBookLoader, expectedCount);
+    }
+}
+
+isolated function assertDispatchCountForUpdateAuthorLoader(int expectedCount) {
+    lock {
+        test:assertEquals(dispatchCountOfUpdateAuthorLoader, expectedCount);
+    }
+}
+
+isolated function assertDispatchCountForAuthorLoader(int expectedCount) {
+    lock {
+        test:assertEquals(dispatchCountOfAuthorLoader, expectedCount);
+    }
+}

ballerina-tests/tests/batch_load_functions.bal

@@ -0,0 +1,73 @@
+// Copyright (c) 2023 WSO2 LLC. (http://www.wso2.com) All Rights Reserved.
+//
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+isolated function authorLoaderFunction(readonly & anydata[] ids) returns AuthorRow[]|error {
+    readonly & int[] keys = check ids.ensureType();
+    // Simulate query: SELECT * FROM authors WHERE id IN (...keys);
+    lock {
+        dispatchCountOfAuthorLoader += 1;
+    }
+    lock {
+        readonly & int[] validKeys = keys.'filter(key => authorTable.hasKey(key)).cloneReadOnly();
+        return keys.length() != validKeys.length() ? error("Invalid keys found for authors")
+            : validKeys.'map(key => authorTable.get(key));
+    }
+};
+
+isolated function bookLoaderFunction(readonly & anydata[] ids) returns BookRow[][]|error {
+    final readonly & int[] keys = check ids.ensureType();
+    // Simulate query: SELECT * FROM books WHERE author IN (...keys);
+    lock {
+        dispatchCountOfBookLoader += 1;
+    }
+    return keys.'map(isolated function(readonly & int key) returns BookRow[] {
+        lock {
+            return bookTable.'filter(book => book.author == key).toArray().clone();
+        }
+    });
+};
+
+isolated function authorUpdateLoaderFunction(readonly & anydata[] idNames) returns AuthorRow[]|error {
+    readonly & [int, string][] idValuePair = check idNames.ensureType();
+    // Simulate batch udpate
+    lock {
+        dispatchCountOfUpdateAuthorLoader += 1;
+    }
+    lock {
+        foreach [int, string] [key, _] in idValuePair {
+            if !authorTable.hasKey(key) {
+                return error(string `Invalid key author key found: ${key}`);
+            }
+        }
+
+        AuthorRow[] updatedAuthorRows = [];
+        foreach [int, string] [key, name] in idValuePair {
+            AuthorRow authorRow = {id: key, name};
+            authorTable.put(authorRow);
+            updatedAuthorRows.push(authorRow.clone());
+        }
+        return updatedAuthorRows.clone();
+    }
+};
+
+isolated function faultyAuthorLoaderFunction(readonly & anydata[] ids) returns AuthorRow[]|error {
+    readonly & int[] keys = check ids.ensureType();
+    lock {
+        readonly & int[] validKeys = keys.'filter(key => authorTable.hasKey(key)).cloneReadOnly();
+        // This method may return an array of size not equal to the input array (ids) size.
+        return validKeys.'map(key => authorTable.get(key));
+    }
+};

ballerina-tests/tests/interceptors.bal

@@ -868,3 +868,23 @@ readonly service class LogSubfields {
         return context.resolve('field);
     }
 }
+
+readonly service class AuthorInterceptor {
+    *graphql:Interceptor;
+
+    isolated remote function execute(graphql:Context context, graphql:Field 'field) returns anydata {
+        var data = context.resolve('field);
+        // Return only the first author
+        return ('field.getName() == "authors" && data is anydata[]) ? [data[0]] : data;
+    }
+}
+
+readonly service class BookInterceptor {
+    *graphql:Interceptor;
+
+    isolated remote function execute(graphql:Context context, graphql:Field 'field) returns anydata {
+        var books = context.resolve('field);
+        // Return only the first book
+        return (books is anydata[]) ? [books[0]] : books;
+    }
+}

ballerina-tests/tests/object_types.bal

@@ -15,6 +15,7 @@
 // under the License.
 
 import ballerina/graphql;
+import ballerina/graphql.dataloader;
 import ballerina/lang.runtime;
 
 public type PeopleService StudentService|TeacherService;
@@ -393,3 +394,71 @@ public distinct isolated service class CustomerAddress {
         return self.city;
     }
 }
+
+public isolated distinct service class AuthorData {
+    private final readonly & AuthorRow author;
+
+    isolated function init(AuthorRow author) {
+        self.author = author.cloneReadOnly();
+    }
+
+    isolated resource function get name() returns string {
+        return self.author.name;
+    }
+
+    isolated function preBooks(graphql:Context ctx) {
+        dataloader:DataLoader bookLoader = ctx.getDataLoader(BOOK_LOADER);
+        bookLoader.add(self.author.id);
+    }
+
+    isolated resource function get books(graphql:Context ctx) returns BookData[]|error {
+        dataloader:DataLoader bookLoader = ctx.getDataLoader(BOOK_LOADER);
+        BookRow[] bookrows = check bookLoader.get(self.author.id);
+        return from BookRow bookRow in bookrows
+            select new BookData(bookRow);
+    }
+}
+
+public isolated distinct service class AuthorDetail {
+    private final readonly & AuthorRow author;
+
+    isolated function init(AuthorRow author) {
+        self.author = author.cloneReadOnly();
+    }
+
+    isolated resource function get name() returns string {
+        return self.author.name;
+    }
+
+    isolated function prefetchBooks(graphql:Context ctx) {
+        dataloader:DataLoader bookLoader = ctx.getDataLoader(BOOK_LOADER);
+        bookLoader.add(self.author.id);
+    }
+
+    @graphql:ResourceConfig {
+        interceptors: new BookInterceptor(),
+        prefetchMethodName: "prefetchBooks"
+    }
+    isolated resource function get books(graphql:Context ctx) returns BookData[]|error {
+        dataloader:DataLoader bookLoader = ctx.getDataLoader(BOOK_LOADER);
+        BookRow[] bookrows = check bookLoader.get(self.author.id);
+        return from BookRow bookRow in bookrows
+            select new BookData(bookRow);
+    }
+}
+
+public isolated distinct service class BookData {
+    private final readonly & BookRow book;
+
+    isolated function init(BookRow book) {
+        self.book = book.cloneReadOnly();
+    }
+
+    isolated resource function get id() returns int {
+        return self.book.id;
+    }
+
+    isolated resource function get title() returns string {
+        return self.book.title;
+    }
+}

ballerina-tests/tests/records.bal

@@ -351,3 +351,14 @@ public type DeprecatedAddress record {|
     # City field is deprecated
     string city;
 |};
+
+type BookRow record {|
+    readonly int id;
+    string title;
+    int author;
+|};
+
+type AuthorRow record {|
+    readonly int id;
+    string name;
+|};

ballerina-tests/tests/resources/documents/batch_function_returing_errors.graphql

@@ -0,0 +1,9 @@
+query {
+  authors(ids: [1, 2, 3, 4, 5, 6]) {
+    name
+    books {
+      id
+      title
+    }
+  }
+}