From 11bb305143ff5331b8c10fb3127cd0e31a79fbf9 Mon Sep 17 00:00:00 2001 From: pmbrull Date: Tue, 28 Jul 2026 16:18:11 +0200 Subject: [PATCH 1/3] fix(search): restore ContextMemory search for owners and shared users MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #29541 kept PRIVATE/SHARED memories out of the search index, but GET /contextCenter/memories serves the ContextCenter listing from search whenever it is given a query, filter, sort or offset — which the UI always sends. Owners stopped finding their own memories and shared principals stopped finding the ones shared with them. Index every memory again and enforce shareConfig visibility per query instead. The request builders now fail closed: a search path that never resolved the caller's subject gets org-wide memories only, so a forgotten call site cannot leak. Paths that build raw requests — field query, aggregations, lineage and entity-relationship traversals, doc-by-id — apply the same restriction explicitly. Restricted memories stay out of the vector index behind a new isVectorEmbeddable hook. Vector chunk documents carry no visibility fields and the vector query path takes no subject, so it cannot be filtered per caller. Also stubs the indexing policy on globally registered mock repositories, which left SearchRepositoryTest failing on main. Refs #29541 Co-Authored-By: Claude --- .../it/tests/ContextMemoryIT.java | 143 +++++++++++------- .../java/org/openmetadata/service/Entity.java | 47 +++--- .../searchIndex/ElasticSearchBulkSink.java | 6 + .../searchIndex/OpenSearchBulkSink.java | 6 + .../jdbi3/ContextMemoryRepository.java | 36 ++--- .../service/jdbi3/EntityRepository.java | 20 ++- .../service/jdbi3/ListFilter.java | 29 ---- .../ElasticSearchAggregationManager.java | 45 ++++-- .../ElasticSearchEntityManager.java | 7 +- .../ElasticSearchRequestBuilder.java | 45 ++++++ .../ElasticSearchSearchManager.java | 81 +++++----- .../service/search/elasticsearch/EsUtils.java | 35 ++++- .../OpenSearchAggregationManager.java | 45 ++++-- .../opensearch/OpenSearchEntityManager.java | 7 +- .../opensearch/OpenSearchRequestBuilder.java | 45 ++++++ .../opensearch/OpenSearchSearchManager.java | 98 ++++++------ .../service/search/opensearch/OsUtils.java | 35 ++++- .../ContextMemorySearchVisibility.java | 51 ++++++- .../vector/ElasticSearchVectorService.java | 24 +++ .../vector/OpenSearchVectorService.java | 24 +++ .../search/vector/VectorEmbeddingHandler.java | 22 ++- .../search/vector/VectorIndexService.java | 8 + .../org/openmetadata/csv/EntityCsvTest.java | 9 +- .../ElasticSearchBulkSinkBehaviorTest.java | 9 ++ .../OpenSearchBulkSinkBehaviorTest.java | 20 +++ .../jdbi3/ContextMemoryRepositoryTest.java | 48 +++--- .../service/jdbi3/ListFilterTest.java | 17 +-- ...rchRequestBuilderMemoryVisibilityTest.java | 114 ++++++++++++++ .../ContextMemorySearchVisibilityTest.java | 73 +++++++++ .../vector/VectorEmbeddingHandlerTest.java | 18 ++- .../policyevaluator/RuleEvaluatorTest.java | 28 +++- 31 files changed, 886 insertions(+), 309 deletions(-) create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRequestBuilderMemoryVisibilityTest.java diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java index 203937308179..a1d0b7b2435f 100644 --- a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java +++ b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContextMemoryIT.java @@ -10,7 +10,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; +import java.util.stream.Collectors; import org.awaitility.Awaitility; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.Execution; @@ -28,6 +30,7 @@ import org.openmetadata.schema.entity.teams.User; import org.openmetadata.schema.type.EntityHistory; import org.openmetadata.schema.type.EntityReference; +import org.openmetadata.sdk.client.OpenMetadataClient; import org.openmetadata.sdk.fluent.Users; import org.openmetadata.sdk.models.ListParams; import org.openmetadata.sdk.models.ListResponse; @@ -59,18 +62,16 @@ public ContextMemoryIT() { // ABSTRACT METHOD IMPLEMENTATIONS (Required by BaseEntityIT) // =================================================================== - // Default to ENTITY (org-wide) visibility so the generic search-index tests have a searchable - // memory: only ENTITY memories are indexed — PRIVATE/SHARED are excluded at index time (see - // ContextMemoryRepository.isSearchIndexable / getReindexFilter). The exclusion is covered by - // privateAndSharedMemoriesAreExcludedFromSearchIndex below. + // No shareConfig, so these fall back to the default (PRIVATE) visibility. The generic + // search-index tests in BaseEntityIT therefore double as the regression guard that a restricted + // memory still reaches the index — visibility is enforced at query time, not at index time. @Override protected CreateContextMemory createMinimalRequest(TestNamespace ns) { return new CreateContextMemory() .withName(ns.prefix("context-memory")) .withDescription("Test context memory") .withQuestion("How do I find certified tables?") - .withAnswer("Filter the Explore page by the Certification tag.") - .withShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY)); + .withAnswer("Filter the Explore page by the Certification tag."); } @Override @@ -79,8 +80,7 @@ protected CreateContextMemory createRequest(String name, TestNamespace ns) { .withName(name) .withDescription("Test context memory") .withQuestion("What is the data quality SLA?") - .withAnswer("Critical tables must pass tests every 24 hours.") - .withShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY)); + .withAnswer("Critical tables must pass tests every 24 hours."); } @Override @@ -422,76 +422,105 @@ void post_contextMemoryWithShareConfigVisibility_200_OK(TestNamespace ns) { assertEquals(MemoryVisibility.PRIVATE, memory.getShareConfig().getVisibility()); } + /** + * The ContextCenter serves its listing from search whenever it passes a query, filter, sort or + * offset — which it always does. Restricted memories must therefore reach the search index and be + * filtered per subject at query time: an owner and a shared principal see their own restricted + * memories, everyone else sees only the org-wide ones. + */ @Test - void privateAndSharedMemoriesAreExcludedFromSearchIndex(TestNamespace ns) throws Exception { - ContextMemory privateMemory = - createEntity(memoryWithVisibility(ns, "private-mem", MemoryVisibility.PRIVATE)); - ContextMemory sharedMemory = - createEntity(memoryWithVisibility(ns, "shared-mem", MemoryVisibility.SHARED)); - // Create the ENTITY memory last. Once it is searchable, keep checking the earlier PRIVATE and - // SHARED memories across a stability window so delayed or out-of-order indexing cannot make the - // test pass before a leaked document appears. - ContextMemory entityMemory = - createEntity(memoryWithVisibility(ns, "entity-mem", MemoryVisibility.ENTITY)); - - Awaitility.await("ENTITY-visibility memory indexed") - .pollDelay(Duration.ofMillis(500)) - .pollInterval(Duration.ofSeconds(2)) - .atMost(Duration.ofSeconds(180)) - .ignoreExceptions() - .untilAsserted( - () -> - assertTrue( - searchForEntity(entityMemory.getId().toString()) - .contains(entityMemory.getId().toString()), - "ENTITY-visibility memory must be searchable")); + void restrictedMemoriesAreSearchableByOwnerAndSharedPrincipalOnly(TestNamespace ns) { + ContextMemory ownedByUser1 = + createEntity( + memoryWithVisibility(ns, "owned-private", MemoryVisibility.PRIVATE) + .withOwners(List.of(testUser1Ref()))); + ContextMemory sharedWithUser1 = + createEntity( + memoryWithVisibility(ns, "shared-with-user1", MemoryVisibility.SHARED) + .withShareConfig( + new MemoryShareConfig() + .withVisibility(MemoryVisibility.SHARED) + .withSharedWith( + List.of(new MemorySharedPrincipal().withPrincipal(testUser1Ref()))))); + ContextMemory ownedByAdmin = + createEntity(memoryWithVisibility(ns, "admin-private", MemoryVisibility.PRIVATE)); + ContextMemory orgWide = + createEntity(memoryWithVisibility(ns, "org-wide", MemoryVisibility.ENTITY)); + + // Admin bypasses the visibility filter, so an admin-visible listing containing all four is the + // barrier proving every document — restricted ones included — reached the index. + awaitSearchBackedListContains( + SdkClients.adminClient(), + List.of(ownedByUser1, sharedWithUser1, ownedByAdmin, orgWide), + "every memory must be indexed regardless of visibility"); + + Set visibleToUser1 = searchBackedListIds(SdkClients.user1Client()); - Awaitility.await("PRIVATE and SHARED memories remain excluded") - .pollDelay(Duration.ofMillis(500)) - .pollInterval(Duration.ofSeconds(2)) - .during(Duration.ofSeconds(10)) - .atMost(Duration.ofSeconds(30)) - .ignoreExceptions() - .untilAsserted( - () -> { - assertFalse( - searchForEntity(privateMemory.getId().toString()) - .contains(privateMemory.getId().toString()), - "PRIVATE memory must not be in the search index"); - assertFalse( - searchForEntity(sharedMemory.getId().toString()) - .contains(sharedMemory.getId().toString()), - "SHARED memory must not be in the search index"); - }); + assertTrue( + visibleToUser1.contains(ownedByUser1.getId().toString()), + "an owner must find their own PRIVATE memory through the search-backed listing"); + assertTrue( + visibleToUser1.contains(sharedWithUser1.getId().toString()), + "a shared principal must find the SHARED memory through the search-backed listing"); + assertTrue( + visibleToUser1.contains(orgWide.getId().toString()), + "org-wide memories stay visible to everyone"); + assertFalse( + visibleToUser1.contains(ownedByAdmin.getId().toString()), + "another user's PRIVATE memory must never surface"); } @Test - void memoryVisibilityFlipAddsAndRemovesFromSearchIndex(TestNamespace ns) throws Exception { + void memoryStaysInSearchIndexAcrossVisibilityFlips(TestNamespace ns) throws Exception { ContextMemory memory = createEntity(memoryWithVisibility(ns, "flip", MemoryVisibility.ENTITY)); String id = memory.getId().toString(); - awaitSearchPresence(id, true, "ENTITY-visibility memory must be searchable"); + awaitSearchPresence(id, "ENTITY-visibility memory must be searchable"); - // Flip ENTITY -> PRIVATE: the live update must delete the doc from the index. ContextMemory toPrivate = getEntity(id); toPrivate.setShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.PRIVATE)); patchEntity(id, toPrivate); - awaitSearchPresence( - id, false, "Memory flipped to PRIVATE must be removed from the search index"); + awaitSearchPresence(id, "a memory flipped to PRIVATE stays indexed for its owner"); - // Flip PRIVATE -> ENTITY: the live update must re-index the doc. ContextMemory toEntity = getEntity(id); toEntity.setShareConfig(new MemoryShareConfig().withVisibility(MemoryVisibility.ENTITY)); patchEntity(id, toEntity); - awaitSearchPresence(id, true, "Memory flipped back to ENTITY must be searchable again"); + awaitSearchPresence(id, "a memory flipped back to ENTITY stays searchable"); + } + + private void awaitSearchPresence(String id, String message) { + Awaitility.await(message) + .pollDelay(Duration.ofMillis(500)) + .pollInterval(Duration.ofSeconds(2)) + .atMost(Duration.ofSeconds(180)) + .ignoreExceptions() + .untilAsserted(() -> assertTrue(searchForEntity(id).contains(id), message)); } - private void awaitSearchPresence(String id, boolean present, String message) { + private void awaitSearchBackedListContains( + OpenMetadataClient client, List memories, String message) { Awaitility.await(message) .pollDelay(Duration.ofMillis(500)) .pollInterval(Duration.ofSeconds(2)) .atMost(Duration.ofSeconds(180)) .ignoreExceptions() - .untilAsserted(() -> assertEquals(present, searchForEntity(id).contains(id), message)); + .untilAsserted( + () -> { + Set listed = searchBackedListIds(client); + for (ContextMemory memory : memories) { + assertTrue(listed.contains(memory.getId().toString()), message); + } + }); + } + + /** Lists memories through the search-backed branch of the endpoint (offset forces it). */ + private Set searchBackedListIds(OpenMetadataClient client) { + ListParams params = new ListParams(); + params.setLimit(1000); + params.setOffset(0); + return new ContextMemoryService(client.getHttpClient()) + .list(params).getData().stream() + .map(memory -> memory.getId().toString()) + .collect(Collectors.toSet()); } private CreateContextMemory memoryWithVisibility( @@ -500,7 +529,7 @@ private CreateContextMemory memoryWithVisibility( .withName(ns.prefix(name)) .withDescription("Visibility indexing test") .withQuestion("Is this memory searchable?") - .withAnswer("Only when visibility is Entity.") + .withAnswer("Visibility is enforced at query time, not at index time.") .withShareConfig(new MemoryShareConfig().withVisibility(visibility)); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java b/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java index 0b30c4a1b5c2..2dca8675e828 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java @@ -39,6 +39,7 @@ import java.util.Set; import java.util.TreeSet; import java.util.UUID; +import java.util.function.BiPredicate; import java.util.stream.Collectors; import lombok.Getter; import lombok.NonNull; @@ -741,26 +742,38 @@ public static boolean hasEntityRepository(@NonNull String entityType) { /** * Whether an entity instance should be written to the search index, per its repository's {@link - * EntityRepository#isSearchIndexable} policy (e.g. {@code ContextMemoryRepository} excludes - * non-org-wide memories). Defaults to {@code true} for entity types with no regular entity - * repository (index-only / time-series sub-entities), so the live index paths never throw on an - * indexable but repository-less type. Returns {@code false} when the entity or its reference is - * missing. + * EntityRepository#isSearchIndexable} policy. Defaults to {@code true} for entity types with no + * regular entity repository (index-only / time-series sub-entities), so the live index paths + * never throw on an indexable but repository-less type. Returns {@code false} when the entity or + * its reference is missing. */ public static boolean isSearchIndexable(EntityInterface entity) { - if (entity == null) { - return false; - } - EntityReference entityReference = entity.getEntityReference(); - if (entityReference == null) { - return false; - } - String entityType = entityReference.getType(); - if (entityType == null) { - return false; + return repositoryPolicyAllows(entity, EntityRepository::isSearchIndexable); + } + + /** + * Whether an entity instance should be embedded into the vector/semantic index, per its + * repository's {@link EntityRepository#isVectorEmbeddable} policy (e.g. {@code + * ContextMemoryRepository} keeps non-org-wide memories out because the vector query path carries + * no per-document visibility filter). Defaults and null-handling match {@link + * #isSearchIndexable}. + */ + public static boolean isVectorEmbeddable(EntityInterface entity) { + return repositoryPolicyAllows(entity, EntityRepository::isVectorEmbeddable); + } + + private static boolean repositoryPolicyAllows( + EntityInterface entity, + BiPredicate, EntityInterface> policy) { + boolean allowed = false; + EntityReference entityReference = entity == null ? null : entity.getEntityReference(); + String entityType = entityReference == null ? null : entityReference.getType(); + if (entityType != null) { + EntityRepository repository = + ENTITY_REPOSITORY_MAP.get(entityType); + allowed = repository == null || policy.test(repository, entity); } - EntityRepository repository = ENTITY_REPOSITORY_MAP.get(entityType); - return repository == null || repository.isSearchIndexable(entity); + return allowed; } /** diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java index b28e373143f0..f53e728b9749 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSink.java @@ -1106,6 +1106,12 @@ private String enrichWithEmbedding( Map existingEmbeddingsById, StageStatsTracker tracker) { try { + // Per-instance gate, mirroring VectorEmbeddingHandler: the entity-type check above cannot + // see that an individual ContextMemory is Private/Shared, and the vector query path carries + // no per-document visibility filter. + if (!Entity.isVectorEmbeddable(entity)) { + return json; + } ElasticSearchVectorService vectorService = ElasticSearchVectorService.getInstance(); if (vectorService == null) { return json; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java index 8b78e3925ea9..e6b626342cc0 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSink.java @@ -1135,6 +1135,12 @@ private String enrichWithEmbedding( StageStatsTracker tracker, String stagedChunkTarget) { try { + // Per-instance gate, mirroring VectorEmbeddingHandler: the entity-type check above cannot + // see that an individual ContextMemory is Private/Shared, and the vector query path carries + // no per-document visibility filter. + if (!Entity.isVectorEmbeddable(entity)) { + return json; + } OpenSearchVectorService vectorService = OpenSearchVectorService.getInstance(); if (vectorService == null) { return json; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContextMemoryRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContextMemoryRepository.java index 04c9cda90048..bc6d8fd5bc0e 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContextMemoryRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContextMemoryRepository.java @@ -42,6 +42,14 @@ import org.openmetadata.service.util.EntityUtil.RelationIncludes; import org.openmetadata.service.util.FullyQualifiedName; +/** + * Every memory is written to the search index regardless of its {@code shareConfig.visibility}, and + * privacy is enforced at query time by {@link + * org.openmetadata.service.search.security.ContextMemorySearchVisibility}. Indexing only org-wide + * memories would hide a user's own PRIVATE memories and the SHARED ones they are a principal of + * from {@code GET /contextCenter/memories}, which serves the ContextCenter listing from search + * whenever it is given a query, filter, sort or offset. + */ @Slf4j @Repository(name = "ContextMemoryRepository") public class ContextMemoryRepository extends EntityRepository { @@ -69,28 +77,20 @@ public ContextMemoryRepository() { } /** - * Only org-wide ({@link MemoryVisibility#ENTITY}) memories are searchable; PRIVATE and SHARED - * memories (and the {@code null}/unset default, which is PRIVATE) are kept out of the search - * index. This live-write check and {@link #getReindexFilter()} (the bulk-reindex equivalent - * expressed as a DB query) are the primary index-time enforcement points. A query-time - * defense-in-depth filter still protects restricted documents written before this policy was - * introduced. Owners and shared principals read restricted memories through the REST {@code - * /contextCenter/memories} endpoints (see {@link - * org.openmetadata.service.resources.context.ContextMemoryVisibility}). + * Only org-wide ({@link MemoryVisibility#ENTITY}) memories are embedded. The vector query path + * has no per-document visibility filter to apply — chunk documents carry neither {@code + * visibility} nor {@code sharedWithIds} — so a restricted memory reachable through semantic + * search could be returned to any caller. Keyword search does not need this restriction because + * {@link org.openmetadata.service.search.security.ContextMemorySearchVisibility} filters it per + * subject. */ @Override - public boolean isSearchIndexable(EntityInterface entity) { - boolean searchable = false; + public boolean isVectorEmbeddable(EntityInterface entity) { + boolean embeddable = false; if (entity instanceof ContextMemory memory && memory.getShareConfig() != null) { - searchable = memory.getShareConfig().getVisibility() == MemoryVisibility.ENTITY; + embeddable = memory.getShareConfig().getVisibility() == MemoryVisibility.ENTITY; } - return searchable; - } - - @Override - public ListFilter getReindexFilter() { - return new ListFilter(Include.ALL) - .addQueryParam(ListFilter.MEMORY_SEARCH_VISIBILITY_PARAM, MemoryVisibility.ENTITY.value()); + return embeddable; } @Override diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java index ce5774d1e80a..d2d8ec50d4c9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java @@ -2360,19 +2360,27 @@ public List findByNames(List entityFQNs, Include include) { /** * Whether a single entity instance should be written to the search index on the live - * create/update path. Defaults to true; override to keep specific instances out of the index - * (e.g. {@link ContextMemoryRepository} excludes non-org-wide memories). The bulk reindex applies - * the same rule at the DB-query level via {@link #getReindexFilter()}. + * create/update path. Defaults to true; override to keep specific instances out of the index. + * The bulk reindex applies the same rule at the DB-query level via {@link #getReindexFilter()}. */ public boolean isSearchIndexable(EntityInterface entity) { return true; } + /** + * Whether a single entity instance should be embedded into the vector/semantic index. Defaults to + * {@link #isSearchIndexable}; override when an instance may be keyword-searchable but must not be + * reachable through the vector path (e.g. {@link ContextMemoryRepository} keeps non-org-wide + * memories out because vector chunks carry no visibility fields to filter on). + */ + public boolean isVectorEmbeddable(EntityInterface entity) { + return isSearchIndexable(entity); + } + /** * Filter the search reindex uses to both list and count this entity's rows. Defaults to all rows; - * override to keep specific instances out of the search index (e.g. {@link ContextMemoryRepository} - * excludes non-org-wide memories). The reader and every entity-count site share this filter so the - * job total matches what actually gets indexed. + * override to keep specific instances out of the search index. The reader and every entity-count + * site share this filter so the job total matches what actually gets indexed. */ public ListFilter getReindexFilter() { return new ListFilter(Include.ALL); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java index 9fbbf767b75d..5fb5fdbbcf8d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java @@ -26,15 +26,6 @@ public class ListFilter extends Filter { public static final String NULL_PARAM = "null"; - /** - * Restricts a ContextMemory listing to a single {@code shareConfig.visibility} value. Set ONLY by - * {@link org.openmetadata.service.jdbi3.ContextMemoryRepository#getReindexFilter()} so the search - * reindex indexes org-wide ({@code Entity}) memories only. Do NOT set this on the REST {@code - * /contextCenter/memories} listing — owners and shared principals must still read their own - * PRIVATE/SHARED memories there. - */ - public static final String MEMORY_SEARCH_VISIBILITY_PARAM = "memorySearchVisibility"; - private static final String TASK_STATUS_GROUP_OPEN = "open"; private static final String TASK_STATUS_GROUP_ACTIVE = "active"; private static final String TASK_STATUS_GROUP_CLOSED = "closed"; @@ -102,7 +93,6 @@ public String getCondition(String tableName) { conditions.add(getTaskCreatedAtRangeCondition(tableName)); conditions.add(getDarSearchCondition()); conditions.add(getEntityStatusCondition(tableName)); - conditions.add(getMemorySearchVisibilityCondition()); conditions.add(getServerIdCondition()); conditions.add(getNameFilterCondition()); conditions.add(getPrimaryEntityCondition()); @@ -404,25 +394,6 @@ private String getAgentTypeCondition() { } } - /** - * Restricts ContextMemory listings to a single {@code shareConfig.visibility}. Used only by the - * search-reindex reader/counter to index org-wide (ENTITY) memories and skip PRIVATE/SHARED ones, - * so the privacy boundary is enforced at index time. NOT applied to the REST {@code - * /contextCenter/memories} listing — owners must still read their own restricted memories there. - */ - private String getMemorySearchVisibilityCondition() { - String visibility = queryParams.get(MEMORY_SEARCH_VISIBILITY_PARAM); - String condition = ""; - if (!nullOrEmpty(visibility)) { - String bindPlaceholder = ":" + MEMORY_SEARCH_VISIBILITY_PARAM; - condition = - Boolean.TRUE.equals(DatasourceConfig.getInstance().isMySQL()) - ? "JSON_UNQUOTE(JSON_EXTRACT(json, '$.shareConfig.visibility')) = " + bindPlaceholder - : "json->'shareConfig'->>'visibility' = " + bindPlaceholder; - } - return condition; - } - public String getProviderCondition(String tableName) { String provider = queryParams.get("provider"); if (provider == null) { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java index 28c693338c1d..6d032c7ae910 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java @@ -40,12 +40,17 @@ import org.openmetadata.service.search.SearchUtils; import org.openmetadata.service.search.elasticsearch.aggregations.ElasticAggregationsBuilder; import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilder; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilderFactory; import org.openmetadata.service.search.queries.OMQueryBuilder; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import org.openmetadata.service.search.security.RBACConditionEvaluator; import org.openmetadata.service.security.policyevaluator.SubjectContext; @Slf4j public class ElasticSearchAggregationManager implements AggregationManagementClient { + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new ElasticQueryBuilderFactory()); + private final ElasticsearchClient client; private final boolean isClientAvailable; private final ObjectMapper mapper; @@ -65,6 +70,20 @@ public ElasticSearchAggregationManager( mapper = new ObjectMapper(); } + /** + * ANDs the org-wide-only ContextMemory filter into an aggregation query. Aggregations run without + * a {@link org.openmetadata.service.security.policyevaluator.SubjectContext}, so they cannot tell + * whose restricted memory a document is and must fail closed: only memories everyone may read are + * aggregated. Non-memory documents are unaffected. + */ + private Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((ElasticQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + private String praseJsonQuery(String jsonQuery) throws JsonProcessingException { JsonNode rootNode = mapper.readTree(jsonQuery); String queryToProcess = jsonQuery; @@ -138,9 +157,7 @@ public Response aggregate(AggregationRequest request) throws IOException { } } - if (query != null) { - searchRequestBuilder.query(query); - } + searchRequestBuilder.query(restrictToOrgWideMemories(query)); String aggregationField = SearchSourceBuilderFactory.resolveFieldForSortOrAggregation(request.getFieldName()); @@ -255,7 +272,7 @@ public DataQualityReport genericAggregation( String indexName = Entity.getSearchRepository().getIndexOrAliasName(index); searchRequestBuilder.index(indexName); - Query parsedQuery; + Query parsedQuery = null; if (query != null) { // Check if query string contains outer "query" wrapper and extract inner query if (query.trim().startsWith("{")) { @@ -264,8 +281,8 @@ public DataQualityReport genericAggregation( } else { parsedQuery = Query.of(q -> q.queryString(qs -> qs.query(query))); } - searchRequestBuilder.query(parsedQuery); } + searchRequestBuilder.query(restrictToOrgWideMemories(parsedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); @@ -366,9 +383,7 @@ public DataQualityReport genericAggregation( } } - if (parsedQuery != null) { - searchRequestBuilder.query(parsedQuery); - } + searchRequestBuilder.query(restrictToOrgWideMemories(parsedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); @@ -465,14 +480,14 @@ public JsonObject aggregate( final Query finalParsedQuery = parsedQuery; final Query finalFilterQuery = filterQuery; - if (finalParsedQuery != null && finalFilterQuery != null) { - searchRequestBuilder.query( - q -> q.bool(b -> b.must(finalParsedQuery).filter(finalFilterQuery))); - } else if (finalParsedQuery != null) { - searchRequestBuilder.query(finalParsedQuery); - } else if (finalFilterQuery != null) { - searchRequestBuilder.query(q -> q.bool(b -> b.filter(finalFilterQuery))); + Query combinedQuery = finalParsedQuery; + if (finalFilterQuery != null) { + combinedQuery = + finalParsedQuery != null + ? Query.of(q -> q.bool(b -> b.must(finalParsedQuery).filter(finalFilterQuery))) + : Query.of(q -> q.bool(b -> b.filter(finalFilterQuery))); } + searchRequestBuilder.query(restrictToOrgWideMemories(combinedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchEntityManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchEntityManager.java index a67c5949c9c3..f8b1e9e06cbd 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchEntityManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchEntityManager.java @@ -74,6 +74,7 @@ import org.openmetadata.service.search.SearchIndexRetryQueue; import org.openmetadata.service.search.SearchRetryUtil; import org.openmetadata.service.search.SearchUtils; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import org.openmetadata.service.workflows.searchIndex.ReindexingUtil; /** @@ -537,7 +538,11 @@ public Response getDocByID(String indexName, String entityId) throws IOException g.index(Entity.getSearchRepository().getIndexOrAliasName(indexName)).id(entityId), Map.class); - if (response != null && response.found()) { + // This path runs no query and has no SubjectContext, so it cannot tell whose restricted + // memory a document is: a non-org-wide memory reads as not found rather than leaking. + if (response != null + && response.found() + && ContextMemorySearchVisibility.isOrgWideReadable(response.source())) { return Response.status(Response.Status.OK).entity(response.source()).build(); } } catch (ElasticsearchException e) { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchRequestBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchRequestBuilder.java index a27901c848c3..d04b4dbef90f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchRequestBuilder.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchRequestBuilder.java @@ -16,6 +16,9 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilder; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilderFactory; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; public class ElasticSearchRequestBuilder { private Query query; @@ -32,6 +35,10 @@ public class ElasticSearchRequestBuilder { private Boolean explain; private List searchAfter; private String preference; + private boolean contextMemoryVisibilityResolved; + + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new ElasticQueryBuilderFactory()); public ElasticSearchRequestBuilder() {} @@ -44,6 +51,34 @@ public Query query() { return this.query; } + /** ANDs {@code filterQuery} into the current query as a non-scoring filter clause. */ + public ElasticSearchRequestBuilder filter(Query filterQuery) { + if (query == null) { + query = filterQuery; + } else { + Query existingQuery = query; + query = + Query.of( + qb -> + qb.bool( + b -> { + b.must(existingQuery); + b.filter(filterQuery); + return b; + })); + } + return this; + } + + /** + * Records that the caller evaluated ContextMemory visibility for this request's subject, so + * {@link #build} leaves the query alone instead of applying its org-wide-only default. + */ + public ElasticSearchRequestBuilder contextMemoryVisibilityResolved() { + this.contextMemoryVisibilityResolved = true; + return this; + } + public ElasticSearchRequestBuilder postFilter(Query postFilterQuery) { this.postFilterQuery = postFilterQuery; return this; @@ -211,7 +246,17 @@ public String preference() { return this.preference; } + /** + * Builds the request, defaulting to org-wide-only ContextMemory visibility unless the caller + * resolved it for a concrete subject via {@link #contextMemoryVisibilityResolved()}. Memory + * privacy is enforced per query rather than by keeping restricted memories out of the index, so a + * search path that forgets to evaluate the subject must fail closed rather than leak. + */ public SearchRequest build(String... indices) { + if (!contextMemoryVisibilityResolved) { + filter(((ElasticQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2()); + contextMemoryVisibilityResolved = true; + } return SearchRequest.of( s -> { s.index(Arrays.asList(indices)); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchSearchManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchSearchManager.java index ff2ed75b2e53..a66f949dcfee 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchSearchManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchSearchManager.java @@ -151,16 +151,13 @@ public Response searchBySourceUrl(String sourceUrl) throws IOException { throw new IOException("Elasticsearch client is not available"); } + Query sourceUrlQuery = + Query.of(q -> q.bool(b -> b.must(m -> m.term(t -> t.field("sourceUrl").value(sourceUrl))))); SearchRequest searchRequest = SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query( - q -> - q.bool( - b -> - b.must( - m -> m.term(t -> t.field("sourceUrl").value(sourceUrl)))))); + .query(restrictToOrgWideMemories(sourceUrlQuery))); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); SearchResponse response; @@ -183,22 +180,20 @@ public Response searchByField( throw new IOException("Elasticsearch client is not available"); } + Query fieldQuery = + Query.of( + q -> + q.bool( + b -> + b.must(m -> m.wildcard(w -> w.field(fieldName).value(fieldValue))) + .filter(f -> f.term(t -> t.field("deleted").value(deleted))))); SearchRequest searchRequest = SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(index)) .from(from) .size(size) - .query( - q -> - q.bool( - b -> - b.must( - m -> - m.wildcard( - w -> w.field(fieldName).value(fieldValue))) - .filter( - f -> f.term(t -> t.field("deleted").value(deleted)))))); + .query(restrictToOrgWideMemories(fieldQuery))); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); SearchResponse response; @@ -398,6 +393,20 @@ private void applyRbacCondition( applyContextMemoryVisibility(subjectContext, requestBuilder); } + /** + * ANDs the org-wide-only ContextMemory filter into a query built outside {@link + * ElasticSearchRequestBuilder} (which applies the same default in its {@code build}). These paths + * take no {@link SubjectContext}, so they cannot tell whose restricted memory a document is and + * must fail closed. Non-memory documents are unaffected. + */ + private Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((ElasticQueryBuilder) contextMemoryVisibility.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + /** * Enforces ContextMemory shareConfig visibility on search results for non-admin subjects. Applied * independently of {@code shouldApplyRbacConditions} because memory visibility is a per-memory @@ -409,22 +418,13 @@ private void applyContextMemoryVisibility( OMQueryBuilder visibilityBuilder = contextMemoryVisibility.buildVisibilityFilter(subjectContext); if (visibilityBuilder != null) { - Query visibilityQuery = ((ElasticQueryBuilder) visibilityBuilder).buildV2(); - Query existingQuery = requestBuilder.query(); - if (existingQuery != null) { - Query combinedQuery = - Query.of( - qb -> - qb.bool( - b -> { - b.must(existingQuery); - b.filter(visibilityQuery); - return b; - })); - requestBuilder.query(combinedQuery); - } else { - requestBuilder.query(visibilityQuery); - } + requestBuilder.filter(((ElasticQueryBuilder) visibilityBuilder).buildV2()); + } + // Admins get no filter but are still resolved. An unidentifiable subject is NOT resolved, so + // ElasticSearchRequestBuilder#build falls back to its org-wide-only default instead of running + // the search unfiltered. + if (contextMemoryVisibility.isSubjectResolvable(subjectContext)) { + requestBuilder.contextMemoryVisibilityResolved(); } } @@ -599,6 +599,7 @@ public Response searchWithNLQ( requestBuilder.from(request.getFrom()); requestBuilder.size(request.getSize()); + applyContextMemoryVisibility(subjectContext, requestBuilder); // Add aggregations for NLQ query addAggregationsToNLQQuery(requestBuilder, request.getIndex()); @@ -910,7 +911,10 @@ private SearchResponse performLineageSearchForDQ( Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); SearchResponse response; try { - response = client.search(s -> s.index(indexName).query(boolQuery).size(1000), JsonData.class); + response = + client.search( + s -> s.index(indexName).query(restrictToOrgWideMemories(boolQuery)).size(1000), + JsonData.class); } finally { if (searchTimerSample != null) { RequestLatencyContext.endSearchOperation(searchTimerSample); @@ -1382,7 +1386,10 @@ private ElasticSearchRequestBuilder buildHierarchyQuery( } baseQueryBuilder.minimumShouldMatch(1); - Query originalQuery = baseQueryBuilder.build(); + // The subject's visibility filter was ANDed into the query this rewrite has just demoted to a + // should clause, where minimumShouldMatch(1) lets a name/displayName phrase match satisfy the + // query without it. Re-AND the memory guard so hierarchy search cannot become a bypass. + Query originalQuery = restrictToOrgWideMemories(baseQueryBuilder.build()); requestBuilder.query(originalQuery); // Add fqnParts aggregation to fetch parent terms @@ -1648,7 +1655,7 @@ private Map searchEntityRelationshipInternal( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(query) + .query(restrictToOrgWideMemories(query)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); @@ -1703,7 +1710,7 @@ private Map searchSchemaEntityRelationshipInternal( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(query) + .query(restrictToOrgWideMemories(query)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); @@ -1814,7 +1821,7 @@ private void getEntityRelationship( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(finalMainQuery) + .query(restrictToOrgWideMemories(finalMainQuery)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/EsUtils.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/EsUtils.java index 60a0844b4995..c9abe201b6ac 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/EsUtils.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/EsUtils.java @@ -40,6 +40,9 @@ import org.openmetadata.sdk.exception.SearchException; import org.openmetadata.service.Entity; import org.openmetadata.service.monitoring.RequestLatencyContext; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilder; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilderFactory; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; @Slf4j public class EsUtils { @@ -47,11 +50,26 @@ public class EsUtils { private static final ObjectMapper mapper; private static final JacksonJsonpMapper jsonpMapper; + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new ElasticQueryBuilderFactory()); + static { mapper = new ObjectMapper(); jsonpMapper = new JacksonJsonpMapper(mapper); } + /** + * ANDs the org-wide-only ContextMemory filter into a query built on a raw {@code + * SearchRequest.Builder}. Non-memory documents are unaffected. + */ + static Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((ElasticQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + public static Map jsonDataToMap(JsonData jsonData) { try { // Convert JsonData to JSON string, then parse it with Jackson @@ -250,7 +268,10 @@ public static SearchRequest getSearchRequest( baseQuery = Query.of(q -> q.bool(b -> b.must(finalBaseQuery).must(deletedQuery))); } - searchRequestBuilder.query(baseQuery); + // Lineage and entity-relationship traversals query the global alias and hand the matched hit's + // raw _source back as the node payload, with no SubjectContext to resolve memory visibility + // against: only org-wide memories may surface as nodes. + searchRequestBuilder.query(restrictToOrgWideMemories(baseQuery)); searchRequestBuilder.from(from); searchRequestBuilder.size(size); @@ -451,7 +472,10 @@ public static SearchRequest getSearchRequest( baseQuery = Query.of(q -> q.bool(b -> b.must(finalBaseQuery).must(deletedQuery))); } - searchRequestBuilder.query(baseQuery); + // Lineage and entity-relationship traversals query the global alias and hand the matched hit's + // raw _source back as the node payload, with no SubjectContext to resolve memory visibility + // against: only org-wide memories may surface as nodes. + searchRequestBuilder.query(restrictToOrgWideMemories(baseQuery)); searchRequestBuilder.from(from); searchRequestBuilder.size(size); @@ -476,8 +500,13 @@ public static SearchResponse searchEntities( Query deletedQuery = Query.of(q -> q.term(t -> t.field("deleted").value(!nullOrEmpty(deleted) && deleted))); + // Platform lineage returns each hit's raw _source and runs without a SubjectContext, so it + // cannot tell whose restricted memory a document is: only org-wide memories may come back. SearchRequest.Builder searchRequestBuilder = - new SearchRequest.Builder().index(indexName).query(deletedQuery).size(10000); + new SearchRequest.Builder() + .index(indexName) + .query(restrictToOrgWideMemories(deletedQuery)) + .size(10000); // Apply query filter buildSearchSourceFilter(queryFilter, searchRequestBuilder); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java index 378743215969..1e58d8c18f5d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java @@ -31,7 +31,9 @@ import org.openmetadata.service.search.SearchUtils; import org.openmetadata.service.search.opensearch.aggregations.OpenAggregationsBuilder; import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilder; +import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilderFactory; import org.openmetadata.service.search.queries.OMQueryBuilder; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import org.openmetadata.service.search.security.RBACConditionEvaluator; import org.openmetadata.service.security.policyevaluator.SubjectContext; import os.org.opensearch.client.json.JsonData; @@ -45,6 +47,9 @@ @Slf4j public class OpenSearchAggregationManager implements AggregationManagementClient { + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new OpenSearchQueryBuilderFactory()); + private final OpenSearchClient client; private final boolean isClientAvailable; private final ObjectMapper mapper; @@ -64,6 +69,20 @@ public OpenSearchAggregationManager( mapper = new ObjectMapper(); } + /** + * ANDs the org-wide-only ContextMemory filter into an aggregation query. Aggregations run without + * a {@link org.openmetadata.service.security.policyevaluator.SubjectContext}, so they cannot tell + * whose restricted memory a document is and must fail closed: only memories everyone may read are + * aggregated. Non-memory documents are unaffected. + */ + private Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((OpenSearchQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + private String praseJsonQuery(String jsonQuery) throws JsonProcessingException { JsonNode rootNode = mapper.readTree(jsonQuery); String queryToProcess = jsonQuery; @@ -138,9 +157,7 @@ public Response aggregate(AggregationRequest request) throws IOException { } } - if (query != null) { - searchRequestBuilder.query(query); - } + searchRequestBuilder.query(restrictToOrgWideMemories(query)); String aggregationField = SearchSourceBuilderFactory.resolveFieldForSortOrAggregation(request.getFieldName()); @@ -253,7 +270,7 @@ public DataQualityReport genericAggregation( String indexName = Entity.getSearchRepository().getIndexOrAliasName(index); searchRequestBuilder.index(indexName); - Query parsedQuery; + Query parsedQuery = null; if (query != null) { // Check if query string contains outer "query" wrapper and extract inner query if (query.trim().startsWith("{")) { @@ -262,8 +279,8 @@ public DataQualityReport genericAggregation( } else { parsedQuery = Query.of(q -> q.queryString(qs -> qs.query(query))); } - searchRequestBuilder.query(parsedQuery); } + searchRequestBuilder.query(restrictToOrgWideMemories(parsedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); @@ -346,9 +363,7 @@ public DataQualityReport genericAggregation( } } - if (parsedQuery != null) { - searchRequestBuilder.query(parsedQuery); - } + searchRequestBuilder.query(restrictToOrgWideMemories(parsedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); @@ -425,14 +440,14 @@ public JsonObject aggregate( final Query finalParsedQuery = parsedQuery; final Query finalFilterQuery = filterQuery; - if (finalParsedQuery != null && finalFilterQuery != null) { - searchRequestBuilder.query( - q -> q.bool(b -> b.must(finalParsedQuery).filter(finalFilterQuery))); - } else if (finalParsedQuery != null) { - searchRequestBuilder.query(finalParsedQuery); - } else if (finalFilterQuery != null) { - searchRequestBuilder.query(q -> q.bool(b -> b.filter(finalFilterQuery))); + Query combinedQuery = finalParsedQuery; + if (finalFilterQuery != null) { + combinedQuery = + finalParsedQuery != null + ? Query.of(q -> q.bool(b -> b.must(finalParsedQuery).filter(finalFilterQuery))) + : Query.of(q -> q.bool(b -> b.filter(finalFilterQuery))); } + searchRequestBuilder.query(restrictToOrgWideMemories(combinedQuery)); searchRequestBuilder.aggregations(aggregations); searchRequestBuilder.size(0); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchEntityManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchEntityManager.java index 64f0656a7ee4..b1fa2fabb285 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchEntityManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchEntityManager.java @@ -53,6 +53,7 @@ import org.openmetadata.service.search.SearchIndexRetryQueue; import org.openmetadata.service.search.SearchRetryUtil; import org.openmetadata.service.search.SearchUtils; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import org.openmetadata.service.workflows.searchIndex.ReindexingUtil; import os.org.opensearch.client.json.JsonData; import os.org.opensearch.client.opensearch.OpenSearchAsyncClient; @@ -587,7 +588,11 @@ public Response getDocByID(String indexName, String entityId) throws IOException g.index(Entity.getSearchRepository().getIndexOrAliasName(indexName)).id(entityId), Map.class); - if (response != null && response.found()) { + // This path runs no query and has no SubjectContext, so it cannot tell whose restricted + // memory a document is: a non-org-wide memory reads as not found rather than leaking. + if (response != null + && response.found() + && ContextMemorySearchVisibility.isOrgWideReadable(response.source())) { return Response.status(Response.Status.OK).entity(response.source()).build(); } } catch (OpenSearchException e) { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchRequestBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchRequestBuilder.java index 08ee772dba7b..3a72cdb7629f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchRequestBuilder.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchRequestBuilder.java @@ -5,6 +5,9 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilder; +import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilderFactory; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import os.org.opensearch.client.opensearch._types.FieldSort; import os.org.opensearch.client.opensearch._types.FieldValue; import os.org.opensearch.client.opensearch._types.NestedSortValue; @@ -34,6 +37,10 @@ public class OpenSearchRequestBuilder { private List searchAfter; private SearchType searchType; private String preference; + private boolean contextMemoryVisibilityResolved; + + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new OpenSearchQueryBuilderFactory()); public OpenSearchRequestBuilder() {} @@ -46,6 +53,34 @@ public Query query() { return this.query; } + /** ANDs {@code filterQuery} into the current query as a non-scoring filter clause. */ + public OpenSearchRequestBuilder filter(Query filterQuery) { + if (query == null) { + query = filterQuery; + } else { + Query existingQuery = query; + query = + Query.of( + qb -> + qb.bool( + b -> { + b.must(existingQuery); + b.filter(filterQuery); + return b; + })); + } + return this; + } + + /** + * Records that the caller evaluated ContextMemory visibility for this request's subject, so + * {@link #build} leaves the query alone instead of applying its org-wide-only default. + */ + public OpenSearchRequestBuilder contextMemoryVisibilityResolved() { + this.contextMemoryVisibilityResolved = true; + return this; + } + public OpenSearchRequestBuilder postFilter(Query postFilterQuery) { this.postFilterQuery = postFilterQuery; return this; @@ -228,7 +263,17 @@ public String preference() { return this.preference; } + /** + * Builds the request, defaulting to org-wide-only ContextMemory visibility unless the caller + * resolved it for a concrete subject via {@link #contextMemoryVisibilityResolved()}. Memory + * privacy is enforced per query rather than by keeping restricted memories out of the index, so a + * search path that forgets to evaluate the subject must fail closed rather than leak. + */ public SearchRequest build(String... indices) { + if (!contextMemoryVisibilityResolved) { + filter(((OpenSearchQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2()); + contextMemoryVisibilityResolved = true; + } return SearchRequest.of( s -> { s.index(Arrays.asList(indices)); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchSearchManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchSearchManager.java index 671c898b54dc..7b154778e045 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchSearchManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchSearchManager.java @@ -168,20 +168,20 @@ public Response searchBySourceUrl(String sourceUrl) throws IOException { throw new IOException("OpenSearch client is not available"); } + Query sourceUrlQuery = + Query.of( + q -> + q.bool( + b -> + b.must( + m -> + m.term( + t -> t.field("sourceUrl").value(FieldValue.of(sourceUrl)))))); SearchRequest searchRequest = SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query( - q -> - q.bool( - b -> - b.must( - m -> - m.term( - t -> - t.field("sourceUrl") - .value(FieldValue.of(sourceUrl))))))); + .query(restrictToOrgWideMemories(sourceUrlQuery))); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); SearchResponse response; @@ -203,26 +203,23 @@ public Response searchByField( throw new IOException("OpenSearch client is not available"); } + Query fieldQuery = + Query.of( + q -> + q.bool( + b -> + b.must(m -> m.wildcard(w -> w.field(fieldName).value(fieldValue))) + .filter( + f -> + f.term( + t -> t.field("deleted").value(FieldValue.of(deleted)))))); SearchRequest searchRequest = SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(index)) .from(from) .size(size) - .query( - q -> - q.bool( - b -> - b.must( - m -> - m.wildcard( - w -> w.field(fieldName).value(fieldValue))) - .filter( - f -> - f.term( - t -> - t.field("deleted") - .value(FieldValue.of(deleted))))))); + .query(restrictToOrgWideMemories(fieldQuery))); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); SearchResponse response; @@ -582,6 +579,7 @@ public Response searchWithNLQ( requestBuilder.from(request.getFrom()); requestBuilder.size(request.getSize()); + applyContextMemoryVisibility(subjectContext, requestBuilder); // Add aggregations for NLQ query addAggregationsToNLQQuery(requestBuilder, request.getIndex()); @@ -866,7 +864,9 @@ private SearchResponse performLineageSearchForDQ( Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); try { - return client.search(s -> s.index(indexName).query(boolQuery).size(1000), JsonData.class); + return client.search( + s -> s.index(indexName).query(restrictToOrgWideMemories(boolQuery)).size(1000), + JsonData.class); } finally { if (searchTimerSample != null) { RequestLatencyContext.endSearchOperation(searchTimerSample); @@ -880,27 +880,32 @@ private SearchResponse performLineageSearchForDQ( * privacy guarantee, not an RBAC policy — disabling RBAC search filtering must not expose private * memories. Non-memory documents are left untouched. */ + /** + * ANDs the org-wide-only ContextMemory filter into a query built outside {@link + * OpenSearchRequestBuilder} (which applies the same default in its {@code build}). These paths + * take no {@link SubjectContext}, so they cannot tell whose restricted memory a document is and + * must fail closed. Non-memory documents are unaffected. + */ + private Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((OpenSearchQueryBuilder) contextMemoryVisibility.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + private void applyContextMemoryVisibility( SubjectContext subjectContext, OpenSearchRequestBuilder requestBuilder) { OMQueryBuilder visibilityBuilder = contextMemoryVisibility.buildVisibilityFilter(subjectContext); if (visibilityBuilder != null) { - Query visibilityQuery = ((OpenSearchQueryBuilder) visibilityBuilder).buildV2(); - Query existingQuery = requestBuilder.query(); - if (existingQuery != null) { - Query combinedQuery = - Query.of( - qb -> - qb.bool( - b -> { - b.must(existingQuery); - b.filter(visibilityQuery); - return b; - })); - requestBuilder.query(combinedQuery); - } else { - requestBuilder.query(visibilityQuery); - } + requestBuilder.filter(((OpenSearchQueryBuilder) visibilityBuilder).buildV2()); + } + // Admins get no filter but are still resolved. An unidentifiable subject is NOT resolved, so + // OpenSearchRequestBuilder#build falls back to its org-wide-only default instead of running the + // search unfiltered. + if (contextMemoryVisibility.isSubjectResolvable(subjectContext)) { + requestBuilder.contextMemoryVisibilityResolved(); } } @@ -1417,7 +1422,10 @@ private OpenSearchRequestBuilder buildHierarchyQuery( } baseQueryBuilder.minimumShouldMatch(1); - Query originalQuery = baseQueryBuilder.build(); + // The subject's visibility filter was ANDed into the query this rewrite has just demoted to a + // should clause, where minimumShouldMatch(1) lets a name/displayName phrase match satisfy the + // query without it. Re-AND the memory guard so hierarchy search cannot become a bypass. + Query originalQuery = restrictToOrgWideMemories(baseQueryBuilder.build()); requestBuilder.query(originalQuery); // Add fqnParts aggregation to fetch parent terms @@ -1673,7 +1681,7 @@ public Map searchEntityRelationshipInternal( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(query) + .query(restrictToOrgWideMemories(query)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); @@ -1728,7 +1736,7 @@ public Map searchSchemaEntityRelationshipInternal( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(query) + .query(restrictToOrgWideMemories(query)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); @@ -1838,7 +1846,7 @@ private void getEntityRelationship( SearchRequest.of( s -> s.index(Entity.getSearchRepository().getIndexOrAliasName(GLOBAL_SEARCH_ALIAS)) - .query(finalBaseQuery) + .query(restrictToOrgWideMemories(finalBaseQuery)) .size(1000)); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OsUtils.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OsUtils.java index 7ed87047f8fb..92e4788fbefb 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OsUtils.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OsUtils.java @@ -30,6 +30,9 @@ import org.openmetadata.service.Entity; import org.openmetadata.service.monitoring.RequestLatencyContext; import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilder; +import org.openmetadata.service.search.opensearch.queries.OpenSearchQueryBuilderFactory; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; import os.org.opensearch.client.json.JsonData; import os.org.opensearch.client.json.jackson.JacksonJsonpMapper; import os.org.opensearch.client.opensearch.OpenSearchClient; @@ -46,6 +49,21 @@ @Slf4j public class OsUtils { + private static final ContextMemorySearchVisibility MEMORY_VISIBILITY = + new ContextMemorySearchVisibility(new OpenSearchQueryBuilderFactory()); + + /** + * ANDs the org-wide-only ContextMemory filter into a query built on a raw {@code + * SearchRequest.Builder}. Non-memory documents are unaffected. + */ + static Query restrictToOrgWideMemories(Query query) { + Query memoryFilter = + ((OpenSearchQueryBuilder) MEMORY_VISIBILITY.buildOrgWideOnlyFilter()).buildV2(); + return query == null + ? memoryFilter + : Query.of(q -> q.bool(b -> b.must(query).filter(memoryFilter))); + } + public static Map jsonDataToMap(JsonData jsonData) { try { // Convert JsonData to JSON string, then parse it with Jackson @@ -259,7 +277,10 @@ public static SearchRequest getSearchRequest( baseQuery = Query.of(q -> q.bool(b -> b.must(finalBaseQuery).must(deletedQuery))); } - searchRequestBuilder.query(baseQuery); + // Lineage and entity-relationship traversals query the global alias and hand the matched hit's + // raw _source back as the node payload, with no SubjectContext to resolve memory visibility + // against: only org-wide memories may surface as nodes. + searchRequestBuilder.query(restrictToOrgWideMemories(baseQuery)); searchRequestBuilder.from(from); searchRequestBuilder.size(size); @@ -461,7 +482,10 @@ public static SearchRequest getSearchRequest( baseQuery = Query.of(q -> q.bool(b -> b.must(finalBaseQuery).must(deletedQuery))); } - searchRequestBuilder.query(baseQuery); + // Lineage and entity-relationship traversals query the global alias and hand the matched hit's + // raw _source back as the node payload, with no SubjectContext to resolve memory visibility + // against: only org-wide memories may surface as nodes. + searchRequestBuilder.query(restrictToOrgWideMemories(baseQuery)); searchRequestBuilder.from(from); searchRequestBuilder.size(size); @@ -490,8 +514,13 @@ public static SearchResponse searchEntities( t -> t.field("deleted").value(FieldValue.of(!nullOrEmpty(deleted) && deleted)))); + // Platform lineage returns each hit's raw _source and runs without a SubjectContext, so it + // cannot tell whose restricted memory a document is: only org-wide memories may come back. SearchRequest.Builder searchRequestBuilder = - new SearchRequest.Builder().index(indexName).query(deletedQuery).size(10000); + new SearchRequest.Builder() + .index(indexName) + .query(restrictToOrgWideMemories(deletedQuery)) + .size(10000); // Apply query filter buildSearchSourceFilter(queryFilter, searchRequestBuilder); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/security/ContextMemorySearchVisibility.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/security/ContextMemorySearchVisibility.java index 03858468843e..5e2260bef2ad 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/security/ContextMemorySearchVisibility.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/security/ContextMemorySearchVisibility.java @@ -17,6 +17,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import org.openmetadata.schema.entity.context.MemoryVisibility; import org.openmetadata.schema.entity.teams.User; import org.openmetadata.schema.type.EntityReference; @@ -37,9 +38,10 @@ * so it is applied for every non-admin subject regardless of the RBAC access-control toggle. * Disabling RBAC search filtering must never expose another user's private memories. * - *

Current live writes and reindexing exclude restricted memories at index time. This filter - * remains as defense-in-depth for documents written before that policy was introduced and for - * mixed-version deployments during an upgrade. + *

Every memory is indexed, whatever its visibility, so that owners and shared principals keep + * finding their own restricted memories through the search-backed {@code /contextCenter/memories} + * listing and global search. That makes this filter the sole enforcement point on the search side: + * every path that can return a ContextMemory document to a caller must apply it. * *

The produced query is engine-agnostic via {@link QueryBuilderFactory}, so the same builder * serves both Elasticsearch and OpenSearch. Each OR-group is a bool with only {@code should} @@ -73,14 +75,51 @@ public OMQueryBuilder buildVisibilityFilter(SubjectContext subjectContext) { return filter; } - private boolean isVisibilityEnforced(SubjectContext subjectContext) { + /** + * Returns a filter admitting only org-wide ({@link MemoryVisibility#ENTITY}) memories. This is the + * fail-closed default for search paths that carry no {@link SubjectContext} and therefore cannot + * decide who a restricted memory belongs to — they get the memories everyone may read and nothing + * else. Like {@link #buildVisibilityFilter}, non-memory documents always pass. + */ + public OMQueryBuilder buildOrgWideOnlyFilter() { + return scopeMemoriesTo( + queryBuilderFactory.termQuery(FIELD_VISIBILITY, MemoryVisibility.ENTITY.value())); + } + + /** + * The document-level equivalent of {@link #buildOrgWideOnlyFilter}, for fetch-by-id paths that + * run no query to filter. Returns false only for a restricted (non-org-wide) context memory; + * every other document passes. + */ + public static boolean isOrgWideReadable(Map document) { + boolean readable = true; + if (document != null && Entity.CONTEXT_MEMORY.equals(document.get(FIELD_ENTITY_TYPE))) { + readable = MemoryVisibility.ENTITY.value().equals(document.get(FIELD_VISIBILITY)); + } + return readable; + } + + /** + * Whether the subject is identifiable enough to decide memory visibility from. Only then may a + * caller mark the request as resolved; an unidentifiable subject must fall back to the + * org-wide-only default rather than search unfiltered. + */ + public boolean isSubjectResolvable(SubjectContext subjectContext) { return subjectContext != null - && !subjectContext.isAdmin() && subjectContext.user() != null && subjectContext.user().getId() != null; } + private boolean isVisibilityEnforced(SubjectContext subjectContext) { + return isSubjectResolvable(subjectContext) && !subjectContext.isAdmin(); + } + private OMQueryBuilder buildFilter(User user) { + return scopeMemoriesTo(buildVisibleToUserClause(user)); + } + + /** Applies {@code memoryClause} to context memory documents only, letting every other type pass. */ + private OMQueryBuilder scopeMemoriesTo(OMQueryBuilder memoryClause) { OMQueryBuilder nonMemory = queryBuilderFactory .boolQuery() @@ -92,7 +131,7 @@ private OMQueryBuilder buildFilter(User user) { .must( List.of( queryBuilderFactory.termQuery(FIELD_ENTITY_TYPE, Entity.CONTEXT_MEMORY), - buildVisibleToUserClause(user))); + memoryClause)); return queryBuilderFactory.boolQuery().should(List.of(nonMemory, memoryVisible)); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/ElasticSearchVectorService.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/ElasticSearchVectorService.java index 70082a769b1a..becc2a9bf4f6 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/ElasticSearchVectorService.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/ElasticSearchVectorService.java @@ -477,6 +477,30 @@ public void partialUpdateEntity( } } + @Override + public void clearEntityEmbedding(String entityIndexName, String entityId) { + // A doc merge cannot delete keys, so removal has to go through a script. + try { + String fieldsJson = MAPPER.writeValueAsString(EMBEDDING_SOURCE_FIELDS); + String updateBody = + "{\"script\":{\"source\":\"for (field in params.fields) { ctx._source.remove(field) }\"," + + "\"params\":{\"fields\":" + + fieldsJson + + "}}}"; + executeGenericRequest( + "POST", + "/" + entityIndexName + "/_update/" + entityId + "?retry_on_conflict=3", + updateBody); + } catch (Exception e) { + LOG.error( + "Failed to clear embedding for entity {} in {}: {}", + entityId, + entityIndexName, + e.getMessage(), + e); + } + } + public void close() { // No-op by design, mirroring OpenSearchVectorService. The elasticsearch-java client stored // here was constructed elsewhere and its Rest5Client transport is shared with diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/OpenSearchVectorService.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/OpenSearchVectorService.java index ae351946e442..7942c3bc0ec3 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/OpenSearchVectorService.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/OpenSearchVectorService.java @@ -1646,6 +1646,30 @@ public void partialUpdateEntity( } } + @Override + public void clearEntityEmbedding(String entityIndexName, String entityId) { + // A doc merge cannot delete keys, so removal has to go through a script. + try { + String fieldsJson = MAPPER.writeValueAsString(EMBEDDING_SOURCE_FIELDS); + String updateBody = + "{\"script\":{\"source\":\"for (field in params.fields) { ctx._source.remove(field) }\"," + + "\"params\":{\"fields\":" + + fieldsJson + + "}}}"; + executeGenericRequest( + "POST", + "/" + entityIndexName + "/_update/" + entityId + "?retry_on_conflict=3", + updateBody); + } catch (Exception e) { + LOG.error( + "Failed to clear embedding for entity {} in {}: {}", + entityId, + entityIndexName, + e.getMessage(), + e); + } + } + @Override public String executeGenericRequest(String method, String endpoint, String body) { try { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorEmbeddingHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorEmbeddingHandler.java index 96237af3be03..02fa7c33804b 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorEmbeddingHandler.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorEmbeddingHandler.java @@ -43,8 +43,9 @@ public void onEntityUpdated( if (entity.getDeleted() != null && entity.getDeleted()) { return; } - if (!Entity.isSearchIndexable(entity)) { + if (!Entity.isVectorEmbeddable(entity)) { deleteChunks(entity); + clearEntityEmbedding(entity); return; } updateEmbedding(entity); @@ -69,6 +70,23 @@ public void onEntitySoftDeletedOrRestored( } } + /** + * Entity-doc embedding writes are partial merges, so an entity that stops being embeddable keeps + * its stale {@code embedding} — and stays kNN-matchable — unless the field is removed explicitly. + */ + private void clearEntityEmbedding(EntityInterface entity) { + if (entity != null && entity.getId() != null && isSupported(entity)) { + try { + String entityIndexName = resolveEntityIndexName(extractEntityType(entity)); + if (entityIndexName != null) { + vectorService.clearEntityEmbedding(entityIndexName, entity.getId().toString()); + } + } catch (Exception e) { + LOG.error("Failed to clear embedding for entity {}: {}", entity.getId(), e.getMessage(), e); + } + } + } + private void deleteChunks(EntityInterface entity) { if (entity != null && entity.getId() != null && isSupported(entity)) { try { @@ -92,7 +110,7 @@ private void updateEmbedding(EntityInterface entity) { if (!isSupported(entity)) { return; } - if (!Entity.isSearchIndexable(entity)) { + if (!Entity.isVectorEmbeddable(entity)) { return; } try { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorIndexService.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorIndexService.java index b7e0826e66b8..721d6403a059 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorIndexService.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/vector/VectorIndexService.java @@ -24,6 +24,14 @@ default void updateEntityEmbeddingChunks(EntityInterface entity) {} /** Remove all chunk documents for the given parent entity (hard/soft delete cleanup). */ default void deleteEntityChunks(String parentId) {} + /** + * Strip the embedding fields from an entity document. Entity-doc embedding writes are partial + * merges, so an entity that becomes non-embeddable (e.g. a ContextMemory flipped from org-wide to + * Private) would otherwise keep a stale {@code embedding} and stay matchable by kNN. Default is a + * no-op for backends without the legacy entity-doc embedding. + */ + default void clearEntityEmbedding(String entityIndexName, String entityId) {} + /** * Combined write: refresh both the legacy entity-doc embedding (read by hybrid search) and the * chunk documents (read by the semantic vector path). Implementations should embed each chunk diff --git a/openmetadata-service/src/test/java/org/openmetadata/csv/EntityCsvTest.java b/openmetadata-service/src/test/java/org/openmetadata/csv/EntityCsvTest.java index dd7782501363..e439749050c9 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/csv/EntityCsvTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/csv/EntityCsvTest.java @@ -36,6 +36,7 @@ import org.apache.commons.lang3.StringEscapeUtils; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; import org.mockito.MockedStatic; import org.mockito.Mockito; import org.openmetadata.common.utils.CommonUtil; @@ -102,7 +103,13 @@ public class EntityCsvTest { @BeforeAll public static void setup() { - Entity.registerEntity(Table.class, Entity.TABLE, Mockito.mock(TableRepository.class)); + TableRepository tableRepository = Mockito.mock(TableRepository.class); + // This registration is global and never torn down, so the stand-in must answer the indexing + // policy hooks the way the real repository does. A bare mock answers false, which would make + // Entity.isSearchIndexable report every table as non-indexable for the rest of the JVM. + Mockito.doReturn(true).when(tableRepository).isSearchIndexable(ArgumentMatchers.any()); + Mockito.doReturn(true).when(tableRepository).isVectorEmbeddable(ArgumentMatchers.any()); + Entity.registerEntity(Table.class, Entity.TABLE, tableRepository); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSinkBehaviorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSinkBehaviorTest.java index b001cf871edb..6ebe0bc3d0e7 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSinkBehaviorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ElasticSearchBulkSinkBehaviorTest.java @@ -708,6 +708,9 @@ void enrichWithEmbeddingReusesCachedFieldsWhenServiceReportsMatch() throws Excep EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); ElasticSearchVectorService vectorService = mock(ElasticSearchVectorService.class); @@ -750,6 +753,9 @@ void enrichWithEmbeddingReusesCachedFieldsWhenServiceReportsMatch() throws Excep void enrichWithEmbeddingRecomputesWhenNoCachedEntryAvailable() throws Exception { EntityInterface entity = mock(EntityInterface.class); when(entity.getId()).thenReturn(UUID.randomUUID()); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); ElasticSearchVectorService vectorService = mock(ElasticSearchVectorService.class); @@ -786,6 +792,9 @@ void enrichWithEmbeddingRecomputesWhenCachedDimensionMismatchesClient() throws E EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); EmbeddingClient embeddingClient = mock(EmbeddingClient.class); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSinkBehaviorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSinkBehaviorTest.java index 9d9bf0de7fd0..52a8e5915044 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSinkBehaviorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/OpenSearchBulkSinkBehaviorTest.java @@ -203,6 +203,8 @@ void relationshipPartialUpdateUsesVectorEnrichedFullUpsertWithoutRewritingOtherF processorConstruction.constructed().getFirst(); entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); entityMock.when(() -> Entity.getEntityTypeFromObject(entity)).thenReturn(Entity.TEST_CASE); + // The bulk sink consults this to keep restricted ContextMemory out of the vector index. + entityMock.when(() -> Entity.isVectorEmbeddable(entity)).thenReturn(true); entityMock .when(() -> Entity.buildSearchIndex(Entity.TEST_CASE, entity)) .thenReturn( @@ -647,6 +649,9 @@ void enrichWithEmbeddingReusesCachedFieldsWhenServiceReportsMatch() throws Excep EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); @@ -783,6 +788,9 @@ void enrichWithEmbeddingRecomputesWhenNoCachedEntryAvailable() throws Exception EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); @@ -838,6 +846,9 @@ void enrichWithEmbeddingRecomputesWhenCachedEntryHasNoEmbedding() throws Excepti EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); @@ -888,6 +899,9 @@ void enrichWithEmbeddingRecomputesWhenCachedNodeIsNotAnObject() throws Exception EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); @@ -941,6 +955,9 @@ void enrichWithEmbeddingRecomputesWhenCachedDimensionMismatchesClient() throws E EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); @@ -1005,6 +1022,9 @@ void enrichWithEmbeddingSkipsGenerationWhenProviderUnavailable() throws Exceptio EntityInterface entity = mock(EntityInterface.class); UUID entityId = UUID.randomUUID(); when(entity.getId()).thenReturn(entityId); + // A real entity always resolves a reference; the bulk sink consults it via + // Entity.isVectorEmbeddable to keep restricted ContextMemory out of the vector index. + when(entity.getEntityReference()).thenReturn(new EntityReference().withType("table")); StageStatsTracker tracker = mock(StageStatsTracker.class); OpenSearchVectorService vectorService = mock(OpenSearchVectorService.class); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ContextMemoryRepositoryTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ContextMemoryRepositoryTest.java index 952f0b0fe2bf..8b6949858939 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ContextMemoryRepositoryTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ContextMemoryRepositoryTest.java @@ -12,7 +12,6 @@ */ package org.openmetadata.service.jdbi3; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; @@ -26,6 +25,8 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.Execution; import org.junit.jupiter.api.parallel.ExecutionMode; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; import org.openmetadata.schema.EntityInterface; import org.openmetadata.schema.entity.context.ContextMemory; import org.openmetadata.schema.entity.context.MemoryShareConfig; @@ -34,11 +35,11 @@ import org.openmetadata.service.Entity; /** - * The ContextMemory index-time search-visibility rule lives in {@link ContextMemoryRepository}: - * only org-wide ({@link MemoryVisibility#ENTITY}) memories are indexed. {@link - * ContextMemoryRepository#isSearchIndexable} gates the live create/update path and {@link - * ContextMemoryRepository#getReindexFilter} carries the same rule into the bulk reindex as a DB - * query. These tests pin both halves across every visibility state. + * ContextMemory is indexed whatever its {@code shareConfig.visibility}; privacy is enforced at + * query time by {@link org.openmetadata.service.search.security.ContextMemorySearchVisibility}. + * Excluding restricted memories at index time hid a user's own PRIVATE memories and the SHARED ones + * they are a principal of from the search-backed {@code GET /contextCenter/memories} listing, so + * these tests pin that no visibility is filtered out of the index. */ @Execution(ExecutionMode.SAME_THREAD) class ContextMemoryRepositoryTest { @@ -60,42 +61,29 @@ void tearDown() { Entity.cleanup(); } - @Test - void isSearchIndexable_trueForEntityVisibility() { - assertTrue(repository.isSearchIndexable(memory(MemoryVisibility.ENTITY))); - } - - @Test - void isSearchIndexable_falseForPrivate() { - assertFalse(repository.isSearchIndexable(memory(MemoryVisibility.PRIVATE))); - } - - @Test - void isSearchIndexable_falseForShared() { - assertFalse(repository.isSearchIndexable(memory(MemoryVisibility.SHARED))); + @ParameterizedTest + @EnumSource(MemoryVisibility.class) + void isSearchIndexable_trueForEveryVisibility(MemoryVisibility visibility) { + assertTrue(repository.isSearchIndexable(memory(visibility))); } @Test - void isSearchIndexable_falseWhenShareConfigMissing() { - assertFalse( + void isSearchIndexable_trueWhenShareConfigMissing() { + assertTrue( repository.isSearchIndexable( new ContextMemory().withId(UUID.randomUUID()).withName("mem"))); } @Test - void getReindexFilter_restrictsToEntityVisibility() { - assertEquals( - MemoryVisibility.ENTITY.value(), - repository - .getReindexFilter() - .getQueryParams() - .get(ListFilter.MEMORY_SEARCH_VISIBILITY_PARAM)); + void getReindexFilter_doesNotRestrictByVisibility() { + assertTrue(repository.getReindexFilter().getQueryParams().isEmpty()); } @Test - void entityFacade_isSearchIndexable_dispatchesToRepository() { + void entityFacade_isSearchIndexable_trueForRestrictedMemories() { assertTrue(Entity.isSearchIndexable(memory(MemoryVisibility.ENTITY))); - assertFalse(Entity.isSearchIndexable(memory(MemoryVisibility.PRIVATE))); + assertTrue(Entity.isSearchIndexable(memory(MemoryVisibility.PRIVATE))); + assertTrue(Entity.isSearchIndexable(memory(MemoryVisibility.SHARED))); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java index a184ceef26b6..ad5de791cd1e 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java @@ -127,20 +127,13 @@ void getCondition() { } @Test - void test_getMemorySearchVisibilityCondition() { - ListFilter filter = new ListFilter(); - String condition = filter.getCondition("context_memory"); + void getCondition_neverFiltersMemoriesByShareConfigVisibility() { + // The search reindex reads memories through this filter, and a user's own PRIVATE/SHARED + // memories have to reach the index to stay findable in the ContextCenter listing. + String condition = new ListFilter().getCondition("context_memory"); + assertFalse(condition.contains("shareConfig.visibility")); assertFalse(condition.contains("shareConfig'->>'visibility")); - - filter = new ListFilter(); - filter.addQueryParam(ListFilter.MEMORY_SEARCH_VISIBILITY_PARAM, "Entity"); - condition = filter.getCondition("context_memory"); - String bindPlaceholder = ":" + ListFilter.MEMORY_SEARCH_VISIBILITY_PARAM; - assertTrue( - condition.contains( - "JSON_UNQUOTE(JSON_EXTRACT(json, '$.shareConfig.visibility')) = " + bindPlaceholder) - || condition.contains("json->'shareConfig'->>'visibility' = " + bindPlaceholder)); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRequestBuilderMemoryVisibilityTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRequestBuilderMemoryVisibilityTest.java new file mode 100644 index 000000000000..e3946a0d8e7c --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRequestBuilderMemoryVisibilityTest.java @@ -0,0 +1,114 @@ +/* + * Copyright 2024 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.openmetadata.service.search; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import es.co.elastic.clients.elasticsearch.core.SearchRequest; +import es.co.elastic.clients.json.jackson.JacksonJsonpMapper; +import jakarta.json.stream.JsonGenerator; +import java.io.StringWriter; +import java.util.UUID; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; +import org.openmetadata.schema.entity.teams.User; +import org.openmetadata.service.Entity; +import org.openmetadata.service.search.elasticsearch.ElasticSearchRequestBuilder; +import org.openmetadata.service.search.elasticsearch.queries.ElasticQueryBuilderFactory; +import org.openmetadata.service.search.opensearch.OpenSearchRequestBuilder; +import org.openmetadata.service.search.security.ContextMemorySearchVisibility; +import org.openmetadata.service.security.policyevaluator.SubjectContext; + +/** + * ContextMemory privacy is enforced per query, not by keeping restricted memories out of the index, + * so the request builders must fail closed: a search path that never evaluated the caller's subject + * gets org-wide (Entity) memories only. Without this default, every path that forgets to call + * {@code applyContextMemoryVisibility} — {@code /search/fieldQuery}, {@code /search/sourceUrl}, + * deep pagination, the entity-relationship traversals — would return other users' private memories. + */ +@Execution(ExecutionMode.CONCURRENT) +class SearchRequestBuilderMemoryVisibilityTest { + + private static final JacksonJsonpMapper JACKSON_JSONP_MAPPER = new JacksonJsonpMapper(); + + @Test + void elasticBuildRestrictsMemoriesWhenVisibilityWasNeverResolved() { + String json = serialize(new ElasticSearchRequestBuilder().build("table_search_index")); + + assertTrue(json.contains(Entity.CONTEXT_MEMORY), "the memory guard must be present"); + assertTrue(json.contains("\"Entity\""), "only org-wide memories may pass"); + } + + @Test + void elasticBuildLeavesQueryAloneWhenVisibilityWasResolved() { + String json = + serialize( + new ElasticSearchRequestBuilder() + .contextMemoryVisibilityResolved() + .build("table_search_index")); + + assertFalse( + json.contains(Entity.CONTEXT_MEMORY), + "a subject-aware caller already applied its own visibility filter"); + } + + @Test + void openSearchBuildRestrictsMemoriesWhenVisibilityWasNeverResolved() { + String json = new OpenSearchRequestBuilder().build("table_search_index").toJsonString(); + + assertTrue(json.contains(Entity.CONTEXT_MEMORY), "the memory guard must be present"); + assertTrue(json.contains("\"Entity\""), "only org-wide memories may pass"); + } + + @Test + void openSearchBuildLeavesQueryAloneWhenVisibilityWasResolved() { + String json = + new OpenSearchRequestBuilder() + .contextMemoryVisibilityResolved() + .build("table_search_index") + .toJsonString(); + + assertFalse( + json.contains(Entity.CONTEXT_MEMORY), + "a subject-aware caller already applied its own visibility filter"); + } + + @Test + void onlyAnIdentifiedSubjectMayMarkVisibilityResolved() { + // buildVisibilityFilter returns null both for admins and for a null/unidentifiable subject. + // Only the former may mark the request resolved — internal callers that pass a null + // SubjectContext must fall back to the org-wide-only default, not search unfiltered. + ContextMemorySearchVisibility visibility = + new ContextMemorySearchVisibility(new ElasticQueryBuilderFactory()); + + assertFalse(visibility.isSubjectResolvable(null), "a null subject is not resolvable"); + assertFalse( + visibility.isSubjectResolvable(new SubjectContext(new User(), null)), + "a subject whose user carries no id is not resolvable"); + assertTrue( + visibility.isSubjectResolvable( + new SubjectContext(new User().withId(UUID.randomUUID()).withIsAdmin(true), null)), + "an identified admin is resolvable and needs no filter"); + } + + private String serialize(SearchRequest request) { + StringWriter writer = new StringWriter(); + JsonGenerator generator = JACKSON_JSONP_MAPPER.jsonProvider().createGenerator(writer); + request.serialize(generator, JACKSON_JSONP_MAPPER); + generator.close(); + return writer.toString(); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/search/security/ContextMemorySearchVisibilityTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/search/security/ContextMemorySearchVisibilityTest.java index 50e5a709c410..9ef030fa228f 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/search/security/ContextMemorySearchVisibilityTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/search/security/ContextMemorySearchVisibilityTest.java @@ -26,10 +26,12 @@ import jakarta.json.stream.JsonGenerator; import java.io.StringWriter; import java.util.List; +import java.util.Map; import java.util.UUID; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.parallel.Execution; import org.junit.jupiter.api.parallel.ExecutionMode; +import org.openmetadata.schema.entity.context.MemoryVisibility; import org.openmetadata.schema.entity.teams.User; import org.openmetadata.schema.type.EntityReference; import org.openmetadata.service.Entity; @@ -190,6 +192,77 @@ void nullSubjectGetsNoVisibilityFilter() { assertNull(filter, "A missing subject must not silently expose memories"); } + @Test + void orgWideOnlyFilterLeavesNonMemoryDocumentsUntouched() { + DocumentContext json = JsonPath.parse(orgWideOnlyJson()); + + assertFieldExists( + json, + "$.bool.should[0].bool.must_not[?(@.term['entityType'].value=='" + + Entity.CONTEXT_MEMORY + + "')]", + "non-memory documents pass through via a must_not entityType branch"); + } + + @Test + void orgWideOnlyFilterAdmitsOnlyEntityVisibilityMemories() { + DocumentContext json = JsonPath.parse(orgWideOnlyJson()); + + assertFieldExists( + json, + "$.bool.should[1].bool.must[?(@.term['visibility'].value=='Entity')]", + "the memory branch admits Entity-visibility memories"); + assertFieldDoesNotExist( + json, "$..term['owners.id']", "a subject-less path must not match by ownership"); + assertFieldDoesNotExist( + json, "$..terms['sharedWithIds']", "a subject-less path must not match by sharedWithIds"); + } + + @Test + void orgWideOnlyFilterIsShouldOnlySoMinimumShouldMatchDefaultsToOne() { + DocumentContext json = JsonPath.parse(orgWideOnlyJson()); + + assertFieldExists(json, "$.bool.should", "the outer access clause must be should-based"); + assertFieldDoesNotExist( + json, "$.bool.must", "a must sibling would drop minimum_should_match to 0 (fail-open)"); + assertFieldDoesNotExist( + json, "$.bool.filter", "a filter sibling would drop minimum_should_match to 0 (fail-open)"); + } + + @Test + void isOrgWideReadableAdmitsEveryNonMemoryDocument() { + assertTrue(ContextMemorySearchVisibility.isOrgWideReadable(null)); + assertTrue(ContextMemorySearchVisibility.isOrgWideReadable(Map.of())); + assertTrue( + ContextMemorySearchVisibility.isOrgWideReadable( + Map.of("entityType", Entity.TABLE, "name", "orders"))); + } + + @Test + void isOrgWideReadableRejectsRestrictedMemories() { + assertTrue( + ContextMemorySearchVisibility.isOrgWideReadable(memoryDocument(MemoryVisibility.ENTITY))); + assertFalse( + ContextMemorySearchVisibility.isOrgWideReadable(memoryDocument(MemoryVisibility.PRIVATE))); + assertFalse( + ContextMemorySearchVisibility.isOrgWideReadable(memoryDocument(MemoryVisibility.SHARED))); + assertFalse( + ContextMemorySearchVisibility.isOrgWideReadable( + Map.of("entityType", Entity.CONTEXT_MEMORY)), + "a memory with no indexed visibility is not org-wide"); + } + + private Map memoryDocument(MemoryVisibility visibility) { + return Map.of("entityType", Entity.CONTEXT_MEMORY, "visibility", visibility.value()); + } + + private String orgWideOnlyJson() { + OMQueryBuilder filter = + new ContextMemorySearchVisibility(new ElasticQueryBuilderFactory()) + .buildOrgWideOnlyFilter(); + return serializeElasticQuery(((ElasticQueryBuilder) filter).build()); + } + @Test void openSearchBuilderProducesEquivalentFilter() { OMQueryBuilder filter = diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorEmbeddingHandlerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorEmbeddingHandlerTest.java index fb5d3ea4d5bd..7c7bd5ac202f 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorEmbeddingHandlerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/search/vector/VectorEmbeddingHandlerTest.java @@ -54,7 +54,7 @@ void testOnEntityCreatedForSupportedType() { SearchRepository searchRepository = mock(SearchRepository.class); IndexMapping indexMapping = mock(IndexMapping.class); entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); - entityMock.when(() -> Entity.isSearchIndexable(entity)).thenReturn(true); + entityMock.when(() -> Entity.isVectorEmbeddable(entity)).thenReturn(true); when(searchRepository.getIndexMapping("table")).thenReturn(indexMapping); when(searchRepository.getClusterAlias()).thenReturn(""); when(indexMapping.getIndexName("")).thenReturn("table_search_index"); @@ -83,7 +83,7 @@ void testOnEntityUpdatedCallsUpdate() { SearchRepository searchRepository = mock(SearchRepository.class); IndexMapping indexMapping = mock(IndexMapping.class); entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); - entityMock.when(() -> Entity.isSearchIndexable(entity)).thenReturn(true); + entityMock.when(() -> Entity.isVectorEmbeddable(entity)).thenReturn(true); when(searchRepository.getIndexMapping("table")).thenReturn(indexMapping); when(searchRepository.getClusterAlias()).thenReturn(""); when(indexMapping.getIndexName("")).thenReturn("table_search_index"); @@ -95,15 +95,25 @@ void testOnEntityUpdatedCallsUpdate() { } @Test - void testOnEntityUpdatedDeletesChunksForNonIndexableEntity() { + void testOnEntityUpdatedDropsChunksAndEmbeddingForNonEmbeddableEntity() { EntityInterface entity = createMockEntity("contextMemory"); + SearchRepository searchRepository = mock(SearchRepository.class); + IndexMapping indexMapping = mock(IndexMapping.class); try (MockedStatic entityMock = mockStatic(Entity.class)) { - entityMock.when(() -> Entity.isSearchIndexable(entity)).thenReturn(false); + entityMock.when(() -> Entity.isVectorEmbeddable(entity)).thenReturn(false); + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); + when(searchRepository.getIndexMapping("contextMemory")).thenReturn(indexMapping); + when(searchRepository.getClusterAlias()).thenReturn(""); + when(indexMapping.getIndexName("")).thenReturn("context_memory_search_index"); handler.onEntityUpdated(entity, null, subjectContext); verify(vectorIndexService).deleteEntityChunks(entity.getId().toString()); + // The entity doc keeps its embedding across partial merges, so it has to be cleared + // explicitly or a memory flipped to Private stays matchable by kNN. + verify(vectorIndexService) + .clearEntityEmbedding("context_memory_search_index", entity.getId().toString()); verify(vectorIndexService, never()).updateEntityEmbeddings(any(), anyString()); } } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/security/policyevaluator/RuleEvaluatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/security/policyevaluator/RuleEvaluatorTest.java index 585304e6ceb4..398e31b6e353 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/security/policyevaluator/RuleEvaluatorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/security/policyevaluator/RuleEvaluatorTest.java @@ -81,7 +81,7 @@ class RuleEvaluatorTest { @BeforeAll public static void setup() { - TeamRepository teamRepository = mock(TeamRepository.class); + TeamRepository teamRepository = stubIndexingPolicy(mock(TeamRepository.class)); Entity.registerEntity(Team.class, Entity.TEAM, teamRepository); Mockito.when(teamRepository.find(any(UUID.class), any(Include.class))) .thenAnswer( @@ -127,24 +127,25 @@ public static void setup() { new ImmutablePair<>(Entity.TEAM, i.getArgument(1))), Team.class)); - tableRepository = mock(TableRepository.class); + tableRepository = stubIndexingPolicy(mock(TableRepository.class)); Entity.registerEntity(Table.class, Entity.TABLE, tableRepository); Mockito.when(tableRepository.getAllTags(any())) .thenAnswer((Answer>) invocationOnMock -> table.getTags()); Mockito.when(tableRepository.getEntityType()).thenReturn(Entity.TABLE); Mockito.when(tableRepository.isSupportsOwners()).thenReturn(Boolean.TRUE); - DatabaseRepository databaseRepository = mock(DatabaseRepository.class); + DatabaseRepository databaseRepository = stubIndexingPolicy(mock(DatabaseRepository.class)); Mockito.when(databaseRepository.getEntityType()).thenReturn(Entity.DATABASE); Mockito.when(databaseRepository.isSupportsOwners()).thenReturn(Boolean.TRUE); Entity.registerEntity(Database.class, Entity.DATABASE, databaseRepository); - DatabaseSchemaRepository databaseSchemaRepository = mock(DatabaseSchemaRepository.class); + DatabaseSchemaRepository databaseSchemaRepository = + stubIndexingPolicy(mock(DatabaseSchemaRepository.class)); Mockito.when(databaseSchemaRepository.getEntityType()).thenReturn(Entity.DATABASE_SCHEMA); Mockito.when(databaseSchemaRepository.isSupportsOwners()).thenReturn(Boolean.TRUE); Entity.registerEntity(DatabaseSchema.class, Entity.DATABASE_SCHEMA, databaseSchemaRepository); - DomainRepository domainRepository = mock(DomainRepository.class); + DomainRepository domainRepository = stubIndexingPolicy(mock(DomainRepository.class)); Mockito.when(domainRepository.getEntityType()).thenReturn(Entity.DOMAIN); Mockito.when(domainRepository.isSupportsOwners()).thenReturn(Boolean.TRUE); Entity.registerEntity(Domain.class, Entity.DOMAIN, domainRepository); @@ -156,7 +157,8 @@ public static void setup() { new ImmutablePair<>(Entity.DOMAIN, i.getArgument(1))), Domain.class)); - DataProductRepository dataProductRepository = mock(DataProductRepository.class); + DataProductRepository dataProductRepository = + stubIndexingPolicy(mock(DataProductRepository.class)); Mockito.when(dataProductRepository.getEntityType()).thenReturn(Entity.DATA_PRODUCT); Mockito.when(dataProductRepository.isSupportsOwners()).thenReturn(Boolean.TRUE); Entity.registerEntity(DataProduct.class, Entity.DATA_PRODUCT, dataProductRepository); @@ -344,7 +346,7 @@ void test_isOwner() { @Test void test_isReviewer() { - GlossaryRepository glossaryRepository = mock(GlossaryRepository.class); + GlossaryRepository glossaryRepository = stubIndexingPolicy(mock(GlossaryRepository.class)); Entity.registerEntity(Glossary.class, Entity.GLOSSARY, glossaryRepository); User reviewer = new User().withId(UUID.randomUUID()).withName("reviewerUser"); @@ -1238,4 +1240,16 @@ void resetContext() { evaluationContext = new StandardEvaluationContext(ruleEvaluator); LOG.info("Context reset to default state after test completion."); } + + /** + * These repository registrations are global and never torn down, so each stand-in has to answer + * the indexing policy hooks the way a real repository does. A bare mock answers false, which + * would make {@link Entity#isSearchIndexable} report every entity of that type as non-indexable + * for the rest of the JVM and break unrelated tests that run later. + */ + private static > T stubIndexingPolicy(T repository) { + Mockito.doReturn(true).when(repository).isSearchIndexable(any()); + Mockito.doReturn(true).when(repository).isVectorEmbeddable(any()); + return repository; + } } From a049f342ed8c9192042b31529ce98b29b7aca690 Mon Sep 17 00:00:00 2001 From: pmbrull Date: Tue, 28 Jul 2026 17:29:22 +0200 Subject: [PATCH 2/3] test(playwright): re-enable ContextCenter memory suites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #30528 marked these describe.fixme after the memory listing stopped returning PRIVATE/SHARED entries. That was the #29541 regression this branch fixes, not a test problem — the suites build Private, Shared and Entity fixtures and assert what a dataConsumer sees. CustomizeDetailPage stays fixme; it is unrelated to this change. Refs #30528 Co-Authored-By: Claude --- .../ui/playwright/e2e/Features/ContextCenterDashboard.spec.ts | 2 +- .../ui/playwright/e2e/Features/ContextCenterMemories.spec.ts | 2 +- .../ui/playwright/e2e/Features/ContextCenterPermission.spec.ts | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterDashboard.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterDashboard.spec.ts index 12a75e628da8..3043e877a03f 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterDashboard.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterDashboard.spec.ts @@ -51,7 +51,7 @@ test.use({ storageState: 'playwright/.auth/admin.json' }); // ─── Suite ──────────────────────────────────────────────────────────────────── -test.describe.fixme('Context Center - Dashboard', () => { +test.describe('Context Center - Dashboard', () => { let dataAsset: TopicClass; test.beforeAll(async ({ browser }) => { diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterMemories.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterMemories.spec.ts index 9fb46667877a..22b3810f1280 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterMemories.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterMemories.spec.ts @@ -49,7 +49,7 @@ const ENTITY_MEMORY_TITLE = `CC Memory Entity ${uuid()}`; // ─── Test suite ─────────────────────────────────────────────────────────────── -test.describe.fixme( +test.describe( 'Context Center - Memories', { tag: ['@Features', '@Governance'] }, () => { diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterPermission.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterPermission.spec.ts index 8793b78bf81c..282fcf081e90 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterPermission.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ContextCenterPermission.spec.ts @@ -1696,7 +1696,7 @@ test.describe('Context Center Permissions', () => { // ─── Memories Permissions (includes isOwner matrix) ───────────────────── - test.describe.fixme('Memories Permissions', () => { + test.describe('Memories Permissions', () => { test('user with view-only permission sees no Add Memory button, no row actions, and no modal action buttons', async ({ viewOnlyPage, }) => { From 3c57cc4d8da8788f1c35afa8a51044e7c186bbbe Mon Sep 17 00:00:00 2001 From: pmbrull Date: Tue, 28 Jul 2026 18:14:27 +0200 Subject: [PATCH 3/3] fix(search): resolve subject for entityTypeCounts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit getEntityTypeCounts built through the request builder without ever resolving the caller, so the fail-closed org-wide-only default always applied — even for an admin or a memory's owner. The contextMemory count then disagreed with the search-backed /contextCenter/memories listing, which does show those callers their PRIVATE and SHARED memories. SearchResource already had the SubjectContext and used it only for domains; thread it down to the aggregation managers and apply the same per-subject filter the listing uses. The subject-less overload stays and still fails closed, so an unidentifiable caller is unaffected. Co-Authored-By: Claude --- .../resources/search/SearchResource.java | 2 +- .../search/AggregationManagementClient.java | 13 +++++++ .../service/search/SearchRepository.java | 5 +++ .../ElasticSearchAggregationManager.java | 35 ++++++++++++++++--- .../elasticsearch/ElasticSearchClient.java | 6 ++++ .../OpenSearchAggregationManager.java | 35 ++++++++++++++++--- .../search/opensearch/OpenSearchClient.java | 6 ++++ .../search/SearchRepositoryBehaviorTest.java | 4 +++ 8 files changed, 97 insertions(+), 9 deletions(-) diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/search/SearchResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/search/SearchResource.java index ce0a2fb8ceaf..3d4afa696e39 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/search/SearchResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/search/SearchResource.java @@ -891,7 +891,7 @@ public Response getEntityTypeCounts( .withPostFilter(postFilter) .withDomains(domains); - return searchRepository.getEntityTypeCounts(request, index); + return searchRepository.getEntityTypeCounts(request, index, subjectContext); } @POST diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/AggregationManagementClient.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/AggregationManagementClient.java index 105e53e2fa70..60ce737ff902 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/AggregationManagementClient.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/AggregationManagementClient.java @@ -77,4 +77,17 @@ JsonObject aggregate( * @throws IOException if the aggregation operation fails */ Response getEntityTypeCounts(SearchRequest request, String index) throws IOException; + + /** + * Same as {@link #getEntityTypeCounts(SearchRequest, String)} but resolves per-subject + * visibility, so a ContextMemory count matches what the same caller sees in the + * {@code /contextCenter/memories} listing instead of being capped at the org-wide memories. + * Defaults to the subject-less overload, which fails closed. + * + * @param subjectContext the caller, used to evaluate per-document visibility rules + */ + default Response getEntityTypeCounts( + SearchRequest request, String index, SubjectContext subjectContext) throws IOException { + return getEntityTypeCounts(request, index); + } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java index bfbeaf7e0a12..78fa77cf962c 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java @@ -4010,6 +4010,11 @@ public Response getEntityTypeCounts(SearchRequest request, String index) throws return searchClient.getEntityTypeCounts(request, index); } + public Response getEntityTypeCounts( + SearchRequest request, String index, SubjectContext subjectContext) throws IOException { + return searchClient.getEntityTypeCounts(request, index, subjectContext); + } + public JsonObject aggregate( String query, String entityType, SearchAggregation searchAggregation, SearchListFilter filter) throws IOException { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java index 6d032c7ae910..0d5a5d36bb0d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchAggregationManager.java @@ -71,10 +71,26 @@ public ElasticSearchAggregationManager( } /** - * ANDs the org-wide-only ContextMemory filter into an aggregation query. Aggregations run without - * a {@link org.openmetadata.service.security.policyevaluator.SubjectContext}, so they cannot tell - * whose restricted memory a document is and must fail closed: only memories everyone may read are - * aggregated. Non-memory documents are unaffected. + * Resolves ContextMemory visibility for {@code subjectContext} so the counts match what that + * caller sees in the listing. Leaving the request unresolved (no identifiable subject) lets + * {@code ElasticSearchRequestBuilder#build} apply its org-wide-only default instead. + */ + private void applyContextMemoryVisibility( + SubjectContext subjectContext, ElasticSearchRequestBuilder requestBuilder) { + OMQueryBuilder visibilityBuilder = MEMORY_VISIBILITY.buildVisibilityFilter(subjectContext); + if (visibilityBuilder != null) { + requestBuilder.filter(((ElasticQueryBuilder) visibilityBuilder).buildV2()); + } + if (MEMORY_VISIBILITY.isSubjectResolvable(subjectContext)) { + requestBuilder.contextMemoryVisibilityResolved(); + } + } + + /** + * ANDs the org-wide-only ContextMemory filter into an aggregation query run without an + * identifiable subject. Such a query cannot tell whose restricted memory a document is and must + * fail closed: only memories everyone may read are aggregated. Non-memory documents are + * unaffected. */ private Query restrictToOrgWideMemories(Query query) { Query memoryFilter = @@ -584,6 +600,15 @@ private String serializeSearchResponseFallback( @Override public Response getEntityTypeCounts( org.openmetadata.schema.search.SearchRequest request, String index) throws IOException { + return getEntityTypeCounts(request, index, null); + } + + @Override + public Response getEntityTypeCounts( + org.openmetadata.schema.search.SearchRequest request, + String index, + SubjectContext subjectContext) + throws IOException { if (!isClientAvailable) { LOG.error("ElasticSearch client is not available. Cannot perform get entity type counts"); throw new IOException("ElasticSearch client is not available"); @@ -653,6 +678,8 @@ public Response getEntityTypeCounts( String resolvedIndex = Entity.getSearchRepository().getIndexOrAliasName(index != null ? index : "all"); + applyContextMemoryVisibility(subjectContext, requestBuilder); + // Build and execute search SearchRequest searchRequest = requestBuilder.build(resolvedIndex); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java index b99502e9261e..31e68be61756 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/elasticsearch/ElasticSearchClient.java @@ -553,6 +553,12 @@ public Response getEntityTypeCounts(SearchRequest request, String index) throws return aggregationManager.getEntityTypeCounts(request, index); } + @Override + public Response getEntityTypeCounts( + SearchRequest request, String index, SubjectContext subjectContext) throws IOException { + return aggregationManager.getEntityTypeCounts(request, index, subjectContext); + } + @Override public Response aggregate(AggregationRequest request) throws IOException { return aggregationManager.aggregate(request); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java index 1e58d8c18f5d..4d63d2c86e62 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchAggregationManager.java @@ -70,10 +70,26 @@ public OpenSearchAggregationManager( } /** - * ANDs the org-wide-only ContextMemory filter into an aggregation query. Aggregations run without - * a {@link org.openmetadata.service.security.policyevaluator.SubjectContext}, so they cannot tell - * whose restricted memory a document is and must fail closed: only memories everyone may read are - * aggregated. Non-memory documents are unaffected. + * Resolves ContextMemory visibility for {@code subjectContext} so the counts match what that + * caller sees in the listing. Leaving the request unresolved (no identifiable subject) lets + * {@code OpenSearchRequestBuilder#build} apply its org-wide-only default instead. + */ + private void applyContextMemoryVisibility( + SubjectContext subjectContext, OpenSearchRequestBuilder requestBuilder) { + OMQueryBuilder visibilityBuilder = MEMORY_VISIBILITY.buildVisibilityFilter(subjectContext); + if (visibilityBuilder != null) { + requestBuilder.filter(((OpenSearchQueryBuilder) visibilityBuilder).buildV2()); + } + if (MEMORY_VISIBILITY.isSubjectResolvable(subjectContext)) { + requestBuilder.contextMemoryVisibilityResolved(); + } + } + + /** + * ANDs the org-wide-only ContextMemory filter into an aggregation query run without an + * identifiable subject. Such a query cannot tell whose restricted memory a document is and must + * fail closed: only memories everyone may read are aggregated. Non-memory documents are + * unaffected. */ private Query restrictToOrgWideMemories(Query query) { Query memoryFilter = @@ -475,6 +491,15 @@ public JsonObject aggregate( @Override public Response getEntityTypeCounts( org.openmetadata.schema.search.SearchRequest request, String index) throws IOException { + return getEntityTypeCounts(request, index, null); + } + + @Override + public Response getEntityTypeCounts( + org.openmetadata.schema.search.SearchRequest request, + String index, + SubjectContext subjectContext) + throws IOException { if (!isClientAvailable) { LOG.error("OpenSearch client is not available. Cannot perform get entity type counts"); throw new IOException("OpenSearch client is not available"); @@ -545,6 +570,8 @@ public Response getEntityTypeCounts( String resolvedIndex = Entity.getSearchRepository().getIndexOrAliasName(index != null ? index : "all"); + applyContextMemoryVisibility(subjectContext, requestBuilder); + // Build and execute search SearchRequest searchRequest = requestBuilder.build(resolvedIndex); Timer.Sample searchTimerSample = RequestLatencyContext.startSearchOperation(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java index ae9d2acab8e1..1545a4da4a8b 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/search/opensearch/OpenSearchClient.java @@ -519,6 +519,12 @@ public Response getEntityTypeCounts(SearchRequest request, String index) throws return aggregationManager.getEntityTypeCounts(request, index); } + @Override + public Response getEntityTypeCounts( + SearchRequest request, String index, SubjectContext subjectContext) throws IOException { + return aggregationManager.getEntityTypeCounts(request, index, subjectContext); + } + @Override public Response aggregate(AggregationRequest request) throws IOException { return aggregationManager.aggregate(request); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRepositoryBehaviorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRepositoryBehaviorTest.java index c6b84f8e9d72..6e3f4ac46e99 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRepositoryBehaviorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/search/SearchRepositoryBehaviorTest.java @@ -3062,6 +3062,7 @@ void requestAndAggregationWrappersDelegateToSearchClient() throws IOException { when(searchClient.searchBySourceUrl("https://src")).thenReturn(response); when(searchClient.aggregate(aggregationRequest)).thenReturn(response); when(searchClient.getEntityTypeCounts(request, "global")).thenReturn(response); + when(searchClient.getEntityTypeCounts(request, "global", subjectContext)).thenReturn(response); when(searchClient.getQueryCostRecords("service")).thenReturn(queryCostResult); assertSame(response, repository.previewSearch(request, subjectContext, null)); @@ -3073,6 +3074,9 @@ void requestAndAggregationWrappersDelegateToSearchClient() throws IOException { assertSame(response, repository.searchBySourceUrl("https://src")); assertSame(response, repository.aggregate(aggregationRequest)); assertSame(response, repository.getEntityTypeCounts(request, "global")); + // The subject-aware overload is what SearchResource calls, so a ContextMemory count matches + // what the same caller sees in the listing instead of being capped at org-wide memories. + assertSame(response, repository.getEntityTypeCounts(request, "global", subjectContext)); assertSame(queryCostResult, repository.getQueryCostRecords("service")); }