Skip to content

Commit 3829053

Browse files
lukasz-antoniaklukasz-antoniak
committed
Return local DC for all profiles
1 parent 101621a commit 3829053

File tree

2 files changed

+80
-28
lines changed

2 files changed

+80
-28
lines changed

‎core/src/main/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilder.java

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
package com.datastax.oss.driver.internal.core.context;
1919

2020
import com.datastax.dse.driver.api.core.config.DseDriverOption;
21-
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
2221
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
2322
import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
2423
import com.datastax.oss.driver.api.core.loadbalancing.LocalDcAwareLoadBalancingPolicy;
@@ -124,9 +123,9 @@ public Map<String, String> build() {
124123
builder.put(APPLICATION_VERSION_KEY, applicationVersion);
125124
}
126125
// do not cache local DC as it can change within LBP implementation
127-
String applicationLocalDc = localDc(config);
128-
if (applicationLocalDc != null) {
129-
builder.put(DRIVER_LOCAL_DC, applicationLocalDc);
126+
String applicationLocalDcs = localDcs();
127+
if (applicationLocalDcs != null) {
128+
builder.put(DRIVER_LOCAL_DC, applicationLocalDcs);
130129
}
131130

132131
return builder.build();
@@ -152,19 +151,28 @@ protected String getDriverVersion() {
152151
return Session.OSS_DRIVER_COORDINATES.getVersion().toString();
153152
}
154153

155-
private String localDc(DriverExecutionProfile profile) {
156-
String dc = context.getLocalDatacenter(profile.getName()); // DC set programmatically
157-
if (dc == null) {
158-
// DC from load balancing policy
159-
LoadBalancingPolicy lbp = context.getLoadBalancingPolicy(DriverExecutionProfile.DEFAULT_NAME);
160-
if (lbp instanceof LocalDcAwareLoadBalancingPolicy) {
161-
dc = ((LocalDcAwareLoadBalancingPolicy) lbp).getLocalDatacenter();
154+
private String localDcs() {
155+
StringBuilder result = new StringBuilder();
156+
boolean first = true;
157+
for (Map.Entry<String, LoadBalancingPolicy> entry :
158+
context.getLoadBalancingPolicies().entrySet()) {
159+
String dc = getLocalDc(entry.getValue());
160+
if (dc != null) {
161+
if (!first) {
162+
result.append(", ");
163+
} else {
164+
first = false;
165+
}
166+
result.append(entry.getKey()).append(": ").append(dc);
162167
}
163168
}
164-
if (dc == null && profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) {
165-
// DC from configuration
166-
dc = profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER);
169+
return first ? null : result.toString();
170+
}
171+
172+
private String getLocalDc(LoadBalancingPolicy loadBalancingPolicy) {
173+
if (loadBalancingPolicy instanceof LocalDcAwareLoadBalancingPolicy) {
174+
return ((LocalDcAwareLoadBalancingPolicy) loadBalancingPolicy).getLocalDatacenter();
167175
}
168-
return dc;
176+
return null;
169177
}
170178
}

‎core/src/test/java/com/datastax/dse/driver/internal/core/context/DseStartupOptionsBuilderTest.java

Lines changed: 57 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import static org.assertj.core.api.Assertions.assertThat;
2121
import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
22+
import static org.mockito.Mockito.mock;
2223
import static org.mockito.Mockito.when;
2324
import static org.mockito.MockitoAnnotations.initMocks;
2425

@@ -28,14 +29,19 @@
2829
import com.datastax.oss.driver.api.core.config.DriverConfig;
2930
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
3031
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
32+
import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
33+
import com.datastax.oss.driver.api.core.loadbalancing.LocalDcAwareLoadBalancingPolicy;
3134
import com.datastax.oss.driver.api.core.session.ProgrammaticArguments;
3235
import com.datastax.oss.driver.api.core.session.Session;
3336
import com.datastax.oss.driver.api.core.uuid.Uuids;
3437
import com.datastax.oss.driver.internal.core.context.DefaultDriverContext;
3538
import com.datastax.oss.driver.internal.core.context.StartupOptionsBuilder;
39+
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
3640
import com.datastax.oss.protocol.internal.request.Startup;
3741
import com.tngtech.java.junit.dataprovider.DataProvider;
3842
import com.tngtech.java.junit.dataprovider.DataProviderRunner;
43+
import edu.umd.cs.findbugs.annotations.NonNull;
44+
import java.util.Map;
3945
import java.util.UUID;
4046
import org.junit.Before;
4147
import org.junit.Test;
@@ -61,16 +67,34 @@ public void before() {
6167
}
6268

6369
private void buildContext(
64-
UUID clientId, String applicationName, String applicationVersion, String localDc) {
70+
UUID clientId,
71+
String applicationName,
72+
String applicationVersion,
73+
Map<String, String> localDcPerProfile) {
6574
ProgrammaticArguments.Builder builder =
6675
ProgrammaticArguments.builder()
6776
.withStartupClientId(clientId)
6877
.withStartupApplicationName(applicationName)
6978
.withStartupApplicationVersion(applicationVersion);
70-
if (localDc != null) {
71-
builder.withLocalDatacenter(DriverExecutionProfile.DEFAULT_NAME, localDc);
72-
}
73-
this.driverContext = new DefaultDriverContext(configLoader, builder.build());
79+
this.driverContext =
80+
new DefaultDriverContext(configLoader, builder.build()) {
81+
@NonNull
82+
@Override
83+
public Map<String, LoadBalancingPolicy> getLoadBalancingPolicies() {
84+
if (localDcPerProfile != null) {
85+
ImmutableMap.Builder<String, LoadBalancingPolicy> map = ImmutableMap.builder();
86+
localDcPerProfile.forEach(
87+
(profile, dc) -> {
88+
LocalDcAwareLoadBalancingPolicy loadBalancingPolicy =
89+
mock(LocalDcAwareLoadBalancingPolicy.class);
90+
when(loadBalancingPolicy.getLocalDatacenter()).thenReturn(dc);
91+
map.put(profile, loadBalancingPolicy);
92+
});
93+
return map.build();
94+
}
95+
return super.getLoadBalancingPolicies();
96+
}
97+
};
7498
}
7599

76100
private void assertDefaultStartupOptions(Startup startup) {
@@ -157,13 +181,14 @@ public void should_build_startup_options_with_all_options() {
157181

158182
UUID customClientId = Uuids.random();
159183

160-
buildContext(customClientId, "Custom_App_Name", "Custom_App_Version", "dc6");
184+
buildContext(
185+
customClientId, "Custom_App_Name", "Custom_App_Version", ImmutableMap.of("default", "dc6"));
161186
Startup startup = new Startup(driverContext.getStartupOptions());
162187
assertThat(startup.options)
163188
.containsEntry(StartupOptionsBuilder.CLIENT_ID_KEY, customClientId.toString())
164189
.containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Custom_App_Name")
165190
.containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version")
166-
.containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "dc6");
191+
.containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "default: dc6");
167192
assertThat(startup.options).containsEntry(Startup.COMPRESSION_KEY, "snappy");
168193
assertDefaultStartupOptions(startup);
169194
}
@@ -176,8 +201,6 @@ public void should_use_configuration_when_no_programmatic_values_provided() {
176201
.thenReturn("Config_App_Version");
177202
when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
178203
.thenReturn("none");
179-
when(defaultProfile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))
180-
.thenReturn("Config_DC_1");
181204
when(defaultProfile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))
182205
.thenReturn(true);
183206
when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
@@ -187,8 +210,7 @@ public void should_use_configuration_when_no_programmatic_values_provided() {
187210

188211
assertThat(startup.options)
189212
.containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Config_App_Name")
190-
.containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Config_App_Version")
191-
.containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "Config_DC_1");
213+
.containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Config_App_Version");
192214
}
193215

194216
@Test
@@ -197,12 +219,34 @@ public void should_ignore_configuration_when_programmatic_values_provided() {
197219
.thenReturn("none");
198220
when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
199221

200-
buildContext(null, "Custom_App_Name", "Custom_App_Version", "us-west-2");
222+
buildContext(
223+
null, "Custom_App_Name", "Custom_App_Version", ImmutableMap.of("default", "us-west-2"));
201224
Startup startup = new Startup(driverContext.getStartupOptions());
202225

203226
assertThat(startup.options)
204227
.containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Custom_App_Name")
205228
.containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version")
206-
.containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "us-west-2");
229+
.containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "default: us-west-2");
230+
}
231+
232+
@Test
233+
public void should_include_all_local_dc_in_startup_message() {
234+
when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
235+
.thenReturn("none");
236+
when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
237+
238+
buildContext(
239+
null,
240+
"Custom_App_Name",
241+
"Custom_App_Version",
242+
ImmutableMap.of("default", "us-west-2", "oltp", "us-east-2", "olap", "eu-central-1"));
243+
Startup startup = new Startup(driverContext.getStartupOptions());
244+
245+
assertThat(startup.options)
246+
.containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Custom_App_Name")
247+
.containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version")
248+
.containsEntry(
249+
StartupOptionsBuilder.DRIVER_LOCAL_DC,
250+
"default: us-west-2, oltp: us-east-2, olap: eu-central-1");
207251
}
208252
}

0 commit comments

Comments
 (0)