Spring transactions in GraphQL Java

May 3, 2021

The example below integrates GraphQL Java with Spring transaction management to ensure each query/mutation is executed in a transaction (similar to request-scoped transactions for RESTful APIs) so as to ensure consistency when different loaders are executed. This is done by providing a custom ExecutionStrategy which starts a transaction before executing the query/mutation and commits it afterwards. In order for this to work properly, it is necessary that for a single query all loaders are executed in sequential order on the same thread.

It was built using version 2.0 of graphql-java-spring-boot-starter-webmvc, but it should be easy to carry over to other setups.

Loaders.kt

package at.sigmoid

import java.util.concurrent.CompletableFuture
import org.dataloader.BatchLoader
import org.dataloader.BatchLoaderEnvironment
import org.dataloader.BatchLoaderWithContext
import org.springframework.core.task.SyncTaskExecutor

// Make sure *all* loaders go through this function
fun <T> runLoader(f: () -> T): CompletableFuture<T> {
    return CompletableFuture.supplyAsync(f, SyncTaskExecutor())
}

fun <K, V> makeBatchLoader(load: (List<K>) -> List<V?>): BatchLoader<K, V> {
    return BatchLoader { keys -> runLoader { load(keys) } }
}

fun <K, V> makeBatchLoaderWithContext(load: (List<K>, BatchLoaderEnvironment) -> List<V?>): BatchLoaderWithContext<K, V> {
    return BatchLoaderWithContext { keys, env -> runLoader { load(keys, env) } }
}

TransactionalExecutionStrategy.kt

package at.sigmoid

import graphql.execution.AsyncExecutionStrategy
import graphql.execution.DataFetcherExceptionHandler
import graphql.execution.ExecutionContext
import graphql.execution.ExecutionStrategyParameters
import java.util.concurrent.CompletableFuture
import org.springframework.transaction.PlatformTransactionManager
import org.springframework.transaction.TransactionDefinition
import org.springframework.transaction.annotation.Isolation
import org.springframework.transaction.annotation.Propagation
import org.springframework.transaction.annotation.Transactional


class TransactionalExecutionStrategy(
    val transactionManager: PlatformTransactionManager,
    val exceptionHandler: DataFetcherExceptionHandler,
) : AsyncExecutionStrategy(exceptionHandler) {
    override fun execute(
        executionContext: ExecutionContext,
        parameters: ExecutionStrategyParameters
    ): CompletableFuture<ExecutionResult> {
        // Create transaction for query execution
        // Adjust transaction settings as required
        val tx = transactionManager.getTransaction(
            object : TransactionDefinition {
                override fun getPropagationBehavior(): Int {
                    return Propagation.REQUIRED
                }

                override fun getIsolationLevel(): Int {
                    return Isolation.SERIALIZABLE
                }

                override fun getTimeout(): Int {
                    return TransactionDefinition.TIMEOUT_DEFAULT
                }

                override fun isReadOnly(): Boolean {
                    return false
                }

                override fun getName(): String? {
                    return super.getName()
                }
            }
        )

        return super.execute(executionContext, parameters)
            // Commit/rollback transaction after execution completes
            .whenComplete { _, throwable ->
                // Roll back for any exception, customize as required
                if (throwable != null) {
                    tx.setRollbackOnly()
                }

                if (tx.isRollbackOnly) {
                    transactionManager.rollback(tx)
                } else {
                    transactionManager.commit(tx)
                }
            }
    }
}

GraphQLConfiguration.kt

package at.sigmoid

import at.sigmoid.TransactionExecutionStrategy
import graphql.GraphQL
import graphql.schema.idl.RuntimeWiring
import graphql.schema.idl.SchemaGenerator
import graphql.schema.idl.SchemaParser
import graphql.execution.DataFetcherExceptionHandler
import org.springframework.core.io.ClassPathResource
import org.springframework.transaction.PlatformTransactionManager

@Configuration
class GraphQLConfiguration {

    @Bean
    fun transactionalExecutionStrategy(
      val transactionManager: PlatformTransactionManager,
      val exceptionHandler: DataFetcherExceptionHandler, // Provide custom exception handler as separate bean
    ) {
      return TransactionalExecutionStrategy(transactionManager, exceptionHandler)
    }

    @Bean
    fun graphql(
        runtimeWiring: RuntimeWiring, // Provide custom runtime wiring as a separate bean
        executionStrategy: TransactionalExecutionStrategy,
    ): GraphQL {
        val sdl = ClassPathResource("schema.graphql").inputStream.use { it.bufferedReader().readText() }
        val typeRegistry = SchemaParser().parse(sdl)
        val schemaGenerator = SchemaGenerator()
        val schema = schemaGenerator.makeExecutableSchema(typeRegistry, runtimeWiring)

        return GraphQL
            .newGraphQL(schema)
            .mutationExecutionStrategy(executionStrategy)
            .queryExecutionStrategy(executionStrategy)
            .build()
    }
}