Avoid unnecessary conversion to ConsistencyLevel in QueryTraceFetcher

olim7t · olim7t · commit 0fb742ac8ee2 · 2018-11-01T14:09:10.000-07:00
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/QueryTraceFetcher.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/QueryTraceFetcher.java
@@ -15,7 +15,6 @@
  */
 package com.datastax.oss.driver.internal.core.cql;
 
-import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
@@ -56,18 +55,12 @@ class QueryTraceFetcher {
     this.tracingId = tracingId;
     this.session = session;
 
-    ConsistencyLevel regularConsistency =
-        context
-            .getConsistencyLevelRegistry()
-            .fromName(config.getString(DefaultDriverOption.REQUEST_CONSISTENCY));
-    ConsistencyLevel traceConsistency =
-        context
-            .getConsistencyLevelRegistry()
-            .fromName(config.getString(DefaultDriverOption.REQUEST_TRACE_CONSISTENCY));
+    String regularConsistency = config.getString(DefaultDriverOption.REQUEST_CONSISTENCY);
+    String traceConsistency = config.getString(DefaultDriverOption.REQUEST_TRACE_CONSISTENCY);
     this.config =
         (traceConsistency.equals(regularConsistency))
             ? config
-            : config.withString(DefaultDriverOption.REQUEST_CONSISTENCY, traceConsistency.name());
+            : config.withString(DefaultDriverOption.REQUEST_CONSISTENCY, traceConsistency);
 
     this.maxAttempts = config.getInt(DefaultDriverOption.REQUEST_TRACE_ATTEMPTS);
     this.intervalNanos = config.getDuration(DefaultDriverOption.REQUEST_TRACE_INTERVAL).toNanos();
diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/QueryTraceFetcherTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/QueryTraceFetcherTest.java
@@ -32,7 +32,6 @@
 import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import com.datastax.oss.driver.api.core.cql.TraceEvent;
 import com.datastax.oss.driver.api.core.uuid.Uuids;
-import com.datastax.oss.driver.internal.core.DefaultConsistencyLevelRegistry;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.context.NettyOptions;
 import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
@@ -105,9 +104,6 @@ public void setup() {
             config.withString(
                 DefaultDriverOption.REQUEST_CONSISTENCY, DefaultConsistencyLevel.ONE.name()))
         .thenReturn(traceConfig);
-
-    Mockito.when(context.getConsistencyLevelRegistry())
-        .thenReturn(new DefaultConsistencyLevelRegistry());
   }
 
   @Test
