Added extra option to add readOnly thrift HMS uri

waggle-dance-api/src/main/java/com/hotels/bdp/waggledance/api/model/AbstractMetaStore.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2023 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -57,6 +57,7 @@ public abstract class AbstractMetaStore {
   private long latency = 0;
   private transient @JsonIgnore HashBiMap<String, String> databaseNameBiMapping = HashBiMap.create();
   private GlueConfig glueConfig;
+  private String readOnlyRemoteMetaStoreUris;
 
   public AbstractMetaStore() {}
 
@@ -124,6 +125,14 @@ public void setRemoteMetaStoreUris(String remoteMetaStoreUris) {
     this.remoteMetaStoreUris = remoteMetaStoreUris;
   }
 
+  public String getReadOnlyRemoteMetaStoreUris() {
+    return readOnlyRemoteMetaStoreUris;
+  }
+
+  public void setReadOnlyRemoteMetaStoreUris(String readOnlyRemoteMetaStoreUris) {
+    this.readOnlyRemoteMetaStoreUris = readOnlyRemoteMetaStoreUris;
+  }
+
   public MetastoreTunnel getMetastoreTunnel() {
     return metastoreTunnel;
   }
@@ -244,6 +253,7 @@ public String toString() {
         .add("databasePrefix", databasePrefix)
         .add("federationType", getFederationType())
         .add("remoteMetaStoreUris", remoteMetaStoreUris)
+        .add("readOnlyRemoteMetaStoreUris", readOnlyRemoteMetaStoreUris)
         .add("metastoreTunnel", metastoreTunnel)
         .add("accessControlType", accessControlType)
         .add("writableDatabaseWhiteList", writableDatabaseWhitelist)
@@ -252,4 +262,5 @@ public String toString() {
         .toString();
   }
 
+
 }

waggle-dance-api/src/test/java/com/hotels/bdp/waggledance/api/model/FederatedMetaStoreTest.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -72,7 +72,7 @@ public void nullDatabasePrefix() {
 
   @Test
   public void toJson() throws Exception {
-    String expected = "{\"accessControlType\":\"READ_ONLY\",\"connectionType\":\"DIRECT\",\"databaseNameMapping\":{},\"databasePrefix\":\"name_\",\"federationType\":\"FEDERATED\",\"glueConfig\":null,\"hiveMetastoreFilterHook\":null,\"latency\":0,\"mappedDatabases\":null,\"mappedTables\":null,\"metastoreTunnel\":null,\"name\":\"name\",\"remoteMetaStoreUris\":\"uri\",\"status\":\"UNKNOWN\",\"writableDatabaseWhiteList\":[]}";
+    String expected = "{\"accessControlType\":\"READ_ONLY\",\"connectionType\":\"DIRECT\",\"databaseNameMapping\":{},\"databasePrefix\":\"name_\",\"federationType\":\"FEDERATED\",\"glueConfig\":null,\"hiveMetastoreFilterHook\":null,\"latency\":0,\"mappedDatabases\":null,\"mappedTables\":null,\"metastoreTunnel\":null,\"name\":\"name\",\"readOnlyRemoteMetaStoreUris\":null,\"remoteMetaStoreUris\":\"uri\",\"status\":\"UNKNOWN\",\"writableDatabaseWhiteList\":[]}";
     ObjectMapper mapper = new ObjectMapper();
     // Sorting to get deterministic test behaviour
     mapper.enable(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY);

waggle-dance-api/src/test/java/com/hotels/bdp/waggledance/api/model/PrimaryMetaStoreTest.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -89,7 +89,7 @@ public void nonEmptyDatabasePrefix() {
 
   @Test
   public void toJson() throws Exception {
-    String expected = "{\"accessControlType\":\"READ_ONLY\",\"connectionType\":\"DIRECT\",\"databaseNameMapping\":{},\"databasePrefix\":\"\",\"federationType\":\"PRIMARY\",\"glueConfig\":null,\"hiveMetastoreFilterHook\":null,\"latency\":0,\"mappedDatabases\":null,\"mappedTables\":null,\"metastoreTunnel\":null,\"name\":\"name\",\"remoteMetaStoreUris\":\"uri\",\"status\":\"UNKNOWN\",\"writableDatabaseWhiteList\":[]}";
+    String expected = "{\"accessControlType\":\"READ_ONLY\",\"connectionType\":\"DIRECT\",\"databaseNameMapping\":{},\"databasePrefix\":\"\",\"federationType\":\"PRIMARY\",\"glueConfig\":null,\"hiveMetastoreFilterHook\":null,\"latency\":0,\"mappedDatabases\":null,\"mappedTables\":null,\"metastoreTunnel\":null,\"name\":\"name\",\"readOnlyRemoteMetaStoreUris\":null,\"remoteMetaStoreUris\":\"uri\",\"status\":\"UNKNOWN\",\"writableDatabaseWhiteList\":[]}";
     ObjectMapper mapper = new ObjectMapper();
     // Sorting to get deterministic test behaviour
     mapper.enable(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY);

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/client/CloseableThriftHiveMetastoreIfaceClientFactory.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,16 +44,19 @@ public class CloseableThriftHiveMetastoreIfaceClientFactory {
   private final int defaultConnectionTimeout = (int) TimeUnit.SECONDS.toMillis(2L);
   private final WaggleDanceConfiguration waggleDanceConfiguration;
   private final GlueClientFactory glueClientFactory;
+  private final SplitTrafficMetastoreClientFactory splitTrafficMetaStoreClientFactory;
 
   public CloseableThriftHiveMetastoreIfaceClientFactory(
       TunnelingMetaStoreClientFactory tunnelingMetaStoreClientFactory,
       DefaultMetaStoreClientFactory defaultMetaStoreClientFactory,
       GlueClientFactory glueClientFactory,
-      WaggleDanceConfiguration waggleDanceConfiguration) {
+      WaggleDanceConfiguration waggleDanceConfiguration,
+      SplitTrafficMetastoreClientFactory splitTrafficMetaStoreClientFactory) {
     this.tunnelingMetaStoreClientFactory = tunnelingMetaStoreClientFactory;
     this.defaultMetaStoreClientFactory = defaultMetaStoreClientFactory;
     this.glueClientFactory = glueClientFactory;
     this.waggleDanceConfiguration = waggleDanceConfiguration;
+    this.splitTrafficMetaStoreClientFactory = splitTrafficMetaStoreClientFactory;
   }
 
   public CloseableThriftHiveMetastoreIface newInstance(AbstractMetaStore metaStore) {
@@ -64,14 +67,24 @@ public CloseableThriftHiveMetastoreIface newInstance(AbstractMetaStore metaStore
     if (metaStore.getGlueConfig() != null) {
       return newGlueInstance(metaStore, properties);
     }
-    return newHiveInstance(metaStore, properties);
+    String name = metaStore.getName().toLowerCase(Locale.ROOT);
+    if (metaStore.getReadOnlyRemoteMetaStoreUris() != null) {
+      CloseableThriftHiveMetastoreIface readWrite = newHiveInstance(metaStore, name, metaStore.getRemoteMetaStoreUris(),
+          properties);
+      CloseableThriftHiveMetastoreIface readOnly = newHiveInstance(metaStore, name+"_ro",
+          metaStore.getReadOnlyRemoteMetaStoreUris(), properties);
+      return splitTrafficMetaStoreClientFactory.newInstance(readWrite, readOnly);
+
+    }
+    return newHiveInstance(metaStore, name, metaStore.getRemoteMetaStoreUris(), properties);
   }
 
   private CloseableThriftHiveMetastoreIface newHiveInstance(
       AbstractMetaStore metaStore,
+      String name,
+      String metaStoreUris,
       Map<String, String> properties) {
-    String uris = MetaStoreUriNormaliser.normaliseMetaStoreUris(metaStore.getRemoteMetaStoreUris());
-    String name = metaStore.getName().toLowerCase(Locale.ROOT);
+    String uris = MetaStoreUriNormaliser.normaliseMetaStoreUris(metaStoreUris);
     // Connection timeout should not be less than 1
     // A timeout of zero is interpreted as an infinite timeout, so this is avoided
     int connectionTimeout = Math.max(1, defaultConnectionTimeout + (int) metaStore.getLatency());

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/client/SplitTrafficMetastoreClientFactory.java

@@ -0,0 +1,87 @@
+/**
+ * Copyright (C) 2016-2024 Expedia, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.hotels.bdp.waggledance.client;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+
+/**
+ * This class splits the traffic for read only calls (get* for instance getTable, getPartition) to the readOnly client
+ * and everything else will go to readWrite client.
+ */
+public class SplitTrafficMetastoreClientFactory {
+
+  static final Class<?>[] INTERFACES = new Class<?>[] { CloseableThriftHiveMetastoreIface.class };
+
+  private static class SplitTrafficClientInvocationHandler implements InvocationHandler {
+
+    private final CloseableThriftHiveMetastoreIface readWrite;
+    private final CloseableThriftHiveMetastoreIface readOnly;
+
+    public SplitTrafficClientInvocationHandler(
+        CloseableThriftHiveMetastoreIface readWrite,
+        CloseableThriftHiveMetastoreIface readOnly) {
+      this.readWrite = readWrite;
+      this.readOnly = readOnly;
+    }
+
+    @Override
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+      switch (method.getName()) {
+      case "isOpen":
+        return readWrite.isOpen() && readOnly.isOpen();
+      case "close":
+        try {
+          readWrite.close();
+        } finally {
+          readOnly.close();
+        }
+        return null;
+      case "set_ugi":
+        Object result = doRealCall(readWrite, method, args);
+        // we skip the result for readOnly (it should always be the same).
+        doRealCall(readOnly, method, args);
+        return result;
+      default:
+        if (method.getName().startsWith("get")) {
+          doRealCall(readOnly, method, args);
+        }
+        return doRealCall(readWrite, method, args);
+      }
+    }
+
+    private Object doRealCall(CloseableThriftHiveMetastoreIface client, Method method, Object[] args)
+      throws IllegalAccessException, Throwable {
+      try {
+        return method.invoke(client, args);
+      } catch (InvocationTargetException e) {
+        Throwable realException = e.getTargetException();
+        throw realException;
+      }
+    }
+  }
+
+  public CloseableThriftHiveMetastoreIface newInstance(
+      CloseableThriftHiveMetastoreIface readWrite,
+      CloseableThriftHiveMetastoreIface readOnly) {
+    return (CloseableThriftHiveMetastoreIface) Proxy
+        .newProxyInstance(getClass().getClassLoader(), INTERFACES,
+            new SplitTrafficClientInvocationHandler(readWrite, readOnly));
+  }
+
+}

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/client/adapter/MetastoreIfaceAdapter.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/client/compatibility/HiveCompatibleThriftHiveMetastoreIfaceFactory.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2019 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/context/CommonBeans.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2023 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,6 +24,7 @@
 import com.hotels.bdp.waggledance.client.CloseableThriftHiveMetastoreIfaceClientFactory;
 import com.hotels.bdp.waggledance.client.DefaultMetaStoreClientFactory;
 import com.hotels.bdp.waggledance.client.GlueClientFactory;
+import com.hotels.bdp.waggledance.client.SplitTrafficMetastoreClientFactory;
 import com.hotels.bdp.waggledance.client.tunnelling.TunnelingMetaStoreClientFactory;
 import com.hotels.bdp.waggledance.conf.WaggleDanceConfiguration;
 import com.hotels.bdp.waggledance.mapping.model.ASTQueryMapping;
@@ -53,11 +54,16 @@ public PrefixNamingStrategy prefixNamingStrategy(WaggleDanceConfiguration waggle
     return new LowerCasePrefixNamingStrategy();
   }
 
+  @Bean
+  public SplitTrafficMetastoreClientFactory splitTrafficMetaStoreClientFactory() {
+    return new SplitTrafficMetastoreClientFactory();
+  }
+
   @Bean
   public CloseableThriftHiveMetastoreIfaceClientFactory metaStoreClientFactory(
-      WaggleDanceConfiguration waggleDanceConfiguration) {
+      WaggleDanceConfiguration waggleDanceConfiguration, SplitTrafficMetastoreClientFactory splitTrafficMetaStoreClientFactory) {
     return new CloseableThriftHiveMetastoreIfaceClientFactory(new TunnelingMetaStoreClientFactory(),
-        new DefaultMetaStoreClientFactory(), new GlueClientFactory(), waggleDanceConfiguration);
+        new DefaultMetaStoreClientFactory(), new GlueClientFactory(), waggleDanceConfiguration, splitTrafficMetaStoreClientFactory);
   }
 
   @Bean

waggle-dance-core/src/main/java/com/hotels/bdp/waggledance/server/FederatedHMSHandler.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/client/CloseableThriftHiveMetastoreIfaceClientFactoryTest.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,12 +18,15 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.isA;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 
 import static com.hotels.bdp.waggledance.api.model.AbstractMetaStore.newFederatedInstance;
+import static com.hotels.bdp.waggledance.api.model.AbstractMetaStore.newPrimaryInstance;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -43,6 +46,7 @@
 import com.hotels.bdp.waggledance.api.model.AbstractMetaStore;
 import com.hotels.bdp.waggledance.api.model.FederatedMetaStore;
 import com.hotels.bdp.waggledance.api.model.GlueConfig;
+import com.hotels.bdp.waggledance.api.model.PrimaryMetaStore;
 import com.hotels.bdp.waggledance.client.adapter.MetastoreIfaceAdapter;
 import com.hotels.bdp.waggledance.client.tunnelling.TunnelingMetaStoreClientFactory;
 import com.hotels.bdp.waggledance.conf.WaggleDanceConfiguration;
@@ -52,6 +56,7 @@
 public class CloseableThriftHiveMetastoreIfaceClientFactoryTest {
 
   private static final String THRIFT_URI = "thrift://host:port";
+  private static final String THRIFT_URI_READ_ONLY = "thrift://host-read-only:port";
 
   private CloseableThriftHiveMetastoreIfaceClientFactory factory;
   private @Mock TunnelingMetaStoreClientFactory tunnelingMetaStoreClientFactory;
@@ -60,6 +65,7 @@ public class CloseableThriftHiveMetastoreIfaceClientFactoryTest {
   private @Mock WaggleDanceConfiguration waggleDanceConfiguration;
   private final Map<String, String> configurationProperties = new HashMap<>();
   private @Mock AWSCatalogMetastoreClient glueClient;
+  private @Mock SplitTrafficMetastoreClientFactory splitTrafficMetaStoreClientFactory;
 
   @Before
   public void setUp() {
@@ -70,7 +76,7 @@ public void setUp() {
     configurationProperties.put(ConfVars.METASTORE_USE_THRIFT_COMPACT_PROTOCOL.varname, "false");
     when(waggleDanceConfiguration.getConfigurationProperties()).thenReturn(configurationProperties);
     factory = new CloseableThriftHiveMetastoreIfaceClientFactory(tunnelingMetaStoreClientFactory,
-        defaultMetaStoreClientFactory, glueClientFactory, waggleDanceConfiguration);
+        defaultMetaStoreClientFactory, glueClientFactory, waggleDanceConfiguration, splitTrafficMetaStoreClientFactory);
   }
 
   @Test
@@ -90,6 +96,25 @@ public void defaultFactory() {
     assertThat(hiveConf.getBoolVar(ConfVars.METASTORE_USE_THRIFT_COMPACT_PROTOCOL), is(false));
   }
 
+  @Test
+  public void splitTrafficFactory() {
+    PrimaryMetaStore metaStore = newPrimaryInstance("hms", THRIFT_URI);
+    metaStore.setReadOnlyRemoteMetaStoreUris(THRIFT_URI_READ_ONLY);   
+    CloseableThriftHiveMetastoreIface readWriteClient = mock(CloseableThriftHiveMetastoreIface.class);
+    //Using 'any(HiveConf.class); generic matcher because HiveConf doesn't implement equals.
+    when(defaultMetaStoreClientFactory
+        .newInstance(any(HiveConf.class), eq("waggledance-hms"), eq(3), eq(2000))).thenReturn(readWriteClient);
+    CloseableThriftHiveMetastoreIface readOnlyclient = mock(CloseableThriftHiveMetastoreIface.class);
+    when(defaultMetaStoreClientFactory
+        .newInstance(any(HiveConf.class), eq("waggledance-hms_ro"), eq(3), eq(2000))).thenReturn(readOnlyclient);
+
+    factory.newInstance(metaStore);
+
+
+    verify(splitTrafficMetaStoreClientFactory).newInstance(readWriteClient, readOnlyclient);
+    verifyNoInteractions(tunnelingMetaStoreClientFactory);
+  }
+
   @Test
   public void tunnelingFactory() {
     MetastoreTunnel metastoreTunnel = new MetastoreTunnel();

waggle-dance-core/src/test/java/com/hotels/bdp/waggledance/mapping/model/MetaStoreMappingFactoryImplTest.java

@@ -1,5 +1,5 @@
 /**
- * Copyright (C) 2016-2022 Expedia, Inc.
+ * Copyright (C) 2016-2024 Expedia, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -43,6 +43,7 @@
 import com.hotels.bdp.waggledance.client.CloseableThriftHiveMetastoreIfaceClientFactory;
 import com.hotels.bdp.waggledance.client.DefaultMetaStoreClientFactory;
 import com.hotels.bdp.waggledance.client.GlueClientFactory;
+import com.hotels.bdp.waggledance.client.SplitTrafficMetastoreClientFactory;
 import com.hotels.bdp.waggledance.client.tunnelling.TunnelingMetaStoreClientFactory;
 import com.hotels.bdp.waggledance.conf.WaggleDanceConfiguration;
 import com.hotels.bdp.waggledance.mapping.service.PrefixNamingStrategy;
@@ -61,7 +62,7 @@ public class MetaStoreMappingFactoryImplTest {
   private @Mock AccessControlHandlerFactory accessControlHandlerFactory;
   private final CloseableThriftHiveMetastoreIfaceClientFactory metaStoreClientFactory = new CloseableThriftHiveMetastoreIfaceClientFactory(
       new TunnelingMetaStoreClientFactory(), new DefaultMetaStoreClientFactory(), new GlueClientFactory(),
-      new WaggleDanceConfiguration());
+      new WaggleDanceConfiguration(), new SplitTrafficMetastoreClientFactory());
 
   private MetaStoreMappingFactoryImpl factory;