test + review

confluentinc · Feb 25, 2022 · 5b543be · 5b543be
1 parent 67b1011
commit 5b543be
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 4 deletions.
diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java
@@ -933,7 +933,7 @@ static KsqlRestApplication buildApplication(
         internalTopicClient,
         commandTopicName,
         metricCollectors.getMetrics(),
-        ksqlConfig.getDouble("ksql.command.topic.rate.limit")
+        ksqlConfig.getDouble(KsqlConfig.KSQL_COMMAND_TOPIC_RATE_LIMIT_CONFIG)
     );
 
     final KsqlResource ksqlResource = new KsqlResource(

diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandRunner.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandRunner.java
@@ -27,7 +27,6 @@
 import io.confluent.ksql.rest.util.PersistentQueryCleanupImpl;
 import io.confluent.ksql.rest.util.TerminateCluster;
 import io.confluent.ksql.services.KafkaTopicClient;
-import io.confluent.ksql.util.KsqlException;
 import io.confluent.ksql.util.Pair;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.QueryMetadata;
@@ -356,8 +355,9 @@ private void executeStatement(final QueuedCommand queuedCommand) {
       } else {
         if (!rateLimiter.tryAcquire()) {
           throw new KsqlRestException(
-            Errors.tooManyRequests("Too many requests to the command topic within a 1 second timeframe")
-          );
+            Errors.tooManyRequests(
+              "Too many requests to the command topic within a 1 second timeframe"
+            ));
         }
         statementExecutor.handleStatement(queuedCommand);
         LOG.info("Executed statement: " + commandStatement);

diff --git a/...b-rest-app/src/test/java/io/confluent/ksql/rest/server/computation/CommandRunnerTest.java b/...b-rest-app/src/test/java/io/confluent/ksql/rest/server/computation/CommandRunnerTest.java
@@ -17,8 +17,10 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -41,7 +43,9 @@
 import io.confluent.ksql.engine.KsqlEngine;
 import io.confluent.ksql.metrics.MetricCollectors;
 import io.confluent.ksql.rest.Errors;
+import io.confluent.ksql.rest.entity.KsqlErrorMessage;
 import io.confluent.ksql.rest.server.resources.IncompatibleKsqlCommandVersionException;
+import io.confluent.ksql.rest.server.resources.KsqlRestException;
 import io.confluent.ksql.rest.server.state.ServerState;
 import io.confluent.ksql.rest.util.ClusterTerminator;
 import io.confluent.ksql.rest.util.PersistentQueryCleanupImpl;
@@ -653,6 +657,38 @@ public void shouldCloseTheCommandRunnerCorrectly() throws Exception {
 
     verify(commandStore).close();
   }
+
+  @Test
+  public void shouldThrowWhenRateLimitHit() {
+    // Given:
+    final CommandRunner rateLimitedCommandRunner = new CommandRunner(
+      statementExecutor,
+      commandStore,
+      3,
+      clusterTerminator,
+      executor,
+      serverState,
+      "ksql-service-id",
+      Duration.ofMillis(COMMAND_RUNNER_HEALTH_TIMEOUT),
+      "",
+      clock,
+      compactor,
+      incompatibleCommandChecker,
+      commandDeserializer,
+      errorHandler,
+      commandTopicExists,
+      new Metrics(),
+      1.0
+    );
+    givenQueuedCommands(queuedCommand1, queuedCommand2, queuedCommand3);
+
+    // When:
+    // Then:
+    final KsqlRestException e = assertThrows(KsqlRestException.class, rateLimitedCommandRunner::fetchAndRunCommands);
+    assertEquals(e.getResponse().getStatus(), 429);
+    final KsqlErrorMessage errorMessage = (KsqlErrorMessage) e.getResponse().getEntity();
+    assertTrue(errorMessage.getMessage().contains("Too many requests to the command topic within a 1 second timeframe"));
+  }
 
   private Runnable getThreadTask() {
     verify(executor).execute(threadTaskCaptor.capture());