Delete sample config when a data stream is deleted

Author: masseyke

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/indices.get_sample/30_with_data_streams.yml

@@ -273,3 +273,13 @@
         index: sample_test
       catch: missing
 
+---
+teardown:
+
+  - do:
+      indices.delete_data_stream:
+        name: sample_test*
+
+  - do:
+      indices.get_all_sample_configuration: {}
+  - length: { $body: 0 }

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/indices.get_sample_stats/30_with_data_streams.yml

@@ -433,3 +433,14 @@
   - match: { samples_rejected_for_size: 0 }
   - match: { samples_accepted: 0 }
   - match: { last_exception.type: "script_exception" }
+
+---
+teardown:
+
+  - do:
+      indices.delete_data_stream:
+        name: sample_test*
+
+  - do:
+      indices.get_all_sample_configuration: {}
+  - length: { $body: 0 }

server/src/internalClusterTest/java/org/elasticsearch/ingest/SamplingServiceIT.java

@@ -17,15 +17,25 @@
 import org.elasticsearch.action.admin.indices.sampling.GetSampleAction;
 import org.elasticsearch.action.admin.indices.sampling.PutSampleConfigurationAction;
 import org.elasticsearch.action.admin.indices.sampling.SamplingConfiguration;
+import org.elasticsearch.action.admin.indices.template.put.TransportPutComposableIndexTemplateAction;
 import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.bulk.TransportBulkAction;
+import org.elasticsearch.action.datastreams.CreateDataStreamAction;
+import org.elasticsearch.action.datastreams.DeleteDataStreamAction;
 import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.cluster.metadata.ComposableIndexTemplate;
 import org.elasticsearch.cluster.metadata.ProjectId;
+import org.elasticsearch.cluster.metadata.Template;
+import org.elasticsearch.common.compress.CompressedXContent;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.datastreams.DataStreamsPlugin;
+import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.junit.After;
 
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -36,6 +46,12 @@
 import static org.hamcrest.Matchers.equalTo;
 
 public class SamplingServiceIT extends ESIntegTestCase {
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return List.of(DataStreamsPlugin.class);
+    }
+
     public void testTTL() throws Exception {
         assumeTrue("Requires sampling feature flag", RANDOM_SAMPLING_FEATURE_FLAG);
         assertAcked(
@@ -102,6 +118,66 @@ public void testDeleteIndex() throws Exception {
         });
     }
 
+    public void testDeleteDataStream() throws Exception {
+        assumeTrue("Requires sampling feature flag", RANDOM_SAMPLING_FEATURE_FLAG);
+        String indexName = randomIdentifier();
+        var template = ComposableIndexTemplate.builder()
+            .indexPatterns(List.of(indexName))
+
+            .template(new Template(Settings.EMPTY, CompressedXContent.fromJSON("""
+                {
+                  "_doc":{
+                    "dynamic":true,
+                    "properties":{
+                      "foo":{
+                        "type":"text"
+                      },
+                      "bar":{
+                        "type":"text"
+                      }
+                    }
+                  }
+                }
+                """), null))
+            .dataStreamTemplate(new ComposableIndexTemplate.DataStreamTemplate(false, false))
+            .build();
+        var request = new TransportPutComposableIndexTemplateAction.Request("logs-template");
+        request.indexTemplate(template);
+        safeGet(client().execute(TransportPutComposableIndexTemplateAction.TYPE, request));
+        client().execute(
+            CreateDataStreamAction.INSTANCE,
+            new CreateDataStreamAction.Request(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS, indexName)
+        ).actionGet();
+        ensureYellow(indexName);
+        PutSampleConfigurationAction.Request putSampleConfigRequest = new PutSampleConfigurationAction.Request(
+            new SamplingConfiguration(1.0d, 10, null, null, null),
+            TimeValue.THIRTY_SECONDS,
+            TimeValue.THIRTY_SECONDS
+        ).indices(indexName);
+        client().execute(PutSampleConfigurationAction.INSTANCE, putSampleConfigRequest).actionGet();
+        for (int i = 0; i < 5; i++) {
+            BulkRequest bulkRequest = new BulkRequest();
+            for (int j = 0; j < 20; j++) {
+                IndexRequest indexRequest = new IndexRequest(indexName);
+                indexRequest.create(true);
+                indexRequest.source(Map.of("@timestamp", 12345, "foo", randomBoolean() ? 3L : randomLong(), "bar", randomBoolean()));
+                bulkRequest.add(indexRequest);
+            }
+            BulkResponse bulkResponse = client().execute(TransportBulkAction.TYPE, bulkRequest).actionGet();
+            assertThat(bulkResponse.hasFailures(), equalTo(false));
+        }
+        GetSampleAction.Response getSampleResponse = client().execute(GetSampleAction.INSTANCE, new GetSampleAction.Request(indexName))
+            .actionGet();
+        assertThat(getSampleResponse.getSample().size(), equalTo(10));
+        client().execute(DeleteDataStreamAction.INSTANCE, new DeleteDataStreamAction.Request(TimeValue.THIRTY_SECONDS, indexName))
+            .actionGet();
+        assertBusy(() -> {
+            for (SamplingService samplingService : internalCluster().getInstances(SamplingService.class)) {
+                assertThat(samplingService.getLocalSample(ProjectId.DEFAULT, indexName), equalTo(List.of()));
+            }
+        });
+    }
+
     @After
     public void cleanup() {
         Map<String, Object> clearedSettings = new HashMap<>();

server/src/main/java/org/elasticsearch/ingest/SamplingService.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.cluster.ClusterStateAckListener;
 import org.elasticsearch.cluster.ClusterStateListener;
 import org.elasticsearch.cluster.SimpleBatchedAckListenerTaskExecutor;
+import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.metadata.Metadata;
@@ -525,6 +526,22 @@ private void maybeDeleteSamplingConfigurations(ClusterChangedEvent event, Projec
                 }
             }
         }
+        if (currentProject.dataStreams() != previousProject.dataStreams()) {
+            for (DataStream dataStream : previousProject.dataStreams().values()) {
+                DataStream current = currentProject.dataStreams().get(dataStream.getName());
+                if (current == null) {
+                    String dataStreamName = dataStream.getName();
+                    SamplingConfiguration samplingConfiguration = getSamplingConfiguration(
+                        event.state().projectState(projectId).metadata(),
+                        dataStreamName
+                    );
+                    if (samplingConfiguration != null) {
+                        logger.debug("Deleting sample configuration for {} because the data stream has been deleted", dataStreamName);
+                        deleteSampleConfiguration(projectId, dataStreamName);
+                    }
+                }
+            }
+        }
     }
 
     private void maybeScheduleJob() {