Provide a way to automatically deserialize non-OK JSON response

core/src/main/java/com/linecorp/armeria/client/AggregatedResponseAs.java

@@ -17,6 +17,7 @@
 package com.linecorp.armeria.client;
 
 import java.io.IOException;
+import java.util.function.Predicate;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -37,27 +38,34 @@ static ResponseAs<AggregatedHttpResponse, ResponseEntity<String>> string() {
         return response -> ResponseEntity.of(response.headers(), response.contentUtf8(), response.trailers());
     }
 
-    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(Class<? extends T> clazz) {
-        return response -> newJsonResponseEntity(response, bytes -> JacksonUtil.readValue(bytes, clazz));
+    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(
+            Class<? extends T> clazz, Predicate<AggregatedHttpResponse> predicate) {
+        return response -> newJsonResponseEntity(
+                response, bytes -> JacksonUtil.readValue(bytes, clazz), predicate);
     }
 
-    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(Class<? extends T> clazz,
-                                                                          ObjectMapper mapper) {
-        return response -> newJsonResponseEntity(response, bytes -> mapper.readValue(bytes, clazz));
+    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(
+            Class<? extends T> clazz, ObjectMapper mapper, Predicate<AggregatedHttpResponse> predicate) {
+        return response -> newJsonResponseEntity(response, bytes -> mapper.readValue(bytes, clazz), predicate);
     }
 
-    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(TypeReference<? extends T> typeRef) {
-        return response -> newJsonResponseEntity(response, bytes -> JacksonUtil.readValue(bytes, typeRef));
+    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(
+            TypeReference<? extends T> typeRef, Predicate<AggregatedHttpResponse> predicate) {
+        return response -> newJsonResponseEntity(
+                response, bytes -> JacksonUtil.readValue(bytes, typeRef), predicate);
     }
 
-    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(TypeReference<? extends T> typeRef,
-                                                                          ObjectMapper mapper) {
-        return response -> newJsonResponseEntity(response, bytes -> mapper.readValue(bytes, typeRef));
+    static <T> ResponseAs<AggregatedHttpResponse, ResponseEntity<T>> json(
+            TypeReference<? extends T> typeRef, ObjectMapper mapper,
+            Predicate<AggregatedHttpResponse> predicate) {
+        return response -> newJsonResponseEntity(
+                response, bytes -> mapper.readValue(bytes, typeRef), predicate);
     }
 
     private static <T> ResponseEntity<T> newJsonResponseEntity(AggregatedHttpResponse response,
-                                                               JsonDecoder<T> decoder) {
-        if (!response.status().isSuccess()) {
+                                                               JsonDecoder<T> decoder,
+                                                               Predicate<AggregatedHttpResponse> predicate) {
+        if (!predicate.test(response)) {
             throw newInvalidHttpResponseException(response);
         }
 

core/src/main/java/com/linecorp/armeria/client/BlockingConditionalResponseAs.java

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2023 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.client;
+
+import java.util.function.Predicate;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import com.linecorp.armeria.common.AggregatedHttpResponse;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.ResponseEntity;
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * Provides a way for users to add {@link ResponseAs} mappings to transform an aggregated response
+ * given that the corresponding {@link Predicate} is satisfied. Note that the conditionals are
+ * invoked in the order in which they are added.
+ *
+ * <pre>{@code
+ * RestClient.of(...)
+ *   .get("/")
+ *   .execute(
+ *     ResponseAs.<MyResponse>json(MyMessage.class, res -> res.status().isError())
+ *       .andThenJson(EmptyMessage.class, res -> res.status().isInformational())
+ *       .orElseJson(MyError.class);
+ * }
+ */
+@UnstableApi
+public final class BlockingConditionalResponseAs<V>
+        extends DefaultConditionalResponseAs<HttpResponse, AggregatedHttpResponse, ResponseEntity<V>> {
+    private static final Predicate<AggregatedHttpResponse> TRUE_PREDICATE = unused -> true;
+
+    BlockingConditionalResponseAs(ResponseAs<HttpResponse, AggregatedHttpResponse> originalResponseAs,
+                                  ResponseAs<AggregatedHttpResponse, ResponseEntity<V>> responseAs,
+                                  Predicate<AggregatedHttpResponse> predicate) {
+        super(originalResponseAs, responseAs, predicate);
+    }
+
+    /**
+     * Adds a mapping such that the response content will be deserialized
+     * to the specified {@link Class} if the {@link Predicate} is satisfied.
+     */
+    public BlockingConditionalResponseAs<V> andThenJson(
+            Class<? extends V> clazz, Predicate<AggregatedHttpResponse> predicate) {
+        return andThen(AggregatedResponseAs.json(clazz, predicate), predicate);
+    }
+
+    /**
+     * Adds a mapping such that the response content will be deserialized
+     * to the specified {@link Class} using the {@link ObjectMapper} if the {@link Predicate} is satisfied.
+     */
+    public BlockingConditionalResponseAs<V> andThenJson(
+            Class<? extends V> clazz, ObjectMapper objectMapper, Predicate<AggregatedHttpResponse> predicate) {
+        return andThen(AggregatedResponseAs.json(clazz, objectMapper, predicate), predicate);
+    }
+
+    /**
+     * Adds a mapping such that the response content will be deserialized
+     * with the specified {@link TypeReference} if the {@link Predicate} is satisfied.
+     */
+    public BlockingConditionalResponseAs<V> andThenJson(
+            TypeReference<? extends V> typeRef, Predicate<AggregatedHttpResponse> predicate) {
+        return andThen(AggregatedResponseAs.json(typeRef, predicate), predicate);
+    }
+
+    /**
+     * Adds a mapping such that the response content will be deserialized
+     * with the specified {@link TypeReference} using the {@link ObjectMapper}
+     * if the {@link Predicate} is satisfied.
+     */
+    public BlockingConditionalResponseAs<V> andThenJson(
+            TypeReference<? extends V> typeRef, ObjectMapper objectMapper,
+            Predicate<AggregatedHttpResponse> predicate) {
+        return andThen(AggregatedResponseAs.json(typeRef, objectMapper, predicate), predicate);
+    }
+
+    /**
+     * Returns {@link ResponseAs} based on the configured {@link ResponseAs} to {@link Predicate}
+     * mappings. If none of the {@link Predicate}s are satisfied, the content will be deserialized
+     * to the specified {@link Class}.
+     */
+    public ResponseAs<HttpResponse, ResponseEntity<V>> orElseJson(Class<? extends V> clazz) {
+        return orElse(AggregatedResponseAs.json(clazz, TRUE_PREDICATE));
+    }
+
+    /**
+     * Returns {@link ResponseAs} based on the configured {@link ResponseAs} to {@link Predicate}
+     * mappings. If none of the {@link Predicate}s are satisfied, the content will be deserialized
+     * to the specified {@link Class} using the {@link ObjectMapper}.
+     */
+    public ResponseAs<HttpResponse, ResponseEntity<V>> orElseJson(
+            Class<? extends V> clazz, ObjectMapper objectMapper) {
+        return orElse(AggregatedResponseAs.json(clazz, objectMapper, TRUE_PREDICATE));
+    }
+
+    /**
+     * Returns {@link ResponseAs} based on the configured {@link ResponseAs} to {@link Predicate}
+     * mappings. If none of the {@link Predicate}s are satisfied, the content will be deserialized
+     * using the specified {@link TypeReference}.
+     */
+    public ResponseAs<HttpResponse, ResponseEntity<V>> orElseJson(TypeReference<? extends V> typeRef) {
+        return orElse(AggregatedResponseAs.json(typeRef, TRUE_PREDICATE));
+    }
+
+    /**
+     * Returns {@link ResponseAs} based on the configured {@link ResponseAs} to {@link Predicate}
+     * mappings. If none of the {@link Predicate}s are satisfied, the content will be deserialized
+     * using the specified {@link TypeReference} and {@link ObjectMapper}.
+     */
+    public ResponseAs<HttpResponse, ResponseEntity<V>> orElseJson(
+            TypeReference<? extends V> typeRef, ObjectMapper objectMapper) {
+        return orElse(AggregatedResponseAs.json(typeRef, objectMapper, TRUE_PREDICATE));
+    }
+
+    @Override
+    public BlockingConditionalResponseAs<V> andThen(
+            ResponseAs<AggregatedHttpResponse, ResponseEntity<V>> responseAs,
+            Predicate<AggregatedHttpResponse> predicate) {
+        return (BlockingConditionalResponseAs<V>) super.andThen(responseAs, predicate);
+    }
+}

core/src/main/java/com/linecorp/armeria/client/BlockingWebClientRequestPreparation.java

@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.function.Predicate;
 
 import org.reactivestreams.Publisher;
 
@@ -54,6 +55,7 @@ public final class BlockingWebClientRequestPreparation
         implements WebRequestPreparationSetters<AggregatedHttpResponse> {
 
     private final WebClientRequestPreparation delegate;
+    private static final Predicate<AggregatedHttpResponse> SUCCESS_PREDICATE = res -> res.status().isSuccess();
 
     BlockingWebClientRequestPreparation(WebClientRequestPreparation delegate) {
         this.delegate = delegate;
@@ -134,7 +136,7 @@ public TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity<Str
     public <T> TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity<T>> asJson(
             Class<? extends T> clazz) {
         requireNonNull(clazz, "clazz");
-        return as(AggregatedResponseAs.json(clazz));
+        return as(AggregatedResponseAs.json(clazz, SUCCESS_PREDICATE));
     }
 
     /**
@@ -162,7 +164,7 @@ public <T> TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity
             Class<? extends T> clazz, ObjectMapper mapper) {
         requireNonNull(clazz, "clazz");
         requireNonNull(mapper, "mapper");
-        return as(AggregatedResponseAs.json(clazz, mapper));
+        return as(AggregatedResponseAs.json(clazz, mapper, SUCCESS_PREDICATE));
     }
 
     /**
@@ -188,7 +190,7 @@ public <T> TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity
     public <T> TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity<T>> asJson(
             TypeReference<? extends T> typeRef) {
         requireNonNull(typeRef, "typeRef");
-        return as(AggregatedResponseAs.json(typeRef));
+        return as(AggregatedResponseAs.json(typeRef, SUCCESS_PREDICATE));
     }
 
     /**
@@ -214,7 +216,7 @@ public <T> TransformingRequestPreparation<AggregatedHttpResponse, ResponseEntity
             TypeReference<? extends T> typeRef, ObjectMapper mapper) {
         requireNonNull(typeRef, "typeRef");
         requireNonNull(mapper, "mapper");
-        return as(AggregatedResponseAs.json(typeRef, mapper));
+        return as(AggregatedResponseAs.json(typeRef, mapper, SUCCESS_PREDICATE));
     }
 
     @Override

core/src/main/java/com/linecorp/armeria/client/ConditionalResponseAs.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2023 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.client;
+
+import java.util.function.Predicate;
+
+/**
+ * Provides a way for users to add {@link ResponseAs} mappings to transform a response
+ * given that the corresponding {@link Predicate} is satisfied. Note that the conditionals are
+ * invoked in the order in which they are added.
+ *
+ * <pre>{@code
+ * WebClient.of(...)
+ *   .prepare()
+ *   .get("/server_error")
+ *   .as(ResponseAs.blocking()
+ *     .<MyResponse>andThen(
+ *       res -> new MyError(res.status().codeAsText(), res.contentUtf8()),
+ *       res -> res.status().isError())
+ *     .andThen(
+ *       res -> new EmptyMessage(),
+ *       res -> !res.headers().contains("x-header"))
+ *     .orElse(res -> new MyMessage(res.contentUtf8())))
+ * }
+ */
+public interface ConditionalResponseAs<T, R, V> {
+
+    /**
+     * Adds a mapping such that {@link ResponseAs} will be applied if the {@link Predicate} is satisfied.
+     */
+    ConditionalResponseAs<T, R, V> andThen(ResponseAs<R, V> responseAs, Predicate<R> predicate);
+
+    /**
+     * Returns {@link ResponseAs} based on the configured {@link ResponseAs} to {@link Predicate}
+     * mappings. If none of the {@link Predicate}s are satisfied, the specified {@link ResponseAs}
+     * is used as a fallback.
+     */
+    ResponseAs<T, V> orElse(ResponseAs<R, V> responseAs);
+}

core/src/main/java/com/linecorp/armeria/client/DefaultConditionalResponseAs.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2023 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.client;
+
+import java.util.AbstractMap.SimpleEntry;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map.Entry;
+import java.util.function.Predicate;
+
+import com.linecorp.armeria.common.annotation.UnstableApi;
+
+/**
+ * Adds the mapping of {@link ResponseAs} and {@link Predicate} in order to return {@link ResponseAs} whose
+ * {@link Predicate} is evaluated as true.
+ */
+@UnstableApi
+class DefaultConditionalResponseAs<T, R, V> implements ConditionalResponseAs<T, R, V> {
+
+    private final ResponseAs<T, R> originalResponseAs;
+    private final List<Entry<ResponseAs<R, V>, Predicate<R>>> predicateMappingList = new ArrayList<>();
+
+    DefaultConditionalResponseAs(
+            ResponseAs<T, R> originalResponseAs, ResponseAs<R, V> responseAs, Predicate<R> predicate) {
+        this.originalResponseAs = originalResponseAs;
+        andThen(responseAs, predicate);
+    }
+
+    @Override
+    public DefaultConditionalResponseAs<T, R, V> andThen(ResponseAs<R, V> responseAs, Predicate<R> predicate) {
+        predicateMappingList.add(new SimpleEntry<>(responseAs, predicate));
+        return this;
+    }
+
+    @Override
+    public ResponseAs<T, V> orElse(ResponseAs<R, V> responseAs) {
+        return new ResponseAs<T, V>() {
+            @Override
+            public V as(T response) {
+                final R r = originalResponseAs.as(response);
+                for (Entry<ResponseAs<R, V>, Predicate<R>> entry: predicateMappingList) {
+                    if (entry.getValue().test(r)) {
+                        return entry.getKey().as(r);
+                    }
+                }
+                return responseAs.as(r);
+            }
+        };
+    }
+}