support to forceOneShardAttributeValue query parameter (DE-541)

Author: rashtao

core/src/main/java/com/arangodb/model/AqlQueryOptions.java

@@ -433,6 +433,27 @@ public AqlQueryOptions shardIds(final String... shardIds) {
         return this;
     }
 
+    public String getForceOneShardAttributeValue() {
+        return options != null ? options.forceOneShardAttributeValue : null;
+    }
+
+    /**
+     * @param forceOneShardAttributeValue This query option can be used in complex queries in case the query optimizer
+     *                                    cannot automatically detect that the query can be limited to only a single
+     *                                    server (e.g. in a disjoint smart graph case).
+     *                                    <p/>
+     *                                    If the option is set incorrectly, i.e. to a wrong shard key value, then the
+     *                                    query may be shipped to a wrong DB server and may not return results (i.e.
+     *                                    empty result set).
+     *                                    <p/>
+     *                                    Use at your own risk.
+     * @return options
+     */
+    public AqlQueryOptions forceOneShardAttributeValue(final String forceOneShardAttributeValue) {
+        getOptions().forceOneShardAttributeValue = forceOneShardAttributeValue;
+        return this;
+    }
+
     public Options getOptions() {
         if (options == null) {
             options = new Options();
@@ -487,6 +508,7 @@ public static final class Options {
         private Collection<String> shardIds;
         private Double maxRuntime;
         private Boolean fillBlockCache;
+        private String forceOneShardAttributeValue;
 
         public Boolean getFailOnWarning() {
             return failOnWarning;

driver/src/test/java/com/arangodb/ArangoDatabaseTest.java

@@ -915,12 +915,38 @@ void queryWithWarning(ArangoDB arangoDB) {
     @ParameterizedTest(name = "{index}")
     @MethodSource("dbs")
     void queryStream(ArangoDatabase db) {
-        if (isAtLeastVersion(3, 4)) {
-            final ArangoCursor<Void> cursor = db
-                    .query("FOR i IN 1..2 RETURN i", Void.class, new AqlQueryOptions().stream(true).count(true));
-            assertThat((Object) cursor).isNotNull();
-            assertThat(cursor.getCount()).isNull();
-        }
+        final ArangoCursor<Void> cursor = db
+                .query("FOR i IN 1..2 RETURN i", Void.class, new AqlQueryOptions().stream(true).count(true));
+        assertThat((Object) cursor).isNotNull();
+        assertThat(cursor.getCount()).isNull();
+    }
+
+    @ParameterizedTest(name = "{index}")
+    @MethodSource("dbs")
+    void queryForceOneShardAttributeValue(ArangoDatabase db) {
+        assumeTrue(isAtLeastVersion(3, 10));
+        assumeTrue(isCluster());
+        assumeTrue(isEnterprise());
+
+        String cname = "forceOneShardAttr-" + UUID.randomUUID();
+        db.createCollection(cname, new CollectionCreateOptions()
+                .shardKeys("foo")
+                .numberOfShards(3));
+        ArangoCollection col = db.collection(cname);
+        BaseDocument doc = new BaseDocument();
+        doc.addAttribute("foo", "bar");
+        col.insertDocument(doc);
+
+        ArangoCursor<BaseDocument> c1 = db
+                .query("FOR d IN @@c RETURN d", BaseDocument.class, Collections.singletonMap("@c", cname),
+                        new AqlQueryOptions().forceOneShardAttributeValue("bar"));
+        assertThat(c1.hasNext()).isTrue();
+        assertThat(c1.next().getAttribute("foo")).isEqualTo("bar");
+
+        ArangoCursor<BaseDocument> c2 = db
+                .query("FOR d IN @@c RETURN d", BaseDocument.class, Collections.singletonMap("@c", cname),
+                        new AqlQueryOptions().forceOneShardAttributeValue("ooo"));
+        assertThat(c2.hasNext()).isFalse();
     }
 
     @ParameterizedTest(name = "{index}")