From b60fa8a83a11c1d75854658b9532c18547fe2dd3 Mon Sep 17 00:00:00 2001
From: Hishida Masato <hishida@nautilus-technologies.com>
Date: Thu, 22 Aug 2024 10:24:20 +0900
Subject: [PATCH] feat: TgSessionOption: Added setKeepAlive()

---
 .../tsurugidb/iceaxe/TsurugiConnector.java    |  4 ++
 .../iceaxe/session/TgSessionOption.java       | 24 ++++++++++++
 .../iceaxe/session/TgSessionOptionTest.java   | 38 ++++++++++++++++++-
 .../example/Example02SessionOption.java       |  5 +++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/TsurugiConnector.java b/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/TsurugiConnector.java
index 90fe385..b7f1a66 100755
--- a/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/TsurugiConnector.java
+++ b/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/TsurugiConnector.java
@@ -344,6 +344,10 @@ protected SessionBuilder createLowSessionBuilder(@Nullable Credential credential
             lowBuilder.withApplicationName(name);
         });
 
+        sessionOption.findKeepAlive().ifPresent(keepAlive -> {
+            lowBuilder.withKeepAlive(keepAlive);
+        });
+
         return lowBuilder;
     }
 
diff --git a/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/session/TgSessionOption.java b/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/session/TgSessionOption.java
index f17bd63..30109e2 100755
--- a/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/session/TgSessionOption.java
+++ b/modules/iceaxe-core/src/main/java/com/tsurugidb/iceaxe/session/TgSessionOption.java
@@ -4,6 +4,7 @@
 import java.util.EnumMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import javax.annotation.Nonnull;
@@ -123,6 +124,7 @@ public enum TgTimeoutKey {
 
     private String sessionLabel;
     private String applicationName;
+    private Optional<Boolean> keepAlive = Optional.empty();
     private final Map<TgTimeoutKey, TgTimeValue> timeoutMap = Collections.synchronizedMap(new EnumMap<>(TgTimeoutKey.class));
     private TgCommitType commitType = TgCommitType.DEFAULT;
     private TgSessionShutdownType closeShutdownType = TgSessionShutdownType.NOTHING;
@@ -176,6 +178,28 @@ public TgSessionOption setApplicationName(@Nullable String name) {
         return this.applicationName;
     }
 
+    /**
+     * set session keep-alive.
+     *
+     * @param enabled {@code true} to enable session keep-alive, or {@code false} to disable it
+     * @return this
+     * @since X.X.X
+     */
+    public TgSessionOption setKeepAlive(@Nullable Boolean enabled) {
+        this.keepAlive = Optional.ofNullable(enabled);
+        return this;
+    }
+
+    /**
+     * get session keep-alive.
+     *
+     * @return session keep-alive
+     * @since X.X.X
+     */
+    public Optional<Boolean> findKeepAlive() {
+        return this.keepAlive;
+    }
+
     /**
      * set timeout.
      *
diff --git a/modules/iceaxe-core/src/test/java/com/tsurugidb/iceaxe/session/TgSessionOptionTest.java b/modules/iceaxe-core/src/test/java/com/tsurugidb/iceaxe/session/TgSessionOptionTest.java
index e96131c..e4e9d07 100755
--- a/modules/iceaxe-core/src/test/java/com/tsurugidb/iceaxe/session/TgSessionOptionTest.java
+++ b/modules/iceaxe-core/src/test/java/com/tsurugidb/iceaxe/session/TgSessionOptionTest.java
@@ -4,6 +4,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 
 import java.net.URI;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import org.junit.jupiter.api.Test;
@@ -51,12 +52,45 @@ void applicationName() {
         assertEquals("test-app", actual);
     }
 
+    @Test
+    void keepAlive() {
+        {
+            var sessionOption = new TgSessionOption();
+            testKeepAlive(Optional.empty(), sessionOption);
+        }
+        {
+            var sessionOption = new TgSessionOption().setKeepAlive(true);
+            testKeepAlive(Optional.of(true), sessionOption);
+        }
+        {
+            var sessionOption = new TgSessionOption().setKeepAlive(false);
+            testKeepAlive(Optional.of(false), sessionOption);
+        }
+        {
+            var sessionOption = new TgSessionOption().setKeepAlive(null);
+            testKeepAlive(Optional.empty(), sessionOption);
+        }
+    }
+
+    private void testKeepAlive(Optional<Boolean> expected, TgSessionOption sessionOption) {
+        assertEquals(expected, sessionOption.findKeepAlive());
+
+        var connector = new SessionOptionTestConnector();
+        var lowBuilder = connector.createLowSessionBuilder(null, sessionOption);
+        boolean actual = getField(lowBuilder, "doKeepAlive");
+        assertEquals(expected.orElseGet(() -> {
+            var defaultBuilder = SessionBuilder.connect("tcp://test:12345");
+            return getField(defaultBuilder, "doKeepAlive");
+        }), actual);
+    }
+
     // FIXME リフレクションを使わずにSessionBuilderから取得したい
-    private static String getField(SessionBuilder lowBuilder, String name) {
+    @SuppressWarnings("unchecked")
+    private static <T> T getField(SessionBuilder lowBuilder, String name) {
         try {
             var field = lowBuilder.getClass().getDeclaredField(name);
             field.setAccessible(true);
-            return (String) field.get(lowBuilder);
+            return (T) field.get(lowBuilder);
         } catch (RuntimeException e) {
             throw e;
         } catch (Exception e) {
diff --git a/modules/iceaxe-examples/src/main/java/com/tsurugidb/iceaxe/example/Example02SessionOption.java b/modules/iceaxe-examples/src/main/java/com/tsurugidb/iceaxe/example/Example02SessionOption.java
index e1f01e1..0e1bb89 100755
--- a/modules/iceaxe-examples/src/main/java/com/tsurugidb/iceaxe/example/Example02SessionOption.java
+++ b/modules/iceaxe-examples/src/main/java/com/tsurugidb/iceaxe/example/Example02SessionOption.java
@@ -22,6 +22,11 @@ void label() {
         // $TSURUGI_HOME/bin/tgctl session list --verbose
     }
 
+    void keepAlive() {
+        var sessionOption = TgSessionOption.of();
+        sessionOption.setKeepAlive(true);
+    }
+
     void timeout() {
         var sessionOption = TgSessionOption.of();
         sessionOption.setTimeout(TgTimeoutKey.DEFAULT, 1, TimeUnit.MINUTES);