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

* fix: live index on test suite creation * fix: make live indexing use entityInterface
open-metadata · Oct 18, 2024 · 781989e · 781989e
1 parent c2929e6
commit 781989e
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 0 deletions.
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java
@@ -8,12 +8,14 @@
 import static org.openmetadata.service.Entity.TEST_CASE;
 import static org.openmetadata.service.Entity.TEST_CASE_RESULT;
 import static org.openmetadata.service.Entity.TEST_SUITE;
+import static org.openmetadata.service.Entity.getEntity;
 import static org.openmetadata.service.Entity.getEntityTimeSeriesRepository;
 import static org.openmetadata.service.util.FullyQualifiedName.quoteName;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import javax.json.JsonArray;
@@ -23,6 +25,7 @@
 import lombok.SneakyThrows;
 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;
@@ -40,6 +43,8 @@
 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;
@@ -240,6 +245,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;");
+ }
+ }
+
  @SneakyThrows
  private List<ResultSummary> getResultSummary(UUID testSuiteId) {
  List<ResultSummary> resultSummaries = new ArrayList<>();

diff --git a/...rvice/src/test/java/org/openmetadata/service/resources/dqtests/TestSuiteResourceTest.java b/...rvice/src/test/java/org/openmetadata/service/resources/dqtests/TestSuiteResourceTest.java
@@ -3,6 +3,7 @@
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
 import static javax.ws.rs.core.Response.Status.NOT_FOUND;
 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.junit.jupiter.api.Assertions.assertTrue;
@@ -13,10 +14,13 @@
 import static org.openmetadata.service.util.TestUtils.assertResponse;
 import static org.openmetadata.service.util.TestUtils.assertResponseContains;
 
+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;
@@ -25,6 +29,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;
@@ -45,6 +50,7 @@
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.resources.EntityResourceTest;
 import org.openmetadata.service.resources.databases.TableResourceTest;
+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;
@@ -733,6 +739,56 @@ void get_listTestSuiteFromSearchWithPagination(TestInfo testInfo) throws IOExcep
  }
  }
 
+ @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());
+ 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,