REST Client: NodeSelector for node attributes

client/rest/src/main/java/org/elasticsearch/client/HasAttributeNodeSelector.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.client;
+
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * A {@link NodeSelector} that selects nodes that have a particular value
+ * for an attribute.
+ */
+public final class HasAttributeNodeSelector implements NodeSelector {
+    private final String key;
+    private final String value;
+
+    public HasAttributeNodeSelector(String key, String value) {
+        this.key = key;
+        this.value = value;
+    }
+
+    @Override
+    public void select(Iterable<Node> nodes) {
+        Iterator<Node> itr = nodes.iterator();
+        while (itr.hasNext()) {
+            Map<String, List<String>> allAttributes = itr.next().getAttributes();
+            if (allAttributes == null) continue;
+            List<String> values = allAttributes.get(key);
+            if (values == null || false == values.contains(value)) {
+                itr.remove();
+            }
+        }
+    }
+
+    @Override
+    public String toString() {
+        return key + "=" + value;
+    }
+}

client/rest/src/main/java/org/elasticsearch/client/Node.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.client;
 
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
@@ -52,13 +54,18 @@ public class Node {
      * if we don't know what roles the node has.
      */
     private final Roles roles;
+    /**
+     * Attributes declared on the node.
+     */
+    private final Map<String, List<String>> attributes;
 
     /**
      * Create a {@linkplain Node} with metadata. All parameters except
      * {@code host} are nullable and implementations of {@link NodeSelector}
      * need to decide what to do in their absence.
      */
-    public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version, Roles roles) {
+    public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version,
+            Roles roles, Map<String, List<String>> attributes) {
         if (host == null) {
             throw new IllegalArgumentException("host cannot be null");
         }
@@ -67,13 +74,14 @@ public Node(HttpHost host, Set<HttpHost> boundHosts, String name, String version
         this.name = name;
         this.version = version;
         this.roles = roles;
+        this.attributes = attributes;
     }
 
     /**
      * Create a {@linkplain Node} without any metadata.
      */
     public Node(HttpHost host) {
-        this(host, null, null, null, null);
+        this(host, null, null, null, null, null);
     }
 
     /**
@@ -115,6 +123,13 @@ public Roles getRoles() {
         return roles;
     }
 
+    /**
+     * Attributes declared on the node.
+     */
+    public Map<String, List<String>> getAttributes() {
+        return attributes;
+    }
+
     @Override
     public String toString() {
         StringBuilder b = new StringBuilder();
@@ -131,6 +146,9 @@ public String toString() {
         if (roles != null) {
             b.append(", roles=").append(roles);
         }
+        if (attributes != null) {
+            b.append(", attributes=").append(attributes);
+        }
         return b.append(']').toString();
     }
 
@@ -144,12 +162,13 @@ public boolean equals(Object obj) {
             && Objects.equals(boundHosts, other.boundHosts)
             && Objects.equals(name, other.name)
             && Objects.equals(version, other.version)
-            && Objects.equals(roles, other.roles);
+            && Objects.equals(roles, other.roles)
+            && Objects.equals(attributes, other.attributes);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(host, boundHosts, name, version, roles);
+        return Objects.hash(host, boundHosts, name, version, roles, attributes);
     }
 
     /**

client/rest/src/test/java/org/elasticsearch/client/HasAttributeNodeSelectorTests.java

@@ -0,0 +1,59 @@
+/*
+ * 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.client;
+
+import org.apache.http.HttpHost;
+import org.elasticsearch.client.Node.Roles;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+import static org.junit.Assert.assertEquals;
+
+public class HasAttributeNodeSelectorTests extends RestClientTestCase {
+    public void testHasAttribute() {
+        Node hasAttributeValue = dummyNode(singletonMap("attr", singletonList("val")));
+        Node hasAttributeButNotValue = dummyNode(singletonMap("attr", singletonList("notval")));
+        Node hasAttributeValueInList = dummyNode(singletonMap("attr", Arrays.asList("val", "notval")));
+        Node notHasAttribute = dummyNode(singletonMap("notattr", singletonList("val")));
+        List<Node> nodes = new ArrayList<>();
+        nodes.add(hasAttributeValue);
+        nodes.add(hasAttributeButNotValue);
+        nodes.add(hasAttributeValueInList);
+        nodes.add(notHasAttribute);
+        List<Node> expected = new ArrayList<>();
+        expected.add(hasAttributeValue);
+        expected.add(hasAttributeValueInList);
+        new HasAttributeNodeSelector("attr", "val").select(nodes);
+        assertEquals(expected, nodes);
+    }
+
+    private Node dummyNode(Map<String, List<String>> attributes) {
+        return new Node(new HttpHost("dummy"), Collections.<HttpHost>emptySet(),
+                randomAsciiAlphanumOfLength(5), randomAsciiAlphanumOfLength(5),
+                new Roles(randomBoolean(), randomBoolean(), randomBoolean()),
+                attributes);
+    }
+}

client/rest/src/test/java/org/elasticsearch/client/NodeSelectorTests.java

@@ -66,6 +66,7 @@ public void testNotMasterOnly() {
     private Node dummyNode(boolean master, boolean data, boolean ingest) {
         return new Node(new HttpHost("dummy"), Collections.<HttpHost>emptySet(),
                 randomAsciiAlphanumOfLength(5), randomAsciiAlphanumOfLength(5),
-                new Roles(master, data, ingest));
+                new Roles(master, data, ingest),
+                Collections.<String, List<String>>emptyMap());
     }
 }

client/rest/src/test/java/org/elasticsearch/client/NodeTests.java

@@ -23,49 +23,67 @@
 import org.elasticsearch.client.Node.Roles;
 
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 
 import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 public class NodeTests extends RestClientTestCase {
     public void testToString() {
+        Map<String, List<String>> attributes = new HashMap<>();
+        attributes.put("foo", singletonList("bar"));
+        attributes.put("baz", Arrays.asList("bort", "zoom"));
         assertEquals("[host=http://1]", new Node(new HttpHost("1")).toString());
+        assertEquals("[host=http://1, attributes={foo=[bar], baz=[bort, zoom]}]",
+                new Node(new HttpHost("1"), null, null, null, null, attributes).toString());
         assertEquals("[host=http://1, roles=mdi]", new Node(new HttpHost("1"),
-                null, null, null, new Roles(true, true, true)).toString());
+                null, null, null, new Roles(true, true, true), null).toString());
         assertEquals("[host=http://1, version=ver]", new Node(new HttpHost("1"),
-                null, null, "ver", null).toString());
+                null, null, "ver", null, null).toString());
         assertEquals("[host=http://1, name=nam]", new Node(new HttpHost("1"),
-                null, "nam", null, null).toString());
+                null, "nam", null, null, null).toString());
         assertEquals("[host=http://1, bound=[http://1, http://2]]", new Node(new HttpHost("1"),
-                new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))), null, null, null).toString());
-        assertEquals("[host=http://1, bound=[http://1, http://2], name=nam, version=ver, roles=m]",
+                new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))), null, null, null, null).toString());
+        assertEquals(
+                "[host=http://1, bound=[http://1, http://2], name=nam, version=ver, roles=m, attributes={foo=[bar], baz=[bort, zoom]}]",
                 new Node(new HttpHost("1"), new HashSet<>(Arrays.asList(new HttpHost("1"), new HttpHost("2"))),
-                    "nam", "ver", new Roles(true, false, false)).toString());
+                    "nam", "ver", new Roles(true, false, false), attributes).toString());
 
     }
 
     public void testEqualsAndHashCode() {
         HttpHost host = new HttpHost(randomAsciiAlphanumOfLength(5));
         Node node = new Node(host,
-            randomBoolean() ? null : singleton(host),
-            randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
-            randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
-            randomBoolean() ? null : new Roles(true, true, true));
+                randomBoolean() ? null : singleton(host),
+                randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
+                randomBoolean() ? null : randomAsciiAlphanumOfLength(5),
+                randomBoolean() ? null : new Roles(true, true, true),
+                randomBoolean() ? null : singletonMap("foo", singletonList("bar")));
         assertFalse(node.equals(null));
         assertTrue(node.equals(node));
         assertEquals(node.hashCode(), node.hashCode());
-        Node copy = new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(), node.getRoles());
+        Node copy = new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(),
+                node.getRoles(), node.getAttributes());
         assertTrue(node.equals(copy));
         assertEquals(node.hashCode(), copy.hashCode());
         assertFalse(node.equals(new Node(new HttpHost(host.toHostString() + "changed"), node.getBoundHosts(),
-                node.getName(), node.getVersion(), node.getRoles())));
+                node.getName(), node.getVersion(), node.getRoles(), node.getAttributes())));
         assertFalse(node.equals(new Node(host, new HashSet<>(Arrays.asList(host, new HttpHost(host.toHostString() + "changed"))),
-                node.getName(), node.getVersion(), node.getRoles())));
-        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName() + "changed", node.getVersion(), node.getRoles())));
-        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(), node.getVersion() + "changed", node.getRoles())));
-        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(), node.getVersion(), new Roles(false, false, false))));
+                node.getName(), node.getVersion(), node.getRoles(), node.getAttributes())));
+        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName() + "changed",
+                node.getVersion(), node.getRoles(), node.getAttributes())));
+        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
+                node.getVersion() + "changed", node.getRoles(), node.getAttributes())));
+        assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
+                node.getVersion(), new Roles(false, false, false), node.getAttributes())));
+                assertFalse(node.equals(new Node(host, node.getBoundHosts(), node.getName(),
+                node.getVersion(), node.getRoles(), singletonMap("bort", singletonList("bing")))));
     }
 }

client/rest/src/test/java/org/elasticsearch/client/RestClientMultipleHostsTests.java

@@ -344,7 +344,7 @@ public void testSetNodes() throws IOException {
         List<Node> newNodes = new ArrayList<>(nodes.size());
         for (int i = 0; i < nodes.size(); i++) {
             Roles roles = i == 0 ? new Roles(false, true, true) : new Roles(true, false, false);
-            newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles));
+            newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles, null));
         }
         restClient.setNodes(newNodes);
         int rounds = between(1, 10);

client/rest/src/test/java/org/elasticsearch/client/RestClientTests.java

@@ -341,9 +341,9 @@ public void testNullPath() throws IOException {
     }
 
     public void testSelectHosts() throws IOException {
-        Node n1 = new Node(new HttpHost("1"), null, null, "1", null);
-        Node n2 = new Node(new HttpHost("2"), null, null, "2", null);
-        Node n3 = new Node(new HttpHost("3"), null, null, "3", null);
+        Node n1 = new Node(new HttpHost("1"), null, null, "1", null, null);
+        Node n2 = new Node(new HttpHost("2"), null, null, "2", null, null);
+        Node n3 = new Node(new HttpHost("3"), null, null, "3", null, null);
 
         NodeSelector not1 = new NodeSelector() {
             @Override

client/rest/src/test/java/org/elasticsearch/client/documentation/RestClientDocumentation.java

@@ -36,6 +36,7 @@
 import org.apache.http.ssl.SSLContextBuilder;
 import org.apache.http.ssl.SSLContexts;
 import org.apache.http.util.EntityUtils;
+import org.elasticsearch.client.HasAttributeNodeSelector;
 import org.elasticsearch.client.HttpAsyncResponseConsumerFactory;
 import org.elasticsearch.client.HttpAsyncResponseConsumerFactory.HeapBufferedResponseConsumerFactory;
 import org.elasticsearch.client.Node;
@@ -191,11 +192,20 @@ public void onFailure(Exception exception) {
             //tag::rest-client-options-set-singleton
             request.setOptions(COMMON_OPTIONS);
             //end::rest-client-options-set-singleton
-            //tag::rest-client-options-customize
-            RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
-            options.addHeader("cats", "knock things off of other things");
-            request.setOptions(options);
-            //end::rest-client-options-customize
+            {
+                //tag::rest-client-options-customize-header
+                RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
+                options.addHeader("cats", "knock things off of other things");
+                request.setOptions(options);
+                //end::rest-client-options-customize-header
+            }
+            {
+                //tag::rest-client-options-customize-attribute
+                RequestOptions.Builder options = COMMON_OPTIONS.toBuilder();
+                options.setNodeSelector(new HasAttributeNodeSelector("rack", "c12")); // <1>
+                request.setOptions(options);
+                //end::rest-client-options-customize-attribute
+            }
         }
         {
             HttpEntity[] documents = new HttpEntity[10];