MINOR - live index on test suite creation (#18317)

* fix: live index on test suite creation

* fix: make live indexing use entityInterface

(cherry picked from commit 781989e)

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java

@@ -7,6 +7,7 @@
 import static org.openmetadata.service.Entity.TABLE;
 import static org.openmetadata.service.Entity.TEST_CASE;
 import static org.openmetadata.service.Entity.TEST_SUITE;
+import static org.openmetadata.service.Entity.getEntity;
 import static org.openmetadata.service.util.FullyQualifiedName.quoteName;
 
 import java.io.IOException;
@@ -21,6 +22,7 @@
 import javax.ws.rs.core.SecurityContext;
 import lombok.extern.slf4j.Slf4j;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
+import org.openmetadata.schema.EntityInterface;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.tests.DataQualityReport;
 import org.openmetadata.schema.tests.ResultSummary;
@@ -35,6 +37,9 @@
 import org.openmetadata.service.resources.feeds.MessageParser;
 import org.openmetadata.service.search.SearchClient;
 import org.openmetadata.service.search.SearchIndexUtils;
+import org.openmetadata.service.search.SearchListFilter;
+import org.openmetadata.service.search.indexes.SearchIndex;
+import org.openmetadata.service.search.models.IndexMapping;
 import org.openmetadata.service.util.EntityUtil;
 import org.openmetadata.service.util.FullyQualifiedName;
 import org.openmetadata.service.util.JsonUtils;
@@ -259,6 +264,30 @@ public TestSummary getTestSummary(UUID testSuiteId) {
  return null;
  }
 
+ @Override
+ protected void postCreate(TestSuite entity) {
+ super.postCreate(entity);
+ if (Boolean.TRUE.equals(entity.getExecutable())
+ && entity.getExecutableEntityReference() != null) {
+ // Update table index with test suite field
+ EntityInterface entityInterface =
+ getEntity(entity.getExecutableEntityReference(), "testSuite", ALL);
+ IndexMapping indexMapping =
+ searchRepository.getIndexMapping(entity.getExecutableEntityReference().getType());
+ SearchClient searchClient = searchRepository.getSearchClient();
+ SearchIndex index =
+ searchRepository
+ .getSearchIndexFactory()
+ .buildIndex(entity.getExecutableEntityReference().getType(), entityInterface);
+ Map<String, Object> doc = index.buildSearchIndexDoc();
+ searchClient.updateEntity(
+ indexMapping.getIndexName(searchRepository.getClusterAlias()),
+ entity.getExecutableEntityReference().getId().toString(),
+ doc,
+ "ctx._source.testSuite = params.testSuite;");
+ }
+ }
+
  @Override
  public void prepare(TestSuite entity, boolean update) {
  /* Nothing to do */

openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestSuiteResourceTest.java

@@ -4,6 +4,7 @@
 import static javax.ws.rs.core.Response.Status.NOT_FOUND;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.openmetadata.service.util.TestUtils.ADMIN_AUTH_HEADERS;
@@ -15,10 +16,13 @@
 
 import es.org.elasticsearch.search.aggregations.AggregationBuilder;
 import es.org.elasticsearch.search.aggregations.AggregationBuilders;
+import es.org.elasticsearch.client.Request;
+import es.org.elasticsearch.client.RestClient;
 import java.io.IOException;
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -28,6 +32,7 @@
 import javax.ws.rs.client.WebTarget;
 import javax.ws.rs.core.Response;
 import org.apache.http.client.HttpResponseException;
+import org.apache.http.util.EntityUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInfo;
@@ -49,6 +54,7 @@
 import org.openmetadata.service.resources.EntityResourceTest;
 import org.openmetadata.service.resources.databases.TableResourceTest;
 import org.openmetadata.service.search.elasticsearch.ElasticSearchClient;
+import org.openmetadata.service.search.models.IndexMapping;
 import org.openmetadata.service.util.JsonUtils;
 import org.openmetadata.service.util.ResultList;
 import org.openmetadata.service.util.TestUtils;
@@ -878,6 +884,56 @@ void get_listTestSuiteFromSearchWithPagination(TestInfo testInfo)
  }
  }
 
+ @Test
+ void create_executableTestSuiteAndCheckSearchClient(TestInfo test) throws IOException {
+ TableResourceTest tableResourceTest = new TableResourceTest();
+ CreateTable tableReq =
+ tableResourceTest
+ .createRequest(test)
+ .withColumns(
+ List.of(
+ new Column()
+ .withName(C1)
+ .withDisplayName("c1")
+ .withDataType(ColumnDataType.VARCHAR)
+ .withDataLength(10)));
+ Table table = tableResourceTest.createEntity(tableReq, ADMIN_AUTH_HEADERS);
+ CreateTestSuite createTestSuite = createRequest(table.getFullyQualifiedName());
+ TestSuite testSuite = createExecutableTestSuite(createTestSuite, ADMIN_AUTH_HEADERS);
+ RestClient searchClient = getSearchClient();
+ IndexMapping index = Entity.getSearchRepository().getIndexMapping(Entity.TABLE);
+ es.org.elasticsearch.client.Response response;
+ Request request =
+ new Request(
+ "GET",
+ String.format(
+ "%s/_search", index.getIndexName(Entity.getSearchRepository().getClusterAlias())));
+ String query =
+ String.format(
+ "{\"size\": 10,\"query\":{\"bool\":{\"must\":[{\"term\":{\"_id\":\"%s\"}}]}}}",
+ table.getId().toString());
+ request.setJsonEntity(query);
+ try {
+ response = searchClient.performRequest(request);
+ } finally {
+ searchClient.close();
+ }
+ String jsonString = EntityUtils.toString(response.getEntity());
+ HashMap<String, Object> map =
+ (HashMap<String, Object>) JsonUtils.readOrConvertValue(jsonString, HashMap.class);
+ LinkedHashMap<String, Object> hits = (LinkedHashMap<String, Object>) map.get("hits");
+ ArrayList<LinkedHashMap<String, Object>> hitsList =
+ (ArrayList<LinkedHashMap<String, Object>>) hits.get("hits");
+ assertNotEquals(0, hitsList.size());
+ Assertions.assertTrue(
+ hitsList.stream()
+ .allMatch(
+ hit ->
+ ((LinkedHashMap<String, Object>) hit.get("_source"))
+ .get("id")
+ .equals(table.getId().toString())));
+ }
+
  public ResultList<TestSuite> getTestSuites(
  Integer limit,
  String fields,