CASSJAVA-92: Local DC provided for nodetool clientstats

Author: lukasz-antoniak

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

@@ -18,6 +18,7 @@
 package com.datastax.oss.driver.internal.core.context;
 
 import com.datastax.dse.driver.api.core.config.DseDriverOption;
+import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.session.Session;
 import com.datastax.oss.driver.api.core.uuid.Uuids;
@@ -33,6 +34,7 @@ public class StartupOptionsBuilder {
 
   public static final String DRIVER_NAME_KEY = "DRIVER_NAME";
   public static final String DRIVER_VERSION_KEY = "DRIVER_VERSION";
+  public static final String DRIVER_LOCAL_DC = "DRIVER_LOCAL_DC";
   public static final String APPLICATION_NAME_KEY = "APPLICATION_NAME";
   public static final String APPLICATION_VERSION_KEY = "APPLICATION_VERSION";
   public static final String CLIENT_ID_KEY = "CLIENT_ID";
@@ -41,6 +43,7 @@ public class StartupOptionsBuilder {
   private UUID clientId;
   private String applicationName;
   private String applicationVersion;
+  private String applicationLocalDc;
 
   public StartupOptionsBuilder(InternalDriverContext context) {
     this.context = context;
@@ -119,6 +122,12 @@ public Map<String, String> build() {
     if (applicationVersion != null) {
       builder.put(APPLICATION_VERSION_KEY, applicationVersion);
     }
+    if (applicationLocalDc == null) {
+      applicationLocalDc = localDc(config);
+    }
+    if (applicationLocalDc != null) {
+      builder.put(DRIVER_LOCAL_DC, applicationLocalDc);
+    }
 
     return builder.build();
   }
@@ -142,4 +151,14 @@ protected String getDriverName() {
   protected String getDriverVersion() {
     return Session.OSS_DRIVER_COORDINATES.getVersion().toString();
   }
+
+  private String localDc(DriverExecutionProfile profile) {
+    String dc = context.getLocalDatacenter(profile.getName()); // DC set programmatically
+    if (dc == null && profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) {
+      dc =
+          profile.getString(
+              DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER); // DC from configuration
+    }
+    return dc;
+  }
 }

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

@@ -60,15 +60,17 @@ public void before() {
     when(defaultProfile.isDefined(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE)).thenReturn(true);
   }
 
-  private void buildContext(UUID clientId, String applicationName, String applicationVersion) {
-    this.driverContext =
-        new DefaultDriverContext(
-            configLoader,
-            ProgrammaticArguments.builder()
-                .withStartupClientId(clientId)
-                .withStartupApplicationName(applicationName)
-                .withStartupApplicationVersion(applicationVersion)
-                .build());
+  private void buildContext(
+      UUID clientId, String applicationName, String applicationVersion, String localDc) {
+    ProgrammaticArguments.Builder builder =
+        ProgrammaticArguments.builder()
+            .withStartupClientId(clientId)
+            .withStartupApplicationName(applicationName)
+            .withStartupApplicationVersion(applicationVersion);
+    if (localDc != null) {
+      builder.withLocalDatacenter(DriverExecutionProfile.DEFAULT_NAME, localDc);
+    }
+    this.driverContext = new DefaultDriverContext(configLoader, builder.build());
   }
 
   private void assertDefaultStartupOptions(Startup startup) {
@@ -86,7 +88,7 @@ private void assertDefaultStartupOptions(Startup startup) {
   public void should_build_startup_options_with_no_compression_if_undefined() {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("none");
-    buildContext(null, null, null);
+    buildContext(null, null, null, null);
     Startup startup = new Startup(driverContext.getStartupOptions());
     assertThat(startup.options).doesNotContainKey(Startup.COMPRESSION_KEY);
     assertDefaultStartupOptions(startup);
@@ -97,7 +99,7 @@ public void should_build_startup_options_with_no_compression_if_undefined() {
   public void should_build_startup_options_with_compression(String compression) {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn(compression);
-    buildContext(null, null, null);
+    buildContext(null, null, null, null);
     Startup startup = new Startup(driverContext.getStartupOptions());
     // assert the compression option is present
     assertThat(startup.options).containsEntry(Startup.COMPRESSION_KEY, compression);
@@ -110,7 +112,7 @@ public void should_build_startup_options_with_compression(String compression) {
   public void should_fail_to_build_startup_options_with_invalid_compression() {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("foobar");
-    buildContext(null, null, null);
+    buildContext(null, null, null, null);
     assertThatIllegalArgumentException()
         .isThrownBy(() -> new Startup(driverContext.getStartupOptions()));
   }
@@ -120,7 +122,7 @@ public void should_build_startup_options_with_client_id() {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("none");
     UUID customClientId = Uuids.random();
-    buildContext(customClientId, null, null);
+    buildContext(customClientId, null, null, null);
     Startup startup = new Startup(driverContext.getStartupOptions());
     // assert the client id is present
     assertThat(startup.options)
@@ -135,7 +137,7 @@ public void should_build_startup_options_with_client_id() {
   public void should_build_startup_options_with_application_version_and_name() {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("none");
-    buildContext(null, "Custom_App_Name", "Custom_App_Version");
+    buildContext(null, "Custom_App_Name", "Custom_App_Version", null);
     Startup startup = new Startup(driverContext.getStartupOptions());
     // assert the app name and version are present
     assertThat(startup.options)
@@ -151,15 +153,17 @@ public void should_build_startup_options_with_all_options() {
     // mock config to specify "snappy" compression
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("snappy");
+    when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
 
     UUID customClientId = Uuids.random();
 
-    buildContext(customClientId, "Custom_App_Name", "Custom_App_Version");
+    buildContext(customClientId, "Custom_App_Name", "Custom_App_Version", "dc6");
     Startup startup = new Startup(driverContext.getStartupOptions());
     assertThat(startup.options)
         .containsEntry(StartupOptionsBuilder.CLIENT_ID_KEY, customClientId.toString())
         .containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Custom_App_Name")
-        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version");
+        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version")
+        .containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "dc6");
     assertThat(startup.options).containsEntry(Startup.COMPRESSION_KEY, "snappy");
     assertDefaultStartupOptions(startup);
   }
@@ -172,25 +176,33 @@ public void should_use_configuration_when_no_programmatic_values_provided() {
         .thenReturn("Config_App_Version");
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("none");
+    when(defaultProfile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))
+        .thenReturn("Config_DC_1");
+    when(defaultProfile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))
+        .thenReturn(true);
+    when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
 
-    buildContext(null, null, null);
+    buildContext(null, null, null, null);
     Startup startup = new Startup(driverContext.getStartupOptions());
 
     assertThat(startup.options)
         .containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Config_App_Name")
-        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Config_App_Version");
+        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Config_App_Version")
+        .containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "Config_DC_1");
   }
 
   @Test
   public void should_ignore_configuration_when_programmatic_values_provided() {
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn("none");
+    when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME);
 
-    buildContext(null, "Custom_App_Name", "Custom_App_Version");
+    buildContext(null, "Custom_App_Name", "Custom_App_Version", "us-west-2");
     Startup startup = new Startup(driverContext.getStartupOptions());
 
     assertThat(startup.options)
         .containsEntry(StartupOptionsBuilder.APPLICATION_NAME_KEY, "Custom_App_Name")
-        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version");
+        .containsEntry(StartupOptionsBuilder.APPLICATION_VERSION_KEY, "Custom_App_Version")
+        .containsEntry(StartupOptionsBuilder.DRIVER_LOCAL_DC, "us-west-2");
   }
 }