diff --git a/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsExecutionResult.kt b/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsExecutionResult.kt index 43915e579..67c83a50c 100644 --- a/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsExecutionResult.kt +++ b/graphql-dgs/src/main/kotlin/com/netflix/graphql/dgs/internal/DgsExecutionResult.kt @@ -31,61 +31,106 @@ import org.springframework.http.HttpStatus import org.springframework.http.ResponseEntity class DgsExecutionResult @JvmOverloads constructor( - val executionResult: ExecutionResult, - val headers: HttpHeaders = HttpHeaders(), - val status: HttpStatus = HttpStatus.OK -) : ExecutionResult by executionResult { - companion object { - private val sentinelObject = Any() - private val logger: Logger = LoggerFactory.getLogger(DgsExecutionResult::class.java) - } + val executionResult: ExecutionResult, + val headers: HttpHeaders = HttpHeaders(), + val status: HttpStatus = HttpStatus.OK + ) : ExecutionResult by executionResult { + companion object { + // defined in here and DgsRestController, for backwards compatibility. + // keep these two variables synced. + const val DGS_RESPONSE_HEADERS_KEY = "dgs-response-headers" + private val sentinelObject = Any() + private val logger: Logger = LoggerFactory.getLogger(DgsExecutionResult::class.java) + } - constructor( - status: HttpStatus = HttpStatus.OK, - headers: HttpHeaders = HttpHeaders.EMPTY, - errors: List = listOf(), - extensions: Map? = null, + init { + addExtensionsHeaderKeyToHeader() + } - // By default, assign data as a sentinel object. - // If we were to default to null here, this constructor - // would be unable to discriminate between an intentionally null - // response and one that the user left default. - data: Any? = sentinelObject - ) : this( - headers = headers, - status = status, - executionResult = ExecutionResultImpl - .newExecutionResult() - .errors(errors) - .extensions(extensions) - .apply { - if (data != sentinelObject) { - data(data) - } - } - .build() - ) + constructor( + status: HttpStatus = HttpStatus.OK, + headers: HttpHeaders = HttpHeaders.EMPTY, + errors: List = listOf(), + extensions: Map? = null, - fun toSpringResponse( - mapper: ObjectMapper = jacksonObjectMapper() - ): ResponseEntity { - val result = try { - TimeTracer.logTime( - { mapper.writeValueAsBytes(this.toSpecification()) }, - logger, - "Serialized JSON result in {}ms" - ) - } catch (ex: InvalidDefinitionException) { - val errorMessage = "Error serializing response: ${ex.message}" - val errorResponse = ExecutionResultImpl(GraphqlErrorBuilder.newError().message(errorMessage).build()) - logger.error(errorMessage, ex) - mapper.writeValueAsBytes(errorResponse.toSpecification()) - } + // By default, assign data as a sentinel object. + // If we were to default to null here, this constructor + // would be unable to discriminate between an intentionally null + // response and one that the user left default. + data: Any? = sentinelObject + ) : this( + headers = headers, + status = status, + executionResult = ExecutionResultImpl + .newExecutionResult() + .errors(errors) + .extensions(extensions) + .apply { + if (data != sentinelObject) { + data(data) + } + } + .build() + ) - return ResponseEntity( - result, - headers, - status - ) - } -} + // for backwards compat with https://github.com/Netflix/dgs-framework/pull/1261. + private fun addExtensionsHeaderKeyToHeader() { + if (executionResult.extensions?.containsKey(DGS_RESPONSE_HEADERS_KEY) == true) { + val dgsResponseHeaders = executionResult.extensions[DGS_RESPONSE_HEADERS_KEY] + + if (dgsResponseHeaders is Map<*, *>) { + dgsResponseHeaders.forEach { + if (it.key != null) { + headers.add(it.key.toString(), it.value?.toString()) + } + } + } else { + logger.warn( + "{} must be of type java.util.Map, but was {}", + DGS_RESPONSE_HEADERS_KEY, + dgsResponseHeaders?.javaClass?.name + ) + } + } + } + + fun toSpringResponse( + mapper: ObjectMapper = jacksonObjectMapper() + ): ResponseEntity { + val result = try { + TimeTracer.logTime( + { mapper.writeValueAsBytes(this.toSpecification()) }, + logger, + "Serialized JSON result in {}ms" + ) + } catch (ex: InvalidDefinitionException) { + val errorMessage = "Error serializing response: ${ex.message}" + val errorResponse = ExecutionResultImpl(GraphqlErrorBuilder.newError().message(errorMessage).build()) + logger.error(errorMessage, ex) + mapper.writeValueAsBytes(errorResponse.toSpecification()) + } + + return ResponseEntity( + result, + headers, + status + ) + } + + // overridden for compatibility with https://github.com/Netflix/dgs-framework/pull/1261. + override fun toSpecification(): MutableMap { + val spec = executionResult.toSpecification() + + if (spec["extensions"] != null && extensions.containsKey(DGS_RESPONSE_HEADERS_KEY)) { + val extensions = spec["extensions"] as Map<*, *> + + if (extensions.size != 1) { + spec["extensions"] = extensions.minus(DGS_RESPONSE_HEADERS_KEY) + } else { + spec.remove("extensions") + } + } + + return spec + } + } \ No newline at end of file