Skip to content

Commit 80a525b

Browse files
mohityadav766mohityadav766
authored
Fix Healthcheck for vector embedding (#27616)
1 parent a8ddb91 commit 80a525b

2 files changed

Lines changed: 134 additions & 1 deletion

File tree

‎openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/SystemRepository.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@
106106
public class SystemRepository {
107107
private static final String FAILED_TO_UPDATE_SETTINGS = "Failed to Update Settings {}";
108108
public static final String INTERNAL_SERVER_ERROR_WITH_REASON = "Internal Server Error. Reason :";
109+
private static final String VECTOR_EMBEDDING_INDEX_KEY = "vectorEmbedding";
109110
private final SystemDAO dao;
110111
private final MigrationValidationClient migrationValidationClient;
111112

@@ -822,12 +823,17 @@ private StepValidation getSearchValidation(OpenMetadataApplicationConfig applica
822823
}
823824
}
824825

825-
private List<String> findMissingIndexes(SearchRepository searchRepository) {
826+
@VisibleForTesting
827+
List<String> findMissingIndexes(SearchRepository searchRepository) {
826828
List<String> missing = new ArrayList<>();
829+
boolean semanticSearchEnabled = searchRepository.isVectorEmbeddingEnabled();
827830
try {
828831
Map<String, org.openmetadata.search.IndexMapping> indexMap =
829832
searchRepository.getEntityIndexMap();
830833
for (Map.Entry<String, org.openmetadata.search.IndexMapping> entry : indexMap.entrySet()) {
834+
if (!semanticSearchEnabled && VECTOR_EMBEDDING_INDEX_KEY.equals(entry.getKey())) {
835+
continue;
836+
}
831837
if (!searchRepository.indexExists(entry.getValue())) {
832838
missing.add(entry.getKey());
833839
}

‎openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/SystemRepositoryMissingIndexesTest.java‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
package org.openmetadata.service.jdbi3;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertFalse;
5+
import static org.junit.jupiter.api.Assertions.assertTrue;
6+
import static org.mockito.Mockito.mock;
7+
import static org.mockito.Mockito.mockStatic;
8+
import static org.mockito.Mockito.when;
9+
10+
import java.util.List;
11+
import java.util.Map;
12+
import org.junit.jupiter.api.AfterEach;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import org.mockito.MockedStatic;
16+
import org.openmetadata.search.IndexMapping;
17+
import org.openmetadata.service.Entity;
18+
import org.openmetadata.service.jdbi3.CollectionDAO.SystemDAO;
19+
import org.openmetadata.service.migration.MigrationValidationClient;
20+
import org.openmetadata.service.search.SearchRepository;
21+
22+
class SystemRepositoryMissingIndexesTest {
23+
24+
private MockedStatic<Entity> entityMock;
25+
private MockedStatic<MigrationValidationClient> migrationMock;
26+
private SearchRepository searchRepository;
27+
private SystemRepository systemRepository;
28+
29+
@BeforeEach
30+
void setup() {
31+
entityMock = mockStatic(Entity.class);
32+
migrationMock = mockStatic(MigrationValidationClient.class);
33+
34+
CollectionDAO collectionDAO = mock(CollectionDAO.class);
35+
SystemDAO systemDAO = mock(SystemDAO.class);
36+
when(collectionDAO.systemDAO()).thenReturn(systemDAO);
37+
entityMock.when(Entity::getCollectionDAO).thenReturn(collectionDAO);
38+
39+
MigrationValidationClient migrationClient = mock(MigrationValidationClient.class);
40+
migrationMock.when(MigrationValidationClient::getInstance).thenReturn(migrationClient);
41+
42+
searchRepository = mock(SearchRepository.class);
43+
entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository);
44+
45+
systemRepository = new SystemRepository();
46+
}
47+
48+
@AfterEach
49+
void tearDown() {
50+
entityMock.close();
51+
migrationMock.close();
52+
}
53+
54+
@Test
55+
void testVectorEmbeddingIndexSkippedWhenSemanticSearchDisabled() {
56+
IndexMapping tableMapping = mock(IndexMapping.class);
57+
IndexMapping vectorMapping = mock(IndexMapping.class);
58+
Map<String, IndexMapping> indexMap =
59+
Map.of("table", tableMapping, "vectorEmbedding", vectorMapping);
60+
61+
when(searchRepository.isVectorEmbeddingEnabled()).thenReturn(false);
62+
when(searchRepository.getEntityIndexMap()).thenReturn(indexMap);
63+
when(searchRepository.indexExists(tableMapping)).thenReturn(true);
64+
when(searchRepository.indexExists(vectorMapping)).thenReturn(false);
65+
66+
List<String> missing = systemRepository.findMissingIndexes(searchRepository);
67+
68+
assertTrue(missing.isEmpty(), "vectorEmbedding should be ignored when semantic search is off");
69+
}
70+
71+
@Test
72+
void testVectorEmbeddingIndexReportedMissingWhenSemanticSearchEnabled() {
73+
IndexMapping tableMapping = mock(IndexMapping.class);
74+
IndexMapping vectorMapping = mock(IndexMapping.class);
75+
Map<String, IndexMapping> indexMap =
76+
Map.of("table", tableMapping, "vectorEmbedding", vectorMapping);
77+
78+
when(searchRepository.isVectorEmbeddingEnabled()).thenReturn(true);
79+
when(searchRepository.getEntityIndexMap()).thenReturn(indexMap);
80+
when(searchRepository.indexExists(tableMapping)).thenReturn(true);
81+
when(searchRepository.indexExists(vectorMapping)).thenReturn(false);
82+
83+
List<String> missing = systemRepository.findMissingIndexes(searchRepository);
84+
85+
assertEquals(1, missing.size());
86+
assertEquals("vectorEmbedding", missing.get(0));
87+
}
88+
89+
@Test
90+
void testNonVectorMissingIndexesAlwaysReported() {
91+
IndexMapping tableMapping = mock(IndexMapping.class);
92+
IndexMapping glossaryMapping = mock(IndexMapping.class);
93+
IndexMapping vectorMapping = mock(IndexMapping.class);
94+
Map<String, IndexMapping> indexMap =
95+
Map.of(
96+
"table", tableMapping, "glossary", glossaryMapping, "vectorEmbedding", vectorMapping);
97+
98+
when(searchRepository.isVectorEmbeddingEnabled()).thenReturn(false);
99+
when(searchRepository.getEntityIndexMap()).thenReturn(indexMap);
100+
when(searchRepository.indexExists(tableMapping)).thenReturn(true);
101+
when(searchRepository.indexExists(glossaryMapping)).thenReturn(false);
102+
when(searchRepository.indexExists(vectorMapping)).thenReturn(false);
103+
104+
List<String> missing = systemRepository.findMissingIndexes(searchRepository);
105+
106+
assertEquals(1, missing.size());
107+
assertEquals("glossary", missing.get(0));
108+
assertFalse(missing.contains("vectorEmbedding"));
109+
}
110+
111+
@Test
112+
void testAllIndexesPresentReturnsEmpty() {
113+
IndexMapping tableMapping = mock(IndexMapping.class);
114+
IndexMapping vectorMapping = mock(IndexMapping.class);
115+
Map<String, IndexMapping> indexMap =
116+
Map.of("table", tableMapping, "vectorEmbedding", vectorMapping);
117+
118+
when(searchRepository.isVectorEmbeddingEnabled()).thenReturn(true);
119+
when(searchRepository.getEntityIndexMap()).thenReturn(indexMap);
120+
when(searchRepository.indexExists(tableMapping)).thenReturn(true);
121+
when(searchRepository.indexExists(vectorMapping)).thenReturn(true);
122+
123+
List<String> missing = systemRepository.findMissingIndexes(searchRepository);
124+
125+
assertTrue(missing.isEmpty());
126+
}
127+
}

0 commit comments

Comments
 (0)