diff --git a/marklogic-data-hub/src/main/java/com/marklogic/hub/HubConfig.java b/marklogic-data-hub/src/main/java/com/marklogic/hub/HubConfig.java index eb6d277685..c2e177e8fd 100644 --- a/marklogic-data-hub/src/main/java/com/marklogic/hub/HubConfig.java +++ b/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); diff --git a/marklogic-data-hub/src/main/java/com/marklogic/hub/flow/impl/FlowRunnerImpl.java b/marklogic-data-hub/src/main/java/com/marklogic/hub/flow/impl/FlowRunnerImpl.java index c67d3dafbd..4fbbbb2204 100644 --- a/marklogic-data-hub/src/main/java/com/marklogic/hub/flow/impl/FlowRunnerImpl.java +++ b/marklogic-data-hub/src/main/java/com/marklogic/hub/flow/impl/FlowRunnerImpl.java @@ -209,7 +209,7 @@ public JobTicket run() { HashMap 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)); diff --git a/marklogic-data-hub/src/main/resources/scaffolding/gradle_properties b/marklogic-data-hub/src/main/resources/scaffolding/gradle_properties index 5238570501..10ef75ac49 100644 --- a/marklogic-data-hub/src/main/resources/scaffolding/gradle_properties +++ b/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 diff --git a/marklogic-data-hub/src/test/java/com/marklogic/hub/HubConfigTest.java b/marklogic-data-hub/src/test/java/com/marklogic/hub/HubConfigTest.java new file mode 100644 index 0000000000..5043f05ab3 --- /dev/null +++ b/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]); + } + +}