From 7eac932030932697ea6f4cac2f6d15c6d5fea826 Mon Sep 17 00:00:00 2001
From: Belonogov Nikolay <nikolai.belonogov@vk.team>
Date: Mon, 22 Jan 2024 13:21:41 +0300
Subject: [PATCH] Add support for "balance" option

Closes #462
---
 .../api/space/options/SelectOptions.java      |  3 +-
 .../crud/OperationWithBalanceOptions.java     | 33 +++++++++++++++++++
 .../space/options/crud/enums/ProxyOption.java |  4 ++-
 .../options/ProxySpaceSelectOptionsIT.java    | 18 ++++++++++
 4 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 src/main/java/io/tarantool/driver/api/space/options/crud/OperationWithBalanceOptions.java

diff --git a/src/main/java/io/tarantool/driver/api/space/options/SelectOptions.java b/src/main/java/io/tarantool/driver/api/space/options/SelectOptions.java
index c7c6815d..92941aa3 100644
--- a/src/main/java/io/tarantool/driver/api/space/options/SelectOptions.java
+++ b/src/main/java/io/tarantool/driver/api/space/options/SelectOptions.java
@@ -1,5 +1,6 @@
 package io.tarantool.driver.api.space.options;
 
+import io.tarantool.driver.api.space.options.crud.OperationWithBalanceOptions;
 import io.tarantool.driver.api.space.options.crud.OperationWithBatchSizeOptions;
 import io.tarantool.driver.api.space.options.crud.OperationWithBucketIdOptions;
 import io.tarantool.driver.api.space.options.crud.OperationWithFieldsOptions;
@@ -16,5 +17,5 @@
  */
 public interface SelectOptions<T extends SelectOptions<T>>
     extends OperationWithBucketIdOptions<T>, OperationWithTimeoutOptions<T>, OperationWithFieldsOptions<T>,
-    OperationWithModeOptions<T>, OperationWithBatchSizeOptions<T> {
+            OperationWithModeOptions<T>, OperationWithBatchSizeOptions<T>, OperationWithBalanceOptions<T> {
 }
diff --git a/src/main/java/io/tarantool/driver/api/space/options/crud/OperationWithBalanceOptions.java b/src/main/java/io/tarantool/driver/api/space/options/crud/OperationWithBalanceOptions.java
new file mode 100644
index 00000000..fb4e2dcb
--- /dev/null
+++ b/src/main/java/io/tarantool/driver/api/space/options/crud/OperationWithBalanceOptions.java
@@ -0,0 +1,33 @@
+package io.tarantool.driver.api.space.options.crud;
+
+import java.util.Optional;
+
+import io.tarantool.driver.api.space.options.Options;
+import io.tarantool.driver.api.space.options.Self;
+import io.tarantool.driver.api.space.options.crud.enums.ProxyOption;
+
+/**
+ * Base interface for all operation options that may have a configurable "balance" option.
+ *
+ * @author Belonogov Nikolay
+ */
+public interface OperationWithBalanceOptions<T extends OperationWithBalanceOptions<T>> extends Options, Self<T> {
+
+    /**
+     * Sets "balance" option value to true. If true then use replica according to vshard load balancing policy. Default
+     * value is false.
+     *
+     * @return this options instance.
+     */
+    default T balance() {
+        addOption(ProxyOption.BALANCE, true);
+        return self();
+    }
+
+    /**
+     * @return "balance" option value.
+     */
+    default Optional<Boolean> getBalance() {
+        return getOption(ProxyOption.BALANCE, Boolean.class);
+    }
+}
diff --git a/src/main/java/io/tarantool/driver/api/space/options/crud/enums/ProxyOption.java b/src/main/java/io/tarantool/driver/api/space/options/crud/enums/ProxyOption.java
index c2b9de5f..d79b83c0 100644
--- a/src/main/java/io/tarantool/driver/api/space/options/crud/enums/ProxyOption.java
+++ b/src/main/java/io/tarantool/driver/api/space/options/crud/enums/ProxyOption.java
@@ -23,7 +23,9 @@ public enum ProxyOption {
 
     AFTER("after"),
 
-    FIRST("first");
+    FIRST("first"),
+
+    BALANCE("balance");
 
     private final String name;
 
diff --git a/src/test/java/io/tarantool/driver/integration/proxy/options/ProxySpaceSelectOptionsIT.java b/src/test/java/io/tarantool/driver/integration/proxy/options/ProxySpaceSelectOptionsIT.java
index 8c023dd7..fb88eabe 100644
--- a/src/test/java/io/tarantool/driver/integration/proxy/options/ProxySpaceSelectOptionsIT.java
+++ b/src/test/java/io/tarantool/driver/integration/proxy/options/ProxySpaceSelectOptionsIT.java
@@ -8,6 +8,7 @@
 import io.tarantool.driver.api.space.options.SelectOptions;
 import io.tarantool.driver.api.space.TarantoolSpaceOperations;
 import io.tarantool.driver.api.space.options.ProxySelectOptions;
+import io.tarantool.driver.api.space.options.crud.enums.ProxyOption;
 import io.tarantool.driver.api.tuple.DefaultTarantoolTupleFactory;
 import io.tarantool.driver.api.tuple.TarantoolTuple;
 import io.tarantool.driver.api.tuple.TarantoolTupleFactory;
@@ -28,6 +29,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
  * @author Artyom Dubinin
@@ -174,4 +176,20 @@ public void withModeTest() throws ExecutionException, InterruptedException {
         crudSelectOpts = client.eval("return crud_select_opts").get();
         assertEquals(Mode.WRITE.value(), ((HashMap<?, ?>) crudSelectOpts.get(0)).get("mode"));
     }
+
+    @Test
+    public void withBalanceTest() {
+        TarantoolSpaceOperations<TarantoolTuple, TarantoolResult<TarantoolTuple>> operations =
+            client.space(TEST_SPACE_NAME);
+
+        SelectOptions<ProxySelectOptions> options = ProxySelectOptions.create().balance();
+
+        assertTrue(options.getBalance().isPresent());
+        assertTrue(options.getBalance().get());
+
+        operations.select(Conditions.any(), options).join();
+        List<?> crudSelectOpts = client.eval("return crud_select_opts").join();
+
+        assertEquals(true, ((HashMap<?, ?>) crudSelectOpts.get(0)).get(ProxyOption.BALANCE.toString()));
+    }
 }