[DE-541] forceOneShardAttributeValue (#491)

Author: rashtao

ChangeLog.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 
 ## [Unreleased]
 
+## [6.22.0] - 2023-03-07
+
+- added support to `forceOneShardAttributeValue` query parameter (DE-541)
+
 ## [6.21.0] - 2023-03-07
 
 - added `x-arango-driver` header (DE-479)

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

@@ -403,6 +403,28 @@ 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;
+    }
+
     private Options getOptions() {
         if (options == null) {
             options = new Options();
@@ -429,6 +451,7 @@ public static class Options implements Serializable {
         private Collection<String> shardIds;
         private Double maxRuntime;
         private Boolean fillBlockCache;
+        private String forceOneShardAttributeValue;
 
         protected Optimizer getOptimizer() {
             if (optimizer == null) {

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

@@ -912,13 +912,41 @@ void queryWithWarning(ArangoDB arangoDB) {
     @ParameterizedTest(name = "{index}")
     @MethodSource("dbs")
     void queryStream(ArangoDatabase db) {
-        if (isAtLeastVersion(3, 4)) {
-            final ArangoCursor<VPackSlice> cursor = db
-                    .query("FOR i IN 1..2 RETURN i", null, new AqlQueryOptions().stream(true).count(true),
-                            VPackSlice.class);
-            assertThat((Object) cursor).isNotNull();
-            assertThat(cursor.getCount()).isNull();
-        }
+        final ArangoCursor<VPackSlice> cursor = db
+                .query("FOR i IN 1..2 RETURN i", null, new AqlQueryOptions().stream(true).count(true),
+                        VPackSlice.class);
+        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", Collections.singletonMap("@c", cname), new AqlQueryOptions()
+                                .forceOneShardAttributeValue("bar"),
+                        BaseDocument.class);
+        assertThat(c1.hasNext()).isTrue();
+        assertThat(c1.next().getAttribute("foo")).isEqualTo("bar");
+
+        ArangoCursor<BaseDocument> c2 = db
+                .query("FOR d IN @@c RETURN d", Collections.singletonMap("@c", cname), new AqlQueryOptions()
+                                .forceOneShardAttributeValue("ooo"),
+                        BaseDocument.class);
+        assertThat(c2.hasNext()).isFalse();
     }
 
     @ParameterizedTest(name = "{index}")