fix: DROP stream for persistent query doesn't always drop underlying …

…query (#7601)

ksqldb-common/src/main/java/io/confluent/ksql/util/KsqlConstants.java

@@ -52,6 +52,10 @@ public enum KsqlQueryType {
     PUSH
   }
 
+  public enum PersistentQueryType {
+    CREATE_AS, INSERT
+  }
+
   public enum KsqlQueryStatus {
     RUNNING,
     ERROR,

ksqldb-engine/src/main/java/io/confluent/ksql/query/QueryExecutor.java

@@ -51,6 +51,7 @@
 import io.confluent.ksql.serde.WindowInfo;
 import io.confluent.ksql.services.ServiceContext;
 import io.confluent.ksql.util.KsqlConfig;
+import io.confluent.ksql.util.KsqlConstants;
 import io.confluent.ksql.util.KsqlException;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.PersistentQueryMetadataImpl;
@@ -215,6 +216,7 @@ private static Optional<ScalablePushRegistry> applyScalablePushProcessor(
   }
 
   PersistentQueryMetadata buildPersistentQuery(
+      final KsqlConstants.PersistentQueryType persistentQueryType,
       final String statementText,
       final QueryId queryId,
       final DataSource sinkDataSource,
@@ -262,6 +264,7 @@ PersistentQueryMetadata buildPersistentQuery(
         .orElse(userErrorClassifiers);
 
     return new PersistentQueryMetadataImpl(
+        persistentQueryType,
         statementText,
         querySchema,
         sources,

ksqldb-engine/src/main/java/io/confluent/ksql/query/QueryRegistryImpl.java

@@ -29,6 +29,7 @@
 import io.confluent.ksql.schema.ksql.LogicalSchema;
 import io.confluent.ksql.serde.WindowInfo;
 import io.confluent.ksql.services.ServiceContext;
+import io.confluent.ksql.util.KsqlConstants;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.PersistentQueryMetadataImpl;
 import io.confluent.ksql.util.QueryMetadata;
@@ -165,6 +166,9 @@ public PersistentQueryMetadata createOrReplacePersistentQuery(
     final QueryExecutor executor =
         executorFactory.create(config, processingLogContext, serviceContext, metaStore);
     final PersistentQueryMetadata query = executor.buildPersistentQuery(
+        createAsQuery
+            ? KsqlConstants.PersistentQueryType.CREATE_AS
+            : KsqlConstants.PersistentQueryType.INSERT,
         statementText,
         queryId,
         sinkDataSource,
@@ -290,16 +294,21 @@ private void unregisterQuery(final QueryMetadata query) {
       final QueryId queryId = persistentQuery.getQueryId();
       persistentQueries.remove(queryId);
 
-      // If query is a INSERT query, then this line should not cause any effect
-      createAsQueries.remove(persistentQuery.getSinkName());
-
-      // If query is a C*AS query, then these lines should not cause any effect
-      sinkAndSources(persistentQuery).forEach(sourceName ->
-          insertQueries.computeIfPresent(sourceName, (s, queries) -> {
-            queries.remove(queryId);
-            return (queries.isEmpty()) ? null : queries;
-          })
-      );
+      switch (persistentQuery.getPersistentQueryType()) {
+        case CREATE_AS:
+          createAsQueries.remove(persistentQuery.getSinkName());
+          break;
+        case INSERT:
+          sinkAndSources(persistentQuery).forEach(sourceName ->
+              insertQueries.computeIfPresent(sourceName, (s, queries) -> {
+                queries.remove(queryId);
+                return (queries.isEmpty()) ? null : queries;
+              })
+          );
+          break;
+        default:
+          // nothing to do with unknown query types
+      }
     }
 
     allLiveQueries.remove(query);

ksqldb-engine/src/main/java/io/confluent/ksql/util/PersistentQueryMetadata.java

@@ -46,6 +46,8 @@ public interface PersistentQueryMetadata extends QueryMetadata {
 
   DataSource getSink();
 
+  KsqlConstants.PersistentQueryType getPersistentQueryType();
+
   ProcessingLogger getProcessingLogger();
 
   Optional<Materialization> getMaterialization(

ksqldb-engine/src/main/java/io/confluent/ksql/util/PersistentQueryMetadataImpl.java

@@ -46,7 +46,7 @@
  */
 public class PersistentQueryMetadataImpl
     extends QueryMetadataImpl implements PersistentQueryMetadata {
-
+  private final KsqlConstants.PersistentQueryType persistentQueryType;
   private final DataSource sinkDataSource;
   private final QuerySchemas schemas;
   private final PhysicalSchema resultSchema;
@@ -60,6 +60,7 @@ public class PersistentQueryMetadataImpl
 
   // CHECKSTYLE_RULES.OFF: ParameterNumberCheck
   public PersistentQueryMetadataImpl(
+      final KsqlConstants.PersistentQueryType persistentQueryType,
       final String statementString,
       final PhysicalSchema schema,
       final Set<SourceName> sourceNames,
@@ -111,6 +112,7 @@ public PersistentQueryMetadataImpl(
         requireNonNull(materializationProviderBuilder, "materializationProviderBuilder");
     this.processingLogger = requireNonNull(processingLogger, "processingLogger");
     this.scalablePushRegistry = requireNonNull(scalablePushRegistry, "scalablePushRegistry");
+    this.persistentQueryType = requireNonNull(persistentQueryType, "persistentQueryType");
   }
 
   // for creating sandbox instances
@@ -127,6 +129,7 @@ protected PersistentQueryMetadataImpl(
     this.materializationProviderBuilder = original.materializationProviderBuilder;
     this.processingLogger = original.processingLogger;
     this.scalablePushRegistry = original.scalablePushRegistry;
+    this.persistentQueryType = original.getPersistentQueryType();
   }
 
   @Override
@@ -179,6 +182,10 @@ public DataSource getSink() {
     return sinkDataSource;
   }
 
+  public KsqlConstants.PersistentQueryType getPersistentQueryType() {
+    return persistentQueryType;
+  }
+
   @VisibleForTesting
   public Optional<MaterializationProvider> getMaterializationProvider() {
     return materializationProvider;

ksqldb-engine/src/test/java/io/confluent/ksql/query/QueryExecutorTest.java

@@ -52,6 +52,7 @@
 import io.confluent.ksql.services.KafkaTopicClient;
 import io.confluent.ksql.services.ServiceContext;
 import io.confluent.ksql.util.KsqlConfig;
+import io.confluent.ksql.util.KsqlConstants;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.QueryMetadata;
 import io.confluent.ksql.util.QueryMetadataImpl;
@@ -261,7 +262,7 @@ public void shouldBuildTransientQueryCorrectly() {
   }
 
   @Test
-  public void shouldBuildPersistentQueryCorrectly() {
+  public void shouldBuildCreateAsPersistentQueryCorrectly() {
     // Given:
     final ProcessingLogger uncaughtProcessingLogger = mock(ProcessingLogger.class);
     when(processingLoggerFactory.getLogger(
@@ -271,6 +272,7 @@ public void shouldBuildPersistentQueryCorrectly() {
 
     // When:
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -295,12 +297,55 @@ public void shouldBuildPersistentQueryCorrectly() {
     assertThat(queryMetadata.getOverriddenProperties(), equalTo(OVERRIDES));
     assertThat(queryMetadata.getStreamsProperties(), equalTo(capturedStreamsProperties()));
     assertThat(queryMetadata.getProcessingLogger(), equalTo(uncaughtProcessingLogger));
+    assertThat(queryMetadata.getPersistentQueryType(),
+        equalTo(KsqlConstants.PersistentQueryType.CREATE_AS));
+  }
+
+  @Test
+  public void shouldBuildInsertPersistentQueryCorrectly() {
+    // Given:
+    final ProcessingLogger uncaughtProcessingLogger = mock(ProcessingLogger.class);
+    when(processingLoggerFactory.getLogger(
+        QueryLoggerUtil.queryLoggerName(QUERY_ID, new QueryContext.Stacker()
+            .push("ksql.logger.thread.exception.uncaught").getQueryContext())
+    )).thenReturn(uncaughtProcessingLogger);
+
+    // When:
+    final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.INSERT,
+        STATEMENT_TEXT,
+        QUERY_ID,
+        sink,
+        SOURCES,
+        physicalPlan,
+        SUMMARY,
+        queryListener,
+        Collections::emptyList
+    );
+    queryMetadata.initialize();
+
+    // Then:
+    assertThat(queryMetadata.getStatementString(), equalTo(STATEMENT_TEXT));
+    assertThat(queryMetadata.getQueryId(), equalTo(QUERY_ID));
+    assertThat(queryMetadata.getSinkName(), equalTo(SINK_NAME));
+    assertThat(queryMetadata.getPhysicalSchema(), equalTo(SINK_PHYSICAL_SCHEMA));
+    assertThat(queryMetadata.getResultTopic(), is(ksqlTopic));
+    assertThat(queryMetadata.getSourceNames(), equalTo(SOURCES));
+    assertThat(queryMetadata.getDataSourceType(), equalTo(DataSourceType.KSTREAM));
+    assertThat(queryMetadata.getExecutionPlan(), equalTo(SUMMARY));
+    assertThat(queryMetadata.getTopology(), is(topology));
+    assertThat(queryMetadata.getOverriddenProperties(), equalTo(OVERRIDES));
+    assertThat(queryMetadata.getStreamsProperties(), equalTo(capturedStreamsProperties()));
+    assertThat(queryMetadata.getProcessingLogger(), equalTo(uncaughtProcessingLogger));
+    assertThat(queryMetadata.getPersistentQueryType(),
+        equalTo(KsqlConstants.PersistentQueryType.INSERT));
   }
 
   @Test
   public void shouldBuildPersistentQueryWithCorrectStreamsApp() {
     // When:
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -321,6 +366,7 @@ public void shouldBuildPersistentQueryWithCorrectStreamsApp() {
   public void shouldStartPersistentQueryWithCorrectMaterializationProvider() {
     // Given:
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -344,6 +390,7 @@ public void shouldStartPersistentQueryWithCorrectMaterializationProvider() {
   public void shouldCreateKSMaterializationCorrectly() {
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -373,6 +420,7 @@ public void shouldCreateKSMaterializationCorrectly() {
   public void shouldMaterializeCorrectlyOnStart() {
     // Given:
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -403,6 +451,7 @@ public void shouldNotIncludeMaterializationProviderIfNoMaterialization() {
     when(tableHolder.getMaterializationBuilder()).thenReturn(Optional.empty());
 
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -425,6 +474,7 @@ public void shouldNotIncludeMaterializationProviderIfNoMaterialization() {
   public void shouldCreateExpectedServiceId() {
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -447,6 +497,7 @@ public void shouldCreateExpectedServiceId() {
   public void shouldAddMetricsInterceptorsToStreamsConfig() {
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -478,6 +529,7 @@ private void shouldUseProvidedOptimizationConfig(final Object value) {
 
     // When:
     final PersistentQueryMetadata queryMetadata = queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -540,6 +592,7 @@ public void shouldAddMetricsInterceptorsToExistingList() {
 
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -566,6 +619,7 @@ public void shouldAddMetricsInterceptorsToExistingString() {
 
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -593,6 +647,7 @@ public void shouldAddMetricsInterceptorsToExistingStringList() {
 
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,
@@ -624,6 +679,7 @@ public void shouldConfigureProducerErrorHandler() {
 
     // When:
     queryBuilder.buildPersistentQuery(
+        KsqlConstants.PersistentQueryType.CREATE_AS,
         STATEMENT_TEXT,
         QUERY_ID,
         sink,

ksqldb-engine/src/test/java/io/confluent/ksql/query/QueryRegistryImplTest.java

@@ -2,6 +2,7 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.ArgumentMatchers.any;
@@ -24,6 +25,7 @@
 import io.confluent.ksql.query.QueryRegistryImpl.QueryExecutorFactory;
 import io.confluent.ksql.schema.ksql.LogicalSchema;
 import io.confluent.ksql.services.ServiceContext;
+import io.confluent.ksql.util.KsqlConstants;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.PersistentQueryMetadataImpl;
 import io.confluent.ksql.util.QueryMetadata;
@@ -196,6 +198,47 @@ public void shouldGetQueriesInsertingIntoOrReadingFromSource() {
     assertThat(queries, contains(new QueryId("q2"), new QueryId("q3"), new QueryId("q4")));
   }
 
+  @Test
+  public void shouldRemoveQueryFromCreateAsQueriesWhenTerminatingCreateAsQuery() {
+    // Given:
+    givenCreate(registry, "q1", "source", "sink1", true);
+    givenCreate(registry, "i1", "source", "sink1", false);
+    final QueryRegistry sandbox = registry.createSandbox();
+    final QueryMetadata q1 = sandbox.getPersistentQuery(new QueryId("q1")).get();
+    final QueryMetadata i1 = sandbox.getPersistentQuery(new QueryId("i1")).get();
+
+    // When:
+    q1.close();
+    final Optional<QueryMetadata> createAsQueries =
+        sandbox.getCreateAsQuery(SourceName.of("sink1"));
+    final Set<QueryId> insertQueries =
+        sandbox.getInsertQueries(SourceName.of("sink1"), (n, q) -> true);
+
+    // Then:
+    assertThat(createAsQueries, equalTo(Optional.empty()));
+    assertThat(insertQueries, contains(i1.getQueryId()));
+  }
+
+  @Test
+  public void shouldRemoveQueryFromInsertQueriesWhenTerminatingInsertQuery() {
+    // Given:
+    givenCreate(registry, "q1", "source", "sink1", true);
+    givenCreate(registry, "i1", "source", "sink1", false);
+    final QueryRegistry sandbox = registry.createSandbox();
+    final QueryMetadata q1 = sandbox.getPersistentQuery(new QueryId("q1")).get();
+    final QueryMetadata i1 = sandbox.getPersistentQuery(new QueryId("i1")).get();
+
+    // When:
+    i1.close();
+    final QueryMetadata createAsQueries = sandbox.getCreateAsQuery(SourceName.of("sink1")).get();
+    final Set<QueryId> insertQueries =
+        sandbox.getInsertQueries(SourceName.of("sink1"), (n, q) -> true);
+
+    // Then:
+    assertThat(createAsQueries, equalTo(q1));
+    assertThat(insertQueries, empty());
+  }
+
   @Test
   public void shouldCopyInsertsOnSandbox() {
     // Given:
@@ -319,7 +362,7 @@ private QueryMetadata.Listener givenCreateGetListener(
   ) {
     givenCreate(registry, id, "source", "sink1", true);
     verify(executor).buildPersistentQuery(
-        any(), any(), any(), any(), any(), any(), queryListenerCaptor.capture(), any());
+        any(), any(), any(), any(), any(), any(), any(), queryListenerCaptor.capture(), any());
     return queryListenerCaptor.getValue();
   }
 
@@ -338,8 +381,11 @@ private PersistentQueryMetadata givenCreate(
     when(query.getSinkName()).thenReturn(SourceName.of(sink));
     when(query.getSink()).thenReturn(sinkSource);
     when(query.getSourceNames()).thenReturn(ImmutableSet.of(SourceName.of(source)));
+    when(query.getPersistentQueryType()).thenReturn(createAs
+        ? KsqlConstants.PersistentQueryType.CREATE_AS
+        : KsqlConstants.PersistentQueryType.INSERT);
     when(executor.buildPersistentQuery(
-        any(), any(), any(), any(), any(), any(), any(), any())
+        any(), any(), any(), any(), any(), any(), any(), any(), any())
     ).thenReturn(query);
     registry.createOrReplacePersistentQuery(
         config,