graphql-java · andimarek · Mar 7, 2024 · Feb 4, 2024 · Feb 4, 2024 · Feb 4, 2024
diff --git a/agent-test/build.gradle b/agent-test/build.gradle
@@ -0,0 +1,57 @@
+plugins {
+    id 'java'
+}
+
+dependencies {
+    implementation(rootProject)
+    implementation("net.bytebuddy:byte-buddy-agent:1.14.11")
+
+    testImplementation 'org.junit.jupiter:junit-jupiter:5.7.1'
+    testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
+
+    testImplementation("org.assertj:assertj-core:3.25.1")
+
+}
+
+java {
+    toolchain {
+        languageVersion = JavaLanguageVersion.of(11)
+    }
+}
+
+tasks.named('test', Test) {
+    dependsOn(':agent:shadowJar')
+    useJUnitPlatform()
+
+    maxHeapSize = '4G'
+
+    testLogging {
+        events "passed"
+    }
+}
+
+
+repositories {
+    mavenCentral()
+    mavenLocal()
+}
+
+
+java {
+    toolchain {
+        languageVersion = JavaLanguageVersion.of(11)
+    }
+}
+
+
+jar {
+    manifest {
+        attributes(
+                'Agent-Class': 'graphql.agent.GraphQLJavaAgent',
+                'Can-Redefine-Classes': 'true',
+                'Can-Retransform-Classes': 'true',
+                'Premain-Class': 'graphql.agent.GraphQLJavaAgent'
+        )
+    }
+}
+
diff --git a/agent-test/src/main/java/graphql/GraphQLApp.java b/agent-test/src/main/java/graphql/GraphQLApp.java
@@ -0,0 +1,40 @@
+package graphql;
+
+import graphql.agent.result.ExecutionTrackingResult;
+import graphql.schema.GraphQLSchema;
+import graphql.schema.idl.RuntimeWiring;
+import graphql.schema.idl.SchemaGenerator;
+import graphql.schema.idl.SchemaParser;
+import graphql.schema.idl.TypeDefinitionRegistry;
+
+/**
+ * Used for testing loading the agent on startup.
+ * See StartAgentOnStartupTest
+ */
+public class GraphQLApp {
+
+    public static void main(String[] args) {
+        String schema = "type Query { hello: String }";
+        TypeDefinitionRegistry typeDefinitionRegistry = new SchemaParser().parse(schema);
+        RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
+                .type("Query", builder -> builder.dataFetcher("hello", environment -> "world"))
+                .build();
+        GraphQLSchema graphQLSchema = new SchemaGenerator().makeExecutableSchema(typeDefinitionRegistry, runtimeWiring);
+        GraphQL graphQL = GraphQL.newGraphQL(graphQLSchema).build();
+        ExecutionInput executionInput = ExecutionInput.newExecutionInput().query("{ hello alias: hello alias2: hello }").build();
+        GraphQLContext graphQLContext = executionInput.getGraphQLContext();
+        ExecutionResult executionResult = graphQL.execute(executionInput);
+        System.out.println(executionResult.getData().toString());
+        ExecutionTrackingResult executionTrackingResult = graphQLContext.get(ExecutionTrackingResult.EXECUTION_TRACKING_KEY);
+        if (executionTrackingResult == null) {
+            System.out.println("No tracking data found");
+            System.exit(1);
+        }
+        if (executionTrackingResult.timePerPath.size() != 3) {
+            System.out.println("Expected 3 paths, got " + executionTrackingResult.timePerPath.size());
+            System.exit(1);
+        }
+        System.out.println("Successfully tracked execution");
+        System.exit(0);
+    }
+}
diff --git a/agent-test/src/test/java/graphql/test/AgentTest.java b/agent-test/src/test/java/graphql/test/AgentTest.java
@@ -0,0 +1,72 @@
+package graphql.test;
+
+import graphql.agent.result.ExecutionTrackingResult;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class AgentTest {
+
+    @BeforeAll
+    static void init() {
+        LoadAgent.loadIntoCurrentJVM();
+    }
+
+    @AfterAll
+    static void cleanup() {
+    }
+
+    @Test
+    void test() {
+        ExecutionTrackingResult executionTrackingResult = TestQuery.executeQuery();
+        assertThat(executionTrackingResult.dataFetcherCount()).isEqualTo(5);
+        assertThat(executionTrackingResult.getTime("/issues")).isGreaterThan(100);
+        assertThat(executionTrackingResult.getDfResultTypes("/issues"))
+                .isEqualTo(ExecutionTrackingResult.DFResultType.DONE_OK);
+
+        verifyAgentDataIsEmpty();
+
+    }
+
+    @Test
+    void testBatchLoader() {
+        ExecutionTrackingResult executionTrackingResult = TestQuery.executeBatchedQuery();
+        assertThat(executionTrackingResult.dataFetcherCount()).isEqualTo(9);
+        assertThat(executionTrackingResult.getTime("/issues")).isGreaterThan(100);
+        assertThat(executionTrackingResult.getDfResultTypes("/issues[0]/author"))
+                .isEqualTo(ExecutionTrackingResult.DFResultType.PENDING);
+        assertThat(executionTrackingResult.getDfResultTypes("/issues[1]/author"))
+                .isEqualTo(ExecutionTrackingResult.DFResultType.PENDING);
+
+        assertThat(executionTrackingResult.getDataLoaderNames()).isEqualTo(Collections.singletonList("userLoader"));
+
+        assertThat(executionTrackingResult.dataLoaderNameToBatchCall).hasSize(1);
+        List<ExecutionTrackingResult.BatchLoadingCall> userLoaderCalls = executionTrackingResult.dataLoaderNameToBatchCall.get("userLoader");
+        assertThat(userLoaderCalls).hasSize(1);
+        ExecutionTrackingResult.BatchLoadingCall batchLoadingCall = userLoaderCalls.get(0);
+
+        assertThat(batchLoadingCall.keyCount).isEqualTo(2);
+
+        verifyAgentDataIsEmpty();
+    }
+
+    private void verifyAgentDataIsEmpty() {
+        try {
+            Class<?> agent = Class.forName("graphql.agent.GraphQLJavaAgent");
+            Map executionIdToData = (Map) agent.getField("executionIdToData").get(null);
+            Map dataLoaderToExecutionId = (Map) agent.getField("dataLoaderToExecutionId").get(null);
+            assertThat(executionIdToData).isEmpty();
+            assertThat(dataLoaderToExecutionId).isEmpty();
+
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+}
diff --git a/agent-test/src/test/java/graphql/test/LoadAgent.java b/agent-test/src/test/java/graphql/test/LoadAgent.java
@@ -0,0 +1,15 @@
+package graphql.test;
+
+import net.bytebuddy.agent.ByteBuddyAgent;
+
+import java.io.File;
+
+
+public class LoadAgent {
+
+
+    public static void loadIntoCurrentJVM() {
+        ByteBuddyAgent.attach(new File("../agent/build/libs/agent.jar"), String.valueOf(ProcessHandle.current().pid()));
+    }
+
+}
diff --git a/agent-test/src/test/java/graphql/test/StartAgentOnStartupTest.java b/agent-test/src/test/java/graphql/test/StartAgentOnStartupTest.java
@@ -0,0 +1,23 @@
+package graphql.test;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class StartAgentOnStartupTest {
+
+
+    @Test
+    void testAgentCanBeLoadedAtStartup() throws IOException, InterruptedException {
+        // we use the classpath of the current test
+        String classPath = System.getProperty("java.class.path");
+        ProcessBuilder processBuilder = new ProcessBuilder("java", "-javaagent:../agent/build/libs/agent.jar", "-classpath", classPath, "graphql.GraphQLApp");
+        Process process = processBuilder.start();
+        process.getErrorStream().transferTo(System.err);
+        process.getInputStream().transferTo(System.out);
+        int i = process.waitFor();
+        assertThat(i).isZero();
+    }
+}
diff --git a/agent-test/src/test/java/graphql/test/TestQuery.java b/agent-test/src/test/java/graphql/test/TestQuery.java
@@ -0,0 +1,102 @@
+package graphql.test;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+import graphql.GraphQL;
+import graphql.agent.result.ExecutionTrackingResult;
+import graphql.schema.DataFetcher;
+import graphql.schema.GraphQLSchema;
+import graphql.schema.idl.RuntimeWiring;
+import graphql.schema.idl.SchemaGenerator;
+import graphql.schema.idl.SchemaParser;
+import graphql.schema.idl.TypeDefinitionRegistry;
+import org.assertj.core.api.Assertions;
+import org.dataloader.BatchLoader;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderFactory;
+import org.dataloader.DataLoaderRegistry;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+
+public class TestQuery {
+
+
+    static ExecutionTrackingResult executeQuery() {
+        String sdl = "type Query{issues: [Issue]} type Issue {id: ID, title: String}";
+        TypeDefinitionRegistry typeDefinitionRegistry = new SchemaParser().parse(sdl);
+        DataFetcher<List> issuesDF = (env) -> {
+            return List.of(
+                Map.of("id", "1", "title", "issue-1"),
+                Map.of("id", "2", "title", "issue-2"));
+        };
+
+        RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
+            .type("Query", builder -> builder.dataFetcher("issues", issuesDF))
+            .build();
+        GraphQLSchema graphQLSchema = new SchemaGenerator().makeExecutableSchema(typeDefinitionRegistry, runtimeWiring);
+
+        GraphQL graphQL = GraphQL.newGraphQL(graphQLSchema).build();
+
+        ExecutionInput executionInput = ExecutionInput.newExecutionInput().query("{issues{id title}}").build();
+        ExecutionResult result = graphQL.execute(executionInput);
+        Assertions.assertThat(result.getErrors()).isEmpty();
+        ExecutionTrackingResult trackingResult = executionInput.getGraphQLContext().get(ExecutionTrackingResult.EXECUTION_TRACKING_KEY);
+        return trackingResult;
+    }
+
+    static ExecutionTrackingResult executeBatchedQuery() {
+        String sdl = "type Query{issues: [Issue]} " +
+            "type Issue {id: ID, author: User}" +
+            "type User {id: ID, name: String}";
+
+        DataFetcher<List> issuesDF = (env) -> List.of(
+            Map.of("id", "1", "title", "issue-1", "authorId", "user-1"),
+            Map.of("id", "2", "title", "issue-2", "authorId", "user-2"));
+
+        BatchLoader<String, Map> userBatchLoader = keys -> {
+            // System.out.println("batch users with keys: " + keys);
+            return CompletableFuture.supplyAsync(() -> {
+                try {
+                    Thread.sleep(100);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+                return List.of(
+                    Map.of("id", "user-1", "name", "Foo-1"),
+                    Map.of("id", "user-2", "name", "Foo-2")
+                );
+            });
+        };
+        DataLoaderRegistry dataLoaderRegistry = new DataLoaderRegistry();
+        dataLoaderRegistry.register("userLoader", DataLoaderFactory.newDataLoader(userBatchLoader));
+
+        DataFetcher<CompletableFuture<Map>> authorDF = (env) -> {
+            DataLoader<String, Map> userLoader = env.getDataLoader("userLoader");
+            // System.out.println("author id: " + (String) ((Map) env.getSource()).get("authorId"));
+            return userLoader.load((String) ((Map) env.getSource()).get("authorId"));
+        };
+        TypeDefinitionRegistry typeDefinitionRegistry = new SchemaParser().parse(sdl);
+
+        RuntimeWiring runtimeWiring = RuntimeWiring.newRuntimeWiring()
+            .type("Query", builder -> builder.dataFetcher("issues", issuesDF))
+            .type("Issue", builder -> builder.dataFetcher("author", authorDF))
+            .build();
+        GraphQLSchema graphQLSchema = new SchemaGenerator().makeExecutableSchema(typeDefinitionRegistry, runtimeWiring);
+
+        GraphQL graphQL = GraphQL.newGraphQL(graphQLSchema).build();
+        String query = "{issues" +
+            "{id author {id name}}" +
+            "}";
+        ExecutionInput executionInput = ExecutionInput.newExecutionInput()
+            .dataLoaderRegistry(dataLoaderRegistry)
+            .query(query).build();
+        ExecutionResult result = graphQL.execute(executionInput);
+        Assertions.assertThat(result.getErrors()).isEmpty();
+        ExecutionTrackingResult trackingResult = executionInput.getGraphQLContext().get(ExecutionTrackingResult.EXECUTION_TRACKING_KEY);
+        return trackingResult;
+    }
+
+
+}
diff --git a/agent/build.gradle b/agent/build.gradle
@@ -0,0 +1,38 @@
+plugins {
+    id 'java'
+    id "com.github.johnrengelman.shadow" version "8.1.1"
+}
+
+dependencies {
+    implementation("net.bytebuddy:byte-buddy:1.14.11")
+    implementation(rootProject)
+}
+
+repositories {
+    mavenCentral()
+    mavenLocal()
+}
+
+
+java {
+    toolchain {
+        languageVersion = JavaLanguageVersion.of(11)
+    }
+}
+
+shadowJar {
+    minimize()
+    archiveClassifier.set('')
+    configurations = [project.configurations.compileClasspath]
+    dependencies {
+        exclude(dependency(rootProject))
+    }
+    manifest {
+        attributes(
+                'Agent-Class': 'graphql.agent.GraphQLJavaAgent',
+                'Premain-Class': 'graphql.agent.GraphQLJavaAgent',
+                'Can-Redefine-Classes': 'true',
+                'Can-Retransform-Classes': 'true',
+        )
+    }
+}