using single try catch block for connection

caraml-dev · Oct 2, 2024 · 3b4f01a · 3b4f01a
1 parent 4355a1a
commit 3b4f01a
Showing 1 changed file with 10 additions and 15 deletions.
diff --git a/...ml-store-serving/src/main/java/dev/caraml/serving/store/bigtable/BigTableStoreConfig.java b/...ml-store-serving/src/main/java/dev/caraml/serving/store/bigtable/BigTableStoreConfig.java
@@ -30,24 +30,18 @@ public class BigTableStoreConfig {
 
   @Bean
   public OnlineRetriever getRetriever() {
-    // Using HBase SDK
-    if (isUsingHBaseSDK) {
-      org.apache.hadoop.conf.Configuration config =
-          BigtableConfiguration.configure(projectId, instanceId);
-      config.set(BigtableOptionsFactory.APP_PROFILE_ID_KEY, appProfileId);
+    try {
+      // Using HBase SDK
+      if (isUsingHBaseSDK) {
+        org.apache.hadoop.conf.Configuration config =
+            BigtableConfiguration.configure(projectId, instanceId);
+        config.set(BigtableOptionsFactory.APP_PROFILE_ID_KEY, appProfileId);
 
-      Connection connection;
-      try {
-        connection = BigtableConfiguration.connect(config);
-      } catch (IllegalStateException e) {
-        throw new RuntimeException(e);
+        Connection connection = BigtableConfiguration.connect(config);
+        return new HBaseOnlineRetriever(connection);
       }
 
-      return new HBaseOnlineRetriever(connection);
-    }
-
-    // Using BigTable SDK
-    try {
+      // Using BigTable SDK
       BigtableDataSettings.Builder builder =
           BigtableDataSettings.newBuilder()
               .setProjectId(projectId)
@@ -66,6 +60,7 @@ public OnlineRetriever getRetriever() {
       }
       BigtableDataClient client = BigtableDataClient.create(settings);
       return new BigTableOnlineRetriever(client);
+
     } catch (IOException e) {
       throw new RuntimeException(e);
     }