OMID-310 Make Zookeeper SASL login context name configurable

stoty · stoty · commit 3f12e66b9997 · 2025-04-09T10:29:05.000+02:00
diff --git a/common/pom.xml b/common/pom.xml
@@ -44,6 +44,13 @@
             <artifactId>curator-framework</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.zookeeper</groupId>
+            <artifactId>zookeeper</artifactId>
+            <version>${zookeeper.version}</version>
+            <scope>provided</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.google.protobuf</groupId>
             <artifactId>protobuf-java</artifactId>
diff --git a/common/src/main/java/org/apache/omid/zk/ZKUtils.java b/common/src/main/java/org/apache/omid/zk/ZKUtils.java
@@ -21,6 +21,7 @@
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.ExponentialBackoffRetry;
+import org.apache.zookeeper.client.ZKClientConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -31,16 +32,25 @@ public class ZKUtils {
 
     private static final Logger LOG = LoggerFactory.getLogger(ZKUtils.class);
 
-    public static CuratorFramework initZKClient(String zkCluster, String namespace, int zkConnectionTimeoutInSec)
+    public static CuratorFramework initZKClient(String zkCluster, String namespace, int zkConnectionTimeoutInSec, String zkLoginContextName)
             throws IOException {
 
         LOG.info("Creating Zookeeper Client connecting to {}", zkCluster);
 
+        ZKClientConfig zkConfig = new ZKClientConfig();
+        if (zkLoginContextName != null) {
+            // TODO should we check if this exists ?
+            // Or just error out with an unsuccessful connection, as we do now ?
+            LOG.info("Using Login Context {} for Zookeeper", zkCluster);
+            zkConfig.setProperty(ZKClientConfig.LOGIN_CONTEXT_NAME_KEY, zkLoginContextName);
+        }
+
         RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
         CuratorFramework zkClient = CuratorFrameworkFactory.builder()
                 .namespace(namespace)
                 .connectString(zkCluster)
                 .retryPolicy(retryPolicy)
+                .zkClientConfig(zkConfig)
                 .build();
 
         zkClient.start();
diff --git a/hbase-client/src/test/java/org/apache/omid/transaction/TestHALeaseManagementModule.java b/hbase-client/src/test/java/org/apache/omid/transaction/TestHALeaseManagementModule.java
@@ -54,7 +54,7 @@ class TestHALeaseManagementModule extends AbstractModule {
 
     @Override
     protected void configure() {
-        install(new ZKModule(zkCluster, zkNamespace));
+        install(new ZKModule(zkCluster, zkNamespace, null));
     }
 
     @Provides
diff --git a/pom.xml b/pom.xml
@@ -168,6 +168,7 @@
         <guava.version>32.1.3-jre</guava.version>
         <!-- 2.12+ shades guava -->
         <curator.version>5.6.0</curator.version>
+        <zookeeper.version>3.8.4</zookeeper.version>
         <snakeyaml.version>2.2</snakeyaml.version>
         <beanutils.version>1.9.4</beanutils.version>
         <commons-io.version>2.18.0</commons-io.version>
diff --git a/timestamp-storage/src/main/java/org/apache/omid/timestamp/storage/DefaultZKTimestampStorageModule.java b/timestamp-storage/src/main/java/org/apache/omid/timestamp/storage/DefaultZKTimestampStorageModule.java
@@ -27,10 +27,11 @@ public class DefaultZKTimestampStorageModule extends AbstractModule {
 
     private String zkCluster = "localhost:2181";
     private String namespace = "omid";
+    private String zkLoginContextName;
 
     @Override
     public void configure() {
-        install(new ZKModule(zkCluster, namespace));
+        install(new ZKModule(zkCluster, namespace, zkLoginContextName));
         install(new ZKTimestampStorageModule());
     }
 
@@ -54,4 +55,12 @@ public void setNamespace(String namespace) {
         this.namespace = namespace;
     }
 
+    public String getZkLoginContextName() {
+        return zkLoginContextName;
+    }
+
+    public void setZkLoginContextName(String zkLoginContextName) {
+        this.zkLoginContextName = zkLoginContextName;
+    }
+
 }
diff --git a/timestamp-storage/src/main/java/org/apache/omid/timestamp/storage/ZKModule.java b/timestamp-storage/src/main/java/org/apache/omid/timestamp/storage/ZKModule.java
@@ -25,15 +25,19 @@
 import javax.inject.Singleton;
 import java.io.IOException;
 
-//TODO:IK: move to common?
+// TODO:IK: move to common?
+// The problem is that common doesn't depend on Guice, and we want to keep that that way.
+// Ideally, this would be in a new module, something like omid-server-common
 public class ZKModule extends AbstractModule {
 
     private final String zkCluster;
     private final String namespace;
+    private final String zkLoginContextName;
 
-    public ZKModule(String zkCluster, String namespace) {
+    public ZKModule(String zkCluster, String namespace, String zkLoginContextName) {
         this.zkCluster = zkCluster;
         this.namespace = namespace;
+        this.zkLoginContextName = zkLoginContextName;
     }
 
     @Override
@@ -43,7 +47,7 @@ public void configure() {
     @Provides
     @Singleton
     CuratorFramework provideInitializedZookeeperClient() throws IOException {
-        return ZKUtils.initZKClient(zkCluster, namespace, 10);
+        return ZKUtils.initZKClient(zkCluster, namespace, 10, zkLoginContextName);
     }
 
     // ----------------------------------------------------------------------------------------------------------------
diff --git a/transaction-client/src/main/java/org/apache/omid/tso/client/OmidClientConfiguration.java b/transaction-client/src/main/java/org/apache/omid/tso/client/OmidClientConfiguration.java
@@ -45,6 +45,7 @@ public enum ConflictDetectionLevel {CELL, ROW}
     private String zkCurrentTsoPath;
     private String zkNamespace;
     private int zkConnectionTimeoutInSecs;
+    private String zkLoginContextName;
 
     // Communication protocol related params
 
@@ -202,6 +203,16 @@ public void setZkNamespace(String zkNamespace) {
         this.zkNamespace = zkNamespace;
     }
 
+    public String getZkLoginContextName() {
+        return zkLoginContextName;
+    }
+
+    @Inject(optional = true)
+    @Named("omid.ha.zkLoginContextName")
+    public void setZkLoginContextName(String zkLoginContextName) {
+        this.zkLoginContextName = zkLoginContextName;
+    }
+
     public PostCommitMode getPostCommitMode() {
         return postCommitMode;
     }
diff --git a/transaction-client/src/main/java/org/apache/omid/tso/client/TSOClient.java b/transaction-client/src/main/java/org/apache/omid/tso/client/TSOClient.java
@@ -135,7 +135,8 @@ private TSOClient(OmidClientConfiguration omidConf) throws IOException {
             case HA:
                 zkClient = ZKUtils.initZKClient(omidConf.getConnectionString(),
                                                 omidConf.getZkNamespace(),
-                                                omidConf.getZkConnectionTimeoutInSecs());
+                                                omidConf.getZkConnectionTimeoutInSecs(),
+                                                omidConf.getZkLoginContextName());
                 zkCurrentTsoPath = omidConf.getZkCurrentTsoPath();
                 configureCurrentTSOServerZNodeCache(zkCurrentTsoPath);
                 String tsoInfo = getCurrentTSOInfoFoundInZK(zkCurrentTsoPath);
diff --git a/tso-server/src/main/java/org/apache/omid/tso/HALeaseManagementModule.java b/tso-server/src/main/java/org/apache/omid/tso/HALeaseManagementModule.java
@@ -40,6 +40,7 @@ public class HALeaseManagementModule extends AbstractModule {
     private String currentTsoPath = "/current-tso";
     private String zkCluster = "localhost:2181";
     private String zkNamespace = "omid";
+    private String zkLoginContextName;
 
     // ----------------------------------------------------------------------------------------------------------------
     // WARNING: Do not remove empty constructor, needed by snake_yaml!
@@ -51,20 +52,20 @@ public HALeaseManagementModule() {
 
     @VisibleForTesting
     public HALeaseManagementModule(long leasePeriodInMs, String tsoLeasePath, String currentTsoPath,
-                                   String zkCluster, String zkNamespace) {
+                                   String zkCluster, String zkNamespace, String zkLoginContextName) {
 
         this.leasePeriodInMs = leasePeriodInMs;
         this.tsoLeasePath = tsoLeasePath;
         this.currentTsoPath = currentTsoPath;
         this.zkCluster = zkCluster;
         this.zkNamespace = zkNamespace;
+        this.zkLoginContextName = zkLoginContextName;
 
     }
 
     @Override
     protected void configure() {
-
-        install(new ZKModule(zkCluster, zkNamespace));
+        install(new ZKModule(zkCluster, zkNamespace, zkLoginContextName));
 
     }
 
@@ -133,4 +134,12 @@ public void setZkNamespace(String zkNamespace) {
         this.zkNamespace = zkNamespace;
     }
 
+    public String getZkLoginContextName() {
+        return zkLoginContextName;
+    }
+
+    public void setZkLoginContextName(String zkLoginContextName) {
+        this.zkLoginContextName = zkLoginContextName;
+    }
+
 }
diff --git a/tso-server/src/test/java/org/apache/omid/tso/client/TestTSOClientConnectionToTSO.java b/tso-server/src/test/java/org/apache/omid/tso/client/TestTSOClientConnectionToTSO.java
@@ -589,7 +589,7 @@ public void testSuccessfulConnectionToTSOThroughZK() throws Exception {
         TSOServerConfig config = new TSOServerConfig();
         config.setConflictMapSize(1000);
         config.setPort(tsoPortForTest);
-        config.setLeaseModule(new HALeaseManagementModule(1000, TSO_LEASE_PATH, CURRENT_TSO_PATH, zkClusterForTest, "omid"));
+        config.setLeaseModule(new HALeaseManagementModule(1000, TSO_LEASE_PATH, CURRENT_TSO_PATH, zkClusterForTest, "omid", null));
         injector = Guice.createInjector(new TSOMockModule(config));
         LOG.info("Starting TSO");
         tsoServer = injector.getInstance(TSOServer.class);
@@ -629,7 +629,7 @@ public void testSuccessOfTSOClientReconnectionsToARestartedTSOWithZKPublishing()
         TSOServerConfig config = new TSOServerConfig();
         config.setConflictMapSize(1000);
         config.setPort(tsoPortForTest);
-        config.setLeaseModule(new HALeaseManagementModule(1000, TSO_LEASE_PATH, CURRENT_TSO_PATH, zkClusterForTest, "omid"));
+        config.setLeaseModule(new HALeaseManagementModule(1000, TSO_LEASE_PATH, CURRENT_TSO_PATH, zkClusterForTest, "omid", null));
         injector = Guice.createInjector(new TSOMockModule(config));
         LOG.info("Starting Initial TSO");
         tsoServer = injector.getInstance(TSOServer.class);

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ class TestHALeaseManagementModule extends AbstractModule {`
`54`	`54`
`55`	`55`	`@Override`
`56`	`56`	`protected void configure() {`
`57`		`- install(new ZKModule(zkCluster, zkNamespace));`
	`57`	`+ install(new ZKModule(zkCluster, zkNamespace, null));`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`@Provides`