1

scylladb · Oct 8, 2024 · d2fcb14 · d2fcb14
1 parent be5f5f8
commit d2fcb14
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 9 deletions.
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java b/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java
@@ -300,7 +300,8 @@ private void init(
                 .withOwnerLogPrefix(logPrefix + "|control")
                 .build();
 
-        Queue<Node> nodes = context.getLoadBalancingPolicyWrapper().newQueryPlan();
+        Queue<Node> nodes =
+            context.getLoadBalancingPolicyWrapper().newControlReconnectionQueryPlan();
 
         connect(
             nodes,

diff --git a/.../main/java/com/datastax/oss/driver/internal/core/metadata/LoadBalancingPolicyWrapper.java b/.../main/java/com/datastax/oss/driver/internal/core/metadata/LoadBalancingPolicyWrapper.java
@@ -161,11 +161,6 @@ public Queue<Node> newQueryPlan(
     }
   }
 
-  @NonNull
-  public Queue<Node> newQueryPlan() {
-    return newQueryPlan(null, DriverExecutionProfile.DEFAULT_NAME, null);
-  }
-
   @NonNull
   public Queue<Node> newControlReconnectionQueryPlan() {
     // First try the original way

diff --git a/...rc/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java b/...rc/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java
@@ -141,7 +141,7 @@ public void setup() {
   }
 
   protected void mockQueryPlan(Node... nodes) {
-    when(loadBalancingPolicyWrapper.newQueryPlan())
+    when(loadBalancingPolicyWrapper.newControlReconnectionQueryPlan())
         .thenAnswer(
             i -> {
               ConcurrentLinkedQueue<Node> queryPlan = new ConcurrentLinkedQueue<>();

diff --git a/...t/java/com/datastax/oss/driver/internal/core/metadata/LoadBalancingPolicyWrapperTest.java b/...t/java/com/datastax/oss/driver/internal/core/metadata/LoadBalancingPolicyWrapperTest.java
@@ -118,10 +118,22 @@ public void setup() {
                 policy3));
   }
 
+  @Test
+  public void should_build_control_connection_query_plan_from_contact_points_before_init() {
+    // When
+    Queue<Node> queryPlan = wrapper.newControlReconnectionQueryPlan();
+
+    // Then
+    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
+      verify(policy, never()).newQueryPlan(null, null);
+    }
+    assertThat(queryPlan).hasSameElementsAs(contactPoints);
+  }
+
   @Test
   public void should_build_query_plan_from_contact_points_before_init() {
     // When
-    Queue<Node> queryPlan = wrapper.newQueryPlan();
+    Queue<Node> queryPlan = wrapper.newQueryPlan(null, DriverExecutionProfile.DEFAULT_NAME, null);
 
     // Then
     for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
@@ -139,7 +151,24 @@ public void should_fetch_query_plan_from_policy_after_init() {
     }
 
     // When
-    Queue<Node> queryPlan = wrapper.newQueryPlan();
+    Queue<Node> queryPlan = wrapper.newControlReconnectionQueryPlan();
+
+    // Then
+    // no-arg newQueryPlan() uses the default profile
+    verify(policy1).newQueryPlan(null, null);
+    assertThat(queryPlan).isEqualTo(defaultPolicyQueryPlan);
+  }
+
+  @Test
+  public void should_fetch_control_connection_query_plan_from_policy_after_init() {
+    // Given
+    wrapper.init();
+    for (LoadBalancingPolicy policy : ImmutableList.of(policy1, policy2, policy3)) {
+      verify(policy).init(anyMap(), any(DistanceReporter.class));
+    }
+
+    // When
+    Queue<Node> queryPlan = wrapper.newQueryPlan(null, DriverExecutionProfile.DEFAULT_NAME, null);
 
     // Then
     // no-arg newQueryPlan() uses the default profile