fix: get rows returned metric working (#8230)

confluentinc · Oct 5, 2021 · da4f71e · da4f71e
1 parent c4a036a
commit da4f71e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/ksqldb-engine/src/main/java/io/confluent/ksql/util/ScalablePushQueryMetadata.java b/ksqldb-engine/src/main/java/io/confluent/ksql/util/ScalablePushQueryMetadata.java
@@ -20,10 +20,10 @@
 import io.confluent.ksql.physical.scalablepush.PushQueryPreparer;
 import io.confluent.ksql.physical.scalablepush.PushQueryQueuePopulator;
 import io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle;
-import io.confluent.ksql.query.BlockingRowQueue;
 import io.confluent.ksql.query.CompletionHandler;
 import io.confluent.ksql.query.LimitHandler;
 import io.confluent.ksql.query.QueryId;
+import io.confluent.ksql.query.TransientQueryQueue;
 import io.confluent.ksql.schema.ksql.LogicalSchema;
 import io.confluent.ksql.util.KsqlConstants.QuerySourceType;
 import io.confluent.ksql.util.KsqlConstants.RoutingNodeType;
@@ -39,7 +39,7 @@ public class ScalablePushQueryMetadata implements PushQueryMetadata {
   private volatile boolean closed = false;
   private final LogicalSchema logicalSchema;
   private final QueryId queryId;
-  private final BlockingRowQueue rowQueue;
+  private final TransientQueryQueue transientQueryQueue;
   private final Optional<ScalablePushQueryMetrics> scalablePushQueryMetrics;
   private final ResultType resultType;
   private final PushQueryQueuePopulator pushQueryQueuePopulator;
@@ -59,7 +59,7 @@ public class ScalablePushQueryMetadata implements PushQueryMetadata {
   public ScalablePushQueryMetadata(
       final LogicalSchema logicalSchema,
       final QueryId queryId,
-      final BlockingRowQueue blockingRowQueue,
+      final TransientQueryQueue transientQueryQueue,
       final Optional<ScalablePushQueryMetrics> scalablePushQueryMetrics,
       final ResultType resultType,
       final PushQueryQueuePopulator pushQueryQueuePopulator,
@@ -70,7 +70,7 @@ public ScalablePushQueryMetadata(
   ) {
     this.logicalSchema = logicalSchema;
     this.queryId = queryId;
-    this.rowQueue = blockingRowQueue;
+    this.transientQueryQueue = transientQueryQueue;
     this.scalablePushQueryMetrics = scalablePushQueryMetrics;
     this.resultType = resultType;
     this.pushQueryQueuePopulator = pushQueryQueuePopulator;
@@ -108,7 +108,7 @@ public void start() {
 
   @Override
   public void close() {
-    rowQueue.close();
+    transientQueryQueue.close();
     startFuture.thenApply(handle -> {
       handle.close();
       return null;
@@ -123,18 +123,18 @@ public boolean isRunning() {
 
   @Override
   @SuppressFBWarnings(value = "EI_EXPOSE_REP")
-  public BlockingRowQueue getRowQueue() {
-    return rowQueue;
+  public TransientQueryQueue getRowQueue() {
+    return transientQueryQueue;
   }
 
   @Override
   public void setLimitHandler(final LimitHandler limitHandler) {
-    rowQueue.setLimitHandler(limitHandler);
+    transientQueryQueue.setLimitHandler(limitHandler);
   }
 
   @Override
   public void setCompletionHandler(final CompletionHandler completionHandler) {
-    rowQueue.setCompletionHandler(completionHandler);
+    transientQueryQueue.setCompletionHandler(completionHandler);
   }
 
   @Override
@@ -186,7 +186,7 @@ public RoutingNodeType getRoutingNodeType() {
   }
 
   public long getTotalRowsReturned() {
-    return rowQueue.size();
+    return transientQueryQueue.getTotalRowsQueued();
   }
 
   public long getTotalRowsProcessed() {

diff --git a/ksqldb-engine/src/test/java/io/confluent/ksql/util/ScalablePushQueryMetadataTest.java b/ksqldb-engine/src/test/java/io/confluent/ksql/util/ScalablePushQueryMetadataTest.java
@@ -24,8 +24,8 @@
 import io.confluent.ksql.internal.ScalablePushQueryMetrics;
 import io.confluent.ksql.physical.scalablepush.PushQueryQueuePopulator;
 import io.confluent.ksql.physical.scalablepush.PushRouting.PushConnectionsHandle;
-import io.confluent.ksql.query.BlockingRowQueue;
 import io.confluent.ksql.query.QueryId;
+import io.confluent.ksql.query.TransientQueryQueue;
 import io.confluent.ksql.schema.ksql.LogicalSchema;
 import io.confluent.ksql.util.PushQueryMetadata.ResultType;
 import java.util.Optional;
@@ -45,7 +45,7 @@ public class ScalablePushQueryMetadataTest {
   @Mock
   private LogicalSchema logicalSchema;
   @Mock
-  private BlockingRowQueue blockingRowQueue;
+  private TransientQueryQueue transientQueryQueue;
   @Mock
   private PushQueryQueuePopulator populator;
   @Mock
@@ -64,7 +64,7 @@ public void setUp()  {
     query = new ScalablePushQueryMetadata(
         logicalSchema,
         new QueryId("queryid"),
-        blockingRowQueue,
+        transientQueryQueue,
         metrics,
         ResultType.STREAM,
         populator,