replaced isAvailable call (#331)

* replaced isAvailable call * Update CHANGELOG.md --------- Co-authored-by: Jay Green-Stevens <[email protected]>
ExpediaGroup · Oct 4, 2024 · 78b20f1 · 78b20f1
1 parent 720a4f5
commit 78b20f1
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,11 @@
+## [3.13.4] - 2024-10-04
+### Fix
+- Removed call to primary and made it only when Kerberos (SASL) is enabled. 
+- Fixed Handler creation for SASL.
+
 ## [3.13.3] - 2024-10-02
 ### Added
-- Metric for monitoring open transports. `<prefix>.open_transports_gauge`
+- Metric for monitoring open transports. `<prefix>.com_hotels_bdp_waggledance_open_transports_gauge`
 ### Changed
 - Removed RetryingHMSHandler. Retries are done in the client there should be no need to wrap everything in retry logic again. 
 

diff --git a/...va/com/hotels/bdp/waggledance/mapping/service/impl/PrefixBasedDatabaseMappingService.java b/...va/com/hotels/bdp/waggledance/mapping/service/impl/PrefixBasedDatabaseMappingService.java
@@ -97,11 +97,6 @@ private void add(AbstractMetaStore metaStore) {
 
     if (metaStore.getFederationType() == PRIMARY) {
       primaryDatabaseMapping = databaseMapping;
-      if (!metaStoreMapping.isAvailable()) {
-        throw new WaggleDanceException(
-                String.format("Primary metastore is unavailable %s", metaStore.getRemoteMetaStoreUris())
-        );
-      }
     }
 
     mappingsByPrefix.put(metaStoreMapping.getDatabasePrefix(), databaseMapping);

diff --git a/...e-core/src/main/java/com/hotels/bdp/waggledance/server/TSetIpAddressProcessorFactory.java b/...e-core/src/main/java/com/hotels/bdp/waggledance/server/TSetIpAddressProcessorFactory.java
@@ -20,6 +20,7 @@
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.metastore.IHMSHandler;
 import org.apache.hadoop.hive.metastore.TSetIpAddressProcessor;
+import org.apache.thrift.TException;
 import org.apache.thrift.TProcessor;
 import org.apache.thrift.TProcessorFactory;
 import org.apache.thrift.transport.TSocket;
@@ -59,7 +60,14 @@ public TProcessor getProcessor(TTransport transport) {
 
       boolean useSASL = hiveConf.getBoolVar(HiveConf.ConfVars.METASTORE_USE_THRIFT_SASL);
       if (useSASL) {
-        IHMSHandler handler = TokenWrappingHMSHandler.newProxyInstance(baseHandler, useSASL);
+        try {
+          baseHandler.getStatus();
+        } catch (TException e) {
+          throw new RuntimeException("Error creating TProcessor. Could not get status.", e);
+        }
+        IHMSHandler tokenHandler = TokenWrappingHMSHandler.newProxyInstance(baseHandler, useSASL);
+        IHMSHandler handler = ExceptionWrappingHMSHandler.newProxyInstance(tokenHandler);
+        transportMonitor.monitor(transport, baseHandler);
         return new TSetIpAddressProcessor<>(handler);
       } else {
         IHMSHandler handler = ExceptionWrappingHMSHandler.newProxyInstance(baseHandler);

diff --git a/...re/src/test/java/com/hotels/bdp/waggledance/server/TSetIpAddressProcessorFactoryTest.java b/...re/src/test/java/com/hotels/bdp/waggledance/server/TSetIpAddressProcessorFactoryTest.java
@@ -42,12 +42,13 @@ public class TSetIpAddressProcessorFactoryTest {
   private @Mock TTransportMonitor transportMonitor;
   private @Mock TTransport transport;
 
-  private final HiveConf hiveConf = new HiveConf();
+  private HiveConf hiveConf;
   private TSetIpAddressProcessorFactory factory;
 
   @Before
   public void init() {
     when(federatedHMSHandlerFactory.create()).thenReturn(federatedHMSHandler);
+    hiveConf = new HiveConf();
     factory = new TSetIpAddressProcessorFactory(hiveConf, federatedHMSHandlerFactory, transportMonitor);
   }
 
@@ -67,4 +68,16 @@ public void connectionIsMonitored() throws Exception {
     assertThat(transportCaptor.getValue(), is(transport));
     assertThat(handlerCaptor.getValue(), is(instanceOf(FederatedHMSHandler.class)));
   }
+
+  @Test
+  public void connectionIsMonitoredSasl() throws Exception {
+    hiveConf.setBoolVar(HiveConf.ConfVars.METASTORE_USE_THRIFT_SASL, Boolean.TRUE);
+    factory.getProcessor(transport);
+
+    ArgumentCaptor<TTransport> transportCaptor = ArgumentCaptor.forClass(TTransport.class);
+    ArgumentCaptor<Closeable> handlerCaptor = ArgumentCaptor.forClass(Closeable.class);
+    verify(transportMonitor).monitor(transportCaptor.capture(), handlerCaptor.capture());
+    assertThat(transportCaptor.getValue(), is(transport));
+    assertThat(handlerCaptor.getValue(), is(instanceOf(FederatedHMSHandler.class)));
+  }
 }