Feature/add whitelist

marklogic-data-hub/src/main/java/com/marklogic/hub/HubConfig.java

@@ -151,6 +151,8 @@ public class HubConfig {
     public String hubRoleName = DEFAULT_ROLE_NAME;
     public String hubUserName = DEFAULT_USER_NAME;
 
+    public String[] loadBalancerHosts;
+
     public String customForestPath = DEFAULT_CUSTOM_FOREST_PATH;
 
     public String modulePermissions = "rest-reader,read,rest-writer,insert,rest-writer,update,rest-extension-user,execute";
@@ -345,6 +347,11 @@ private void loadConfigurationFromProperties(Properties environmentProperties) {
             hubRoleName = getEnvPropString(environmentProperties, "mlHubUserRole", hubRoleName);
             hubUserName = getEnvPropString(environmentProperties, "mlHubUserName", hubUserName);
 
+            String lbh = getEnvPropString(environmentProperties, "mlLoadBalancerHosts", null);
+            if (lbh != null && lbh.length() > 0) {
+                loadBalancerHosts = lbh.split(",");
+            }
+
             projectDir = getEnvPropString(environmentProperties, "hubProjectDir", projectDir);
 
             logger.info("Hub Project Dir: " + projectDir);

marklogic-data-hub/src/main/java/com/marklogic/hub/flow/impl/FlowRunnerImpl.java

@@ -209,7 +209,7 @@ public JobTicket run() {
 
         HashMap<String, JobTicket> ticketWrapper = new HashMap<>();
 
-        QueryBatcher queryBatcher = dataMovementManager.newQueryBatcher(uris.iterator())
+        QueryBatcher tempQueryBatcher = dataMovementManager.newQueryBatcher(uris.iterator())
             .withBatchSize(batchSize)
             .withThreadCount(threadCount)
             .onUrisReady((QueryBatch batch) -> {
@@ -275,6 +275,17 @@ public JobTicket run() {
                 failedEvents.addAndGet(batchSize);
             });
 
+
+        if (hubConfig.loadBalancerHosts != null && hubConfig.loadBalancerHosts.length > 0){
+            tempQueryBatcher = tempQueryBatcher.withForestConfig(
+                new FilteredForestConfiguration(
+                    dataMovementManager.readForestConfig()
+                ).withWhiteList(hubConfig.loadBalancerHosts)
+            );
+        }
+        QueryBatcher queryBatcher = tempQueryBatcher;
+
+
         JobTicket jobTicket = dataMovementManager.startJob(queryBatcher);
         ticketWrapper.put("jobTicket", jobTicket);
         jobManager.saveJob(job.withStatus(JobStatus.RUNNING_HARMONIZE));

marklogic-data-hub/src/main/resources/scaffolding/gradle_properties

@@ -106,3 +106,9 @@ mlHubUserRole=%%mlHubUserRole%%
 mlHubUserName=%%mlHubUserName%%
 # this password is autogenerated for you via the 'gradle hubInit' task
 mlHubUserPassword=%%mlHubUserPassword%%
+
+# If you are working with a load balancer you can
+# specify the hostname(s) of the load balancer here separated by commas
+# see: https://docs.marklogic.com/guide/java/data-movement#id_26583
+# for more info about what's happening under the hood
+# mlLoadBalancerHosts=your-load-balancer-hostnames,separated,by-commas

marklogic-data-hub/src/test/java/com/marklogic/hub/HubConfigTest.java

@@ -0,0 +1,79 @@
+package com.marklogic.hub;
+
+import org.apache.commons.io.FileUtils;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.Properties;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class HubConfigTest extends HubTestBase {
+
+    private static File projectPath = new File(PROJECT_PATH);
+
+    @Before
+    public void setup() throws IOException {
+        FileUtils.deleteDirectory(projectPath);
+        HubConfig config = getHubConfig();
+        config.initHubProject();
+    }
+
+    @After
+    public void teardown() throws IOException {
+        FileUtils.deleteDirectory(projectPath);
+    }
+
+    private void deleteProp(String key) {
+        try {
+            File gradleProperties = new File(projectPath, "gradle.properties");
+            Properties props = new Properties();
+            props.load(new FileInputStream(gradleProperties));
+            props.remove(key);
+            props.store(new FileOutputStream(gradleProperties), "");
+        }
+        catch(IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void writeProp(String key, String value) {
+        try {
+            File gradleProperties = new File(projectPath, "gradle.properties");
+            Properties props = new Properties();
+            props.load(new FileInputStream(gradleProperties));
+            props.put(key, value);
+            props.store(new FileOutputStream(gradleProperties), "");
+        }
+        catch(IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    public void testLoadBalancerProps() {
+        deleteProp("mlLoadBalancerHosts");
+        assertNull(getHubConfig().loadBalancerHosts);
+
+        writeProp("mlLoadBalancerHosts", "");
+        assertNull(getHubConfig().loadBalancerHosts);
+
+        writeProp("mlLoadBalancerHosts", "host1,host2");
+        HubConfig config = getHubConfig();
+        assertEquals(2, config.loadBalancerHosts.length);
+        assertEquals("host1", config.loadBalancerHosts[0]);
+        assertEquals("host2", config.loadBalancerHosts[1]);
+
+        writeProp("mlLoadBalancerHosts", "host1");
+        config = getHubConfig();
+        assertEquals(1, config.loadBalancerHosts.length);
+        assertEquals("host1", config.loadBalancerHosts[0]);
+    }
+
+}