diff --git a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/BookQueryService.kt b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/BookQueryService.kt
index 5425e96b46..0bdff2957f 100644
--- a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/BookQueryService.kt
+++ b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/BookQueryService.kt
@@ -19,6 +19,7 @@ package com.expediagroup.graphql.examples.server.ktor.schema
 import com.expediagroup.graphql.examples.server.ktor.schema.dataloaders.BookDataLoader
 import com.expediagroup.graphql.examples.server.ktor.schema.models.Book
 import com.expediagroup.graphql.generator.annotations.GraphQLDescription
+import com.expediagroup.graphql.server.extensions.getValuesFromDataLoader
 import com.expediagroup.graphql.server.operations.Query
 import graphql.schema.DataFetchingEnvironment
 import java.util.concurrent.CompletableFuture
@@ -30,8 +31,7 @@ class BookQueryService : Query {
     @GraphQLDescription("Return list of books based on BookSearchParameter options")
     @Suppress("unused")
     fun searchBooks(params: BookSearchParameters, dfe: DataFetchingEnvironment): CompletableFuture<List<Book>> =
-        dfe.getDataLoader<Int, Book>(BookDataLoader.dataLoaderName)
-            .loadMany(params.ids)
+        dfe.getValuesFromDataLoader(BookDataLoader.dataLoaderName, params.ids)
 }
 
 data class BookSearchParameters(val ids: List<Int>)
diff --git a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/CourseQueryService.kt b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/CourseQueryService.kt
index a237820550..bcde4d7e1e 100644
--- a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/CourseQueryService.kt
+++ b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/CourseQueryService.kt
@@ -18,14 +18,14 @@ package com.expediagroup.graphql.examples.server.ktor.schema
 
 import com.expediagroup.graphql.examples.server.ktor.schema.dataloaders.CourseDataLoader
 import com.expediagroup.graphql.examples.server.ktor.schema.models.Course
+import com.expediagroup.graphql.server.extensions.getValuesFromDataLoader
 import com.expediagroup.graphql.server.operations.Query
 import graphql.schema.DataFetchingEnvironment
 import java.util.concurrent.CompletableFuture
 
 class CourseQueryService : Query {
     fun searchCourses(params: CourseSearchParameters, dfe: DataFetchingEnvironment): CompletableFuture<List<Course>> =
-        dfe.getDataLoader<Int, Course>(CourseDataLoader.dataLoaderName)
-            .loadMany(params.ids)
+        dfe.getValuesFromDataLoader(CourseDataLoader.dataLoaderName, params.ids)
 }
 
 data class CourseSearchParameters(val ids: List<Int>)
diff --git a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/UniversityQueryService.kt b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/UniversityQueryService.kt
index 58e939339d..ffaab784dc 100644
--- a/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/UniversityQueryService.kt
+++ b/examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/schema/UniversityQueryService.kt
@@ -18,14 +18,14 @@ package com.expediagroup.graphql.examples.server.ktor.schema
 
 import com.expediagroup.graphql.examples.server.ktor.schema.dataloaders.UniversityDataLoader
 import com.expediagroup.graphql.examples.server.ktor.schema.models.University
+import com.expediagroup.graphql.server.extensions.getValuesFromDataLoader
 import com.expediagroup.graphql.server.operations.Query
 import graphql.schema.DataFetchingEnvironment
 import java.util.concurrent.CompletableFuture
 
 class UniversityQueryService : Query {
     fun searchUniversities(params: UniversitySearchParameters, dfe: DataFetchingEnvironment): CompletableFuture<List<University>> =
-        dfe.getDataLoader<Int, University>(UniversityDataLoader.dataLoaderName)
-            .loadMany(params.ids)
+        dfe.getValuesFromDataLoader(UniversityDataLoader.dataLoaderName, params.ids)
 }
 
 data class UniversitySearchParameters(val ids: List<Int>)
diff --git a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/query/DataLoaderQuery.kt b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/query/DataLoaderQuery.kt
index f30ed260a7..0e8c93af0b 100644
--- a/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/query/DataLoaderQuery.kt
+++ b/examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/query/DataLoaderQuery.kt
@@ -51,7 +51,7 @@ class DataLoaderQuery : Query {
 class CompanyDataFetcher : DataFetcher<CompletableFuture<Company>> {
 
     override fun get(environment: DataFetchingEnvironment): CompletableFuture<Company> {
-        val companyId = environment.getSource<Employee>().companyId
+        val companyId = environment.getSource<Employee>()?.companyId ?: throw IllegalArgumentException("companyId is null")
         return environment.getValueFromDataLoader(CompanyDataLoader.name, companyId)
     }
 }
diff --git a/executions/graphql-kotlin-automatic-persisted-queries/src/test/kotlin/com/expediagroup/graphql/apq/fixture/ProductGraphQL.kt b/executions/graphql-kotlin-automatic-persisted-queries/src/test/kotlin/com/expediagroup/graphql/apq/fixture/ProductGraphQL.kt
index 53b1603aa6..17c7fa0435 100644
--- a/executions/graphql-kotlin-automatic-persisted-queries/src/test/kotlin/com/expediagroup/graphql/apq/fixture/ProductGraphQL.kt
+++ b/executions/graphql-kotlin-automatic-persisted-queries/src/test/kotlin/com/expediagroup/graphql/apq/fixture/ProductGraphQL.kt
@@ -52,7 +52,9 @@ object ProductGraphQL {
     private val runtimeWiring = RuntimeWiring.newRuntimeWiring().apply {
         type(
             TypeRuntimeWiring.newTypeWiring("Query")
-                .dataFetcher("product") { products[it.getArgument<String>("id").toInt()] }
+                .dataFetcher("product") {
+                    it.getArgument<String>("id")?.let { id -> products[id.toInt()] }
+                }
         )
     }.build()
 
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt
index 0d0ae4d06d..5927716b38 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/AstronautGraphQL.kt
@@ -93,18 +93,16 @@ object AstronautGraphQL {
 
     private val astronautService = AstronautService()
     private val astronautDataFetcher = DataFetcher { environment ->
+        val astronautId = environment.getArgument<String>("id")?.toInt() ?: throw IllegalArgumentException("Astronaut ID is null")
         astronautService.getAstronaut(
-            AstronautServiceRequest(
-                environment.getArgument<String>("id").toInt()
-            ),
+            AstronautServiceRequest(astronautId),
             environment
         )
     }
     private val createAstronautDataFetcher = DataFetcher { environment ->
+        val astronautName = environment.getArgument<String>("name") ?: throw IllegalArgumentException("Astronaut name is null")
         astronautService.createAstronaut(
-            CreateAstronautServiceRequest(
-                environment.getArgument("name")
-            )
+            CreateAstronautServiceRequest(astronautName)
         )
     }
     private val astronautsDataFetcher = DataFetcher { environment ->
@@ -118,10 +116,9 @@ object AstronautGraphQL {
 
     private val missionService = MissionService()
     private val missionDataFetcher = DataFetcher { environment ->
+        val missionId = environment.getArgument<String>("id")?.toInt() ?: throw IllegalArgumentException("Mission ID is null")
         missionService.getMission(
-            MissionServiceRequest(
-                environment.getArgument<String>("id").toInt()
-            ),
+            MissionServiceRequest(missionId),
             environment
         )
     }
@@ -134,26 +131,26 @@ object AstronautGraphQL {
         )
     }
     private val missionsByAstronautDataFetcher = DataFetcher { environment ->
-        val astronaut = environment.getSource<Astronaut>()
+        val astronautId = environment.getSource<Astronaut>()?.id ?: throw IllegalArgumentException("Astronaut ID is null")
         missionService
             .getMissionsByAstronaut(
-                MissionServiceRequest(0, astronaut.id),
+                MissionServiceRequest(0, astronautId),
                 environment
             )
     }
 
     private val planetService = PlanetService()
     private val planetsByMissionDataFetcher = DataFetcher { environment ->
-        val mission = environment.getSource<Mission>()
+        val missionId = environment.getSource<Mission>()?.id ?: throw IllegalArgumentException("Mission ID is null")
         planetService.getPlanets(
-            PlanetServiceRequest(0, mission.id),
+            PlanetServiceRequest(0, missionId),
             environment
         )
     }
     private val planetsByAstronautDataFetcher = DataFetcher { environment ->
-        val astronaut = environment.getSource<Astronaut>()
+        val astronautId = environment.getSource<Astronaut>()?.id ?: throw IllegalArgumentException("Astronaut ID is null")
         astronautService.getPlanets(
-            AstronautServiceRequest(astronaut.id),
+            AstronautServiceRequest(astronautId),
             environment
         )
     }
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt
index 50aef25a12..89f6e020c1 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/ProductGraphQL.kt
@@ -64,7 +64,7 @@ object ProductGraphQL {
     private val productService = ProductService()
 
     private val productDataFetcher = DataFetcher<CompletableFuture<Product>> { environment ->
-        val productId = environment.getArgument<String>("id").toInt()
+        val productId = environment.getArgument<String>("id")?.toInt() ?: throw IllegalArgumentException("Product ID is null")
         val selectionFields = environment.selectionSet.immediateFields.map(SelectedField::getName)
         productService.getProduct(
             ProductServiceRequest(productId, selectionFields),
@@ -73,7 +73,7 @@ object ProductGraphQL {
     }
 
     private val productSummaryDataFetcher = DataFetcher<CompletableFuture<ProductSummary>> { environment ->
-        val productId = environment.getArgument<String>("productId").toInt()
+        val productId = environment.getArgument<String>("productId")?.toInt() ?: throw IllegalArgumentException("Product ID is null")
         val selectionFields = listOf("summary")
         productService.getProduct(
             ProductServiceRequest(productId, selectionFields),
@@ -82,7 +82,7 @@ object ProductGraphQL {
     }
 
     private val productDetailsDataFetcher = DataFetcher<CompletableFuture<ProductDetails>> { environment ->
-        val productId = environment.getArgument<String>("productId").toInt()
+        val productId = environment.getArgument<String>("productId")?.toInt() ?: throw IllegalArgumentException("Product ID is null")
         val selectionFields = listOf("details")
         productService.getProduct(
             ProductServiceRequest(productId, selectionFields),
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt
index 0ec154dddc..4e7e814191 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/AstronautService.kt
@@ -59,7 +59,7 @@ class AstronautService {
     ): CompletableFuture<Astronaut> =
         environment
             .getDataLoader<AstronautServiceRequest, Astronaut>("AstronautDataLoader")
-            .load(request)
+            ?.load(request) ?: throw IllegalArgumentException("No data loader called AstronautDataLoader was found")
 
     fun createAstronaut(
         request: CreateAstronautServiceRequest
@@ -73,7 +73,7 @@ class AstronautService {
         requests.isNotEmpty() -> {
             environment
                 .getDataLoader<AstronautServiceRequest, Astronaut>("AstronautDataLoader")
-                .loadMany(requests)
+                ?.loadMany(requests) ?: throw IllegalArgumentException("No data loader called AstronautDataLoader was found")
         }
         else -> {
             AstronautRepository
@@ -89,7 +89,9 @@ class AstronautService {
         environment: DataFetchingEnvironment
     ): CompletableFuture<List<Planet>> {
         val missionsByAstronautDataLoader = environment.getDataLoader<MissionServiceRequest, List<Mission>>("MissionsByAstronautDataLoader")
+            ?: throw IllegalArgumentException("No data loader called MissionsByAstronautDataLoader was found")
         val planetsByMissionDataLoader = environment.getDataLoader<PlanetServiceRequest, List<Planet>>("PlanetsByMissionDataLoader")
+            ?: throw IllegalArgumentException("No data loader called PlanetsByMissionDataLoader was found")
         return missionsByAstronautDataLoader
             .load(MissionServiceRequest(0, astronautId = request.id))
             .thenCompose { missions ->
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/MissionService.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/MissionService.kt
index a5ab03a21c..7e02343cd4 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/MissionService.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/MissionService.kt
@@ -66,7 +66,7 @@ class MissionService {
     ): CompletableFuture<Mission> =
         environment
             .getDataLoader<MissionServiceRequest, Mission>("MissionDataLoader")
-            .load(request)
+            ?.load(request) ?: throw IllegalArgumentException("No data loader called MissionDataLoader was found")
 
     fun getMissions(
         requests: List<MissionServiceRequest>,
@@ -75,7 +75,7 @@ class MissionService {
         requests.isNotEmpty() -> {
             environment
                 .getDataLoader<MissionServiceRequest, Mission>("MissionDataLoader")
-                .loadMany(requests)
+                ?.loadMany(requests) ?: throw IllegalArgumentException("No data loader called MissionDataLoader was found")
         }
         else -> {
             MissionRepository
@@ -92,5 +92,5 @@ class MissionService {
     ): CompletableFuture<List<Mission>> =
         environment
             .getDataLoader<MissionServiceRequest, List<Mission>>("MissionsByAstronautDataLoader")
-            .load(request)
+            ?.load(request) ?: throw IllegalArgumentException("No data loader called MissionsByAstronautDataLoader was found")
 }
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/PlanetService.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/PlanetService.kt
index cd281d333a..945af39a14 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/PlanetService.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/PlanetService.kt
@@ -50,5 +50,5 @@ class PlanetService {
     ): CompletableFuture<List<Planet>> =
         environment
             .getDataLoader<PlanetServiceRequest, List<Planet>>("PlanetsByMissionDataLoader")
-            .load(request)
+            ?.load(request) ?: throw IllegalStateException("No data loader called PlanetsByMissionDataLoader was found")
 }
diff --git a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/ProductService.kt b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/ProductService.kt
index 5ec977b6cd..ff4e9588ea 100644
--- a/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/ProductService.kt
+++ b/executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/fixture/datafetcher/ProductService.kt
@@ -53,5 +53,5 @@ class ProductService {
     ): CompletableFuture<Product> =
         environment
             .getDataLoader<ProductServiceRequest, Product>("ProductDataLoader")
-            .load(request)
+            ?.load(request) ?: throw IllegalStateException("No data loader called ProductDataLoader was found")
 }
diff --git a/generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcher.kt b/generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcher.kt
index fbd345d11d..b0843e31a5 100644
--- a/generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcher.kt
+++ b/generator/graphql-kotlin-federation/src/main/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcher.kt
@@ -54,7 +54,7 @@ open class EntitiesDataFetcher(
      * @return list of resolved nullable entities
      */
     override fun get(env: DataFetchingEnvironment): CompletableFuture<DataFetcherResult<List<Any?>>> {
-        val representations: List<Map<String, Any>> = env.getArgument(REPRESENTATIONS)
+        val representations: List<Map<String, Any>> = env.getArgumentOrDefault(REPRESENTATIONS, listOf(emptyMap()))
 
         val representationsWithoutResolver = mutableListOf<IndexedValue<Map<String, Any>>>()
         val entitiesWithPromiseResolver = mutableListOf<ResolvableEntity<FederatedTypePromiseResolver<*>>>()
diff --git a/generator/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcherTest.kt b/generator/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcherTest.kt
index a08c45aa09..c97387c6de 100644
--- a/generator/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcherTest.kt
+++ b/generator/graphql-kotlin-federation/src/test/kotlin/com/expediagroup/graphql/generator/federation/execution/EntitiesDataFetcherTest.kt
@@ -41,7 +41,7 @@ class EntitiesDataFetcherTest {
             mapOf<String, Any>("__typename" to "User", "userId" to 123, "name" to "testName")
         )
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -57,7 +57,7 @@ class EntitiesDataFetcherTest {
             mapOf<String, Any>("__typename" to "Author", "authorId" to 1)
         )
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
         }
         val result = resolver.get(env).get()
         verifyData(result.data, Author(1, "Author 1"))
@@ -72,7 +72,7 @@ class EntitiesDataFetcherTest {
             mapOf<String, Any>("__typename" to "User", "userId" to 123, "name" to "testName")
         )
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
         }
 
         val result = resolver.get(env).get()
@@ -86,7 +86,7 @@ class EntitiesDataFetcherTest {
         val mockUserResolver = mockk<FederatedTypeSuspendResolver<*>> { every { typeName } returns "User" }
         val resolver = EntitiesDataFetcher(mockBookResolver, mockUserResolver)
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns listOf(emptyMap<String, Any>())
+            every { getArgumentOrDefault<Any>(any(), any()) } returns listOf(emptyMap<String, Any>())
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -100,7 +100,7 @@ class EntitiesDataFetcherTest {
         val resolver = EntitiesDataFetcher(emptyList())
         val representations = listOf(mapOf<String, Any>("__typename" to "User", "userId" to 123, "name" to "testName"))
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -118,7 +118,7 @@ class EntitiesDataFetcherTest {
         val resolver = EntitiesDataFetcher(listOf(mockUserResolver))
         val representations = listOf(mapOf<String, Any>("__typename" to "User", "userId" to 123, "name" to "testName"))
         val env: DataFetchingEnvironment = mockk {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -136,7 +136,7 @@ class EntitiesDataFetcherTest {
         }
         val representations = listOf(user1.toRepresentation(), book.toRepresentation(), user2.toRepresentation())
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -163,7 +163,7 @@ class EntitiesDataFetcherTest {
         }
         val representations = listOf(user.toRepresentation(), book.toRepresentation())
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
@@ -195,7 +195,7 @@ class EntitiesDataFetcherTest {
         )
 
         val env = mockk<DataFetchingEnvironment> {
-            every { getArgument<Any>(any()) } returns representations
+            every { getArgumentOrDefault<Any>(any(), any()) } returns representations
             every { graphQlContext } returns GraphQLContext.newContext().build()
         }
 
diff --git a/generator/graphql-kotlin-schema-generator/src/test/kotlin/com/expediagroup/graphql/generator/directives/DirectiveTests.kt b/generator/graphql-kotlin-schema-generator/src/test/kotlin/com/expediagroup/graphql/generator/directives/DirectiveTests.kt
index 62c6098b53..12364e5cea 100644
--- a/generator/graphql-kotlin-schema-generator/src/test/kotlin/com/expediagroup/graphql/generator/directives/DirectiveTests.kt
+++ b/generator/graphql-kotlin-schema-generator/src/test/kotlin/com/expediagroup/graphql/generator/directives/DirectiveTests.kt
@@ -160,7 +160,7 @@ class QueryWithDeprecatedFields {
 data class ClassWithDeprecatedField(
     val something: String,
     @Deprecated("this field is deprecated")
-    val deprecatedField: String,
+    val deprecatedField: String?,
     @GraphQLDeprecated("this field is also deprecated")
-    val graphqlDeprecatedField: String = ""
+    val graphqlDeprecatedField: String? = ""
 )
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 2835826d59..e460a39e6e 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -3,7 +3,7 @@ android-plugin = "8.1.1"
 classgraph = "4.8.170"
 dataloader = "3.3.0"
 federation = "4.4.1"
-graphql-java = "0.0.0-2024-04-23-expedia-branch-1"
+graphql-java = "22.1"
 graalvm = "0.10.1"
 jackson = "2.15.4"
 # kotlin version has to match the compile-testing compiler version
diff --git a/integration/graalvm/common-graalvm-server/src/main/kotlin/com/expediagroup/graalvm/schema/AsyncQuery.kt b/integration/graalvm/common-graalvm-server/src/main/kotlin/com/expediagroup/graalvm/schema/AsyncQuery.kt
index e40323060a..9a5b68cc8b 100644
--- a/integration/graalvm/common-graalvm-server/src/main/kotlin/com/expediagroup/graalvm/schema/AsyncQuery.kt
+++ b/integration/graalvm/common-graalvm-server/src/main/kotlin/com/expediagroup/graalvm/schema/AsyncQuery.kt
@@ -18,6 +18,7 @@ package com.expediagroup.graalvm.schema
 
 import com.expediagroup.graalvm.schema.dataloader.EXAMPLE_LOADER
 import com.expediagroup.graphql.generator.scalars.ID
+import com.expediagroup.graphql.server.extensions.getValueFromDataLoader
 import com.expediagroup.graphql.server.operations.Query
 import graphql.schema.DataFetchingEnvironment
 import kotlinx.coroutines.coroutineScope
@@ -33,7 +34,7 @@ class AsyncQuery : Query {
 
     fun future(): CompletableFuture<Int> = CompletableFuture.completedFuture(random.nextInt())
 
-    fun dataLoader(env: DataFetchingEnvironment, id: ID): CompletableFuture<String> = env.getDataLoader<ID, String>(EXAMPLE_LOADER).load(id)
+    fun dataLoader(env: DataFetchingEnvironment, id: ID): CompletableFuture<String> = env.getValueFromDataLoader(EXAMPLE_LOADER, id)
 
     suspend fun coroutine(): Int = coroutineScope {
         delay(10)
diff --git a/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandlerTest.kt b/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandlerTest.kt
index b2813c32b6..c796973b50 100644
--- a/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandlerTest.kt
+++ b/servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/execution/GraphQLRequestHandlerTest.kt
@@ -25,6 +25,7 @@ import com.expediagroup.graphql.generator.execution.FlowSubscriptionExecutionStr
 import com.expediagroup.graphql.generator.extensions.toGraphQLContext
 import com.expediagroup.graphql.generator.hooks.FlowSubscriptionSchemaGeneratorHooks
 import com.expediagroup.graphql.generator.toSchema
+import com.expediagroup.graphql.server.extensions.getValueFromDataLoader
 import com.expediagroup.graphql.server.types.GraphQLBatchRequest
 import com.expediagroup.graphql.server.types.GraphQLBatchResponse
 import com.expediagroup.graphql.server.types.GraphQLRequest
@@ -389,7 +390,7 @@ class GraphQLRequestHandlerTest {
             id: Int,
             dataFetchingEnvironment: DataFetchingEnvironment
         ): CompletableFuture<User> =
-            dataFetchingEnvironment.getDataLoader<Int, User>("UserDataLoader").load(id)
+            dataFetchingEnvironment.getValueFromDataLoader("UserDataLoader", id)
 
         fun hello(name: String): String = "Hello $name!"