add finalize to NadelTransform

lib/src/main/java/graphql/nadel/engine/plan/NadelExecutionPlan.kt

@@ -2,6 +2,7 @@ package graphql.nadel.engine.plan
 
 import graphql.nadel.Service
 import graphql.nadel.engine.transform.NadelTransform
+import graphql.nadel.engine.transform.NadelTransformServiceExecutionContext
 import graphql.nadel.instrumentation.parameters.NadelInstrumentationTimingParameters
 import graphql.normalized.ExecutableNormalizedField
 
@@ -18,6 +19,7 @@ data class NadelExecutionPlan(
         val queryTransformTimingStep: NadelInstrumentationTimingParameters.Step,
         val resultTransformTimingStep: NadelInstrumentationTimingParameters.Step,
         val state: T,
+        val transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     )
 
     /**

lib/src/main/java/graphql/nadel/engine/plan/NadelExecutionPlanFactory.kt

@@ -11,6 +11,7 @@ import graphql.nadel.engine.transform.NadelRenameArgumentInputTypesTransform
 import graphql.nadel.engine.transform.NadelRenameTransform
 import graphql.nadel.engine.transform.NadelServiceTypeFilterTransform
 import graphql.nadel.engine.transform.NadelTransform
+import graphql.nadel.engine.transform.NadelTransformServiceExecutionContext
 import graphql.nadel.engine.transform.NadelTypeRenameResultTransform
 import graphql.nadel.engine.transform.hydration.NadelHydrationTransform
 import graphql.nadel.engine.transform.hydration.batch.NadelBatchHydrationTransform
@@ -64,7 +65,8 @@ internal class NadelExecutionPlanFactory(
     ): NadelExecutionPlan {
         val executionSteps: MutableMap<ExecutableNormalizedField, List<NadelExecutionPlan.Step<Any>>> =
             mutableMapOf()
-
+        val transformContexts: MutableMap<NadelTransform<Any>, NadelTransformServiceExecutionContext?> =
+            mutableMapOf()
         executionContext.timer.batch { timer ->
             traverseQuery(rootField) { field ->
                 val steps = transformsWithTimingStepInfo.mapNotNull { transformWithTimingInfo ->
@@ -75,6 +77,21 @@ internal class NadelExecutionPlanFactory(
                     if (isSkipIncludeSpecialField(field) && ((transform as NadelTransform<*>) !is NadelSkipIncludeTransform)) {
                         null
                     } else {
+                        val executionTransformContext = if (transformContexts.containsKey(transform)) {
+                            transformContexts[transform]
+                        } else {
+                            transformContexts.getOrPut(transform) {
+                                transform.buildContext(
+                                    executionContext,
+                                    serviceExecutionContext,
+                                    executionBlueprint,
+                                    services,
+                                    service,
+                                    rootField,
+                                    serviceHydrationDetails
+                                )
+                            }
+                        }
                         val state = timer.time(step = transformWithTimingInfo.executionPlanTimingStep) {
                             transform.isApplicable(
                                 executionContext,
@@ -83,6 +100,7 @@ internal class NadelExecutionPlanFactory(
                                 services,
                                 service,
                                 field,
+                                executionTransformContext,
                                 serviceHydrationDetails,
                             )
                         }
@@ -97,6 +115,7 @@ internal class NadelExecutionPlanFactory(
                                 queryTransformTimingStep = transformWithTimingInfo.queryTransformTimingStep,
                                 resultTransformTimingStep = transformWithTimingInfo.resultTransformTimingStep,
                                 state = state,
+                                executionTransformContext
                             )
                         }
                     }

lib/src/main/java/graphql/nadel/engine/transform/NadelDeepRenameTransform.kt

@@ -88,6 +88,7 @@ internal class NadelDeepRenameTransform : NadelTransform<NadelDeepRenameTransfor
         services: Map<String, Service>,
         service: Service,
         overallField: ExecutableNormalizedField,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
         hydrationDetails: ServiceExecutionHydrationDetails?,
     ): State? {
         val deepRenameInstructions = executionBlueprint
@@ -161,6 +162,7 @@ internal class NadelDeepRenameTransform : NadelTransform<NadelDeepRenameTransfor
         service: Service,
         field: ExecutableNormalizedField,
         state: State,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): NadelTransformFieldResult {
         val objectTypesNoRenames = field.objectTypeNames.filterNot { it in state.instructionsByObjectTypeNames }
 
@@ -299,6 +301,7 @@ internal class NadelDeepRenameTransform : NadelTransform<NadelDeepRenameTransfor
         result: ServiceExecutionResult,
         state: State,
         nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): List<NadelResultInstruction> {
         val parentNodes = nodes.getNodesAt(
             queryPath = underlyingParentField?.queryPath ?: NadelQueryPath.root,

lib/src/main/java/graphql/nadel/engine/transform/NadelRenameArgumentInputTypesTransform.kt

@@ -35,6 +35,7 @@ internal class NadelRenameArgumentInputTypesTransform : NadelTransform<State> {
         services: Map<String, Service>,
         service: Service,
         overallField: ExecutableNormalizedField,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
         hydrationDetails: ServiceExecutionHydrationDetails?,
     ): State? {
         // Transform if there's any arguments at all
@@ -55,6 +56,7 @@ internal class NadelRenameArgumentInputTypesTransform : NadelTransform<State> {
         service: Service,
         field: ExecutableNormalizedField,
         state: State,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): NadelTransformFieldResult {
         return NadelTransformFieldResult(
             newField = field
@@ -74,6 +76,7 @@ internal class NadelRenameArgumentInputTypesTransform : NadelTransform<State> {
         result: ServiceExecutionResult,
         state: State,
         nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): List<NadelResultInstruction> {
         return emptyList()
     }

lib/src/main/java/graphql/nadel/engine/transform/NadelRenameTransform.kt

@@ -43,6 +43,7 @@ internal class NadelRenameTransform : NadelTransform<State> {
         services: Map<String, Service>,
         service: Service,
         overallField: ExecutableNormalizedField,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
         hydrationDetails: ServiceExecutionHydrationDetails?,
     ): State? {
         val renameInstructions = executionBlueprint
@@ -73,6 +74,7 @@ internal class NadelRenameTransform : NadelTransform<State> {
         service: Service,
         field: ExecutableNormalizedField,
         state: State,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): NadelTransformFieldResult {
         return NadelTransformFieldResult(
             newField = if (state.objectTypesWithoutRename.isNotEmpty()) {
@@ -185,6 +187,7 @@ internal class NadelRenameTransform : NadelTransform<State> {
         result: ServiceExecutionResult,
         state: State,
         nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): List<NadelResultInstruction> {
         val parentNodes = nodes.getNodesAt(
             queryPath = underlyingParentField?.queryPath ?: NadelQueryPath.root,

lib/src/main/java/graphql/nadel/engine/transform/NadelServiceTypeFilterTransform.kt

@@ -77,6 +77,7 @@ class NadelServiceTypeFilterTransform : NadelTransform<State> {
         services: Map<String, Service>,
         service: Service,
         overallField: ExecutableNormalizedField,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
         hydrationDetails: ServiceExecutionHydrationDetails?,
     ): State? {
         when {
@@ -100,7 +101,9 @@ class NadelServiceTypeFilterTransform : NadelTransform<State> {
             .all { objectTypeName ->
                 objectTypeName in typeNamesOwnedByService
                     || objectTypeName in underlyingTypeNamesOwnedByService
-                    || (executionContext.hints.sharedTypeRenames(service) && executionBlueprint.getUnderlyingTypeName(objectTypeName) in underlyingTypeNamesOwnedByService)
+                    || (executionContext.hints.sharedTypeRenames(service) && executionBlueprint.getUnderlyingTypeName(
+                    objectTypeName
+                ) in underlyingTypeNamesOwnedByService)
             }
 
         if (noForeignTypes) {
@@ -110,7 +113,9 @@ class NadelServiceTypeFilterTransform : NadelTransform<State> {
         val fieldObjectTypeNamesOwnedByService = overallField.objectTypeNames.filter { objectTypeName ->
             objectTypeName in typeNamesOwnedByService
                 || objectTypeName in underlyingTypeNamesOwnedByService
-                || (executionContext.hints.sharedTypeRenames(service) && executionBlueprint.getUnderlyingTypeName(objectTypeName) in underlyingTypeNamesOwnedByService)
+                || (executionContext.hints.sharedTypeRenames(service) && executionBlueprint.getUnderlyingTypeName(
+                objectTypeName
+            ) in underlyingTypeNamesOwnedByService)
         }
 
         return State(
@@ -132,6 +137,7 @@ class NadelServiceTypeFilterTransform : NadelTransform<State> {
         service: Service,
         field: ExecutableNormalizedField,
         state: State,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): NadelTransformFieldResult {
         // Nothing to query if there are no fields, we need to add selection
         if (state.fieldObjectTypeNamesOwnedByService.isEmpty()) {
@@ -190,6 +196,7 @@ class NadelServiceTypeFilterTransform : NadelTransform<State> {
         result: ServiceExecutionResult,
         state: State,
         nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): List<NadelResultInstruction> {
         return emptyList()
     }

lib/src/main/java/graphql/nadel/engine/transform/NadelTransform.kt

@@ -19,6 +19,32 @@ interface NadelTransform<State : Any> {
     val name: String
         get() = javaClass.simpleName.ifBlank { "UnknownTransform" }
 
+    /**
+     * This method is called once before execution of the transform starts.
+     * Override it to create a common object that is shared between all invocations of all other methods
+     * of the transform on all the fields.
+     *
+     * @param executionBlueprint the [NadelOverallExecutionBlueprint] of the Nadel instance being operated on
+     * @param rootField the root [ExecutableNormalizedField] of the operation this [NadelTransform] runs on
+     * @param hydrationDetails the [ServiceExecutionHydrationDetails] when the [NadelTransform] is applied to fields inside
+     * hydrations, `null` otherwise
+     *
+     * @return a common [NadelTransformServiceExecutionContext] that will be fed into all other methods of this transform
+     * when they run or [null] if you don't need such functionality
+     *
+     */
+    suspend fun buildContext(
+        executionContext: NadelExecutionContext,
+        serviceExecutionContext: NadelServiceExecutionContext,
+        executionBlueprint: NadelOverallExecutionBlueprint,
+        services: Map<String, Service>,
+        service: Service,
+        rootField: ExecutableNormalizedField,
+        hydrationDetails: ServiceExecutionHydrationDetails?,
+    ): NadelTransformServiceExecutionContext? {
+        return null
+    }
+
     /**
      * Determines whether the [NadelTransform] should run. If it should run return a [State].
      *
@@ -48,6 +74,7 @@ interface NadelTransform<State : Any> {
         services: Map<String, Service>,
         service: Service,
         overallField: ExecutableNormalizedField,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
         hydrationDetails: ServiceExecutionHydrationDetails? = null,
     ): State?
 
@@ -66,6 +93,7 @@ interface NadelTransform<State : Any> {
         service: Service,
         field: ExecutableNormalizedField,
         state: State,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): NadelTransformFieldResult
 
     /**
@@ -84,7 +112,24 @@ interface NadelTransform<State : Any> {
         result: ServiceExecutionResult,
         state: State,
         nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
     ): List<NadelResultInstruction>
+
+    /**
+     * Called once after all other functions of a transform ran on all fields in the query.
+     * Override this function to perform cleanup or finalization tasks.
+     * This method is optional for implementing classes.
+     */
+    suspend fun onComplete(
+        executionContext: NadelExecutionContext,
+        serviceExecutionContext: NadelServiceExecutionContext,
+        executionBlueprint: NadelOverallExecutionBlueprint,
+        service: Service,
+        result: ServiceExecutionResult,
+        nodes: JsonNodes,
+        transformServiceExecutionContext: NadelTransformServiceExecutionContext?,
+    ) {
+    }
 }
 
 data class NadelTransformFieldResult(