From f4a7793f89b6ffd40d27f50c230c1361d1f9884f Mon Sep 17 00:00:00 2001 From: uboness Date: Thu, 14 Aug 2014 02:55:09 +0200 Subject: [PATCH] Introduced a new elasticsearch exception family that can hold headers - These heades will be copied as response header on the rest response --- .../elasticsearch/ElasticsearchException.java | 62 +++++++++++++++++++ .../elasticsearch/rest/BytesRestResponse.java | 3 + .../elasticsearch/rest/HasRestHeaders.java | 38 ++++++++++++ .../org/elasticsearch/rest/RestResponse.java | 18 +++++- .../rest/BytesRestResponseTests.java | 56 +++++++++++++++++ 5 files changed, 176 insertions(+), 1 deletion(-) create mode 100644 src/main/java/org/elasticsearch/rest/HasRestHeaders.java create mode 100644 src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java diff --git a/src/main/java/org/elasticsearch/ElasticsearchException.java b/src/main/java/org/elasticsearch/ElasticsearchException.java index dc1ab99424a..2da9e46f1d6 100644 --- a/src/main/java/org/elasticsearch/ElasticsearchException.java +++ b/src/main/java/org/elasticsearch/ElasticsearchException.java @@ -19,8 +19,17 @@ package org.elasticsearch; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; +import org.elasticsearch.common.Nullable; +import org.elasticsearch.common.collect.Tuple; +import org.elasticsearch.rest.HasRestHeaders; import org.elasticsearch.rest.RestStatus; +import java.util.List; +import java.util.Map; + /** * A base class for all elasticsearch exceptions. */ @@ -152,4 +161,57 @@ public class ElasticsearchException extends RuntimeException { return false; } } + + /** + * A base class for exceptions that should carry rest headers + */ + @SuppressWarnings("unchecked") + public static class WithRestHeaders extends ElasticsearchException implements HasRestHeaders { + + private final ImmutableMap> headers; + + public WithRestHeaders(String msg, Tuple... headers) { + super(msg); + this.headers = headers(headers); + } + + public WithRestHeaders(String msg, @Nullable ImmutableMap> headers) { + super(msg); + this.headers = headers != null ? headers : ImmutableMap.>of(); + } + + public WithRestHeaders(String msg, Throwable cause, Tuple... headers) { + super(msg, cause); + this.headers = headers(headers); + } + + public WithRestHeaders(String msg, Throwable cause, @Nullable ImmutableMap> headers) { + super(msg, cause); + this.headers = headers != null ? headers : ImmutableMap.>of(); + } + + public ImmutableMap> getHeaders() { + return headers; + } + + protected static Tuple header(String name, String... values) { + return Tuple.tuple(name, values); + } + + private static ImmutableMap> headers(Tuple... headers) { + Map> map = Maps.newHashMap(); + for (Tuple header : headers) { + List list = map.get(header.v1()); + if (list == null) { + list = Lists.newArrayList(header.v2()); + map.put(header.v1(), list); + } else { + for (String value : header.v2()) { + list.add(value); + } + } + } + return ImmutableMap.copyOf(map); + } + } } diff --git a/src/main/java/org/elasticsearch/rest/BytesRestResponse.java b/src/main/java/org/elasticsearch/rest/BytesRestResponse.java index 5143db8f580..9a24099680f 100644 --- a/src/main/java/org/elasticsearch/rest/BytesRestResponse.java +++ b/src/main/java/org/elasticsearch/rest/BytesRestResponse.java @@ -93,6 +93,9 @@ public class BytesRestResponse extends RestResponse { this.content = builder.bytes(); this.contentType = builder.contentType().restContentType(); } + if (t instanceof HasRestHeaders) { + addHeaders(((HasRestHeaders) t).getHeaders()); + } this.contentThreadSafe = true; } diff --git a/src/main/java/org/elasticsearch/rest/HasRestHeaders.java b/src/main/java/org/elasticsearch/rest/HasRestHeaders.java new file mode 100644 index 00000000000..e76a7bed68d --- /dev/null +++ b/src/main/java/org/elasticsearch/rest/HasRestHeaders.java @@ -0,0 +1,38 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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 + * + * http://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 org.elasticsearch.rest; + +import java.util.List; +import java.util.Map; + +/** + * Classes that carry rest headers should implement this interface. Specifically, exceptions that + * get translated to a rest response, can implement this interface and the headers will be added + * the the response. + * + * @see org.elasticsearch.ElasticsearchException.WithRestHeaders + */ +public interface HasRestHeaders { + + /** + * @return The rest headers + */ + Map> getHeaders(); +} diff --git a/src/main/java/org/elasticsearch/rest/RestResponse.java b/src/main/java/org/elasticsearch/rest/RestResponse.java index 1673b2c2190..5d839015894 100644 --- a/src/main/java/org/elasticsearch/rest/RestResponse.java +++ b/src/main/java/org/elasticsearch/rest/RestResponse.java @@ -19,6 +19,8 @@ package org.elasticsearch.rest; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; import org.elasticsearch.common.Nullable; import org.elasticsearch.common.bytes.BytesReference; @@ -30,7 +32,7 @@ import java.util.Map; /** * */ -public abstract class RestResponse { +public abstract class RestResponse implements HasRestHeaders { protected Map> customHeaders; @@ -56,6 +58,20 @@ public abstract class RestResponse { */ public abstract RestStatus status(); + public void addHeaders(Map> headers) { + if (customHeaders == null) { + customHeaders = new HashMap<>(headers.size()); + } + for (Map.Entry> entry : headers.entrySet()) { + List values = customHeaders.get(entry.getKey()); + if (values == null) { + values = Lists.newArrayList(); + customHeaders.put(entry.getKey(), values); + } + values.addAll(entry.getValue()); + } + } + /** * Add a custom header. */ diff --git a/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java b/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java new file mode 100644 index 00000000000..79c0f295ba4 --- /dev/null +++ b/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java @@ -0,0 +1,56 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch 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 + * + * http://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 org.elasticsearch.rest; + +import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.test.ElasticsearchTestCase; +import org.junit.Test; + +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.notNullValue; + +/** + * + */ +public class BytesRestResponseTests extends ElasticsearchTestCase { + + @Test + public void testWithHeaders() throws Exception { + RestRequest request = new FakeRestRequest(); + RestChannel channel = new RestChannel(request) { + @Override + public void sendResponse(RestResponse response) { + } + }; + BytesRestResponse response = new BytesRestResponse(channel, new ExceptionWithHeaders()); + assertThat(response.getHeaders().get("n1"), notNullValue()); + assertThat(response.getHeaders().get("n1"), contains("v11", "v12")); + assertThat(response.getHeaders().get("n2"), notNullValue()); + assertThat(response.getHeaders().get("n2"), contains("v21", "v22")); + } + + + private static class ExceptionWithHeaders extends ElasticsearchException.WithRestHeaders { + + ExceptionWithHeaders() { + super("", header("n1", "v11", "v12"), header("n2", "v21", "v22")); + } + } +}