remove ZooKeeperHiveClientException from method signature of chooseSe…

…rver
apache · Oct 23, 2024 · 84c24e8 · 84c24e8
1 parent c02ec68
commit 84c24e8
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 20 deletions.
diff --git a/.../src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala b/.../src/test/scala/org/apache/kyuubi/ha/client/zookeeper/ZookeeperDiscoveryClientSuite.scala
@@ -262,7 +262,5 @@ object CustomSelectStrategy extends ServerSelectStrategy {
   override def chooseServer(
       serverHosts: util.List[String],
       zkClient: CuratorFramework,
-      namespace: String): String = {
-    serverHosts.get(0)
-  }
+      namespace: String): String = serverHosts.get(0)
 }
diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/ZooKeeperHiveClientHelper.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/ZooKeeperHiveClientHelper.java
@@ -123,16 +123,19 @@ static void configureConnParams(JdbcConnectionParams connParams)
   }
 
   private static String chooseServer(
-      JdbcConnectionParams connParams, List<String> serverHosts, CuratorFramework zkClient)
-      throws ZooKeeperHiveClientException {
+      JdbcConnectionParams connParams, List<String> serverHosts, CuratorFramework zkClient) {
     String zooKeeperNamespace = getZooKeeperNamespace(connParams);
     String strategyName =
         connParams
             .getSessionVars()
             .getOrDefault(
                 JdbcConnectionParams.SERVER_SELECT_STRATEGY, RandomSelectStrategy.strategyName);
-    ServerSelectStrategy strategy = StrategyFactory.createStrategy(strategyName);
-    return strategy.chooseServer(serverHosts, zkClient, zooKeeperNamespace);
+    try {
+      ServerSelectStrategy strategy = StrategyFactory.createStrategy(strategyName);
+      return strategy.chooseServer(serverHosts, zkClient, zooKeeperNamespace);
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to choose server with strategy " + strategyName, e);
+    }
   }
 
   static List<JdbcConnectionParams> getDirectParamsList(JdbcConnectionParams connParams)

diff --git a/...bi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/ServerSelectStrategy.java b/...bi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/ServerSelectStrategy.java
@@ -18,10 +18,8 @@
 package org.apache.kyuubi.jdbc.hive.strategy;
 
 import java.util.List;
-import org.apache.kyuubi.jdbc.hive.ZooKeeperHiveClientException;
 import org.apache.kyuubi.shaded.curator.framework.CuratorFramework;
 
 public interface ServerSelectStrategy {
-  String chooseServer(List<String> serverHosts, CuratorFramework zkClient, String namespace)
-      throws ZooKeeperHiveClientException;
+  String chooseServer(List<String> serverHosts, CuratorFramework zkClient, String namespace);
 }
diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/StrategyFactory.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/StrategyFactory.java
@@ -23,8 +23,7 @@
 import org.apache.kyuubi.jdbc.hive.strategy.zk.RandomSelectStrategy;
 
 public class StrategyFactory {
-  public static ServerSelectStrategy createStrategy(String strategyName)
-      throws ZooKeeperHiveClientException {
+  public static ServerSelectStrategy createStrategy(String strategyName) {
     try {
       switch (strategyName) {
         case PollingSelectStrategy.strategyName:
@@ -43,8 +42,7 @@ public static ServerSelectStrategy createStrategy(String strategyName)
           }
       }
     } catch (Exception e) {
-      throw new ZooKeeperHiveClientException(
-          "Oops, load the chooseStrategy is wrong, please check your connection params", e);
+      throw new RuntimeException("Failed to init server select strategy", e);
     }
   }
 }
diff --git a/...ive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/zk/PollingSelectStrategy.java b/...ive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/strategy/zk/PollingSelectStrategy.java
@@ -18,7 +18,6 @@
 package org.apache.kyuubi.jdbc.hive.strategy.zk;
 
 import java.util.List;
-import org.apache.kyuubi.jdbc.hive.ZooKeeperHiveClientException;
 import org.apache.kyuubi.jdbc.hive.strategy.ServerSelectStrategy;
 import org.apache.kyuubi.shaded.curator.framework.CuratorFramework;
 import org.apache.kyuubi.shaded.curator.framework.recipes.atomic.AtomicValue;
@@ -32,20 +31,19 @@ public class PollingSelectStrategy implements ServerSelectStrategy {
   private static final String COUNTER_PATH_SUFFIX = "-counter";
 
   @Override
-  public String chooseServer(List<String> serverHosts, CuratorFramework zkClient, String namespace)
-      throws ZooKeeperHiveClientException {
+  public String chooseServer(
+      List<String> serverHosts, CuratorFramework zkClient, String namespace) {
     String counterPath = COUNTER_PATH_PREFIX + namespace + COUNTER_PATH_SUFFIX;
     try {
       return serverHosts.get(getAndIncrement(zkClient, counterPath) % serverHosts.size());
     } catch (Exception e) {
-      throw new ZooKeeperHiveClientException(
-          "Oops, PollingSelectStrategy get the server is wrong!", e);
+      throw new RuntimeException("Failed to choose server by polling select strategy", e);
     }
   }
 
   private int getAndIncrement(CuratorFramework zkClient, String path) throws Exception {
     DistributedAtomicInteger dai =
-        new DistributedAtomicInteger(zkClient, path, new RetryForever(1000));
+        new DistributedAtomicInteger(zkClient, path, new RetryForever(3000));
     AtomicValue<Integer> atomicVal;
     do {
       atomicVal = dai.add(1);