Skip to content

Fixes #27796: enforce entityType filter in MCP search_metadata #27821

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ecerulm wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fixes #27796: enforce entityType filter in MCP search_metadata #27821

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 40 additions & 1 deletion ...a-integration-tests/src/test/java/org/openmetadata/it/tests/mcp/McpToolsValidationIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,34 @@ void testSearchMetadataTool() throws Exception {
Map<String, Object> toolCall =
McpTestUtils.createSearchMetadataToolCall("mcp_val_table", 5, Entity.TABLE);
JsonNode result = executeToolCall(toolCall);
validateSearchMetadataResponse(result, "mcp_val_table");
validateSearchMetadataResponse(result, "mcp_val_table", Entity.TABLE);
}

@Test
@Order(1)
void testSearchMetadataEntityTypeFilterIsHonored() throws Exception {
// Regression test for https://github.com/open-metadata/OpenMetadata/issues/27796
// Searching with a specific entityType must only return results of that type, not leak
// other types from the default dataAsset alias.
Map<String, Object> tableSearch =
McpTestUtils.createSearchMetadataToolCall("test", 10, Entity.TABLE);
JsonNode tableResult = executeToolCall(tableSearch);
validateSearchMetadataResponse(tableResult, "test", Entity.TABLE);

Map<String, Object> dashboardSearch =
McpTestUtils.createSearchMetadataToolCall("test", 10, Entity.DASHBOARD);
JsonNode dashboardResult = executeToolCall(dashboardSearch);
JsonNode dashboardResponse =
OBJECT_MAPPER.readTree(dashboardResult.get("content").get(0).get("text").asText());
dashboardResponse
.get("results")
.forEach(
r ->
assertThat(r.get("entityType").asText())
.withFailMessage(
"Expected only %s results but got %s for entity %s",
Entity.DASHBOARD, r.get("entityType").asText(), r.get("name").asText())
.isEqualTo(Entity.DASHBOARD));
}

@Test
Expand Down Expand Up @@ -444,6 +471,11 @@ private Map<String, Object> createSearchToolCallWithDeletedParam(

private void validateSearchMetadataResponse(JsonNode result, String expectedQuery)
throws Exception {
validateSearchMetadataResponse(result, expectedQuery, null);
}

private void validateSearchMetadataResponse(
JsonNode result, String expectedQuery, String expectedEntityType) throws Exception {
assertThat(result.has("content")).isTrue();
JsonNode content = result.get("content");
assertThat(content.isArray()).isTrue();
Expand All @@ -468,6 +500,13 @@ private void validateSearchMetadataResponse(JsonNode result, String expectedQuer
.withFailMessage(
"Missing 'deleted' field in search result for: " + r.get("name"))
.isTrue();
if (expectedEntityType != null) {
assertThat(r.get("entityType").asText())
.withFailMessage(
"Expected entityType '%s' but got '%s' for result '%s'",
expectedEntityType, r.get("entityType").asText(), r.get("name").asText())
.isEqualTo(expectedEntityType);
}
matchingEntities.add(r.get("name").asText());
});

Expand Down
100 changes: 76 additions & 24 deletions openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/SearchMetadataTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import static org.openmetadata.service.security.DefaultAuthorizer.getSubjectContext;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.annotations.VisibleForTesting;
import jakarta.ws.rs.core.Response;
Expand Down Expand Up @@ -192,44 +193,35 @@ public Map<String, Object> execute(
LOG.debug("Applied query filter to query: {}", queryFilter);
}

queryFilter = addEntityTypeFilter(queryFilter, entityType);

LOG.info(
"Search query: {}, index: {}, limit: {}, includeDeleted: {}",
queryFilter,
index,
size,
includeDeleted);

SearchRequest searchRequest;
boolean userProvidedQueryFilter = params.containsKey("queryFilter");
SearchRequest searchRequest =
new SearchRequest()
.withIndex(Entity.getSearchRepository().getIndexOrAliasName(index))
.withSize(size)
.withFrom(from)
.withFetchSource(true)
.withDeleted(includeDeleted);
if (!nullOrEmpty(queryFilter)) {
// When queryFilter is provided, use it directly as it's already a transformed OpenSearch
// query
searchRequest =
new SearchRequest()
.withIndex(Entity.getSearchRepository().getIndexOrAliasName(index))
.withQueryFilter(queryFilter)
.withSize(size)
.withFrom(from)
.withFetchSource(true)
.withDeleted(includeDeleted);
} else {
// Fallback to basic query when no queryFilter is provided
searchRequest =
new SearchRequest()
.withQuery(query)
.withIndex(Entity.getSearchRepository().getIndexOrAliasName(index))
.withSize(size)
.withFrom(from)
.withFetchSource(true)
.withDeleted(includeDeleted);
searchRequest.withQueryFilter(queryFilter);
}
if (!userProvidedQueryFilter) {
searchRequest.withQuery(query);
}

SubjectContext subjectContext = getSubjectContext(securityContext);
Response response;
if (!nullOrEmpty(queryFilter)) {
// Use direct query method when queryFilter is provided since it's already a transformed query
if (userProvidedQueryFilter) {
response = Entity.getSearchRepository().searchWithDirectQuery(searchRequest, subjectContext);
} else {
// Use regular search for basic queries
response = Entity.getSearchRepository().search(searchRequest, subjectContext);
}

Expand Down Expand Up @@ -383,6 +375,66 @@ public static Map<String, Object> cleanSearchResponseObject(Map<String, Object>
return object;
}

/**
* Ensures results are constrained to the requested entityType by injecting an explicit
* {@code term} filter on the {@code entityType} field. Targeting an alias by itself is not
* always sufficient — for example, when the alias resolves to a multi-entity index or fans
* out to {@code dataAsset} — so the request can leak documents of other types. Adding the
* term filter guarantees correctness regardless of how the index alias resolves.
*
* @param existingFilter user-provided OpenSearch query JSON, already wrapped under "query", or
* {@code null}
* @param entityType requested entity type, or {@code null}/blank to leave the filter untouched
* @return a JSON string containing the merged query filter, or {@code existingFilter} if no
* entityType was provided
*/
@VisibleForTesting
static String addEntityTypeFilter(String existingFilter, String entityType) {
if (entityType == null || entityType.isBlank()) {
return existingFilter;
}
ObjectNode termFilter = JsonUtils.getObjectMapper().createObjectNode();
termFilter.putObject("term").put("entityType", entityType);
if (nullOrEmpty(existingFilter)) {
ObjectNode bool = JsonUtils.getObjectMapper().createObjectNode();
ArrayNode filterArray = bool.putObject("bool").putArray("filter");
filterArray.add(termFilter);
ObjectNode wrapper = JsonUtils.getObjectMapper().createObjectNode();
wrapper.set("query", bool);
return JsonUtils.pojoToJson(wrapper);
}
Comment on lines +391 to +405
Copy link
Copy Markdown

@gitar-bot gitar-bot Bot Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Quality: addEntityTypeFilter exceeds 15-line method limit

The new addEntityTypeFilter method is ~45 lines with 4 levels of nesting, exceeding the project's complexity limits (max 15 lines/method, max 3 nesting levels). The method handles three distinct concerns — creating a fresh filter, merging into an existing bool query, and wrapping a non-bool query — each of which could be a separate private helper.

Suggested fix:

Extract into three small methods:
- `createEntityTypeFilter(entityType)` for the null/empty case
- `mergeBoolFilter(boolNode, termFilter)` for adding to existing bool
- `wrapNonBoolQuery(queryObject, termFilter)` for wrapping non-bool

The public method becomes a short dispatcher:
  if (nullOrEmpty(existingFilter)) return createEntityTypeFilter(entityType);
  // parse, then dispatch to merge or wrap

Was this helpful? React with 👍 / 👎 | Reply gitar fix to apply this suggestion

try {
JsonNode root = JsonUtils.getObjectMapper().readTree(existingFilter);
JsonNode queryNode = root.get("query");
if (queryNode == null || !queryNode.isObject()) {
return existingFilter;
}
ObjectNode queryObject = (ObjectNode) queryNode;
ObjectNode boolNode;
if (queryObject.has("bool") && queryObject.get("bool").isObject()) {
boolNode = (ObjectNode) queryObject.get("bool");
} else {
ObjectNode originalCopy = queryObject.deepCopy();
queryObject.removeAll();
boolNode = queryObject.putObject("bool");
boolNode.putArray("must").add(originalCopy);
}
ArrayNode filterArray;
if (boolNode.has("filter") && boolNode.get("filter").isArray()) {
filterArray = (ArrayNode) boolNode.get("filter");
} else {
filterArray = boolNode.putArray("filter");
}
filterArray.add(termFilter);
Comment on lines +422 to +428
Copy link
Copy Markdown

@gitar-bot gitar-bot Bot Apr 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Edge Case: Duplicate entityType filter if user's queryFilter already has one

If a caller passes a queryFilter that already includes a term filter on entityType, addEntityTypeFilter will blindly append a second identical term filter. While this won't produce incorrect results (two identical AND'd terms are logically equivalent to one), it adds unnecessary noise to the OpenSearch query. A defensive check for an existing entityType term in the filter array would be cleaner.

Was this helpful? React with 👍 / 👎 | Reply gitar fix to apply this suggestion

return JsonUtils.pojoToJson(root);
} catch (IOException e) {
LOG.warn(
"Unable to merge entityType filter into provided queryFilter, leaving it unchanged: {}",
e.getMessage());
return existingFilter;
}
}

/**
* Truncates aggregation buckets to prevent excessive response size that could overwhelm LLM
* context windows. Based on industry best practices, LLM performance degrades when context
Expand Down
127 changes: 127 additions & 0 deletions openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/SearchMetadataToolTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,14 @@

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.fasterxml.jackson.databind.JsonNode;
import jakarta.ws.rs.core.Response;
import java.security.Principal;
import java.util.HashMap;
Expand All @@ -15,9 +18,14 @@
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.MockedStatic;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import org.openmetadata.schema.entity.teams.User;
import org.openmetadata.schema.search.SearchRequest;
import org.openmetadata.schema.utils.JsonUtils;
import org.openmetadata.service.Entity;
import org.openmetadata.service.search.SearchRepository;
import org.openmetadata.service.security.Authorizer;
Expand All @@ -35,6 +43,7 @@
* - Response formatting and structure
*/
@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.LENIENT)
class SearchMetadataToolTest {

private SearchMetadataTool searchMetadataTool;
Expand Down Expand Up @@ -135,4 +144,122 @@ void testSearchWithSpecificEntityType() throws Exception {
assertNotNull(result.get("results"));
}
}

@Test
void testEntityTypeIsAlwaysAppliedAsExplicitFilter() throws Exception {
try (MockedStatic<SubjectCache> subjectCacheMock = mockStatic(SubjectCache.class)) {
subjectCacheMock.when(() -> SubjectCache.getUserContext("test-user")).thenReturn(mockUser);

Map<String, Object> params = new HashMap<>();
params.put("query", "test");
params.put("entityType", "metric");

when(searchRepository.getIndexOrAliasName("metric")).thenReturn("metric");

Response mockResponse = mock(Response.class);
when(mockResponse.getEntity()).thenReturn("{\"hits\":{\"hits\":[],\"total\":{\"value\":0}}}");
when(searchRepository.search(any(), any(SubjectContext.class))).thenReturn(mockResponse);

searchMetadataTool.execute(authorizer, securityContext, params);

ArgumentCaptor<SearchRequest> captor = ArgumentCaptor.forClass(SearchRequest.class);
verify(searchRepository).search(captor.capture(), any(SubjectContext.class));

SearchRequest sent = captor.getValue();
assertEquals("test", sent.getQuery());
assertNotNull(sent.getQueryFilter());
JsonNode filter = JsonUtils.readTree(sent.getQueryFilter());
assertEquals("metric", filter.at("/query/bool/filter/0/term/entityType").asText());
}
}

@Test
void testNoEntityTypeLeavesQueryFilterEmpty() throws Exception {
try (MockedStatic<SubjectCache> subjectCacheMock = mockStatic(SubjectCache.class)) {
subjectCacheMock.when(() -> SubjectCache.getUserContext("test-user")).thenReturn(mockUser);

Map<String, Object> params = new HashMap<>();
params.put("query", "test");

when(searchRepository.getIndexOrAliasName("dataAsset")).thenReturn("dataAsset");

Response mockResponse = mock(Response.class);
when(mockResponse.getEntity()).thenReturn("{\"hits\":{\"hits\":[],\"total\":{\"value\":0}}}");
when(searchRepository.search(any(), any(SubjectContext.class))).thenReturn(mockResponse);

searchMetadataTool.execute(authorizer, securityContext, params);

ArgumentCaptor<SearchRequest> captor = ArgumentCaptor.forClass(SearchRequest.class);
verify(searchRepository).search(captor.capture(), any(SubjectContext.class));

SearchRequest sent = captor.getValue();
assertEquals("test", sent.getQuery());
assertNull(sent.getQueryFilter());
}
}

@Test
void testAddEntityTypeFilterWithoutExistingFilter() throws Exception {
String result = SearchMetadataTool.addEntityTypeFilter(null, "metric");
assertNotNull(result);
JsonNode root = JsonUtils.readTree(result);
assertEquals("metric", root.at("/query/bool/filter/0/term/entityType").asText());
}

@Test
void testAddEntityTypeFilterWrapsExistingBoolQuery() throws Exception {
String existing = "{\"query\":{\"bool\":{\"must\":[{\"term\":{\"owners.name\":\"team\"}}]}}}";
String result = SearchMetadataTool.addEntityTypeFilter(existing, "topic");
JsonNode root = JsonUtils.readTree(result);
assertEquals("topic", root.at("/query/bool/filter/0/term/entityType").asText());
assertEquals("team", root.at("/query/bool/must/0/term/owners.name").asText());
}

@Test
void testAddEntityTypeFilterWrapsNonBoolQuery() throws Exception {
String existing = "{\"query\":{\"term\":{\"owners.name\":\"team\"}}}";
String result = SearchMetadataTool.addEntityTypeFilter(existing, "pipeline");
JsonNode root = JsonUtils.readTree(result);
assertEquals("pipeline", root.at("/query/bool/filter/0/term/entityType").asText());
assertEquals("team", root.at("/query/bool/must/0/term/owners.name").asText());
}

@Test
void testAddEntityTypeFilterIsNoopWhenEntityTypeMissing() {
String existing = "{\"query\":{\"term\":{\"owners.name\":\"team\"}}}";
assertEquals(existing, SearchMetadataTool.addEntityTypeFilter(existing, null));
assertEquals(existing, SearchMetadataTool.addEntityTypeFilter(existing, ""));
assertEquals(existing, SearchMetadataTool.addEntityTypeFilter(existing, " "));
assertNull(SearchMetadataTool.addEntityTypeFilter(null, null));
}

@Test
void testEntityTypeFilterMergesWithUserQueryFilter() throws Exception {
try (MockedStatic<SubjectCache> subjectCacheMock = mockStatic(SubjectCache.class)) {
subjectCacheMock.when(() -> SubjectCache.getUserContext("test-user")).thenReturn(mockUser);

Map<String, Object> params = new HashMap<>();
params.put("entityType", "metric");
params.put(
"queryFilter",
"{\"query\":{\"bool\":{\"must\":[{\"term\":{\"owners.name\":\"finance\"}}]}}}");

when(searchRepository.getIndexOrAliasName("metric")).thenReturn("metric");

Response mockResponse = mock(Response.class);
when(mockResponse.getEntity()).thenReturn("{\"hits\":{\"hits\":[],\"total\":{\"value\":0}}}");
when(searchRepository.searchWithDirectQuery(any(), any(SubjectContext.class)))
.thenReturn(mockResponse);

searchMetadataTool.execute(authorizer, securityContext, params);

ArgumentCaptor<SearchRequest> captor = ArgumentCaptor.forClass(SearchRequest.class);
verify(searchRepository).searchWithDirectQuery(captor.capture(), any(SubjectContext.class));

SearchRequest sent = captor.getValue();
JsonNode filter = JsonUtils.readTree(sent.getQueryFilter());
assertEquals("metric", filter.at("/query/bool/filter/0/term/entityType").asText());
assertEquals("finance", filter.at("/query/bool/must/0/term/owners.name").asText());
}
}
}
Loading