[executions] do not apply data loader instrumentation if execution is…

… a Mutation (#1468) * feat: do not apply instrumentation if execution is a mutation * feat: update docs * feat: fix ktlint check * feat: add mutation tests Co-authored-by: samvazquez <[email protected]>
ExpediaGroup · Jun 15, 2022 · 5721ea3 · 5721ea3
1 parent ccb5449
commit 5721ea3
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 17 deletions.
diff --git a/.../expediagroup/graphql/dataloader/instrumentation/extensions/ExecutionContextExtensions.kt b/.../expediagroup/graphql/dataloader/instrumentation/extensions/ExecutionContextExtensions.kt
@@ -19,6 +19,7 @@ package com.expediagroup.graphql.dataloader.instrumentation.extensions
 import graphql.analysis.QueryTraverser
 import graphql.analysis.QueryVisitorFieldEnvironment
 import graphql.execution.ExecutionContext
+import graphql.language.OperationDefinition
 import kotlin.math.max
 
 /**
@@ -41,3 +42,10 @@ internal fun ExecutionContext.getDocumentHeight(): Int {
             0
         )
 }
+
+/**
+ * Checks if the [ExecutionContext] is a [OperationDefinition.Operation.MUTATION]
+ * @return Boolean indicating if GraphQL Operation is a Mutation
+ */
+internal fun ExecutionContext.isMutation(): Boolean =
+    this.operationDefinition.operation == OperationDefinition.Operation.MUTATION
diff --git a/...loader/instrumentation/level/execution/AbstractExecutionLevelDispatchedInstrumentation.kt b/...loader/instrumentation/level/execution/AbstractExecutionLevelDispatchedInstrumentation.kt
@@ -17,6 +17,7 @@
 package com.expediagroup.graphql.dataloader.instrumentation.level.execution
 
 import com.expediagroup.graphql.dataloader.instrumentation.NoOpExecutionStrategyInstrumentationContext
+import com.expediagroup.graphql.dataloader.instrumentation.extensions.isMutation
 import com.expediagroup.graphql.dataloader.instrumentation.level.state.ExecutionLevelDispatchedState
 import com.expediagroup.graphql.dataloader.instrumentation.level.state.Level
 import graphql.ExecutionInput
@@ -54,16 +55,16 @@ abstract class AbstractExecutionLevelDispatchedInstrumentation : SimpleInstrumen
     override fun beginExecuteOperation(
         parameters: InstrumentationExecuteOperationParameters
     ): InstrumentationContext<ExecutionResult> =
-        parameters.executionContext
-            .graphQLContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
             ?.beginExecuteOperation(parameters)
             ?: SimpleInstrumentationContext.noOp()
 
     override fun beginExecutionStrategy(
         parameters: InstrumentationExecutionStrategyParameters
     ): ExecutionStrategyInstrumentationContext =
-        parameters.executionContext
-            .graphQLContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
             ?.beginExecutionStrategy(
                 parameters,
                 this.getOnLevelDispatchedCallback(
@@ -78,8 +79,8 @@ abstract class AbstractExecutionLevelDispatchedInstrumentation : SimpleInstrumen
     override fun beginFieldFetch(
         parameters: InstrumentationFieldFetchParameters
     ): InstrumentationContext<Any> =
-        parameters.executionContext
-            .graphQLContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
             ?.beginFieldFetch(
                 parameters,
                 this.getOnLevelDispatchedCallback(
@@ -95,8 +96,8 @@ abstract class AbstractExecutionLevelDispatchedInstrumentation : SimpleInstrumen
         dataFetcher: DataFetcher<*>,
         parameters: InstrumentationFieldFetchParameters
     ): DataFetcher<*> =
-        parameters.executionContext
-            .graphQLContext.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<ExecutionLevelDispatchedState>(ExecutionLevelDispatchedState::class)
             ?.instrumentDataFetcher(dataFetcher, parameters)
             ?: dataFetcher
 }
diff --git a/...instrumentation/syncexhaustion/execution/AbstractSyncExecutionExhaustedInstrumentation.kt b/...instrumentation/syncexhaustion/execution/AbstractSyncExecutionExhaustedInstrumentation.kt
@@ -17,6 +17,7 @@
 package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.execution
 
 import com.expediagroup.graphql.dataloader.instrumentation.NoOpExecutionStrategyInstrumentationContext
+import com.expediagroup.graphql.dataloader.instrumentation.extensions.isMutation
 import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
 import graphql.ExecutionInput
 import graphql.ExecutionResult
@@ -54,24 +55,24 @@ abstract class AbstractSyncExecutionExhaustedInstrumentation : SimpleInstrumenta
     override fun beginExecuteOperation(
         parameters: InstrumentationExecuteOperationParameters
     ): InstrumentationContext<ExecutionResult> =
-        parameters.executionContext
-            .graphQLContext.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
             ?.beginExecuteOperation(parameters)
             ?: SimpleInstrumentationContext.noOp()
 
     override fun beginExecutionStrategy(
         parameters: InstrumentationExecutionStrategyParameters
     ): ExecutionStrategyInstrumentationContext =
-        parameters.executionContext
-            .graphQLContext.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
             ?.beginExecutionStrategy(parameters)
             ?: NoOpExecutionStrategyInstrumentationContext
 
     override fun beginFieldFetch(
         parameters: InstrumentationFieldFetchParameters
     ): InstrumentationContext<Any> =
-        parameters.executionContext
-            .graphQLContext.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
+        parameters.executionContext.takeIf { !it.isMutation() }
+            ?.graphQLContext?.get<SyncExecutionExhaustedState>(SyncExecutionExhaustedState::class)
             ?.beginFieldFetch(
                 parameters,
                 this.getOnSyncExecutionExhaustedCallback(

diff --git a/...st/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt b/...st/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt
@@ -22,6 +22,7 @@ import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.A
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Astronaut
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.AstronautService
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.AstronautServiceRequest
+import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.CreateAstronautServiceRequest
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.MissionDataLoader
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.MissionService
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.MissionServiceRequest
@@ -56,6 +57,9 @@ object AstronautGraphQL {
             missions(ids: [ID!]): [Mission]!
             nasa: Nasa!
         }
+        type Mutation {
+            createAstronaut(name: String!): Astronaut!
+        }
         type Nasa {
             astronaut(id: ID!): Astronaut
             mission(id: ID!): Mission
@@ -95,6 +99,14 @@ object AstronautGraphQL {
             environment
         )
     }
+    private val createAstronautDataFetcher = DataFetcher { environment ->
+        astronautService.createAstronaut(
+            CreateAstronautServiceRequest(
+                environment.getArgument("name")
+            ),
+            environment
+        )
+    }
     private val astronautsDataFetcher = DataFetcher { environment ->
         astronautService.getAstronauts(
             environment.getArgument<List<String>>("ids")
@@ -155,6 +167,10 @@ object AstronautGraphQL {
                 .dataFetcher("missions", missionsDataFetcher)
                 .dataFetcher("nasa") { Nasa() }
         )
+        type(
+            TypeRuntimeWiring.newTypeWiring("Mutation")
+                .dataFetcher("createAstronaut", createAstronautDataFetcher)
+        )
         type(
             TypeRuntimeWiring.newTypeWiring("Astronaut")
                 .dataFetcher("missions", missionsByAstronautDataFetcher)

diff --git a/...m/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt b/...m/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt
@@ -28,10 +28,13 @@ import org.dataloader.DataLoader
 import org.dataloader.DataLoaderFactory
 import org.dataloader.DataLoaderOptions
 import org.dataloader.stats.SimpleStatisticsCollector
+import reactor.kotlin.core.publisher.toMono
+import java.time.Duration
 import java.util.Optional
 import java.util.concurrent.CompletableFuture
 
 data class AstronautServiceRequest(val id: Int)
+data class CreateAstronautServiceRequest(val name: String)
 
 class AstronautDataLoader : KotlinDataLoader<AstronautServiceRequest, Astronaut?> {
     override val dataLoaderName: String = "AstronautDataLoader"
@@ -57,6 +60,12 @@ class AstronautService {
             .getDataLoader<AstronautServiceRequest, Astronaut>("AstronautDataLoader")
             .load(request)
 
+    fun createAstronaut(
+        request: CreateAstronautServiceRequest,
+        environment: DataFetchingEnvironment
+    ): CompletableFuture<Astronaut> =
+        Astronaut(100, request.name).toMono().delayElement(Duration.ofMillis(100)).toFuture()
+
     fun getAstronauts(
         requests: List<AstronautServiceRequest>,
         environment: DataFetchingEnvironment

diff --git a/.../graphql/dataloader/instrumentation/level/DataLoaderLevelDispatchedInstrumentationTest.kt b/.../graphql/dataloader/instrumentation/level/DataLoaderLevelDispatchedInstrumentationTest.kt
@@ -16,19 +16,29 @@
 
 package com.expediagroup.graphql.dataloader.instrumentation.level
 
-import com.expediagroup.graphql.dataloader.instrumentation.fixture.DataLoaderInstrumentationStrategy
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.AstronautGraphQL
+import com.expediagroup.graphql.dataloader.instrumentation.fixture.DataLoaderInstrumentationStrategy
+import io.mockk.Called
+import io.mockk.clearAllMocks
+import io.mockk.spyk
 import io.mockk.verify
+import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 
 class DataLoaderLevelDispatchedInstrumentationTest {
+    private val dataLoaderLevelDispatchedInstrumentation = spyk(DataLoaderLevelDispatchedInstrumentation())
     private val graphQL = AstronautGraphQL.builder
-        .instrumentation(DataLoaderLevelDispatchedInstrumentation())
+        .instrumentation(dataLoaderLevelDispatchedInstrumentation)
         // graphql java adds DataLoaderDispatcherInstrumentation by default
         .doNotAddDefaultInstrumentations()
         .build()
 
+    @BeforeEach
+    fun clear() {
+        clearAllMocks()
+    }
+
     @Test
     fun `Instrumentation should batch transactions on async top level fields`() {
         val queries = listOf(
@@ -132,4 +142,22 @@ class DataLoaderLevelDispatchedInstrumentationTest {
             kotlinDataLoaderRegistry.dispatchAll()
         }
     }
+
+    @Test
+    fun `Instrumentation should not apply to mutations`() {
+        val queries = listOf(
+            """mutation { createAstronaut(name: "spaceMan") { id name } }"""
+        )
+
+        val (results, _) = AstronautGraphQL.execute(
+            graphQL,
+            queries,
+            DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED
+        )
+
+        assertEquals(1, results.size)
+        verify {
+            dataLoaderLevelDispatchedInstrumentation.getOnLevelDispatchedCallback(ofType()) wasNot Called
+        }
+    }
 }
diff --git a/...der/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentationTest.kt b/...der/instrumentation/syncexhaustion/DataLoaderSyncExecutionExhaustedInstrumentationTest.kt
@@ -19,13 +19,18 @@ package com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.DataLoaderInstrumentationStrategy
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.AstronautGraphQL
 import com.expediagroup.graphql.dataloader.instrumentation.fixture.ProductGraphQL
+import io.mockk.Called
+import io.mockk.clearAllMocks
+import io.mockk.spyk
 import io.mockk.verify
+import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 
 class DataLoaderSyncExecutionExhaustedInstrumentationTest {
+    private val dataLoaderSyncExecutionExhaustedInstrumentation = spyk(DataLoaderSyncExecutionExhaustedInstrumentation())
     private val graphQL = AstronautGraphQL.builder
-        .instrumentation(DataLoaderSyncExecutionExhaustedInstrumentation())
+        .instrumentation(dataLoaderSyncExecutionExhaustedInstrumentation)
         // graphql java adds DataLoaderDispatcherInstrumentation by default
         .doNotAddDefaultInstrumentations()
         .build()
@@ -36,6 +41,11 @@ class DataLoaderSyncExecutionExhaustedInstrumentationTest {
         .doNotAddDefaultInstrumentations()
         .build()
 
+    @BeforeEach
+    fun clear() {
+        clearAllMocks()
+    }
+
     @Test
     fun `Instrumentation should batch transactions on async top level fields`() {
         val queries = listOf(
@@ -498,4 +508,22 @@ class DataLoaderSyncExecutionExhaustedInstrumentationTest {
         assertEquals(1, productStatistics?.batchInvokeCount)
         assertEquals(2, productStatistics?.batchLoadCount)
     }
+
+    @Test
+    fun `Instrumentation should not apply to mutations`() {
+        val queries = listOf(
+            """mutation { createAstronaut(name: "spaceMan") { id name } }"""
+        )
+
+        val (results, _) = AstronautGraphQL.execute(
+            graphQL,
+            queries,
+            DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED
+        )
+
+        assertEquals(1, results.size)
+        verify {
+            dataLoaderSyncExecutionExhaustedInstrumentation.getOnSyncExecutionExhaustedCallback(ofType()) wasNot Called
+        }
+    }
 }