Skip to content

Commit fe88f37

Browse files
chore: tests for batching and deduplication by field selection (ExpediaGroup#1449)
1 parent aec9a74 commit fe88f37

File tree

10 files changed

+369
-20
lines changed

10 files changed

+369
-20
lines changed
Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,7 @@ import kotlinx.coroutines.awaitAll
4747
import kotlinx.coroutines.future.await
4848
import kotlinx.coroutines.runBlocking
4949

50-
enum class DataLoaderInstrumentationStrategy { LEVEL_DISPATCHED, SYNC_EXHAUSTION }
51-
52-
object TestGraphQL {
50+
object AstronautGraphQL {
5351
private val schema = """
5452
type Query {
5553
astronaut(id: ID!): Astronaut
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture
2+
3+
enum class DataLoaderInstrumentationStrategy { LEVEL_DISPATCHED, SYNC_EXHAUSTION }
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright 2022 Expedia, Inc
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.expediagroup.graphql.dataloader.instrumentation.fixture
18+
19+
import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistry
20+
import com.expediagroup.graphql.dataloader.KotlinDataLoaderRegistryFactory
21+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.ProductDataLoader
22+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.ProductService
23+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.ProductServiceRequest
24+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Product
25+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductDetails
26+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductSummary
27+
import com.expediagroup.graphql.dataloader.instrumentation.level.state.ExecutionLevelDispatchedState
28+
import com.expediagroup.graphql.dataloader.instrumentation.syncexhaustion.state.SyncExecutionExhaustedState
29+
import graphql.ExecutionInput
30+
import graphql.ExecutionResult
31+
import graphql.GraphQL
32+
import graphql.schema.DataFetcher
33+
import graphql.schema.SelectedField
34+
import graphql.schema.idl.RuntimeWiring
35+
import graphql.schema.idl.SchemaGenerator
36+
import graphql.schema.idl.SchemaParser
37+
import graphql.schema.idl.TypeRuntimeWiring
38+
import io.mockk.spyk
39+
import kotlinx.coroutines.async
40+
import kotlinx.coroutines.awaitAll
41+
import kotlinx.coroutines.future.await
42+
import kotlinx.coroutines.runBlocking
43+
import java.util.concurrent.CompletableFuture
44+
45+
object ProductGraphQL {
46+
private val schema = """
47+
type Query {
48+
product(id: ID!): Product
49+
productSummary(productId: ID!): ProductSummary
50+
productDetails(productId: ID!): ProductDetails
51+
}
52+
type Product {
53+
summary: ProductSummary
54+
details: ProductDetails
55+
}
56+
type ProductSummary {
57+
name: String!
58+
}
59+
type ProductDetails {
60+
rating: String!
61+
}
62+
""".trimIndent()
63+
64+
private val productService = ProductService()
65+
66+
private val productDataFetcher = DataFetcher<CompletableFuture<Product>> { environment ->
67+
val productId = environment.getArgument<String>("id").toInt()
68+
val selectionFields = environment.selectionSet.immediateFields.map(SelectedField::getName)
69+
productService.getProduct(
70+
ProductServiceRequest(productId, selectionFields),
71+
environment
72+
)
73+
}
74+
75+
private val productSummaryDataFetcher = DataFetcher<CompletableFuture<ProductSummary>> { environment ->
76+
val productId = environment.getArgument<String>("productId").toInt()
77+
val selectionFields = listOf("summary")
78+
productService.getProduct(
79+
ProductServiceRequest(productId, selectionFields),
80+
environment
81+
).thenApply(Product::summary)
82+
}
83+
84+
private val productDetailsDataFetcher = DataFetcher<CompletableFuture<ProductDetails>> { environment ->
85+
val productId = environment.getArgument<String>("productId").toInt()
86+
val selectionFields = listOf("details")
87+
productService.getProduct(
88+
ProductServiceRequest(productId, selectionFields),
89+
environment
90+
).thenApply(Product::details)
91+
}
92+
93+
private val runtimeWiring = RuntimeWiring.newRuntimeWiring().apply {
94+
type(
95+
TypeRuntimeWiring.newTypeWiring("Query")
96+
.dataFetcher("product", productDataFetcher)
97+
.dataFetcher("productSummary", productSummaryDataFetcher)
98+
.dataFetcher("productDetails", productDetailsDataFetcher)
99+
)
100+
}.build()
101+
102+
val builder: GraphQL.Builder = GraphQL.newGraphQL(
103+
SchemaGenerator().makeExecutableSchema(
104+
SchemaParser().parse(schema),
105+
runtimeWiring
106+
)
107+
)
108+
109+
fun execute(
110+
graphQL: GraphQL,
111+
queries: List<String>,
112+
dataLoaderInstrumentationStrategy: DataLoaderInstrumentationStrategy
113+
): Pair<List<ExecutionResult>, KotlinDataLoaderRegistry> {
114+
val kotlinDataLoaderRegistry = spyk(
115+
KotlinDataLoaderRegistryFactory(
116+
ProductDataLoader()
117+
).generate()
118+
)
119+
120+
val graphQLContext = mapOf(
121+
when (dataLoaderInstrumentationStrategy) {
122+
DataLoaderInstrumentationStrategy.SYNC_EXHAUSTION ->
123+
SyncExecutionExhaustedState::class to SyncExecutionExhaustedState(
124+
queries.size,
125+
kotlinDataLoaderRegistry
126+
)
127+
DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED ->
128+
ExecutionLevelDispatchedState::class to ExecutionLevelDispatchedState(
129+
queries.size
130+
)
131+
}
132+
)
133+
134+
val results = runBlocking {
135+
queries.map { query ->
136+
async {
137+
graphQL.executeAsync(
138+
ExecutionInput
139+
.newExecutionInput(query)
140+
.dataLoaderRegistry(kotlinDataLoaderRegistry)
141+
.graphQLContext(graphQLContext)
142+
.build()
143+
).await()
144+
}
145+
}.awaitAll()
146+
}
147+
148+
return Pair(results, kotlinDataLoaderRegistry)
149+
}
150+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher
2+
3+
import com.expediagroup.graphql.dataloader.KotlinDataLoader
4+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Product
5+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.extensions.toListOfNullables
6+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.repository.ProductRepository
7+
import graphql.schema.DataFetchingEnvironment
8+
import org.dataloader.BatchLoader
9+
import java.util.Optional
10+
import java.util.concurrent.CompletableFuture
11+
12+
class ProductDataLoader : KotlinDataLoader<ProductServiceRequest, Product?> {
13+
override val dataLoaderName: String = "ProductDataLoader"
14+
override fun getBatchLoader(): BatchLoader<ProductServiceRequest, Product?> =
15+
BatchLoader<ProductServiceRequest, Product?> { requests ->
16+
ProductRepository
17+
.getProducts(requests)
18+
.collectList()
19+
.map(List<Optional<Product>>::toListOfNullables)
20+
.toFuture()
21+
}
22+
}
23+
24+
data class ProductServiceRequest(val id: Int, val fields: List<String>)
25+
26+
class ProductService {
27+
fun getProduct(
28+
request: ProductServiceRequest,
29+
environment: DataFetchingEnvironment
30+
): CompletableFuture<Product> =
31+
environment
32+
.getDataLoader<ProductServiceRequest, Product>("ProductDataLoader")
33+
.load(request)
34+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture.domain
2+
3+
data class Product(
4+
val id: Int,
5+
val summary: ProductSummary?,
6+
val details: ProductDetails?
7+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture.domain
2+
3+
data class ProductDetails(
4+
val rating: String
5+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture.domain
2+
3+
data class ProductSummary(
4+
val name: String
5+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package com.expediagroup.graphql.dataloader.instrumentation.fixture.repository
2+
3+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.datafetcher.ProductServiceRequest
4+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.Product
5+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductDetails
6+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.domain.ProductSummary
7+
import reactor.core.publisher.Flux
8+
import reactor.kotlin.core.publisher.toFlux
9+
import reactor.kotlin.core.publisher.toMono
10+
import java.time.Duration
11+
import java.util.Optional
12+
13+
object ProductRepository {
14+
private val products = listOf(
15+
Product(
16+
1,
17+
ProductSummary("Product 1"),
18+
ProductDetails("5 out of 10")
19+
),
20+
Product(
21+
2,
22+
ProductSummary("Product 2"),
23+
ProductDetails("10 out of 10")
24+
)
25+
)
26+
27+
/**
28+
* let's assume data batch loader provides 4 requests "keys" to getProducts
29+
* - 2 for productId 1 fetching summary and details respectively
30+
* - 2 for productId 2 fetching summary and details respectively
31+
*
32+
* here we would need to aggregate 2 requests for each productId into 1 like this
33+
* - 1 request for productId 1 fetching summary and details
34+
* - 1 request for productId 1 fetching summary and details
35+
*/
36+
fun getProducts(requests: List<ProductServiceRequest>): Flux<Optional<Product>> {
37+
val reducedRequests = requests
38+
.groupBy(ProductServiceRequest::id)
39+
.mapValues { (productId, requests) ->
40+
ProductServiceRequest(
41+
productId,
42+
requests.map(ProductServiceRequest::fields).flatten().distinct()
43+
)
44+
}.values.toList()
45+
46+
val results = reducedRequests.mapNotNull { productRequest ->
47+
products
48+
.firstOrNull { it.id == productRequest.id }
49+
?.let { property ->
50+
Product(
51+
property.id,
52+
when {
53+
productRequest.fields.contains("summary") -> property.summary
54+
else -> null
55+
},
56+
when {
57+
productRequest.fields.contains("details") -> property.details
58+
else -> null
59+
}
60+
)
61+
}
62+
}.associateBy(Product::id)
63+
64+
return requests
65+
.toFlux()
66+
.flatMap { request ->
67+
Optional.ofNullable(results[request.id]).toMono().delayElement(Duration.ofMillis(200))
68+
}
69+
}
70+
}

executions/graphql-kotlin-dataloader-instrumentation/src/test/kotlin/com/expediagroup/graphql/dataloader/instrumentation/level/DataLoaderLevelDispatchedInstrumentationTest.kt

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,13 @@
1717
package com.expediagroup.graphql.dataloader.instrumentation.level
1818

1919
import com.expediagroup.graphql.dataloader.instrumentation.fixture.DataLoaderInstrumentationStrategy
20-
import com.expediagroup.graphql.dataloader.instrumentation.fixture.TestGraphQL
20+
import com.expediagroup.graphql.dataloader.instrumentation.fixture.AstronautGraphQL
2121
import io.mockk.verify
2222
import org.junit.jupiter.api.Test
2323
import kotlin.test.assertEquals
2424

2525
class DataLoaderLevelDispatchedInstrumentationTest {
26-
private val graphQL = TestGraphQL.builder
26+
private val graphQL = AstronautGraphQL.builder
2727
.instrumentation(DataLoaderLevelDispatchedInstrumentation())
2828
// graphql java adds DataLoaderDispatcherInstrumentation by default
2929
.doNotAddDefaultInstrumentations()
@@ -38,7 +38,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
3838
"{ mission(id: 4) { designation } }"
3939
)
4040

41-
val (results, kotlinDataLoaderRegistry) = TestGraphQL.execute(
41+
val (results, kotlinDataLoaderRegistry) = AstronautGraphQL.execute(
4242
graphQL,
4343
queries,
4444
DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED
@@ -69,7 +69,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
6969
"{ nasa { mission(id: 4) { id designation } } }"
7070
)
7171

72-
val (results, kotlinDataLoaderRegistry) = TestGraphQL.execute(
72+
val (results, kotlinDataLoaderRegistry) = AstronautGraphQL.execute(
7373
graphQL,
7474
queries,
7575
DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED
@@ -104,7 +104,7 @@ class DataLoaderLevelDispatchedInstrumentationTest {
104104
"{ mission(id: 4) { designation } }"
105105
)
106106

107-
val (results, kotlinDataLoaderRegistry) = TestGraphQL.execute(
107+
val (results, kotlinDataLoaderRegistry) = AstronautGraphQL.execute(
108108
graphQL,
109109
queries,
110110
DataLoaderInstrumentationStrategy.LEVEL_DISPATCHED

0 commit comments

Comments
 (0)