diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
index 8f9874bc850495..c4ab026cdff470 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -947,6 +947,10 @@ private void configureContainerResolvers(final RuntimeWiring.Builder builder) {
typeWiring ->
typeWiring
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService))
.dataFetcher("entities", new ContainerEntitiesResolver(entityClient))
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
@@ -1749,6 +1753,10 @@ private void configureDatasetResolvers(final RuntimeWiring.Builder builder) {
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.datasetType))
.dataFetcher(
"lineage",
@@ -1975,7 +1983,11 @@ private void configureGlossaryTermResolvers(final RuntimeWiring.Builder builder)
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
"aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
- .dataFetcher("exists", new EntityExistsResolver(entityService)));
+ .dataFetcher("exists", new EntityExistsResolver(entityService))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService)));
}
private void configureGlossaryNodeResolvers(final RuntimeWiring.Builder builder) {
@@ -1991,7 +2003,11 @@ private void configureGlossaryNodeResolvers(final RuntimeWiring.Builder builder)
"glossaryChildrenSearch",
new GlossaryChildrenSearchResolver(this.entityClient, this.viewService))
.dataFetcher(
- "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)));
+ "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService)));
}
private void configureSchemaFieldResolvers(final RuntimeWiring.Builder builder) {
@@ -2216,6 +2232,10 @@ private void configureDashboardResolvers(final RuntimeWiring.Builder builder) {
typeWiring ->
typeWiring
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dashboardType))
.dataFetcher(
"lineage",
@@ -2347,6 +2367,10 @@ private void configureChartResolvers(final RuntimeWiring.Builder builder) {
typeWiring ->
typeWiring
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.chartType))
.dataFetcher(
"aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
@@ -2547,6 +2571,10 @@ private void configureDataJobResolvers(final RuntimeWiring.Builder builder) {
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dataJobType))
.dataFetcher(
"lineage",
@@ -2650,6 +2678,10 @@ private void configureDataFlowResolvers(final RuntimeWiring.Builder builder) {
typeWiring ->
typeWiring
.dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.dataFlowType))
.dataFetcher(
"lineage",
@@ -2706,6 +2738,10 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher(
"browsePaths", new EntityBrowsePathsResolver(this.mlFeatureTableType))
.dataFetcher(
@@ -2798,6 +2834,10 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher("browsePaths", new EntityBrowsePathsResolver(this.mlModelType))
.dataFetcher(
"lineage",
@@ -2846,6 +2886,10 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher(
"browsePaths", new EntityBrowsePathsResolver(this.mlModelGroupType))
.dataFetcher(
@@ -2879,6 +2923,10 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher(
"lineage",
new EntityLineageResultResolver(
@@ -2905,6 +2953,10 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
typeWiring
.dataFetcher(
"relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge
+ .RelatedDocumentsResolver(documentService, entityClient, groupService))
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
"lineage",
@@ -2951,7 +3003,11 @@ private void configureDomainResolvers(final RuntimeWiring.Builder builder) {
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
"aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
- .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)));
+ .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService)));
builder.type(
"DomainAssociation",
typeWiring ->
@@ -3043,7 +3099,11 @@ private void configureDataProductResolvers(final RuntimeWiring.Builder builder)
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
"aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
- .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)));
+ .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService)));
}
private void configureApplicationResolvers(final RuntimeWiring.Builder builder) {
@@ -3054,7 +3114,11 @@ private void configureApplicationResolvers(final RuntimeWiring.Builder builder)
.dataFetcher("privileges", new EntityPrivilegesResolver(entityClient))
.dataFetcher(
"aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry))
- .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient)));
+ .dataFetcher("relationships", new EntityRelationshipsResultResolver(graphClient))
+ .dataFetcher(
+ "relatedDocuments",
+ new com.linkedin.datahub.graphql.resolvers.knowledge.RelatedDocumentsResolver(
+ documentService, entityClient, groupService)));
builder.type(
"ApplicationAssociation",
typeWiring ->
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtils.java
new file mode 100644
index 00000000000000..e4db48621f0585
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtils.java
@@ -0,0 +1,72 @@
+package com.linkedin.datahub.graphql.resolvers.knowledge;
+
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
+import com.linkedin.metadata.query.filter.Criterion;
+import com.linkedin.metadata.query.filter.CriterionArray;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.utils.CriterionUtils;
+import java.util.ArrayList;
+import java.util.List;
+import javax.annotation.Nonnull;
+
+/**
+ * Utility class for building document search filters with ownership constraints. Shared logic for
+ * SearchDocumentsResolver and ContextDocumentsResolver.
+ */
+public class DocumentSearchFilterUtils {
+
+ private DocumentSearchFilterUtils() {
+ // Utility class - prevent instantiation
+ }
+
+ /**
+ * Builds a combined filter with ownership constraints. The filter structure is: (user-filters AND
+ * PUBLISHED AND NOT-DRAFT) OR (user-filters AND UNPUBLISHED AND owned-by-user-or-groups AND
+ * NOT-DRAFT)
+ *
+ *
Drafts (documents with draftOf field set) are excluded by default from search results. They
+ * should only be accessed directly by URN or through the DocumentDrafts resolver.
+ *
+ * @param baseCriteria The base user criteria (without state filtering)
+ * @param userAndGroupUrns List of URNs for the current user and their groups
+ * @return The combined filter
+ */
+ @Nonnull
+ public static Filter buildCombinedFilter(
+ @Nonnull List baseCriteria, @Nonnull List userAndGroupUrns) {
+ // Build two conjunctive clauses:
+ // 1. Base filters AND PUBLISHED AND NOT-DRAFT
+ // 2. Base filters AND UNPUBLISHED AND owned-by-user-or-groups AND NOT-DRAFT
+
+ List orClauses = new ArrayList<>();
+
+ // Create criterion to exclude drafts (draftOf field must be null/not set)
+ Criterion notDraftCriterion = new Criterion();
+ notDraftCriterion.setField("draftOf");
+ notDraftCriterion.setCondition(Condition.IS_NULL);
+
+ // Clause 1: Published documents (with user filters, excluding drafts)
+ List publishedCriteria = new ArrayList<>(baseCriteria);
+ publishedCriteria.add(CriterionUtils.buildCriterion("state", Condition.EQUAL, "PUBLISHED"));
+ publishedCriteria.add(notDraftCriterion);
+ orClauses.add(new ConjunctiveCriterion().setAnd(new CriterionArray(publishedCriteria)));
+
+ // Clause 2: Unpublished documents owned by user or their groups (with user filters, excluding
+ // drafts)
+ List unpublishedOwnedCriteria = new ArrayList<>(baseCriteria);
+ unpublishedOwnedCriteria.add(
+ CriterionUtils.buildCriterion("state", Condition.EQUAL, "UNPUBLISHED"));
+ unpublishedOwnedCriteria.add(
+ CriterionUtils.buildCriterion("owners", Condition.EQUAL, userAndGroupUrns));
+ // Create a new criterion instance for the second clause to avoid sharing mutable state
+ Criterion notDraftCriterion2 = new Criterion();
+ notDraftCriterion2.setField("draftOf");
+ notDraftCriterion2.setCondition(Condition.IS_NULL);
+ unpublishedOwnedCriteria.add(notDraftCriterion2);
+ orClauses.add(new ConjunctiveCriterion().setAnd(new CriterionArray(unpublishedOwnedCriteria)));
+
+ return new Filter().setOr(new ConjunctiveCriterionArray(orClauses));
+ }
+}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolver.java
new file mode 100644
index 00000000000000..fd605b8f4e5f05
--- /dev/null
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolver.java
@@ -0,0 +1,239 @@
+package com.linkedin.datahub.graphql.resolvers.knowledge;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+
+import com.datahub.authentication.group.GroupService;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
+import com.linkedin.datahub.graphql.generated.Document;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.RelatedDocumentsInput;
+import com.linkedin.datahub.graphql.generated.RelatedDocumentsResult;
+import com.linkedin.datahub.graphql.types.knowledge.DocumentMapper;
+import com.linkedin.datahub.graphql.types.mappers.MapperUtils;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.Criterion;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
+import com.linkedin.metadata.query.filter.SortOrder;
+import com.linkedin.metadata.search.SearchEntity;
+import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.service.DocumentService;
+import com.linkedin.metadata.utils.CriterionUtils;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Resolver for fetching context documents related to an entity. Returns documents where the
+ * entity's URN appears in the relatedAssets field.
+ *
+ *
Filtering behavior: - PUBLISHED documents are shown to all users - UNPUBLISHED documents are
+ * only shown if owned by the current user or a group they belong to
+ */
+@Slf4j
+@RequiredArgsConstructor
+public class RelatedDocumentsResolver
+ implements DataFetcher> {
+
+ private static final Integer DEFAULT_START = 0;
+ private static final Integer DEFAULT_COUNT = 100; // Limit to 100 most recently updated documents
+ private static final String DEFAULT_QUERY = "*";
+
+ private final DocumentService _documentService;
+ private final EntityClient _entityClient;
+ private final GroupService _groupService;
+
+ @Override
+ public CompletableFuture get(final DataFetchingEnvironment environment)
+ throws Exception {
+ final QueryContext context = environment.getContext();
+
+ return GraphQLConcurrencyUtils.supplyAsync(
+ () -> {
+ log.debug("RelatedDocumentsResolver.supplyAsync lambda executing");
+ try {
+ // Get the parent entity from the source
+ final Entity parentEntity = (Entity) environment.getSource();
+ if (parentEntity == null || parentEntity.getUrn() == null) {
+ log.error("Parent entity is null or missing URN");
+ throw new RuntimeException("Parent entity URN is required");
+ }
+
+ final String parentEntityUrn = parentEntity.getUrn();
+ log.debug("Fetching related documents for entity: {}", parentEntityUrn);
+
+ final RelatedDocumentsInput input =
+ bindArgument(environment.getArgument("input"), RelatedDocumentsInput.class);
+ final Integer start = input.getStart() == null ? DEFAULT_START : input.getStart();
+ final Integer count = input.getCount() == null ? DEFAULT_COUNT : input.getCount();
+ // Get current user and their groups for ownership filtering
+ final Urn currentUserUrn = Urn.createFromString(context.getActorUrn());
+ final List userGroupUrns =
+ _groupService.getGroupsForUser(context.getOperationContext(), currentUserUrn);
+ final List userAndGroupUrns = new ArrayList<>();
+ userAndGroupUrns.add(currentUserUrn.toString());
+ userGroupUrns.forEach(groupUrn -> userAndGroupUrns.add(groupUrn.toString()));
+
+ // Build base user criteria from input
+ List baseUserCriteria = buildBaseUserCriteria(input, parentEntityUrn);
+
+ // Build filter that combines user filters with ownership constraints
+ // Filter logic: (PUBLISHED) OR (UNPUBLISHED AND owned-by-user-or-groups)
+ Filter filter =
+ DocumentSearchFilterUtils.buildCombinedFilter(baseUserCriteria, userAndGroupUrns);
+
+ // Step 1: Search using service to get URNs
+ // Sort by lastModifiedAt descending to show most recently updated documents first
+ final SortCriterion sortCriterion =
+ new SortCriterion().setField("lastModifiedAt").setOrder(SortOrder.DESCENDING);
+ final SearchResult gmsResult;
+ try {
+ gmsResult =
+ _documentService.searchDocuments(
+ context.getOperationContext(),
+ DEFAULT_QUERY,
+ filter,
+ sortCriterion,
+ start,
+ count);
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to search context documents", e);
+ }
+
+ // Step 2: Extract URNs from search results
+ final List documentUrns =
+ gmsResult.getEntities().stream()
+ .map(SearchEntity::getEntity)
+ .collect(Collectors.toList());
+
+ // Step 3: Batch hydrate/resolve the Document entities
+ final Map entities =
+ _entityClient.batchGetV2(
+ context.getOperationContext(),
+ Constants.DOCUMENT_ENTITY_NAME,
+ new HashSet<>(documentUrns),
+ com.linkedin.datahub.graphql.types.knowledge.DocumentType.ASPECTS_TO_FETCH);
+
+ // Step 4: Map entities in the same order as search results
+ final List orderedEntityResponses = new ArrayList<>();
+ for (Urn urn : documentUrns) {
+ orderedEntityResponses.add(entities.getOrDefault(urn, null));
+ }
+
+ // Step 5: Convert to GraphQL Document objects
+ final List documents =
+ orderedEntityResponses.stream()
+ .filter(entityResponse -> entityResponse != null)
+ .map(entityResponse -> DocumentMapper.map(context, entityResponse))
+ .collect(Collectors.toList());
+
+ // Step 6: Build the result
+ final RelatedDocumentsResult result = new RelatedDocumentsResult();
+ result.setStart(gmsResult.getFrom());
+ result.setCount(gmsResult.getPageSize());
+ result.setTotal(gmsResult.getNumEntities());
+ result.setDocuments(documents);
+
+ // Map facets
+ if (gmsResult.getMetadata() != null
+ && gmsResult.getMetadata().getAggregations() != null) {
+ result.setFacets(
+ gmsResult.getMetadata().getAggregations().stream()
+ .map(facet -> MapperUtils.mapFacet(context, facet))
+ .collect(Collectors.toList()));
+ } else {
+ result.setFacets(Collections.emptyList());
+ }
+
+ return result;
+ } catch (Exception e) {
+ String entityUrn = null;
+ try {
+ final Entity parentEntity = (Entity) environment.getSource();
+ entityUrn = parentEntity != null ? parentEntity.getUrn() : "unknown";
+ } catch (Exception ignored) {
+ // Ignore errors getting entity URN for logging
+ }
+ log.error(
+ "Failed to fetch related documents for entity {}: {}",
+ entityUrn,
+ e.getMessage(),
+ e);
+ throw new RuntimeException("Failed to fetch related documents", e);
+ }
+ },
+ this.getClass().getSimpleName(),
+ "get");
+ }
+
+ /**
+ * Builds the base user criteria from the input (excludes state filtering). These criteria are
+ * common to both published and unpublished document searches. Automatically adds the
+ * relatedAssets filter with the parent entity's URN.
+ *
+ * @param input The context documents input
+ * @param parentEntityUrn The URN of the parent entity (automatically added to relatedAssets
+ * filter)
+ * @return List of criteria
+ */
+ private List buildBaseUserCriteria(
+ RelatedDocumentsInput input, String parentEntityUrn) {
+ List criteria = new ArrayList<>();
+
+ // Automatically add relatedAssets filter with parent entity URN
+ criteria.add(
+ CriterionUtils.buildCriterion(
+ "relatedAssets", Condition.EQUAL, Collections.singletonList(parentEntityUrn)));
+
+ // Add parent documents filter if provided
+ if (input.getParentDocuments() != null && !input.getParentDocuments().isEmpty()) {
+ criteria.add(
+ CriterionUtils.buildCriterion(
+ "parentDocument", Condition.EQUAL, input.getParentDocuments()));
+ } else if (input.getRootOnly() != null && input.getRootOnly()) {
+ // Filter for root-level documents only (no parent)
+ Criterion noParentCriterion = new Criterion();
+ noParentCriterion.setField("parentDocument");
+ noParentCriterion.setCondition(Condition.IS_NULL);
+ criteria.add(noParentCriterion);
+ }
+
+ // Add types filter if provided (now using subTypes aspect)
+ if (input.getTypes() != null && !input.getTypes().isEmpty()) {
+ criteria.add(CriterionUtils.buildCriterion("subTypes", Condition.EQUAL, input.getTypes()));
+ }
+
+ // Add domains filter if provided
+ if (input.getDomains() != null && !input.getDomains().isEmpty()) {
+ criteria.add(CriterionUtils.buildCriterion("domains", Condition.EQUAL, input.getDomains()));
+ }
+
+ // NOTE: State filtering is handled in DocumentSearchFilterUtils.buildCombinedFilter with
+ // ownership logic
+ // Do not add state filters here
+
+ // Add source type filter if provided (if null, search all)
+ if (input.getSourceType() != null) {
+ criteria.add(
+ CriterionUtils.buildCriterion(
+ "sourceType",
+ Condition.EQUAL,
+ Collections.singletonList(input.getSourceType().toString())));
+ }
+
+ return criteria;
+ }
+}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolver.java
index a575905d2e9ce5..a4478325ac5bb1 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolver.java
@@ -9,7 +9,6 @@
import com.linkedin.datahub.graphql.generated.Document;
import com.linkedin.datahub.graphql.generated.SearchDocumentsInput;
import com.linkedin.datahub.graphql.generated.SearchDocumentsResult;
-import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
import com.linkedin.datahub.graphql.types.knowledge.DocumentMapper;
import com.linkedin.datahub.graphql.types.mappers.MapperUtils;
import com.linkedin.entity.EntityResponse;
@@ -79,7 +78,9 @@ public CompletableFuture get(final DataFetchingEnvironmen
// Build filter that combines user filters with ownership constraints
// Filter logic: (PUBLISHED) OR (UNPUBLISHED AND owned-by-user-or-groups)
- Filter filter = buildCombinedFilter(input, userAndGroupUrns);
+ List baseUserCriteria = buildBaseUserCriteria(input);
+ Filter filter =
+ DocumentSearchFilterUtils.buildCombinedFilter(baseUserCriteria, userAndGroupUrns);
// Step 1: Search using service to get URNs
final SearchResult gmsResult;
@@ -146,46 +147,6 @@ public CompletableFuture get(final DataFetchingEnvironmen
"get");
}
- /**
- * Builds a combined filter with ownership constraints. The filter structure is: (user-filters AND
- * PUBLISHED) OR (user-filters AND UNPUBLISHED AND owned-by-user-or-groups)
- *
- * @param input The search input from the user
- * @param userAndGroupUrns List of URNs for the current user and their groups
- * @return The combined filter
- */
- private Filter buildCombinedFilter(SearchDocumentsInput input, List userAndGroupUrns) {
- // Build the base user filters (without state)
- List baseUserCriteria = buildBaseUserCriteria(input);
-
- // Build two conjunctive clauses:
- // 1. Base filters AND PUBLISHED
- // 2. Base filters AND UNPUBLISHED AND owned-by-user-or-groups
-
- List orClauses = new ArrayList<>();
-
- // Clause 1: Published documents (with user filters)
- List publishedCriteria = new ArrayList<>(baseUserCriteria);
- publishedCriteria.add(CriterionUtils.buildCriterion("state", Condition.EQUAL, "PUBLISHED"));
- orClauses.add(
- new com.linkedin.metadata.query.filter.ConjunctiveCriterion()
- .setAnd(new com.linkedin.metadata.query.filter.CriterionArray(publishedCriteria)));
-
- // Clause 2: Unpublished documents owned by user or their groups (with user filters)
- List unpublishedOwnedCriteria = new ArrayList<>(baseUserCriteria);
- unpublishedOwnedCriteria.add(
- CriterionUtils.buildCriterion("state", Condition.EQUAL, "UNPUBLISHED"));
- unpublishedOwnedCriteria.add(
- CriterionUtils.buildCriterion("owners", Condition.EQUAL, userAndGroupUrns));
- orClauses.add(
- new com.linkedin.metadata.query.filter.ConjunctiveCriterion()
- .setAnd(
- new com.linkedin.metadata.query.filter.CriterionArray(unpublishedOwnedCriteria)));
-
- return new Filter()
- .setOr(new com.linkedin.metadata.query.filter.ConjunctiveCriterionArray(orClauses));
- }
-
/**
* Builds the base user criteria from the search input (excludes state filtering). These criteria
* are common to both published and unpublished document searches.
@@ -193,16 +154,11 @@ private Filter buildCombinedFilter(SearchDocumentsInput input, List user
private List buildBaseUserCriteria(SearchDocumentsInput input) {
List criteria = new ArrayList<>();
- // Add parent document filter if provided
- // If parentDocuments (plural) is provided, use it; otherwise fall back to single parentDocument
+ // Add parent documents filter if provided
if (input.getParentDocuments() != null && !input.getParentDocuments().isEmpty()) {
criteria.add(
CriterionUtils.buildCriterion(
"parentDocument", Condition.EQUAL, input.getParentDocuments()));
- } else if (input.getParentDocument() != null) {
- criteria.add(
- CriterionUtils.buildCriterion(
- "parentDocument", Condition.EQUAL, input.getParentDocument()));
} else if (input.getRootOnly() != null && input.getRootOnly()) {
// Filter for root-level documents only (no parent)
Criterion noParentCriterion = new Criterion();
@@ -228,9 +184,6 @@ private List buildBaseUserCriteria(SearchDocumentsInput input) {
"relatedAssets", Condition.EQUAL, input.getRelatedAssets()));
}
- // NOTE: State filtering is handled in buildCombinedFilter with ownership logic
- // Do not add state filters here
-
// Add source type filter if provided (if null, search all)
if (input.getSourceType() != null) {
criteria.add(
@@ -240,35 +193,6 @@ private List buildBaseUserCriteria(SearchDocumentsInput input) {
Collections.singletonList(input.getSourceType().toString())));
}
- // Exclude documents that are drafts by default, unless explicitly requested
- if (input.getIncludeDrafts() == null || !input.getIncludeDrafts()) {
- Criterion notDraftCriterion = new Criterion();
- notDraftCriterion.setField("draftOf");
- notDraftCriterion.setCondition(Condition.IS_NULL);
- criteria.add(notDraftCriterion);
- }
-
- // Add custom facet filters if provided - convert to AndFilterInput format
- if (input.getFilters() != null && !input.getFilters().isEmpty()) {
- final List orFilters =
- new ArrayList<>();
- final com.linkedin.datahub.graphql.generated.AndFilterInput andFilter =
- new com.linkedin.datahub.graphql.generated.AndFilterInput();
- andFilter.setAnd(input.getFilters());
- orFilters.add(andFilter);
- Filter additionalFilter = ResolverUtils.buildFilter(null, orFilters);
- if (additionalFilter != null && additionalFilter.getOr() != null) {
- additionalFilter
- .getOr()
- .forEach(
- conj -> {
- if (conj.getAnd() != null) {
- criteria.addAll(conj.getAnd());
- }
- });
- }
- }
-
return criteria;
}
}
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapper.java
index 64852152770223..3197b46a42a435 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapper.java
@@ -19,13 +19,13 @@
import com.linkedin.datahub.graphql.generated.DocumentRelatedAsset;
import com.linkedin.datahub.graphql.generated.DocumentRelatedDocument;
import com.linkedin.datahub.graphql.generated.EntityType;
-import com.linkedin.datahub.graphql.types.common.mappers.AuditStampMapper;
import com.linkedin.datahub.graphql.types.common.mappers.BrowsePathsV2Mapper;
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
import com.linkedin.datahub.graphql.types.common.mappers.DataPlatformInstanceAspectMapper;
import com.linkedin.datahub.graphql.types.common.mappers.OwnershipMapper;
import com.linkedin.datahub.graphql.types.domain.DomainAssociationMapper;
import com.linkedin.datahub.graphql.types.glossary.mappers.GlossaryTermsMapper;
+import com.linkedin.datahub.graphql.types.mappers.MapperUtils;
import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertiesMapper;
import com.linkedin.datahub.graphql.types.tag.mappers.GlobalTagsMapper;
import com.linkedin.domain.Domains;
@@ -179,10 +179,10 @@ private static DocumentInfo mapDocumentInfo(
result.setContents(graphqlContent);
// Map created audit stamp
- result.setCreated(AuditStampMapper.map(null, info.getCreated()));
+ result.setCreated(MapperUtils.createResolvedAuditStamp(info.getCreated()));
// Map lastModified audit stamp
- result.setLastModified(AuditStampMapper.map(null, info.getLastModified()));
+ result.setLastModified(MapperUtils.createResolvedAuditStamp(info.getLastModified()));
// Map related assets - create stubs that will be resolved by GraphQL batch loaders
if (info.hasRelatedAssets()) {
diff --git a/datahub-graphql-core/src/main/resources/documents.graphql b/datahub-graphql-core/src/main/resources/documents.graphql
index 753c79febaecb1..7cc79bbe8d130a 100644
--- a/datahub-graphql-core/src/main/resources/documents.graphql
+++ b/datahub-graphql-core/src/main/resources/documents.graphql
@@ -66,11 +66,11 @@ extend type Query {
"""
Search Documents with hybrid semantic search and filtering support.
- Supports filtering by parent document, types, domains, and semantic query.
+ Supports filtering by parent document, types, domains, and query.
- Results are automatically filtered based on document visibility:
+ Results are automatically filtered based on document status:
- PUBLISHED documents are shown to all users
- - UNPUBLISHED documents are only shown if owned by the current user or a group they belong to
+ - UNPUBLISHED documents are only shown if owned by the current calling user or a group they belong to
This ensures that unpublished documents remain private to their owners while published
documents are discoverable by everyone.
@@ -230,12 +230,12 @@ type DocumentInfo {
"""
The audit stamp for when the document was created
"""
- created: AuditStamp!
+ created: ResolvedAuditStamp!
"""
The audit stamp for when the document was last modified (any field)
"""
- lastModified: AuditStamp!
+ lastModified: ResolvedAuditStamp!
"""
Assets referenced by or related to this Document
@@ -603,13 +603,7 @@ input SearchDocumentsInput {
query: String
"""
- Optional parent document URN to filter by (for hierarchical browsing)
- """
- parentDocument: String
-
- """
- Optional list of parent document URNs to filter by (for batch child lookups).
- If both parentDocument and parentDocuments are provided, parentDocuments takes precedence.
+ Optional list of parent document URNs to filter by (for batch child lookups)
"""
parentDocuments: [String!]
@@ -636,44 +630,89 @@ input SearchDocumentsInput {
relatedAssets: [String!]
"""
- DEPRECATED: This field is no longer used for filtering.
+ Optional document source type to filter by.
+ If not provided, searches all documents regardless of source.
+ """
+ sourceType: DocumentSourceType
+}
- Document visibility is now automatically determined by ownership:
- - PUBLISHED documents are shown to all users
- - UNPUBLISHED documents are only shown if owned by the current user or their groups
+"""
+The result obtained when searching Documents
+"""
+type SearchDocumentsResult {
+ """
+ The starting offset of the result set returned
+ """
+ start: Int!
- This field is maintained for backward compatibility but has no effect on search results.
"""
- states: [DocumentState!]
+ The number of Documents in the returned result set
+ """
+ count: Int!
"""
- Optional document source type to filter by.
- If not provided, searches all documents regardless of source.
+ The total number of Documents in the result set
"""
- sourceType: DocumentSourceType
+ total: Int!
"""
- Whether to include draft documents in the search results.
- Draft documents have draftOf set and are hidden from normal browsing by default.
- Defaults to false (excludes drafts).
+ The Documents themselves
"""
- includeDrafts: Boolean
+ documents: [Document!]!
"""
- Optional facet filters to apply
+ Facets for filtering search results
"""
- filters: [FacetFilterInput!]
+ facets: [FacetMetadata!]
+}
+"""
+Input for querying context documents related to an entity.
+The relatedAssets filter is automatically set to the parent entity's URN.
+"""
+input RelatedDocumentsInput {
"""
- Optional flags controlling search options
+ The starting offset of the result set
"""
- searchFlags: SearchFlags
+ start: Int
+
+ """
+ The maximum number of Documents to return
+ """
+ count: Int
+
+ """
+ Optional list of parent document URNs to filter by (for batch child lookups)
+ """
+ parentDocuments: [String!]
+
+ """
+ If true, only returns documents with no parent (root-level documents).
+ If false or not provided, returns all documents regardless of parent.
+ """
+ rootOnly: Boolean
+
+ """
+ Optional list of document types to filter by (ANDed with other filters)
+ """
+ types: [String!]
+
+ """
+ Optional list of domain URNs to filter by (ANDed with other filters)
+ """
+ domains: [String!]
+
+ """
+ Optional document source type to filter by.
+ """
+ sourceType: DocumentSourceType
}
"""
-The result obtained when searching Documents
+Result containing context documents related to an entity.
+Same structure as SearchDocumentsResult.
"""
-type SearchDocumentsResult {
+type RelatedDocumentsResult {
"""
The starting offset of the result set returned
"""
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
index 464e8ce66a0be4..cbd983a8d4a85b 100644
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -1840,6 +1840,11 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -2291,7 +2296,6 @@ type VersionedDataset implements Entity {
No-op, has to be included due to model
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
- @deprecated
}
"""
@@ -2530,6 +2534,11 @@ type GlossaryTerm implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Schema metadata of the dataset
"""
@@ -2693,6 +2702,11 @@ type GlossaryNode implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Recursively get the lineage of glossary nodes for this entity
"""
@@ -3190,6 +3204,11 @@ type Container implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Status metadata of the container
"""
@@ -3646,6 +3665,11 @@ type SchemaFieldEntity implements EntityWithRelationships & Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -6051,6 +6075,11 @@ type Dashboard implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -6408,6 +6437,11 @@ type Chart implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -6822,6 +6856,11 @@ type DataFlow implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -7076,6 +7115,11 @@ type DataJob implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -10431,6 +10475,11 @@ type MLModel implements EntityWithRelationships & Entity & BrowsableEntity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -10575,6 +10624,11 @@ type MLModelGroup implements EntityWithRelationships & Entity & BrowsableEntity
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -10779,6 +10833,11 @@ type MLFeature implements EntityWithRelationships & Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -11045,6 +11104,11 @@ type MLPrimaryKey implements EntityWithRelationships & Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -11196,6 +11260,11 @@ type MLFeatureTable implements EntityWithRelationships & Entity & BrowsableEntit
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Edges extending from this entity grouped by direction in the lineage graph
"""
@@ -11670,6 +11739,11 @@ type Domain implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Experimental API.
For fetching extra entities that do not have custom UI code yet
@@ -13231,6 +13305,11 @@ type DataProduct implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
Children entities inside of the DataProduct
"""
@@ -13446,6 +13525,11 @@ type Application implements Entity {
"""
relationships(input: RelationshipsInput!): EntityRelationshipsResult
+ """
+ Get context documents related to this entity
+ """
+ relatedDocuments(input: RelatedDocumentsInput!): RelatedDocumentsResult
+
"""
The structured glossary terms associated with the Application
"""
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtilsTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtilsTest.java
new file mode 100644
index 00000000000000..289bb7d9835239
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/DocumentSearchFilterUtilsTest.java
@@ -0,0 +1,243 @@
+package com.linkedin.datahub.graphql.resolvers.knowledge;
+
+import static org.testng.Assert.*;
+
+import com.google.common.collect.ImmutableList;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
+import com.linkedin.metadata.query.filter.Criterion;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.utils.CriterionUtils;
+import java.util.ArrayList;
+import java.util.List;
+import org.testng.annotations.Test;
+
+public class DocumentSearchFilterUtilsTest {
+
+ private static final String TEST_USER_URN = "urn:li:corpuser:test";
+ private static final String TEST_GROUP_URN = "urn:li:corpGroup:testGroup";
+
+ @Test
+ public void testBuildCombinedFilterWithEmptyBaseCriteria() {
+ List baseCriteria = new ArrayList<>();
+ List userAndGroupUrns = ImmutableList.of(TEST_USER_URN, TEST_GROUP_URN);
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2); // PUBLISHED OR UNPUBLISHED owned
+
+ // Check first clause: PUBLISHED AND NOT-DRAFT
+ ConjunctiveCriterion publishedClause = filter.getOr().get(0);
+ assertNotNull(publishedClause.getAnd());
+ assertEquals(publishedClause.getAnd().size(), 2); // state + draftOf IS_NULL filter
+ Criterion stateCriterion = publishedClause.getAnd().get(0);
+ assertEquals(stateCriterion.getField(), "state");
+ assertEquals(stateCriterion.getCondition(), Condition.EQUAL);
+ assertEquals(stateCriterion.getValues().get(0), "PUBLISHED");
+ // Verify draftOf IS_NULL filter
+ Criterion draftOfCriterion = publishedClause.getAnd().get(1);
+ assertEquals(draftOfCriterion.getField(), "draftOf");
+ assertEquals(draftOfCriterion.getCondition(), Condition.IS_NULL);
+
+ // Check second clause: UNPUBLISHED AND owned AND NOT-DRAFT
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ assertNotNull(unpublishedClause.getAnd());
+ assertEquals(unpublishedClause.getAnd().size(), 3); // State + owners + draftOf IS_NULL filter
+ Criterion unpublishedStateCriterion = unpublishedClause.getAnd().get(0);
+ assertEquals(unpublishedStateCriterion.getField(), "state");
+ assertEquals(unpublishedStateCriterion.getCondition(), Condition.EQUAL);
+ assertEquals(unpublishedStateCriterion.getValues().get(0), "UNPUBLISHED");
+ Criterion ownersCriterion = unpublishedClause.getAnd().get(1);
+ assertEquals(ownersCriterion.getField(), "owners");
+ assertEquals(ownersCriterion.getCondition(), Condition.EQUAL);
+ assertTrue(ownersCriterion.getValues().contains(TEST_USER_URN));
+ assertTrue(ownersCriterion.getValues().contains(TEST_GROUP_URN));
+ // Verify draftOf IS_NULL filter in unpublished clause
+ Criterion unpublishedDraftOfCriterion = unpublishedClause.getAnd().get(2);
+ assertEquals(unpublishedDraftOfCriterion.getField(), "draftOf");
+ assertEquals(unpublishedDraftOfCriterion.getCondition(), Condition.IS_NULL);
+ }
+
+ @Test
+ public void testBuildCombinedFilterWithBaseCriteria() {
+ List baseCriteria = new ArrayList<>();
+ baseCriteria.add(
+ CriterionUtils.buildCriterion("types", Condition.EQUAL, ImmutableList.of("tutorial")));
+ baseCriteria.add(
+ CriterionUtils.buildCriterion(
+ "domains", Condition.EQUAL, ImmutableList.of("urn:li:domain:test")));
+
+ List userAndGroupUrns = ImmutableList.of(TEST_USER_URN);
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2);
+
+ // Check first clause: base criteria AND PUBLISHED AND NOT-DRAFT
+ ConjunctiveCriterion publishedClause = filter.getOr().get(0);
+ assertNotNull(publishedClause.getAnd());
+ assertEquals(publishedClause.getAnd().size(), 4); // types + domains + state + draftOf IS_NULL
+
+ // Check second clause: base criteria AND UNPUBLISHED AND owned AND NOT-DRAFT
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ assertNotNull(unpublishedClause.getAnd());
+ assertEquals(
+ unpublishedClause.getAnd().size(), 5); // types + domains + state + owners + draftOf IS_NULL
+
+ // Verify base criteria are in both clauses
+ boolean foundTypesInPublished = false;
+ boolean foundDomainsInPublished = false;
+ boolean foundTypesInUnpublished = false;
+ boolean foundDomainsInUnpublished = false;
+ boolean foundDraftOfInPublished = false;
+ boolean foundDraftOfInUnpublished = false;
+
+ for (Criterion criterion : publishedClause.getAnd()) {
+ if ("types".equals(criterion.getField())) {
+ foundTypesInPublished = true;
+ }
+ if ("domains".equals(criterion.getField())) {
+ foundDomainsInPublished = true;
+ }
+ if ("draftOf".equals(criterion.getField()) && criterion.getCondition() == Condition.IS_NULL) {
+ foundDraftOfInPublished = true;
+ }
+ }
+
+ for (Criterion criterion : unpublishedClause.getAnd()) {
+ if ("types".equals(criterion.getField())) {
+ foundTypesInUnpublished = true;
+ }
+ if ("domains".equals(criterion.getField())) {
+ foundDomainsInUnpublished = true;
+ }
+ if ("draftOf".equals(criterion.getField()) && criterion.getCondition() == Condition.IS_NULL) {
+ foundDraftOfInUnpublished = true;
+ }
+ }
+
+ assertTrue(foundTypesInPublished, "Types filter should be in published clause");
+ assertTrue(foundDomainsInPublished, "Domains filter should be in published clause");
+ assertTrue(foundDraftOfInPublished, "DraftOf IS_NULL filter should be in published clause");
+ assertTrue(foundTypesInUnpublished, "Types filter should be in unpublished clause");
+ assertTrue(foundDomainsInUnpublished, "Domains filter should be in unpublished clause");
+ assertTrue(foundDraftOfInUnpublished, "DraftOf IS_NULL filter should be in unpublished clause");
+ }
+
+ @Test
+ public void testBuildCombinedFilterWithSingleUser() {
+ List baseCriteria = new ArrayList<>();
+ List userAndGroupUrns = ImmutableList.of(TEST_USER_URN);
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2);
+
+ // Check owners filter in unpublished clause (state, owners, draftOf)
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ assertEquals(unpublishedClause.getAnd().size(), 3); // state + owners + draftOf IS_NULL
+ Criterion ownersCriterion = unpublishedClause.getAnd().get(1);
+ assertEquals(ownersCriterion.getField(), "owners");
+ assertEquals(ownersCriterion.getValues().size(), 1);
+ assertEquals(ownersCriterion.getValues().get(0), TEST_USER_URN);
+ }
+
+ @Test
+ public void testBuildCombinedFilterWithMultipleGroups() {
+ List baseCriteria = new ArrayList<>();
+ List userAndGroupUrns =
+ ImmutableList.of(TEST_USER_URN, TEST_GROUP_URN, "urn:li:corpGroup:anotherGroup");
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2);
+
+ // Check owners filter includes all URNs (state, owners, draftOf)
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ assertEquals(unpublishedClause.getAnd().size(), 3); // state + owners + draftOf IS_NULL
+ Criterion ownersCriterion = unpublishedClause.getAnd().get(1);
+ assertEquals(ownersCriterion.getField(), "owners");
+ assertEquals(ownersCriterion.getValues().size(), 3);
+ assertTrue(ownersCriterion.getValues().contains(TEST_USER_URN));
+ assertTrue(ownersCriterion.getValues().contains(TEST_GROUP_URN));
+ assertTrue(ownersCriterion.getValues().contains("urn:li:corpGroup:anotherGroup"));
+ }
+
+ @Test
+ public void testBuildCombinedFilterStructure() {
+ // Verify the filter structure: (base AND PUBLISHED AND NOT-DRAFT) OR (base AND UNPUBLISHED AND
+ // owners AND NOT-DRAFT)
+ List baseCriteria = new ArrayList<>();
+ baseCriteria.add(
+ CriterionUtils.buildCriterion(
+ "relatedAssets", Condition.EQUAL, ImmutableList.of("urn:li:dataset:test")));
+
+ List userAndGroupUrns = ImmutableList.of(TEST_USER_URN);
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2);
+
+ // Published clause should have: relatedAssets + state + draftOf IS_NULL
+ ConjunctiveCriterion publishedClause = filter.getOr().get(0);
+ assertEquals(publishedClause.getAnd().size(), 3);
+ assertEquals(publishedClause.getAnd().get(1).getField(), "state");
+ assertEquals(publishedClause.getAnd().get(1).getValues().get(0), "PUBLISHED");
+ assertEquals(publishedClause.getAnd().get(2).getField(), "draftOf");
+ assertEquals(publishedClause.getAnd().get(2).getCondition(), Condition.IS_NULL);
+
+ // Unpublished clause should have: relatedAssets + state + owners + draftOf IS_NULL
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ assertEquals(unpublishedClause.getAnd().size(), 4);
+ assertEquals(unpublishedClause.getAnd().get(1).getField(), "state");
+ assertEquals(unpublishedClause.getAnd().get(1).getValues().get(0), "UNPUBLISHED");
+ assertEquals(unpublishedClause.getAnd().get(2).getField(), "owners");
+ assertEquals(unpublishedClause.getAnd().get(3).getField(), "draftOf");
+ assertEquals(unpublishedClause.getAnd().get(3).getCondition(), Condition.IS_NULL);
+ }
+
+ @Test
+ public void testBuildCombinedFilterExcludesDrafts() {
+ // Verify that drafts are excluded from both published and unpublished clauses
+ List baseCriteria = new ArrayList<>();
+ List userAndGroupUrns = ImmutableList.of(TEST_USER_URN);
+
+ Filter filter = DocumentSearchFilterUtils.buildCombinedFilter(baseCriteria, userAndGroupUrns);
+
+ assertNotNull(filter);
+ assertNotNull(filter.getOr());
+ assertEquals(filter.getOr().size(), 2);
+
+ // Both clauses should have draftOf IS_NULL criterion
+ boolean foundDraftOfInPublished = false;
+ boolean foundDraftOfInUnpublished = false;
+
+ ConjunctiveCriterion publishedClause = filter.getOr().get(0);
+ for (Criterion criterion : publishedClause.getAnd()) {
+ if ("draftOf".equals(criterion.getField()) && criterion.getCondition() == Condition.IS_NULL) {
+ foundDraftOfInPublished = true;
+ }
+ }
+
+ ConjunctiveCriterion unpublishedClause = filter.getOr().get(1);
+ for (Criterion criterion : unpublishedClause.getAnd()) {
+ if ("draftOf".equals(criterion.getField()) && criterion.getCondition() == Condition.IS_NULL) {
+ foundDraftOfInUnpublished = true;
+ }
+ }
+
+ assertTrue(foundDraftOfInPublished, "Published clause should exclude drafts (draftOf IS_NULL)");
+ assertTrue(
+ foundDraftOfInUnpublished, "Unpublished clause should exclude drafts (draftOf IS_NULL)");
+ }
+}
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolverTest.java
new file mode 100644
index 00000000000000..3e5fdba80b1cc7
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/RelatedDocumentsResolverTest.java
@@ -0,0 +1,479 @@
+package com.linkedin.datahub.graphql.resolvers.knowledge;
+
+import static com.linkedin.datahub.graphql.TestUtils.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.*;
+import static org.testng.Assert.*;
+
+import com.datahub.authentication.group.GroupService;
+import com.google.common.collect.ImmutableList;
+import com.linkedin.common.Owner;
+import com.linkedin.common.OwnerArray;
+import com.linkedin.common.Ownership;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.RelatedDocumentsInput;
+import com.linkedin.datahub.graphql.generated.RelatedDocumentsResult;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.EnvelopedAspectMap;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.knowledge.DocumentInfo;
+import com.linkedin.knowledge.DocumentStatus;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.Criterion;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.search.SearchEntity;
+import com.linkedin.metadata.search.SearchEntityArray;
+import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.search.SearchResultMetadata;
+import com.linkedin.metadata.service.DocumentService;
+import graphql.schema.DataFetchingEnvironment;
+import io.datahubproject.metadata.context.OperationContext;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class RelatedDocumentsResolverTest {
+
+ private static final String TEST_DOCUMENT_URN = "urn:li:document:test-document";
+ private static final String TEST_PARENT_ENTITY_URN = "urn:li:dataset:test-dataset";
+ private static final String TEST_USER_URN = "urn:li:corpuser:test";
+ private static final String TEST_GROUP_URN = "urn:li:corpGroup:testGroup";
+ private static final String OTHER_USER_URN = "urn:li:corpuser:other";
+
+ private DocumentService mockService;
+ private EntityClient mockEntityClient;
+ private GroupService mockGroupService;
+ private RelatedDocumentsResolver resolver;
+ private DataFetchingEnvironment mockEnv;
+ private RelatedDocumentsInput input;
+ private Entity mockParentEntity;
+
+ @BeforeMethod
+ public void setupTest() throws Exception {
+ mockService = mock(DocumentService.class);
+ mockEntityClient = mock(EntityClient.class);
+ mockGroupService = mock(GroupService.class);
+ mockEnv = mock(DataFetchingEnvironment.class);
+
+ // Setup mock parent entity
+ mockParentEntity = mock(Entity.class);
+ when(mockParentEntity.getUrn()).thenReturn(TEST_PARENT_ENTITY_URN);
+ when(mockParentEntity.getType()).thenReturn(EntityType.DATASET);
+ when(mockEnv.getSource()).thenReturn(mockParentEntity);
+
+ // Setup default input
+ input = new RelatedDocumentsInput();
+ input.setStart(0);
+ input.setCount(10);
+
+ // Setup mock search result
+ SearchResult searchResult = new SearchResult();
+ searchResult.setFrom(0);
+ searchResult.setPageSize(10);
+ searchResult.setNumEntities(1);
+ searchResult.setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(new SearchEntity().setEntity(UrnUtils.getUrn(TEST_DOCUMENT_URN)))));
+ searchResult.setMetadata(new SearchResultMetadata());
+
+ when(mockService.searchDocuments(
+ any(OperationContext.class),
+ any(String.class),
+ any(),
+ any(),
+ any(Integer.class),
+ any(Integer.class)))
+ .thenReturn(searchResult);
+
+ // Mock EntityClient.batchGetV2 to return a hydrated PUBLISHED entity by default
+ Map entityResponseMap = new HashMap<>();
+ EntityResponse entityResponse =
+ createPublishedDocumentResponse(TEST_DOCUMENT_URN, TEST_USER_URN);
+ entityResponseMap.put(UrnUtils.getUrn(TEST_DOCUMENT_URN), entityResponse);
+
+ when(mockEntityClient.batchGetV2(any(OperationContext.class), any(String.class), any(), any()))
+ .thenReturn(entityResponseMap);
+
+ // Mock GroupService to return empty groups by default
+ when(mockGroupService.getGroupsForUser(any(OperationContext.class), any(Urn.class)))
+ .thenReturn(Collections.emptyList());
+
+ resolver = new RelatedDocumentsResolver(mockService, mockEntityClient, mockGroupService);
+ }
+
+ private EntityResponse createPublishedDocumentResponse(String documentUrn, String ownerUrn) {
+ EntityResponse entityResponse = new EntityResponse();
+ entityResponse.setUrn(UrnUtils.getUrn(documentUrn));
+ entityResponse.setEntityName("document");
+
+ EnvelopedAspectMap aspects = new EnvelopedAspectMap();
+
+ // Add DocumentInfo with PUBLISHED state
+ DocumentInfo docInfo = new DocumentInfo();
+ DocumentStatus status = new DocumentStatus();
+ status.setState(com.linkedin.knowledge.DocumentState.PUBLISHED);
+ docInfo.setStatus(status);
+ // Add required contents field
+ com.linkedin.knowledge.DocumentContents contents =
+ new com.linkedin.knowledge.DocumentContents();
+ contents.setText("Test content");
+ docInfo.setContents(contents);
+ // Add required created and lastModified fields
+ com.linkedin.common.AuditStamp created = new com.linkedin.common.AuditStamp();
+ created.setTime(System.currentTimeMillis());
+ created.setActor(UrnUtils.getUrn(ownerUrn));
+ docInfo.setCreated(created);
+ com.linkedin.common.AuditStamp lastModified = new com.linkedin.common.AuditStamp();
+ lastModified.setTime(System.currentTimeMillis());
+ lastModified.setActor(UrnUtils.getUrn(ownerUrn));
+ docInfo.setLastModified(lastModified);
+ aspects.put(
+ Constants.DOCUMENT_INFO_ASPECT_NAME,
+ new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(docInfo.data())));
+
+ // Add Ownership
+ Ownership ownership = new Ownership();
+ Owner owner = new Owner();
+ owner.setOwner(UrnUtils.getUrn(ownerUrn));
+ owner.setType(com.linkedin.common.OwnershipType.TECHNICAL_OWNER);
+ ownership.setOwners(new OwnerArray(owner));
+ aspects.put(
+ Constants.OWNERSHIP_ASPECT_NAME,
+ new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(ownership.data())));
+
+ entityResponse.setAspects(aspects);
+ return entityResponse;
+ }
+
+ private EntityResponse createUnpublishedDocumentResponse(String documentUrn, String ownerUrn) {
+ EntityResponse entityResponse = new EntityResponse();
+ entityResponse.setUrn(UrnUtils.getUrn(documentUrn));
+ entityResponse.setEntityName("document");
+
+ EnvelopedAspectMap aspects = new EnvelopedAspectMap();
+
+ // Add DocumentInfo with UNPUBLISHED state
+ DocumentInfo docInfo = new DocumentInfo();
+ DocumentStatus status = new DocumentStatus();
+ status.setState(com.linkedin.knowledge.DocumentState.UNPUBLISHED);
+ docInfo.setStatus(status);
+ // Add required contents field
+ com.linkedin.knowledge.DocumentContents contents =
+ new com.linkedin.knowledge.DocumentContents();
+ contents.setText("Test content");
+ docInfo.setContents(contents);
+ // Add required created and lastModified fields
+ com.linkedin.common.AuditStamp created = new com.linkedin.common.AuditStamp();
+ created.setTime(System.currentTimeMillis());
+ created.setActor(UrnUtils.getUrn(ownerUrn));
+ docInfo.setCreated(created);
+ com.linkedin.common.AuditStamp lastModified = new com.linkedin.common.AuditStamp();
+ lastModified.setTime(System.currentTimeMillis());
+ lastModified.setActor(UrnUtils.getUrn(ownerUrn));
+ docInfo.setLastModified(lastModified);
+ aspects.put(
+ Constants.DOCUMENT_INFO_ASPECT_NAME,
+ new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(docInfo.data())));
+
+ // Add Ownership
+ Ownership ownership = new Ownership();
+ Owner owner = new Owner();
+ owner.setOwner(UrnUtils.getUrn(ownerUrn));
+ owner.setType(com.linkedin.common.OwnershipType.TECHNICAL_OWNER);
+ ownership.setOwners(new OwnerArray(owner));
+ aspects.put(
+ Constants.OWNERSHIP_ASPECT_NAME,
+ new EnvelopedAspect().setValue(new com.linkedin.entity.Aspect(ownership.data())));
+
+ entityResponse.setAspects(aspects);
+ return entityResponse;
+ }
+
+ @Test
+ public void testContextDocumentsSuccess() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(result.getStart(), 0);
+ assertEquals(result.getCount(), 10);
+ assertEquals(result.getTotal(), 1);
+ assertEquals(result.getDocuments().size(), 1);
+
+ // Verify service was called with "*" query (no semantic search)
+ ArgumentCaptor filterCaptor = ArgumentCaptor.forClass(Filter.class);
+ verify(mockService, times(1))
+ .searchDocuments(
+ any(OperationContext.class), eq("*"), filterCaptor.capture(), any(), eq(0), eq(10));
+
+ // Verify that relatedAssets filter was automatically added with parent entity URN
+ Filter capturedFilter = filterCaptor.getValue();
+ assertNotNull(capturedFilter);
+ assertNotNull(capturedFilter.getOr());
+ assertEquals(capturedFilter.getOr().size(), 2); // PUBLISHED OR UNPUBLISHED owned
+
+ // Check that relatedAssets filter is in both clauses
+ boolean foundRelatedAssetsFilter = false;
+ for (com.linkedin.metadata.query.filter.ConjunctiveCriterion conj : capturedFilter.getOr()) {
+ if (conj.getAnd() != null) {
+ for (Criterion criterion : conj.getAnd()) {
+ if ("relatedAssets".equals(criterion.getField())
+ && Condition.EQUAL.equals(criterion.getCondition())
+ && criterion.getValues() != null
+ && criterion.getValues().contains(TEST_PARENT_ENTITY_URN)) {
+ foundRelatedAssetsFilter = true;
+ break;
+ }
+ }
+ }
+ }
+ assertTrue(foundRelatedAssetsFilter, "relatedAssets filter with parent URN should be present");
+ }
+
+ @Test
+ public void testContextDocumentsWithAdditionalFilters() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ input.setTypes(ImmutableList.of("tutorial", "guide"));
+ input.setParentDocuments(ImmutableList.of("urn:li:document:parent"));
+ input.setDomains(ImmutableList.of("urn:li:domain:test"));
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+
+ // Verify service was called
+ verify(mockService, times(1))
+ .searchDocuments(any(OperationContext.class), eq("*"), any(), any(), eq(0), eq(10));
+ }
+
+ @Test
+ public void testContextDocumentsDefaultValues() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ // Don't set start/count - should use defaults
+ input.setStart(null);
+ input.setCount(null);
+
+ // Mock should return SearchResult with pageSize matching the default count (100)
+ SearchResult defaultSearchResult = new SearchResult();
+ defaultSearchResult.setFrom(0);
+ defaultSearchResult.setPageSize(100); // Match DEFAULT_COUNT
+ defaultSearchResult.setNumEntities(1);
+ defaultSearchResult.setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(new SearchEntity().setEntity(UrnUtils.getUrn(TEST_DOCUMENT_URN)))));
+ defaultSearchResult.setMetadata(new SearchResultMetadata());
+
+ when(mockService.searchDocuments(
+ any(OperationContext.class),
+ any(String.class),
+ any(),
+ any(),
+ eq(0),
+ eq(100))) // Default count
+ .thenReturn(defaultSearchResult);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(result.getStart(), 0); // Default start
+ assertEquals(result.getCount(), 100); // Default count
+
+ // Verify service was called with defaults
+ verify(mockService, times(1))
+ .searchDocuments(any(OperationContext.class), eq("*"), any(), any(), eq(0), eq(100));
+ }
+
+ @Test
+ public void testContextDocumentsIncludesPublishedDocuments() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(OTHER_USER_URN); // Different user
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ // Create published document owned by different user
+ Map entityResponseMap = new HashMap<>();
+ EntityResponse entityResponse =
+ createPublishedDocumentResponse(TEST_DOCUMENT_URN, OTHER_USER_URN);
+ entityResponseMap.put(UrnUtils.getUrn(TEST_DOCUMENT_URN), entityResponse);
+ when(mockEntityClient.batchGetV2(any(OperationContext.class), any(String.class), any(), any()))
+ .thenReturn(entityResponseMap);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(result.getDocuments().size(), 1); // Published docs visible to all users
+ }
+
+ @Test
+ public void testContextDocumentsIncludesOwnedUnpublishedDocuments() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ // Create unpublished document owned by current user
+ Map entityResponseMap = new HashMap<>();
+ EntityResponse entityResponse =
+ createUnpublishedDocumentResponse(TEST_DOCUMENT_URN, TEST_USER_URN);
+ entityResponseMap.put(UrnUtils.getUrn(TEST_DOCUMENT_URN), entityResponse);
+ when(mockEntityClient.batchGetV2(any(OperationContext.class), any(String.class), any(), any()))
+ .thenReturn(entityResponseMap);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(result.getDocuments().size(), 1); // Owned unpublished docs visible to owner
+ }
+
+ @Test
+ public void testContextDocumentsExcludesOtherUsersUnpublishedDocuments() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ // Create unpublished document owned by different user
+ Map entityResponseMap = new HashMap<>();
+ EntityResponse entityResponse =
+ createUnpublishedDocumentResponse(TEST_DOCUMENT_URN, OTHER_USER_URN);
+ entityResponseMap.put(UrnUtils.getUrn(TEST_DOCUMENT_URN), entityResponse);
+ when(mockEntityClient.batchGetV2(any(OperationContext.class), any(String.class), any(), any()))
+ .thenReturn(entityResponseMap);
+
+ // Mock search to return empty results (filter excludes unpublished docs not owned by user)
+ SearchResult emptySearchResult = new SearchResult();
+ emptySearchResult.setFrom(0);
+ emptySearchResult.setPageSize(10);
+ emptySearchResult.setNumEntities(0);
+ emptySearchResult.setEntities(new SearchEntityArray(Collections.emptyList()));
+ emptySearchResult.setMetadata(new SearchResultMetadata());
+
+ when(mockService.searchDocuments(
+ any(OperationContext.class),
+ any(String.class),
+ any(),
+ any(),
+ any(Integer.class),
+ any(Integer.class)))
+ .thenReturn(emptySearchResult);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(
+ result.getDocuments().size(), 0); // Unpublished docs not owned by user are excluded
+ }
+
+ @Test
+ public void testContextDocumentsWithGroupOwnership() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ // Mock user belongs to a group
+ when(mockGroupService.getGroupsForUser(any(OperationContext.class), any(Urn.class)))
+ .thenReturn(ImmutableList.of(UrnUtils.getUrn(TEST_GROUP_URN)));
+
+ // Create unpublished document owned by user's group
+ Map entityResponseMap = new HashMap<>();
+ EntityResponse entityResponse =
+ createUnpublishedDocumentResponse(TEST_DOCUMENT_URN, TEST_GROUP_URN);
+ entityResponseMap.put(UrnUtils.getUrn(TEST_DOCUMENT_URN), entityResponse);
+ when(mockEntityClient.batchGetV2(any(OperationContext.class), any(String.class), any(), any()))
+ .thenReturn(entityResponseMap);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ assertEquals(
+ result.getDocuments().size(), 1); // Unpublished docs owned by user's group are visible
+ }
+
+ @Test(expectedExceptions = ExecutionException.class)
+ public void testContextDocumentsMissingParentEntity() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+ when(mockEnv.getSource()).thenReturn(null); // No parent entity
+
+ try {
+ resolver.get(mockEnv).get(); // Should throw ExecutionException wrapping RuntimeException
+ } catch (ExecutionException e) {
+ // Verify the cause is RuntimeException (wrapped by catch block)
+ assertTrue(e.getCause() instanceof RuntimeException);
+ // The original exception is wrapped: RuntimeException("Failed to fetch related documents",
+ // originalException)
+ // So we need to check the nested cause
+ assertTrue(e.getCause().getCause() instanceof RuntimeException);
+ assertTrue(e.getCause().getCause().getMessage().contains("Parent entity URN is required"));
+ throw e; // Re-throw to satisfy expectedExceptions
+ }
+ }
+
+ @Test(expectedExceptions = ExecutionException.class)
+ public void testContextDocumentsMissingParentEntityUrn() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ Entity entityWithoutUrn = mock(Entity.class);
+ when(entityWithoutUrn.getUrn()).thenReturn(null);
+ when(mockEnv.getSource()).thenReturn(entityWithoutUrn);
+
+ try {
+ resolver.get(mockEnv).get(); // Should throw ExecutionException wrapping RuntimeException
+ } catch (ExecutionException e) {
+ // Verify the cause is RuntimeException (wrapped by catch block)
+ assertTrue(e.getCause() instanceof RuntimeException);
+ // The original exception is wrapped: RuntimeException("Failed to fetch related documents",
+ // originalException)
+ // So we need to check the nested cause
+ assertTrue(e.getCause().getCause() instanceof RuntimeException);
+ assertTrue(e.getCause().getCause().getMessage().contains("Parent entity URN is required"));
+ throw e; // Re-throw to satisfy expectedExceptions
+ }
+ }
+
+ @Test
+ public void testContextDocumentsWithRootOnlyFilter() throws Exception {
+ QueryContext mockContext = getMockAllowContext();
+ when(mockContext.getActorUrn()).thenReturn(TEST_USER_URN);
+ when(mockEnv.getContext()).thenReturn(mockContext);
+ when(mockEnv.getArgument(eq("input"))).thenReturn(input);
+
+ input.setRootOnly(true);
+
+ RelatedDocumentsResult result = resolver.get(mockEnv).get();
+
+ assertNotNull(result);
+ // Verify service was called
+ verify(mockService, times(1))
+ .searchDocuments(any(OperationContext.class), eq("*"), any(), any(), eq(0), eq(10));
+ }
+}
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolverTest.java
index e5de2a1726c77a..ccd992254dfcb8 100644
--- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolverTest.java
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/knowledge/SearchDocumentsResolverTest.java
@@ -15,7 +15,6 @@
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.DocumentSourceType;
-import com.linkedin.datahub.graphql.generated.DocumentState;
import com.linkedin.datahub.graphql.generated.SearchDocumentsInput;
import com.linkedin.datahub.graphql.generated.SearchDocumentsResult;
import com.linkedin.entity.EntityResponse;
@@ -219,7 +218,7 @@ public void testSearchDocumentsWithFilters() throws Exception {
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
input.setTypes(ImmutableList.of("tutorial", "guide"));
- input.setParentDocument("urn:li:document:parent");
+ input.setParentDocuments(ImmutableList.of("urn:li:document:parent"));
SearchDocumentsResult result = resolver.get(mockEnv).get();
@@ -232,133 +231,92 @@ public void testSearchDocumentsWithFilters() throws Exception {
}
@Test
- public void testSearchDocumentsEmptyQuery() throws Exception {
+ public void testSearchDocumentsWithRootOnly() throws Exception {
QueryContext mockContext = getMockAllowContext();
when(mockEnv.getContext()).thenReturn(mockContext);
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
- input.setQuery(null); // Empty query should default to "*"
+ input.setRootOnly(true);
SearchDocumentsResult result = resolver.get(mockEnv).get();
assertNotNull(result);
- // Verify service was called with "*" query
- verify(mockService, times(1))
- .searchDocuments(any(OperationContext.class), eq("*"), any(), any(), eq(0), eq(10));
- }
-
- // Note: Search operations don't require special authorization like other entity searches
- // Authorization is applied at the entity level when viewing individual documents
-
- @Test
- public void testSearchDocumentsServiceThrowsException() throws Exception {
- QueryContext mockContext = getMockAllowContext();
- when(mockEnv.getContext()).thenReturn(mockContext);
- when(mockEnv.getArgument(eq("input"))).thenReturn(input);
-
- when(mockService.searchDocuments(
- any(OperationContext.class),
- any(),
- any(),
- any(),
- any(Integer.class),
- any(Integer.class)))
- .thenThrow(new RuntimeException("Service error"));
-
- assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
- }
-
- @Test
- public void testSearchDocumentsDefaultsToPublishedState() throws Exception {
- QueryContext mockContext = getMockAllowContext();
- when(mockEnv.getContext()).thenReturn(mockContext);
- when(mockEnv.getArgument(eq("input"))).thenReturn(input);
-
- // Don't set any states - should default to PUBLISHED
- input.setStates(null);
-
- SearchDocumentsResult result = resolver.get(mockEnv).get();
-
- assertNotNull(result);
-
- // Verify service was called (the filter will contain state=PUBLISHED by default)
+ // Verify service was called with rootOnly filter
verify(mockService, times(1))
.searchDocuments(
any(OperationContext.class), eq("test query"), any(), any(), eq(0), eq(10));
}
@Test
- public void testSearchDocumentsWithSingleState() throws Exception {
+ public void testSearchDocumentsWithParentDocumentsAndRootOnly() throws Exception {
QueryContext mockContext = getMockAllowContext();
when(mockEnv.getContext()).thenReturn(mockContext);
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
- // Set to only search UNPUBLISHED documents
- input.setStates(ImmutableList.of(DocumentState.UNPUBLISHED));
+ // When both are set, parentDocuments takes precedence
+ input.setParentDocuments(ImmutableList.of("urn:li:document:parent"));
+ input.setRootOnly(true);
SearchDocumentsResult result = resolver.get(mockEnv).get();
assertNotNull(result);
- // Verify service was called with UNPUBLISHED state filter
+ // Verify service was called (parentDocuments filter should be used, not rootOnly)
verify(mockService, times(1))
.searchDocuments(
any(OperationContext.class), eq("test query"), any(), any(), eq(0), eq(10));
}
@Test
- public void testSearchDocumentsWithMultipleStates() throws Exception {
+ public void testSearchDocumentsEmptyQuery() throws Exception {
QueryContext mockContext = getMockAllowContext();
when(mockEnv.getContext()).thenReturn(mockContext);
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
- // Set to search both PUBLISHED and UNPUBLISHED documents
- input.setStates(ImmutableList.of(DocumentState.PUBLISHED, DocumentState.UNPUBLISHED));
+ input.setQuery(null); // Empty query should default to "*"
SearchDocumentsResult result = resolver.get(mockEnv).get();
assertNotNull(result);
- // Verify service was called with both states in filter
+ // Verify service was called with "*" query
verify(mockService, times(1))
- .searchDocuments(
- any(OperationContext.class), eq("test query"), any(), any(), eq(0), eq(10));
+ .searchDocuments(any(OperationContext.class), eq("*"), any(), any(), eq(0), eq(10));
}
+ // Note: Search operations don't require special authorization like other entity searches
+ // Authorization is applied at the entity level when viewing individual documents
+
@Test
- public void testSearchDocumentsExcludesDraftsByDefault() throws Exception {
+ public void testSearchDocumentsServiceThrowsException() throws Exception {
QueryContext mockContext = getMockAllowContext();
when(mockEnv.getContext()).thenReturn(mockContext);
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
- // Don't set includeDrafts - should exclude drafts by default
- input.setIncludeDrafts(null);
-
- SearchDocumentsResult result = resolver.get(mockEnv).get();
-
- assertNotNull(result);
+ when(mockService.searchDocuments(
+ any(OperationContext.class),
+ any(),
+ any(),
+ any(),
+ any(Integer.class),
+ any(Integer.class)))
+ .thenThrow(new RuntimeException("Service error"));
- // Verify service was called (the filter will exclude draftOf != null by default)
- verify(mockService, times(1))
- .searchDocuments(
- any(OperationContext.class), eq("test query"), any(), any(), eq(0), eq(10));
+ assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
}
@Test
- public void testSearchDocumentsIncludeDrafts() throws Exception {
+ public void testSearchDocumentsDefaultsToPublishedState() throws Exception {
QueryContext mockContext = getMockAllowContext();
when(mockEnv.getContext()).thenReturn(mockContext);
when(mockEnv.getArgument(eq("input"))).thenReturn(input);
- // Explicitly include drafts
- input.setIncludeDrafts(true);
-
SearchDocumentsResult result = resolver.get(mockEnv).get();
assertNotNull(result);
- // Verify service was called without draftOf filter
+ // Verify service was called (the filter will contain state=PUBLISHED by default)
verify(mockService, times(1))
.searchDocuments(
any(OperationContext.class), eq("test query"), any(), any(), eq(0), eq(10));
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapperTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapperTest.java
index b2f3d8f1dbc73c..e7ddac178a4f84 100644
--- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapperTest.java
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/knowledge/DocumentMapperTest.java
@@ -154,6 +154,13 @@ public void testMapDocumentWithAllAspects() throws URISyntaxException {
assertEquals(result.getInfo().getContents().getText(), TEST_CONTENT);
assertNotNull(result.getInfo().getCreated());
assertEquals(result.getInfo().getCreated().getTime(), TEST_TIMESTAMP);
+ // Verify actor is set as CorpUser in ResolvedAuditStamp
+ assertNotNull(result.getInfo().getCreated().getActor());
+ assertEquals(result.getInfo().getCreated().getActor().getUrn(), TEST_ACTOR_URN);
+ assertNotNull(result.getInfo().getLastModified());
+ assertEquals(result.getInfo().getLastModified().getTime(), TEST_TIMESTAMP);
+ assertNotNull(result.getInfo().getLastModified().getActor());
+ assertEquals(result.getInfo().getLastModified().getActor().getUrn(), TEST_ACTOR_URN);
// Relationships are present inside info and constructed as unresolved stubs
assertNotNull(result.getInfo().getParentDocument());
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/README.md b/datahub-web-react/src/app/document/hooks/__tests__/README.md
deleted file mode 100644
index 3d6598bb81d2c3..00000000000000
--- a/datahub-web-react/src/app/document/hooks/__tests__/README.md
+++ /dev/null
@@ -1,92 +0,0 @@
-# DocumentV2 Hooks Test Suite
-
-This directory contains comprehensive unit tests for all hooks in the `documentV2/hooks` directory.
-
-## Test Coverage
-
-### ✅ Fully Tested Hooks
-
-1. **useUpdateDocument** - Document mutation operations
-
- - Update contents, title, status, sub-type
- - Update related entities (assets, documents)
- - Error handling and rollback
- - Loading states
-
-2. **useSearchDocuments** - Document search functionality
-
- - Default and custom queries
- - Filtering by parent, root-only, types, states
- - Pagination support
- - Various fetch policies
-
-3. **useDocumentTreeMutations** - Tree mutation with optimistic updates
-
- - Create documents with optimistic UI
- - Update document titles
- - Move documents in tree
- - Delete documents with rollback on error
-
-4. **useDocumentChildren** - Fetching child documents
-
- - Batch checking for children
- - Fetching children for parent
- - Error handling
-
-5. **useDocumentPermissions** - Permission checks
-
- - Create, edit, delete, move permissions
- - Platform vs entity-level privileges
- - Memoization tests
-
-6. **useDocumentTreeExpansion** - Tree expansion state
-
- - Expand/collapse nodes
- - Loading children on demand
- - Caching loaded children
- - External state control
-
-7. **useExtractMentions** - URN extraction from markdown
-
- - Document and asset URN extraction
- - Duplicate handling
- - Edge cases (empty, malformed, special characters)
-
-8. **useLoadDocumentTree** - Initial tree loading
- - Load root documents
- - Check for children
- - Document sorting
- - Error handling
-
-## Running Tests
-
-```bash
-# Run all documentV2 hook tests
-yarn test src/app/documentV2/hooks/__tests__/
-
-# Run specific test file
-yarn test src/app/documentV2/hooks/__tests__/useExtractMentions.test.ts
-
-# Run with coverage
-yarn test-coverage src/app/documentV2/hooks/__tests__/
-```
-
-## Test Structure
-
-All tests follow these principles:
-
-- **Comprehensive**: Test happy paths, error cases, and edge conditions
-- **Isolated**: Each test is independent with proper setup/teardown
-- **Clear**: Descriptive test names and organized test suites
-- **Realistic**: Use proper mocks for GraphQL, contexts, and dependencies
-
-## Technologies Used
-
-- **Vitest**: Test runner
-- **React Testing Library**: Component and hook testing utilities
-- **Apollo MockedProvider**: GraphQL mock responses
-- **Vi Mocks**: Module and function mocking
-
-## Known Issues
-
-Some tests involving Apollo MockedProvider may have timing issues with async mock resolution. These are being addressed but do not affect the structural correctness of the tests.
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/useDocumentChildren.test.tsx b/datahub-web-react/src/app/document/hooks/__tests__/useDocumentChildren.test.tsx
deleted file mode 100644
index cd93bae2fa8700..00000000000000
--- a/datahub-web-react/src/app/document/hooks/__tests__/useDocumentChildren.test.tsx
+++ /dev/null
@@ -1,316 +0,0 @@
-import { ApolloClient, ApolloProvider, InMemoryCache } from '@apollo/client';
-import { renderHook } from '@testing-library/react-hooks';
-import React from 'react';
-import { beforeEach, describe, expect, it, vi } from 'vitest';
-
-import { useDocumentChildren } from '@app/document/hooks/useDocumentChildren';
-
-import { SearchDocumentsDocument } from '@graphql/document.generated';
-import { DocumentState } from '@types';
-
-describe('useDocumentChildren', () => {
- let mockClient: ApolloClient;
-
- beforeEach(() => {
- vi.clearAllMocks();
- console.log = vi.fn(); // Suppress console.log in tests
- console.error = vi.fn(); // Suppress console.error in tests
-
- mockClient = new ApolloClient({
- cache: new InMemoryCache(),
- defaultOptions: {
- query: {
- fetchPolicy: 'no-cache',
- },
- },
- });
- });
-
- const wrapper = ({ children }: { children: React.ReactNode }) => (
- {children}
- );
-
- describe('checkForChildren', () => {
- it('should return empty map for empty parent URNs array', async () => {
- const mockQuery = vi.spyOn(mockClient, 'query');
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- // Mock useApolloClient
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const childrenMap = await result.current.checkForChildren([]);
-
- expect(childrenMap).toEqual({});
- expect(mockQuery).not.toHaveBeenCalled();
- });
-
- it('should check for children of multiple parents', async () => {
- const parentUrns = ['urn:li:document:parent1', 'urn:li:document:parent2'];
-
- const mockDocuments = [
- {
- urn: 'urn:li:document:child1',
- info: {
- title: 'Child 1',
- parentDocument: {
- document: {
- urn: 'urn:li:document:parent1',
- },
- },
- },
- },
- {
- urn: 'urn:li:document:child2',
- info: {
- title: 'Child 2',
- parentDocument: {
- document: {
- urn: 'urn:li:document:parent1',
- },
- },
- },
- },
- ];
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: {
- searchDocuments: {
- documents: mockDocuments,
- total: 2,
- },
- },
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const childrenMap = await result.current.checkForChildren(parentUrns);
-
- expect(childrenMap).toEqual({
- 'urn:li:document:parent1': true,
- 'urn:li:document:parent2': false,
- });
-
- expect(mockClient.query).toHaveBeenCalledWith({
- query: SearchDocumentsDocument,
- variables: {
- input: {
- query: '*',
- parentDocuments: parentUrns,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
- start: 0,
- count: 200, // 100 * 2 parents
- },
- },
- fetchPolicy: 'network-only',
- });
- });
-
- it('should handle all parents having children', async () => {
- const parentUrns = ['urn:li:document:parent1', 'urn:li:document:parent2'];
-
- const mockDocuments = [
- {
- urn: 'urn:li:document:child1',
- info: {
- parentDocument: {
- document: {
- urn: 'urn:li:document:parent1',
- },
- },
- },
- },
- {
- urn: 'urn:li:document:child2',
- info: {
- parentDocument: {
- document: {
- urn: 'urn:li:document:parent2',
- },
- },
- },
- },
- ];
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: {
- searchDocuments: {
- documents: mockDocuments,
- },
- },
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const childrenMap = await result.current.checkForChildren(parentUrns);
-
- expect(childrenMap).toEqual({
- 'urn:li:document:parent1': true,
- 'urn:li:document:parent2': true,
- });
- });
-
- it('should handle errors gracefully', async () => {
- const parentUrns = ['urn:li:document:parent1'];
-
- mockClient.query = vi.fn().mockRejectedValue(new Error('Network error'));
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const childrenMap = await result.current.checkForChildren(parentUrns);
-
- expect(childrenMap).toEqual({});
- });
- });
-
- describe('fetchChildren', () => {
- it('should fetch children for a parent document', async () => {
- const parentUrn = 'urn:li:document:parent1';
-
- const mockDocuments = [
- {
- urn: 'urn:li:document:child1',
- info: {
- title: 'Child 1',
- },
- },
- {
- urn: 'urn:li:document:child2',
- info: {
- title: 'Child 2',
- },
- },
- ];
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: {
- searchDocuments: {
- documents: mockDocuments,
- },
- },
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const children = await result.current.fetchChildren(parentUrn);
-
- expect(children).toEqual([
- { urn: 'urn:li:document:child1', title: 'Child 1' },
- { urn: 'urn:li:document:child2', title: 'Child 2' },
- ]);
-
- expect(mockClient.query).toHaveBeenCalledWith({
- query: SearchDocumentsDocument,
- variables: {
- input: {
- query: '*',
- parentDocument: parentUrn,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
- start: 0,
- count: 100,
- },
- },
- fetchPolicy: 'cache-first',
- });
- });
-
- it('should use "New Document" as default title when title is missing', async () => {
- const parentUrn = 'urn:li:document:parent1';
-
- const mockDocuments = [
- {
- urn: 'urn:li:document:child1',
- info: {
- title: null,
- },
- },
- ];
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: {
- searchDocuments: {
- documents: mockDocuments,
- },
- },
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const children = await result.current.fetchChildren(parentUrn);
-
- expect(children).toEqual([{ urn: 'urn:li:document:child1', title: 'New Document' }]);
- });
-
- it('should return empty array when no children found', async () => {
- const parentUrn = 'urn:li:document:parent1';
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: {
- searchDocuments: {
- documents: [],
- },
- },
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const children = await result.current.fetchChildren(parentUrn);
-
- expect(children).toEqual([]);
- });
-
- it('should handle errors gracefully', async () => {
- const parentUrn = 'urn:li:document:parent1';
-
- mockClient.query = vi.fn().mockRejectedValue(new Error('Network error'));
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const children = await result.current.fetchChildren(parentUrn);
-
- expect(children).toEqual([]);
- });
-
- it('should handle GraphQL errors', async () => {
- const parentUrn = 'urn:li:document:parent1';
-
- mockClient.query = vi.fn().mockResolvedValue({
- data: null,
- error: new Error('GraphQL error'),
- errors: [{ message: 'GraphQL error' }],
- });
-
- // Note: useApolloClient is mocked via module mock in real implementation
-
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- const children = await result.current.fetchChildren(parentUrn);
-
- expect(children).toEqual([]);
- });
- });
-
- describe('loading state', () => {
- it('should always return loading as false', () => {
- const { result } = renderHook(() => useDocumentChildren(), { wrapper });
-
- expect(result.current.loading).toBe(false);
- });
- });
-});
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/useDocumentTreeMutations.test.tsx b/datahub-web-react/src/app/document/hooks/__tests__/useDocumentTreeMutations.test.tsx
index 42cd968683b042..752ca856e30568 100644
--- a/datahub-web-react/src/app/document/hooks/__tests__/useDocumentTreeMutations.test.tsx
+++ b/datahub-web-react/src/app/document/hooks/__tests__/useDocumentTreeMutations.test.tsx
@@ -3,9 +3,10 @@ import { waitFor } from '@testing-library/react';
import { renderHook } from '@testing-library/react-hooks';
import { message } from 'antd';
import React from 'react';
-import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-import { DocumentTreeContext } from '@app/document/DocumentTreeContext';
+import analytics, { DocumentEditType, EventType } from '@app/analytics';
+import { DocumentTreeNode, DocumentTreeProvider, useDocumentTree } from '@app/document/DocumentTreeContext';
import {
useCreateDocumentTreeMutation,
useDeleteDocumentTreeMutation,
@@ -29,32 +30,49 @@ vi.mock('antd', () => ({
},
}));
+// Mock analytics
+vi.mock('@app/analytics', () => ({
+ default: {
+ event: vi.fn(),
+ },
+ EventType: {
+ CreateDocumentEvent: 'CreateDocumentEvent',
+ EditDocumentEvent: 'EditDocumentEvent',
+ MoveDocumentEvent: 'MoveDocumentEvent',
+ DeleteDocumentEvent: 'DeleteDocumentEvent',
+ },
+ DocumentEditType: {
+ Title: 'Title',
+ },
+}));
+
+const createWrapper = (mocks: any[]) => {
+ return ({ children }: { children: React.ReactNode }) => (
+
+ {children}
+
+ );
+};
+
describe('useDocumentTreeMutations', () => {
beforeEach(() => {
vi.clearAllMocks();
- console.error = vi.fn(); // Suppress console.error in tests
+ console.error = vi.fn();
+ });
+
+ afterEach(() => {
+ vi.restoreAllMocks();
});
describe('useCreateDocumentTreeMutation', () => {
- it('should create a document and update tree optimistically', async () => {
- const mockAddNode = vi.fn();
- const mockDeleteNode = vi.fn();
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: mockAddNode,
- deleteNode: mockDeleteNode,
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: vi.fn(),
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
+ it('should successfully create a root document', async () => {
+ const mockUrn = 'urn:li:document:123';
+ const input = {
+ title: 'New Document',
+ parentDocument: null,
+ subType: 'guide',
};
- const newUrn = 'urn:li:document:new123';
const mocks = [
{
request: {
@@ -63,80 +81,104 @@ describe('useDocumentTreeMutations', () => {
input: {
title: 'New Document',
parentDocument: null,
- subType: undefined,
- contents: { text: '' },
+ subType: 'guide',
state: DocumentState.Published,
+ contents: { text: '' },
settings: { showInGlobalContext: true },
},
},
},
result: {
data: {
- createDocument: newUrn,
+ createDocument: mockUrn,
},
},
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ createMutation: useCreateDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useCreateDocumentTreeMutation(), { wrapper });
+ const newUrn = await result.current.createMutation.createDocument(input);
- const createdUrnPromise = result.current.createDocument({
- title: 'New Document',
- parentDocument: null,
+ await waitFor(() => {
+ expect(newUrn).toBe(mockUrn);
+ // Verify analytics was called
+ expect(analytics.event).toHaveBeenCalledWith({
+ type: EventType.CreateDocumentEvent,
+ documentUrn: mockUrn,
+ documentType: 'guide',
+ hasParent: false,
+ parentDocumentUrn: undefined,
+ });
+ // Verify node was added to tree
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node).toBeTruthy();
+ expect(node?.title).toBe('New Document');
});
+ });
- // Should add optimistic node immediately
- await waitFor(() => {
- expect(mockAddNode).toHaveBeenCalledWith(
- expect.objectContaining({
- title: 'New Document',
- parentUrn: null,
- hasChildren: false,
- }),
- );
+ it('should successfully create a child document', async () => {
+ const parentUrn = 'urn:li:document:parent';
+ const mockUrn = 'urn:li:document:child';
+ const input = {
+ title: 'Child Document',
+ parentDocument: parentUrn,
+ subType: 'reference',
+ };
+
+ const mocks = [
+ {
+ request: {
+ query: CreateDocumentDocument,
+ variables: {
+ input: {
+ title: 'Child Document',
+ parentDocument: parentUrn,
+ subType: 'reference',
+ state: DocumentState.Published,
+ contents: { text: '' },
+ settings: { showInGlobalContext: true },
+ },
+ },
+ },
+ result: {
+ data: {
+ createDocument: mockUrn,
+ },
+ },
+ },
+ ];
+
+ const { result } = renderHook(() => useCreateDocumentTreeMutation(), {
+ wrapper: createWrapper(mocks),
});
- const urn = await createdUrnPromise;
+ const newUrn = await result.current.createDocument(input);
- // Should replace temp node with real node
await waitFor(() => {
- expect(mockDeleteNode).toHaveBeenCalled();
- expect(mockAddNode).toHaveBeenCalledWith(
- expect.objectContaining({
- urn: newUrn,
- title: 'New Document',
- parentUrn: null,
- }),
- );
+ expect(newUrn).toBe(mockUrn);
+ expect(analytics.event).toHaveBeenCalledWith({
+ type: EventType.CreateDocumentEvent,
+ documentUrn: mockUrn,
+ documentType: 'reference',
+ hasParent: true,
+ parentDocumentUrn: parentUrn,
+ });
});
-
- expect(urn).toBe(newUrn);
});
- it('should rollback on error', async () => {
- const mockAddNode = vi.fn();
- const mockDeleteNode = vi.fn();
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: mockAddNode,
- deleteNode: mockDeleteNode,
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: vi.fn(),
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
+ it('should rollback on mutation error', async () => {
+ const input = {
+ title: 'Failed Document',
+ parentDocument: null,
};
const mocks = [
@@ -148,140 +190,89 @@ describe('useDocumentTreeMutations', () => {
title: 'Failed Document',
parentDocument: null,
subType: undefined,
- contents: { text: '' },
state: DocumentState.Published,
+ contents: { text: '' },
settings: { showInGlobalContext: true },
},
},
},
- error: new Error('Server error'),
+ error: new Error('Network error'),
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ createMutation: useCreateDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useCreateDocumentTreeMutation(), { wrapper });
+ const newUrn = await result.current.createMutation.createDocument(input);
- const createdUrn = await result.current.createDocument({
- title: 'Failed Document',
- parentDocument: null,
+ await waitFor(() => {
+ expect(newUrn).toBe(null);
+ expect(message.error).toHaveBeenCalledWith('Failed to create document');
+ // Verify no nodes remain in tree (temp node was rolled back)
+ const rootNodes = result.current.tree.getRootNodes();
+ expect(rootNodes).toHaveLength(0);
});
-
- expect(createdUrn).toBeNull();
- expect(message.error).toHaveBeenCalledWith('Failed to create document');
- // Should have called deleteNode twice: once for the temp node after success, once for rollback
- expect(mockDeleteNode).toHaveBeenCalled();
});
- it('should handle creating document with parent and subType', async () => {
- const mockAddNode = vi.fn();
- const mockDeleteNode = vi.fn();
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: mockAddNode,
- deleteNode: mockDeleteNode,
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: vi.fn(),
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
+ it('should rollback when mutation returns no URN', async () => {
+ const input = {
+ title: 'Failed Document',
+ parentDocument: null,
};
- const newUrn = 'urn:li:document:child123';
- const parentUrn = 'urn:li:document:parent456';
const mocks = [
{
request: {
query: CreateDocumentDocument,
variables: {
input: {
- title: 'Child Document',
- parentDocument: parentUrn,
- subType: 'guide',
- contents: { text: '' },
+ title: 'Failed Document',
+ parentDocument: null,
+ subType: undefined,
state: DocumentState.Published,
+ contents: { text: '' },
settings: { showInGlobalContext: true },
},
},
},
result: {
data: {
- createDocument: newUrn,
+ createDocument: null,
},
},
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ createMutation: useCreateDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useCreateDocumentTreeMutation(), { wrapper });
+ const newUrn = await result.current.createMutation.createDocument(input);
- const createdUrnPromise = result.current.createDocument({
- title: 'Child Document',
- parentDocument: parentUrn,
- subType: 'guide',
- });
-
- // Wait for optimistic node to be added
await waitFor(() => {
- expect(mockAddNode).toHaveBeenCalled();
+ expect(newUrn).toBe(null);
+ expect(message.error).toHaveBeenCalledWith('Failed to create document');
});
-
- const createdUrn = await createdUrnPromise;
-
- expect(createdUrn).toBe(newUrn);
- // Should have been called twice: once with temp URN, once with real URN
- expect(mockAddNode).toHaveBeenCalledTimes(2);
- expect(mockAddNode).toHaveBeenLastCalledWith(
- expect.objectContaining({
- urn: newUrn,
- title: 'Child Document',
- parentUrn,
- }),
- );
});
});
describe('useUpdateDocumentTitleMutation', () => {
- it('should update title optimistically and persist', async () => {
- const mockUpdateNodeTitle = vi.fn();
- const mockGetNode = vi.fn(() => ({
- urn: 'urn:li:document:123',
- title: 'Old Title',
- parentUrn: null,
- hasChildren: false,
- children: [],
- }));
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: mockUpdateNodeTitle,
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
- };
+ it('should successfully update document title', async () => {
+ const mockUrn = 'urn:li:document:123';
+ const newTitle = 'Updated Title';
const mocks = [
{
@@ -289,8 +280,8 @@ describe('useDocumentTreeMutations', () => {
query: UpdateDocumentContentsDocument,
variables: {
input: {
- urn: 'urn:li:document:123',
- title: 'New Title',
+ urn: mockUrn,
+ title: newTitle,
},
},
},
@@ -302,45 +293,95 @@ describe('useDocumentTreeMutations', () => {
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ updateMutation: useUpdateDocumentTitleMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useUpdateDocumentTitleMutation(), { wrapper });
+ // Add node to tree first
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Old Title',
+ parentUrn: null,
+ hasChildren: false,
+ children: [],
+ };
+ result.current.tree.addNode(initialNode);
- const success = await result.current.updateTitle('urn:li:document:123', 'New Title');
+ const success = await result.current.updateMutation.updateTitle(mockUrn, newTitle);
- expect(success).toBe(true);
- expect(mockUpdateNodeTitle).toHaveBeenCalledWith('urn:li:document:123', 'New Title');
+ await waitFor(() => {
+ expect(success).toBe(true);
+ expect(analytics.event).toHaveBeenCalledWith({
+ type: EventType.EditDocumentEvent,
+ documentUrn: mockUrn,
+ editType: DocumentEditType.Title,
+ });
+ // Verify title was updated in tree
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node?.title).toBe(newTitle);
+ });
});
- it('should rollback title on error', async () => {
- const mockUpdateNodeTitle = vi.fn();
- const mockGetNode = vi.fn(() => ({
- urn: 'urn:li:document:123',
- title: 'Old Title',
+ it('should rollback on mutation error', async () => {
+ const mockUrn = 'urn:li:document:123';
+ const oldTitle = 'Old Title';
+ const newTitle = 'Failed Title';
+
+ const mocks = [
+ {
+ request: {
+ query: UpdateDocumentContentsDocument,
+ variables: {
+ input: {
+ urn: mockUrn,
+ title: newTitle,
+ },
+ },
+ },
+ error: new Error('Update failed'),
+ },
+ ];
+
+ const { result } = renderHook(
+ () => ({
+ updateMutation: useUpdateDocumentTitleMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
+ );
+
+ // Add node to tree first
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: oldTitle,
parentUrn: null,
hasChildren: false,
children: [],
- }));
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: mockUpdateNodeTitle,
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
};
+ result.current.tree.addNode(initialNode);
+
+ const success = await result.current.updateMutation.updateTitle(mockUrn, newTitle);
+
+ await waitFor(() => {
+ expect(success).toBe(false);
+ expect(message.error).toHaveBeenCalledWith('Failed to update title');
+ // Verify title was rolled back
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node?.title).toBe(oldTitle);
+ });
+ });
+
+ it('should handle updating non-existent document gracefully', async () => {
+ const mockUrn = 'urn:li:document:nonexistent';
+ const newTitle = 'New Title';
const mocks = [
{
@@ -348,59 +389,36 @@ describe('useDocumentTreeMutations', () => {
query: UpdateDocumentContentsDocument,
variables: {
input: {
- urn: 'urn:li:document:123',
- title: 'New Title',
+ urn: mockUrn,
+ title: newTitle,
},
},
},
- error: new Error('Server error'),
+ result: {
+ data: {
+ updateDocumentContents: true,
+ },
+ },
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
- );
-
- const { result } = renderHook(() => useUpdateDocumentTitleMutation(), { wrapper });
+ const { result } = renderHook(() => useUpdateDocumentTitleMutation(), {
+ wrapper: createWrapper(mocks),
+ });
- const success = await result.current.updateTitle('urn:li:document:123', 'New Title');
+ const success = await result.current.updateTitle(mockUrn, newTitle);
- expect(success).toBe(false);
- expect(message.error).toHaveBeenCalledWith('Failed to update title');
- // Should have called updateNodeTitle twice: once optimistically, once to rollback
- expect(mockUpdateNodeTitle).toHaveBeenCalledTimes(2);
- expect(mockUpdateNodeTitle).toHaveBeenLastCalledWith('urn:li:document:123', 'Old Title');
+ await waitFor(() => {
+ expect(success).toBe(true);
+ });
});
});
describe('useMoveDocumentTreeMutation', () => {
- it('should move document optimistically and persist', async () => {
- const mockMoveNode = vi.fn();
- const mockGetNode = vi.fn(() => ({
- urn: 'urn:li:document:123',
- title: 'Document',
- parentUrn: 'urn:li:document:oldParent',
- hasChildren: false,
- children: [],
- }));
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: vi.fn(),
- moveNode: mockMoveNode,
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
- };
+ it('should successfully move document to new parent', async () => {
+ const mockUrn = 'urn:li:document:child';
+ const oldParentUrn = 'urn:li:document:oldParent';
+ const newParentUrn = 'urn:li:document:newParent';
const mocks = [
{
@@ -408,8 +426,8 @@ describe('useDocumentTreeMutations', () => {
query: MoveDocumentDocument,
variables: {
input: {
- urn: 'urn:li:document:123',
- parentDocument: 'urn:li:document:newParent',
+ urn: mockUrn,
+ parentDocument: newParentUrn,
},
},
},
@@ -421,46 +439,98 @@ describe('useDocumentTreeMutations', () => {
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ moveMutation: useMoveDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useMoveDocumentTreeMutation(), { wrapper });
+ // Add node to tree first
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Child Document',
+ parentUrn: oldParentUrn,
+ hasChildren: false,
+ children: [],
+ };
+ result.current.tree.addNode(initialNode);
- const success = await result.current.moveDocument('urn:li:document:123', 'urn:li:document:newParent');
+ const success = await result.current.moveMutation.moveDocument(mockUrn, newParentUrn);
- expect(success).toBe(true);
- expect(mockMoveNode).toHaveBeenCalledWith('urn:li:document:123', 'urn:li:document:newParent');
- expect(message.success).toHaveBeenCalledWith('Document moved successfully');
+ await waitFor(() => {
+ expect(success).toBe(true);
+ expect(message.success).toHaveBeenCalledWith('Document moved successfully');
+ expect(analytics.event).toHaveBeenCalledWith({
+ type: EventType.MoveDocumentEvent,
+ documentUrn: mockUrn,
+ oldParentDocumentUrn: oldParentUrn,
+ newParentDocumentUrn: newParentUrn,
+ });
+ // Verify parent was updated in tree
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node?.parentUrn).toBe(newParentUrn);
+ });
});
- it('should rollback move on error', async () => {
- const mockMoveNode = vi.fn();
- const mockGetNode = vi.fn(() => ({
- urn: 'urn:li:document:123',
- title: 'Document',
- parentUrn: 'urn:li:document:oldParent',
+ it('should successfully move document to root level', async () => {
+ const mockUrn = 'urn:li:document:child';
+ const oldParentUrn = 'urn:li:document:parent';
+
+ const mocks = [
+ {
+ request: {
+ query: MoveDocumentDocument,
+ variables: {
+ input: {
+ urn: mockUrn,
+ parentDocument: null,
+ },
+ },
+ },
+ result: {
+ data: {
+ moveDocument: true,
+ },
+ },
+ },
+ ];
+
+ const { result } = renderHook(
+ () => ({
+ moveMutation: useMoveDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
+ );
+
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Child Document',
+ parentUrn: oldParentUrn,
hasChildren: false,
children: [],
- }));
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: vi.fn(),
- moveNode: mockMoveNode,
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
};
+ result.current.tree.addNode(initialNode);
+
+ const success = await result.current.moveMutation.moveDocument(mockUrn, null);
+
+ await waitFor(() => {
+ expect(success).toBe(true);
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node?.parentUrn).toBe(null);
+ });
+ });
+
+ it('should rollback on mutation error', async () => {
+ const mockUrn = 'urn:li:document:child';
+ const oldParentUrn = 'urn:li:document:oldParent';
+ const newParentUrn = 'urn:li:document:newParent';
const mocks = [
{
@@ -468,99 +538,121 @@ describe('useDocumentTreeMutations', () => {
query: MoveDocumentDocument,
variables: {
input: {
- urn: 'urn:li:document:123',
- parentDocument: 'urn:li:document:newParent',
+ urn: mockUrn,
+ parentDocument: newParentUrn,
},
},
},
- error: new Error('Server error'),
+ error: new Error('Move failed'),
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ moveMutation: useMoveDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useMoveDocumentTreeMutation(), { wrapper });
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Child Document',
+ parentUrn: oldParentUrn,
+ hasChildren: false,
+ children: [],
+ };
+ result.current.tree.addNode(initialNode);
- const success = await result.current.moveDocument('urn:li:document:123', 'urn:li:document:newParent');
+ const success = await result.current.moveMutation.moveDocument(mockUrn, newParentUrn);
- expect(success).toBe(false);
- expect(message.error).toHaveBeenCalledWith('Failed to move document');
- // Should have called moveNode twice: once optimistically, once to rollback
- expect(mockMoveNode).toHaveBeenCalledTimes(2);
- expect(mockMoveNode).toHaveBeenLastCalledWith('urn:li:document:123', 'urn:li:document:oldParent');
+ await waitFor(() => {
+ expect(success).toBe(false);
+ expect(message.error).toHaveBeenCalledWith('Failed to move document');
+ // Verify parent was rolled back
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node?.parentUrn).toBe(oldParentUrn);
+ });
});
it('should return false when document not found in tree', async () => {
- const mockGetNode = vi.fn(() => undefined);
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
- };
+ const mockUrn = 'urn:li:document:nonexistent';
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
- );
+ const mocks: any[] = [];
- const { result } = renderHook(() => useMoveDocumentTreeMutation(), { wrapper });
+ const { result } = renderHook(() => useMoveDocumentTreeMutation(), {
+ wrapper: createWrapper(mocks),
+ });
- const success = await result.current.moveDocument('urn:li:document:nonexistent', 'urn:li:document:parent');
+ const success = await result.current.moveDocument(mockUrn, null);
expect(success).toBe(false);
});
});
describe('useDeleteDocumentTreeMutation', () => {
- it('should delete document optimistically and persist', async () => {
- const mockDeleteNode = vi.fn();
- const mockGetNode = vi.fn(() => ({
- urn: 'urn:li:document:123',
+ it('should successfully delete document that exists in tree', async () => {
+ const mockUrn = 'urn:li:document:123';
+
+ const mocks = [
+ {
+ request: {
+ query: DeleteDocumentDocument,
+ variables: { urn: mockUrn },
+ },
+ result: {
+ data: {
+ deleteDocument: true,
+ },
+ },
+ },
+ ];
+
+ const { result } = renderHook(
+ () => ({
+ deleteMutation: useDeleteDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
+ );
+
+ // Add node to tree first
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
title: 'Document to Delete',
parentUrn: null,
hasChildren: false,
children: [],
- }));
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: mockDeleteNode,
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
};
+ result.current.tree.addNode(initialNode);
+
+ const success = await result.current.deleteMutation.deleteDocument(mockUrn);
+
+ await waitFor(() => {
+ expect(success).toBe(true);
+ expect(message.success).toHaveBeenCalledWith('Document deleted');
+ expect(analytics.event).toHaveBeenCalledWith({
+ type: EventType.DeleteDocumentEvent,
+ documentUrn: mockUrn,
+ });
+ // Verify node was removed from tree
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node).toBe(undefined);
+ });
+ });
+
+ it('should successfully delete document not in tree (e.g., modal)', async () => {
+ const mockUrn = 'urn:li:document:notInTree';
const mocks = [
{
request: {
query: DeleteDocumentDocument,
- variables: {
- urn: 'urn:li:document:123',
- },
+ variables: { urn: mockUrn },
},
result: {
data: {
@@ -570,109 +662,107 @@ describe('useDocumentTreeMutations', () => {
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
- );
-
- const { result } = renderHook(() => useDeleteDocumentTreeMutation(), { wrapper });
+ const { result } = renderHook(() => useDeleteDocumentTreeMutation(), {
+ wrapper: createWrapper(mocks),
+ });
- const success = await result.current.deleteDocument('urn:li:document:123');
+ const success = await result.current.deleteDocument(mockUrn);
- expect(success).toBe(true);
- expect(mockDeleteNode).toHaveBeenCalledWith('urn:li:document:123');
- expect(message.success).toHaveBeenCalledWith('Document deleted');
+ await waitFor(() => {
+ expect(success).toBe(true);
+ expect(message.success).toHaveBeenCalledWith('Document deleted');
+ });
});
- it('should rollback delete on error', async () => {
- const mockDeleteNode = vi.fn();
- const mockAddNode = vi.fn();
- const savedNode = {
- urn: 'urn:li:document:123',
- title: 'Document to Delete',
- parentUrn: null,
- hasChildren: false,
- children: [],
- };
- const mockGetNode = vi.fn(() => savedNode);
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: mockAddNode,
- deleteNode: mockDeleteNode,
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
- };
+ it('should rollback on mutation error', async () => {
+ const mockUrn = 'urn:li:document:123';
const mocks = [
{
request: {
query: DeleteDocumentDocument,
- variables: {
- urn: 'urn:li:document:123',
- },
+ variables: { urn: mockUrn },
},
- error: new Error('Server error'),
+ error: new Error('Delete failed'),
},
];
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const { result } = renderHook(
+ () => ({
+ deleteMutation: useDeleteDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useDeleteDocumentTreeMutation(), { wrapper });
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Document to Delete',
+ parentUrn: null,
+ hasChildren: false,
+ children: [],
+ };
+ result.current.tree.addNode(initialNode);
- const success = await result.current.deleteDocument('urn:li:document:123');
+ const success = await result.current.deleteMutation.deleteDocument(mockUrn);
- expect(success).toBe(false);
- expect(message.error).toHaveBeenCalledWith('Failed to delete document');
- expect(mockDeleteNode).toHaveBeenCalledWith('urn:li:document:123');
- expect(mockAddNode).toHaveBeenCalledWith(savedNode);
+ await waitFor(() => {
+ expect(success).toBe(false);
+ expect(message.error).toHaveBeenCalledWith('Failed to delete document');
+ // Verify node was restored in tree
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node).toBeTruthy();
+ expect(node?.title).toBe('Document to Delete');
+ });
});
- it('should return false when document not found in tree', async () => {
- const mockGetNode = vi.fn(() => undefined);
-
- const mockContextValue = {
- nodes: new Map(),
- rootUrns: [],
- addNode: vi.fn(),
- deleteNode: vi.fn(),
- updateNodeTitle: vi.fn(),
- moveNode: vi.fn(),
- getNode: mockGetNode,
- getRootNodes: vi.fn(),
- getChildren: vi.fn(),
- initializeTree: vi.fn(),
- setNodeChildren: vi.fn(),
- };
+ it('should rollback when mutation returns false', async () => {
+ const mockUrn = 'urn:li:document:123';
- const wrapper = ({ children }: any) => (
-
-
- {children}
-
-
+ const mocks = [
+ {
+ request: {
+ query: DeleteDocumentDocument,
+ variables: { urn: mockUrn },
+ },
+ result: {
+ data: {
+ deleteDocument: false,
+ },
+ },
+ },
+ ];
+
+ const { result } = renderHook(
+ () => ({
+ deleteMutation: useDeleteDocumentTreeMutation(),
+ tree: useDocumentTree(),
+ }),
+ {
+ wrapper: createWrapper(mocks),
+ },
);
- const { result } = renderHook(() => useDeleteDocumentTreeMutation(), { wrapper });
+ const initialNode: DocumentTreeNode = {
+ urn: mockUrn,
+ title: 'Document to Delete',
+ parentUrn: null,
+ hasChildren: false,
+ children: [],
+ };
+ result.current.tree.addNode(initialNode);
- const success = await result.current.deleteDocument('urn:li:document:nonexistent');
+ const success = await result.current.deleteMutation.deleteDocument(mockUrn);
- expect(success).toBe(false);
+ await waitFor(() => {
+ expect(success).toBe(false);
+ expect(message.error).toHaveBeenCalledWith('Failed to delete document');
+ // Verify node was restored
+ const node = result.current.tree.getNode(mockUrn);
+ expect(node).toBeTruthy();
+ });
});
});
});
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/useLoadDocumentTree.test.tsx b/datahub-web-react/src/app/document/hooks/__tests__/useLoadDocumentTree.test.tsx
index d685c354820769..c623177a78a371 100644
--- a/datahub-web-react/src/app/document/hooks/__tests__/useLoadDocumentTree.test.tsx
+++ b/datahub-web-react/src/app/document/hooks/__tests__/useLoadDocumentTree.test.tsx
@@ -9,7 +9,6 @@ import { useLoadDocumentTree } from '@app/document/hooks/useLoadDocumentTree';
import * as useSearchDocumentsModule from '@app/document/hooks/useSearchDocuments';
import { SearchDocumentsDocument } from '@graphql/document.generated';
-import { DocumentState } from '@types';
vi.mock('../useSearchDocuments');
@@ -305,8 +304,6 @@ describe('useLoadDocumentTree', () => {
input: {
query: '*',
parentDocuments: urns,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
start: 0,
count: 200, // 2 * 100
},
@@ -441,6 +438,22 @@ describe('useLoadDocumentTree', () => {
expect(children[0].urn).toBe('urn:li:document:child1'); // Sorted by time DESC
expect(children[1].urn).toBe('urn:li:document:child2');
expect(mockSetNodeChildren).toHaveBeenCalledWith(parentUrn, children);
+
+ // Verify the first query call (loadChildren) matches the implementation
+ const queryMock = vi.mocked(mockClient.query);
+ const firstCall = queryMock.mock.calls[0];
+ expect(firstCall[0]).toMatchObject({
+ query: SearchDocumentsDocument,
+ variables: {
+ input: {
+ query: '*',
+ parentDocuments: [parentUrn],
+ start: 0,
+ count: 100,
+ },
+ },
+ fetchPolicy: 'cache-first',
+ });
});
it('should handle errors in loadChildren', async () => {
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/useRelatedDocuments.test.tsx b/datahub-web-react/src/app/document/hooks/__tests__/useRelatedDocuments.test.tsx
new file mode 100644
index 00000000000000..18590f2bb6cc51
--- /dev/null
+++ b/datahub-web-react/src/app/document/hooks/__tests__/useRelatedDocuments.test.tsx
@@ -0,0 +1,184 @@
+import { describe, expect, it } from 'vitest';
+
+import { GetRelatedDocumentsQuery, RelatedDocumentsFieldsFragment } from '@graphql/document.generated';
+
+describe('useRelatedDocuments - type extraction', () => {
+ const mockRelatedDocumentsResult: RelatedDocumentsFieldsFragment = {
+ __typename: 'RelatedDocumentsResult',
+ start: 0,
+ count: 2,
+ total: 2,
+ documents: [
+ {
+ __typename: 'Document',
+ urn: 'urn:li:document:1',
+ type: 'DOCUMENT' as any,
+ info: {
+ __typename: 'DocumentInfo',
+ title: 'Test Document 1',
+ lastModified: {
+ __typename: 'ResolvedAuditStamp',
+ time: 1000,
+ actor: {
+ __typename: 'CorpUser',
+ urn: 'urn:li:corpuser:test',
+ type: 'CORP_USER' as any,
+ username: 'testuser',
+ } as any,
+ },
+ },
+ },
+ ],
+ };
+
+ // Test the type extraction logic by creating mock entity structures
+ // that match the GraphQL query response types
+
+ it('should extract relatedDocuments from Dataset entity', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'Dataset',
+ urn: 'urn:li:dataset:test',
+ type: 'DATASET' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ };
+
+ // The type guard should work - entity has relatedDocuments
+ expect(entity).toBeDefined();
+ if (entity && 'relatedDocuments' in entity && entity.relatedDocuments) {
+ expect(entity.relatedDocuments).toEqual(mockRelatedDocumentsResult);
+ expect(entity.relatedDocuments.total).toBe(2);
+ } else {
+ throw new Error('Type guard failed - relatedDocuments should be accessible');
+ }
+ });
+
+ it('should extract relatedDocuments from Dashboard entity', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'Dashboard',
+ urn: 'urn:li:dashboard:test',
+ type: 'DASHBOARD' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ };
+
+ if (entity && 'relatedDocuments' in entity && entity.relatedDocuments) {
+ expect(entity.relatedDocuments).toEqual(mockRelatedDocumentsResult);
+ } else {
+ throw new Error('Type guard failed');
+ }
+ });
+
+ it('should return null when entity does not have relatedDocuments field', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'CorpUser',
+ urn: 'urn:li:corpuser:test',
+ type: 'CORP_USER' as any,
+ // No relatedDocuments field
+ };
+
+ // Type guard should fail - CorpUser doesn't have relatedDocuments
+ if (entity && 'relatedDocuments' in entity) {
+ throw new Error('Type guard should not pass for CorpUser');
+ }
+ if (entity) {
+ expect('relatedDocuments' in entity).toBe(false);
+ }
+ });
+
+ it('should handle null entity', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = null;
+
+ expect(entity).toBeNull();
+ if (entity) {
+ throw new Error('Entity should be null');
+ }
+ });
+
+ it('should handle entity with null relatedDocuments', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'Dataset',
+ urn: 'urn:li:dataset:test',
+ type: 'DATASET' as any,
+ relatedDocuments: null,
+ };
+
+ // Type guard should pass (field exists) but value is null
+ if (entity && 'relatedDocuments' in entity) {
+ expect(entity.relatedDocuments).toBeNull();
+ } else {
+ throw new Error('Type guard should pass - field exists');
+ }
+ });
+
+ it('should handle entity with undefined relatedDocuments', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'Dataset',
+ urn: 'urn:li:dataset:test',
+ type: 'DATASET' as any,
+ relatedDocuments: undefined,
+ };
+
+ if (entity && 'relatedDocuments' in entity) {
+ expect(entity.relatedDocuments).toBeUndefined();
+ } else {
+ throw new Error('Type guard should pass - field exists');
+ }
+ });
+
+ it('should work with multiple entity types that support relatedDocuments', () => {
+ const entityTypes: Array = [
+ {
+ __typename: 'Chart',
+ urn: 'urn:li:chart:test',
+ type: 'CHART' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ },
+ {
+ __typename: 'DataJob',
+ urn: 'urn:li:datajob:test',
+ type: 'DATA_JOB' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ },
+ {
+ __typename: 'Container',
+ urn: 'urn:li:container:test',
+ type: 'CONTAINER' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ },
+ ];
+
+ entityTypes.forEach((entity) => {
+ if (entity && 'relatedDocuments' in entity && entity.relatedDocuments) {
+ expect(entity.relatedDocuments.total).toBe(2);
+ } else {
+ throw new Error(`Type guard failed for ${entity?.__typename}`);
+ }
+ });
+ });
+
+ it('should properly narrow types after type guard check', () => {
+ const entity: GetRelatedDocumentsQuery['entity'] = {
+ __typename: 'Dataset',
+ urn: 'urn:li:dataset:test',
+ type: 'DATASET' as any,
+ relatedDocuments: mockRelatedDocumentsResult,
+ };
+
+ // Before type guard - TypeScript doesn't know if relatedDocuments exists
+ if (!entity) {
+ throw new Error('Entity should exist');
+ }
+
+ // After checking 'relatedDocuments' in entity, TypeScript should narrow the type
+ if ('relatedDocuments' in entity) {
+ // At this point, TypeScript knows entity has relatedDocuments property
+ // We can safely access it without type assertions
+ const result = entity.relatedDocuments;
+ if (result) {
+ // TypeScript knows result is RelatedDocumentsFieldsFragment
+ expect(result.__typename).toBe('RelatedDocumentsResult');
+ expect(result.total).toBe(2);
+ expect(result.documents).toHaveLength(1);
+ }
+ }
+ });
+});
diff --git a/datahub-web-react/src/app/document/hooks/__tests__/useSearchDocuments.test.tsx b/datahub-web-react/src/app/document/hooks/__tests__/useSearchDocuments.test.tsx
index f91c049de454ce..3c1fd542252991 100644
--- a/datahub-web-react/src/app/document/hooks/__tests__/useSearchDocuments.test.tsx
+++ b/datahub-web-react/src/app/document/hooks/__tests__/useSearchDocuments.test.tsx
@@ -41,12 +41,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -96,12 +94,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: 'test query',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -148,12 +144,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: 'urn:li:document:parent',
+ parentDocuments: ['urn:li:document:parent'],
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -199,12 +193,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: true,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -250,12 +242,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: ['guide', 'tutorial'],
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -301,12 +291,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -340,7 +328,6 @@ describe('useSearchDocuments', () => {
it('should include drafts when specified', async () => {
const input: SearchDocumentsInput = {
fetchPolicy: 'network-only',
- includeDrafts: true,
};
const mocks = [
@@ -352,12 +339,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: true,
},
includeParentDocuments: false,
},
@@ -404,12 +389,10 @@ describe('useSearchDocuments', () => {
start: 10,
count: 50,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -456,12 +439,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: true,
},
@@ -506,12 +487,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -557,12 +536,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: 'nonexistent',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -608,12 +585,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
@@ -652,12 +627,10 @@ describe('useSearchDocuments', () => {
start: 0,
count: 100,
query: '*',
- parentDocument: undefined,
+ parentDocuments: undefined,
rootOnly: undefined,
types: undefined,
- states: [DocumentState.Published, DocumentState.Unpublished],
sourceType: 'NATIVE',
- includeDrafts: false,
},
includeParentDocuments: false,
},
diff --git a/datahub-web-react/src/app/document/hooks/useDocumentChildren.ts b/datahub-web-react/src/app/document/hooks/useDocumentChildren.ts
deleted file mode 100644
index 2eecf636cc5baf..00000000000000
--- a/datahub-web-react/src/app/document/hooks/useDocumentChildren.ts
+++ /dev/null
@@ -1,123 +0,0 @@
-import { useApolloClient } from '@apollo/client';
-import { useCallback } from 'react';
-
-import { SearchDocumentsDocument } from '@graphql/document.generated';
-import { DocumentState } from '@types';
-
-export interface DocumentChild {
- urn: string;
- title: string;
-}
-
-// Pagination constants
-const CHILD_PAGE_SIZE = 100; // Fetch up to 100 children per parent level
-
-export function useDocumentChildren() {
- const client = useApolloClient();
-
- /**
- * Check if any of the given parent documents have children
- * Returns a map of parentUrn -> hasChildren
- *
- * TODO: Consider refactoring to use useLazyQuery for better type safety
- * once Apollo is updated to support returning data directly from lazy queries
- */
- const checkForChildren = useCallback(
- async (parentUrns: string[]): Promise> => {
- if (parentUrns.length === 0) {
- return {};
- }
-
- try {
- // Initialize all parents as having no children
- const childrenMap: Record = {};
- parentUrns.forEach((urn) => {
- childrenMap[urn] = false;
- });
-
- // Make ONE batch query for all children of all parents
- const result = await client.query({
- query: SearchDocumentsDocument,
- variables: {
- input: {
- query: '*',
- parentDocuments: parentUrns, // Batch query with all parents
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
- start: 0,
- count: CHILD_PAGE_SIZE * parentUrns.length, // Fetch enough for all parents
- },
- },
- fetchPolicy: 'network-only',
- });
-
- // Group children by their parent URN
- const children = result.data?.searchDocuments?.documents || [];
- children.forEach((child) => {
- const parentUrn = child.info?.parentDocument?.document?.urn;
- if (parentUrn && childrenMap.hasOwnProperty(parentUrn)) {
- childrenMap[parentUrn] = true;
- }
- });
-
- return childrenMap;
- } catch (error) {
- console.error('Failed to check for children:', error);
- return {};
- }
- },
- [client],
- );
-
- /**
- * Fetch all children for a specific parent document
- *
- * TODO: Consider refactoring to use useLazyQuery for better type safety
- * once Apollo is updated to support returning data directly from lazy queries
- */
- const fetchChildren = useCallback(
- async (parentUrn: string): Promise => {
- try {
- const result = await client.query({
- query: SearchDocumentsDocument,
- variables: {
- input: {
- query: '*',
- parentDocument: parentUrn,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
- start: 0,
- count: CHILD_PAGE_SIZE, // Limit children per level
- },
- },
- // Use cache-first to return cached data if available (instant!), otherwise fetch from network
- // This makes folder expansion feel instant when we've updated the cache (e.g., after moves/creates)
- fetchPolicy: 'cache-first',
- });
-
- if (!result || result.error || result.errors) {
- console.error('Failed to fetch children:', result?.error || result?.errors);
- return [];
- }
-
- const { data } = result;
-
- const documents = data?.searchDocuments?.documents || [];
- return documents.map((doc) => ({
- urn: doc.urn,
- title: doc.info?.title || 'New Document',
- }));
- } catch (error) {
- console.error('Failed to fetch children:', error);
- return [];
- }
- },
- [client],
- );
-
- return {
- checkForChildren,
- fetchChildren,
- loading: false, // We manage loading state in the tree component
- };
-}
diff --git a/datahub-web-react/src/app/document/hooks/useDocumentTreeMutations.ts b/datahub-web-react/src/app/document/hooks/useDocumentTreeMutations.ts
index 9e31bf3647aeba..b91722a71aa873 100644
--- a/datahub-web-react/src/app/document/hooks/useDocumentTreeMutations.ts
+++ b/datahub-web-react/src/app/document/hooks/useDocumentTreeMutations.ts
@@ -221,19 +221,16 @@ export function useDeleteDocumentTreeMutation() {
const deleteDocument = useCallback(
async (urn: string) => {
- // Get node for rollback
+ // Get node for rollback (may be null if document isn't in tree, e.g., opened in modal)
const node = getNode(urn);
- if (!node) {
- console.error('Document not found in tree:', urn);
- return false;
+ // 1. Optimistically update tree state (only if node exists in tree)
+ if (node) {
+ deleteNode(urn);
}
- // 1. Optimistically update tree state
- deleteNode(urn);
-
try {
- // 2. Call backend mutation
+ // 2. Call backend mutation (always call, even if not in tree)
const result = await deleteDocumentMutation({
variables: { urn },
});
@@ -254,8 +251,10 @@ export function useDeleteDocumentTreeMutation() {
console.error('Failed to delete document:', error);
message.error('Failed to delete document');
- // 3. Rollback on error
- addNode(node);
+ // 3. Rollback on error (only if node was in tree)
+ if (node) {
+ addNode(node);
+ }
return false;
}
diff --git a/datahub-web-react/src/app/document/hooks/useExtractMentions.ts b/datahub-web-react/src/app/document/hooks/useExtractMentions.ts
index 785a8ff2caee37..9a799b640074ca 100644
--- a/datahub-web-react/src/app/document/hooks/useExtractMentions.ts
+++ b/datahub-web-react/src/app/document/hooks/useExtractMentions.ts
@@ -1,39 +1,13 @@
import { useMemo } from 'react';
+import { extractMentions } from '@app/document/utils/extractMentions';
+
/**
* Hook to extract @ mentions (URNs) from markdown text.
* Searches for markdown link patterns like [@Entity](urn:li:entityType:id)
+ *
+ * This hook memoizes the result of extractMentions to avoid recalculating on every render.
*/
export const useExtractMentions = (content: string) => {
- const mentions = useMemo(() => {
- if (!content) return { documentUrns: [], assetUrns: [] };
-
- // Match markdown link syntax: [text](urn:li:entityType:id)
- // Handle URNs with nested parentheses by matching everything between the markdown link's parens
- // The pattern matches: [text](urn:li:entityType:...) where ... can include nested parens
- // We match the URN prefix, then allow nested paren groups or non-paren characters (one or more)
- const urnPattern = /\[([^\]]+)\]\((urn:li:[a-zA-Z]+:(?:[^)(]+|\([^)]*\))+)\)/g;
- const matches = Array.from(content.matchAll(urnPattern));
-
- const documentUrns: string[] = [];
- const assetUrns: string[] = [];
-
- matches.forEach((match) => {
- const urn = match[2]; // URN is in the second capture group (inside parentheses)
-
- // Check if it's a document URN
- if (urn.includes(':document:')) {
- if (!documentUrns.includes(urn)) {
- documentUrns.push(urn);
- }
- } else if (!assetUrns.includes(urn)) {
- // Everything else is considered an asset
- assetUrns.push(urn);
- }
- });
-
- return { documentUrns, assetUrns };
- }, [content]);
-
- return mentions;
+ return useMemo(() => extractMentions(content), [content]);
};
diff --git a/datahub-web-react/src/app/document/hooks/useLoadDocumentTree.ts b/datahub-web-react/src/app/document/hooks/useLoadDocumentTree.ts
index bfdc6db4edbdd2..b6648c2105fad4 100644
--- a/datahub-web-react/src/app/document/hooks/useLoadDocumentTree.ts
+++ b/datahub-web-react/src/app/document/hooks/useLoadDocumentTree.ts
@@ -3,9 +3,9 @@ import { useCallback, useEffect } from 'react';
import { DocumentTreeNode, useDocumentTree } from '@app/document/DocumentTreeContext';
import { useSearchDocuments } from '@app/document/hooks/useSearchDocuments';
+import { documentToTreeNode, sortDocumentsByCreationTime } from '@app/document/utils/documentUtils';
import { SearchDocumentsDocument } from '@graphql/document.generated';
-import { Document, DocumentState } from '@types';
/**
* Hook to load and populate the document tree from backend queries.
@@ -16,16 +16,6 @@ import { Document, DocumentState } from '@types';
* - Loading children on demand
*/
-function documentToTreeNode(doc: Document, hasChildren: boolean): DocumentTreeNode {
- return {
- urn: doc.urn,
- title: doc.info?.title || 'Untitled',
- parentUrn: doc.info?.parentDocument?.document?.urn || null,
- hasChildren,
- children: undefined, // Not loaded yet
- };
-}
-
export function useLoadDocumentTree() {
const { initializeTree, setNodeChildren, getRootNodes } = useDocumentTree();
const apolloClient = useApolloClient();
@@ -34,8 +24,6 @@ export function useLoadDocumentTree() {
const { documents: rootDocuments, loading: loadingRoot } = useSearchDocuments({
query: '*',
rootOnly: true,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
start: 0,
count: 100,
fetchPolicy: 'cache-first',
@@ -55,8 +43,6 @@ export function useLoadDocumentTree() {
input: {
query: '*',
parentDocuments: urns, // Batch query
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
start: 0,
count: urns.length * 100,
},
@@ -98,9 +84,7 @@ export function useLoadDocumentTree() {
variables: {
input: {
query: '*',
- parentDocument: parentUrn,
- states: [DocumentState.Published, DocumentState.Unpublished],
- includeDrafts: false,
+ parentDocuments: parentUrn ? [parentUrn] : undefined,
start: 0,
count: 100,
},
@@ -111,11 +95,7 @@ export function useLoadDocumentTree() {
const documents = result.data?.searchDocuments?.documents || [];
// Sort by creation time (most recent first)
- const sortedDocuments = [...documents].sort((a, b) => {
- const timeA = a.info?.created?.time || 0;
- const timeB = b.info?.created?.time || 0;
- return timeB - timeA; // DESC order
- });
+ const sortedDocuments = sortDocumentsByCreationTime(documents);
// Check if these documents have children
const childUrns = sortedDocuments.map((doc) => doc.urn);
@@ -148,11 +128,7 @@ export function useLoadDocumentTree() {
if (isTreeEmpty) {
// Sort root documents by creation time (most recent first)
- const sortedRootDocuments = [...rootDocuments].sort((a, b) => {
- const timeA = a.info?.created?.time || 0;
- const timeB = b.info?.created?.time || 0;
- return timeB - timeA; // DESC order
- });
+ const sortedRootDocuments = sortDocumentsByCreationTime(rootDocuments);
// Check which root documents have children
const rootDocUrns = sortedRootDocuments.map((doc) => doc.urn);
diff --git a/datahub-web-react/src/app/document/hooks/useRelatedDocuments.ts b/datahub-web-react/src/app/document/hooks/useRelatedDocuments.ts
new file mode 100644
index 00000000000000..32a32eddea6cbc
--- /dev/null
+++ b/datahub-web-react/src/app/document/hooks/useRelatedDocuments.ts
@@ -0,0 +1,89 @@
+import { useMemo } from 'react';
+
+import {
+ GetRelatedDocumentsQuery,
+ RelatedDocumentsFieldsFragment,
+ useGetRelatedDocumentsQuery,
+} from '@graphql/document.generated';
+import { Document, DocumentSourceType } from '@types';
+
+export interface RelatedDocumentsInput {
+ start?: number;
+ count?: number;
+ parentDocuments?: string[];
+ rootOnly?: boolean;
+ types?: string[];
+ domains?: string[];
+ sourceType?: DocumentSourceType;
+}
+
+/**
+ * Type guard that checks if an entity has the relatedDocuments field.
+ * This properly narrows the union type to entities that support relatedDocuments.
+ */
+function hasRelatedDocuments(
+ entity: NonNullable,
+): entity is Extract, { relatedDocuments?: unknown }> {
+ return 'relatedDocuments' in entity;
+}
+
+/**
+ * Extracts the relatedDocuments result from the entity union type.
+ * Returns null if the entity doesn't have relatedDocuments or if it's null/undefined.
+ *
+ * This function safely extracts relatedDocuments by using a type guard to check
+ * if the property exists on the entity object, which works for all entity types
+ * that support relatedDocuments (Dataset, Dashboard, Chart, DataJob, DataFlow,
+ * Container, MLModel, etc.)
+ */
+function extractRelatedDocuments(entity: GetRelatedDocumentsQuery['entity']): RelatedDocumentsFieldsFragment | null {
+ if (!entity) {
+ return null;
+ }
+
+ // Use type guard to narrow the type - no cast needed!
+ if (hasRelatedDocuments(entity) && entity.relatedDocuments) {
+ return entity.relatedDocuments;
+ }
+
+ return null;
+}
+
+export function useRelatedDocuments(entityUrn: string, input?: RelatedDocumentsInput) {
+ const { data, loading, error, refetch } = useGetRelatedDocumentsQuery({
+ variables: {
+ urn: entityUrn,
+ input: {
+ start: input?.start ?? 0,
+ count: input?.count ?? 100, // Default to 100 most recently updated
+ parentDocuments: input?.parentDocuments,
+ rootOnly: input?.rootOnly,
+ types: input?.types,
+ domains: input?.domains,
+ sourceType: input?.sourceType,
+ },
+ },
+ skip: !entityUrn,
+ fetchPolicy: 'cache-first',
+ });
+
+ const relatedDocumentsResult = useMemo(() => {
+ return extractRelatedDocuments(data?.entity ?? null);
+ }, [data?.entity]);
+
+ const documents = useMemo(() => {
+ return (relatedDocumentsResult?.documents || []) as Document[];
+ }, [relatedDocumentsResult]);
+
+ const total = useMemo(() => {
+ return relatedDocumentsResult?.total ?? 0;
+ }, [relatedDocumentsResult]);
+
+ return {
+ documents,
+ total,
+ loading,
+ error,
+ refetch,
+ };
+}
diff --git a/datahub-web-react/src/app/document/hooks/useSearchDocuments.ts b/datahub-web-react/src/app/document/hooks/useSearchDocuments.ts
index da61ee5cc2ebfa..a21ea13074db0f 100644
--- a/datahub-web-react/src/app/document/hooks/useSearchDocuments.ts
+++ b/datahub-web-react/src/app/document/hooks/useSearchDocuments.ts
@@ -9,7 +9,6 @@ export interface SearchDocumentsInput {
rootOnly?: boolean;
types?: string[];
states?: DocumentState[];
- includeDrafts?: boolean;
start?: number;
count?: number;
fetchPolicy?: 'cache-first' | 'cache-and-network' | 'network-only';
@@ -23,12 +22,10 @@ export function useSearchDocuments(input: SearchDocumentsInput) {
start: input.start || 0,
count: input.count || 100,
query: input.query || '*',
- parentDocument: input.parentDocument,
+ parentDocuments: input.parentDocument ? [input.parentDocument] : undefined,
rootOnly: input.rootOnly,
types: input.types,
- states: input.states || [DocumentState.Published, DocumentState.Unpublished],
sourceType: DocumentSourceType.Native,
- includeDrafts: input.includeDrafts || false,
},
includeParentDocuments: input.includeParentDocuments || false,
},
diff --git a/datahub-web-react/src/app/document/utils/__tests__/documentUtils.test.ts b/datahub-web-react/src/app/document/utils/__tests__/documentUtils.test.ts
new file mode 100644
index 00000000000000..53f5ad4ab36d9b
--- /dev/null
+++ b/datahub-web-react/src/app/document/utils/__tests__/documentUtils.test.ts
@@ -0,0 +1,284 @@
+import { describe, expect, it } from 'vitest';
+
+import { documentToTreeNode, sortDocumentsByCreationTime } from '@app/document/utils/documentUtils';
+
+import { Document } from '@types';
+
+// Helper to create minimal valid Document for testing
+// Uses 'as any' to allow partial documents in tests (consistent with other test files)
+const createTestDocument = (overrides: Partial = {}): Document => {
+ const base = {
+ urn: 'urn:li:document:test',
+ info: {
+ title: 'Test Document',
+ created: { time: 1000 },
+ contents: { text: '' },
+ lastModified: { time: 1000 },
+ },
+ };
+ return {
+ ...base,
+ ...overrides,
+ info: {
+ ...base.info,
+ ...overrides.info,
+ },
+ } as Document;
+};
+
+describe('documentUtils', () => {
+ describe('documentToTreeNode', () => {
+ it('should convert a document with all fields to a tree node', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: 'Test Document',
+ parentDocument: {
+ document: {
+ urn: 'urn:li:document:parent',
+ },
+ },
+ } as any,
+ });
+
+ const result = documentToTreeNode(doc, true);
+
+ expect(result).toEqual({
+ urn: 'urn:li:document:123',
+ title: 'Test Document',
+ parentUrn: 'urn:li:document:parent',
+ hasChildren: true,
+ children: undefined,
+ });
+ });
+
+ it('should use "Untitled" as default title when title is missing', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: null,
+ } as any,
+ });
+
+ const result = documentToTreeNode(doc, false);
+
+ expect(result.title).toBe('Untitled');
+ });
+
+ it('should use "Untitled" as default title when info is missing', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ });
+ // Override to null for this test case
+ (doc as any).info = null;
+
+ const result = documentToTreeNode(doc, false);
+
+ expect(result.title).toBe('Untitled');
+ });
+
+ it('should set parentUrn to null when parentDocument is missing', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: 'Root Document',
+ parentDocument: null,
+ } as any,
+ });
+
+ const result = documentToTreeNode(doc, false);
+
+ expect(result.parentUrn).toBe(null);
+ });
+
+ it('should set parentUrn to null when parentDocument.document is missing', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: 'Test Document',
+ parentDocument: {
+ document: null,
+ },
+ } as any,
+ });
+
+ const result = documentToTreeNode(doc, false);
+
+ expect(result.parentUrn).toBe(null);
+ });
+
+ it('should correctly set hasChildren flag', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: 'Test Document',
+ } as any,
+ });
+
+ const resultWithChildren = documentToTreeNode(doc, true);
+ expect(resultWithChildren.hasChildren).toBe(true);
+
+ const resultWithoutChildren = documentToTreeNode(doc, false);
+ expect(resultWithoutChildren.hasChildren).toBe(false);
+ });
+
+ it('should always set children to undefined', () => {
+ const doc = createTestDocument({
+ urn: 'urn:li:document:123',
+ info: {
+ title: 'Test Document',
+ } as any,
+ });
+
+ const result = documentToTreeNode(doc, true);
+
+ expect(result.children).toBeUndefined();
+ });
+ });
+
+ describe('sortDocumentsByCreationTime', () => {
+ it('should sort documents by creation time in descending order (newest first)', () => {
+ const documents: Document[] = [
+ createTestDocument({
+ urn: 'urn:li:document:1',
+ info: {
+ title: 'Oldest',
+ created: { time: 1000 },
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:2',
+ info: {
+ title: 'Newest',
+ created: { time: 3000 },
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:3',
+ info: {
+ title: 'Middle',
+ created: { time: 2000 },
+ } as any,
+ }),
+ ];
+
+ const result = sortDocumentsByCreationTime(documents);
+
+ expect(result[0].urn).toBe('urn:li:document:2'); // Newest first
+ expect(result[1].urn).toBe('urn:li:document:3');
+ expect(result[2].urn).toBe('urn:li:document:1'); // Oldest last
+ });
+
+ it('should handle documents with missing creation time (treat as 0)', () => {
+ const documents: Document[] = [
+ createTestDocument({
+ urn: 'urn:li:document:1',
+ info: {
+ title: 'Has time',
+ created: { time: 2000 },
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:2',
+ info: {
+ title: 'No time',
+ created: null as any,
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:3',
+ info: {
+ title: 'No created field',
+ created: undefined as any,
+ } as any,
+ }),
+ ];
+
+ const result = sortDocumentsByCreationTime(documents);
+
+ // Documents with time should come first
+ expect(result[0].urn).toBe('urn:li:document:1');
+ // Documents without time should be sorted together (both treated as 0)
+ expect(result.slice(1).map((d) => d.urn)).toContain('urn:li:document:2');
+ expect(result.slice(1).map((d) => d.urn)).toContain('urn:li:document:3');
+ });
+
+ it('should not mutate the original array', () => {
+ const documents: Document[] = [
+ createTestDocument({
+ urn: 'urn:li:document:1',
+ info: {
+ title: 'Doc 1',
+ created: { time: 1000 },
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:2',
+ info: {
+ title: 'Doc 2',
+ created: { time: 2000 },
+ } as any,
+ }),
+ ];
+
+ const originalOrder = documents.map((d) => d.urn);
+ const result = sortDocumentsByCreationTime(documents);
+
+ // Original array should be unchanged
+ expect(documents.map((d) => d.urn)).toEqual(originalOrder);
+ // Result should be sorted
+ expect(result.map((d) => d.urn)).toEqual(['urn:li:document:2', 'urn:li:document:1']);
+ });
+
+ it('should handle empty array', () => {
+ const documents: Document[] = [];
+
+ const result = sortDocumentsByCreationTime(documents);
+
+ expect(result).toEqual([]);
+ });
+
+ it('should handle single document', () => {
+ const documents: Document[] = [
+ createTestDocument({
+ urn: 'urn:li:document:1',
+ info: {
+ title: 'Single Doc',
+ created: { time: 1000 },
+ } as any,
+ }),
+ ];
+
+ const result = sortDocumentsByCreationTime(documents);
+
+ expect(result).toHaveLength(1);
+ expect(result[0].urn).toBe('urn:li:document:1');
+ });
+
+ it('should handle documents with same creation time', () => {
+ const documents: Document[] = [
+ createTestDocument({
+ urn: 'urn:li:document:1',
+ info: {
+ title: 'Doc 1',
+ created: { time: 1000 },
+ } as any,
+ }),
+ createTestDocument({
+ urn: 'urn:li:document:2',
+ info: {
+ title: 'Doc 2',
+ created: { time: 1000 },
+ } as any,
+ }),
+ ];
+
+ const result = sortDocumentsByCreationTime(documents);
+
+ // Both should be present, order may vary but both should be there
+ expect(result).toHaveLength(2);
+ expect(result.map((d) => d.urn)).toContain('urn:li:document:1');
+ expect(result.map((d) => d.urn)).toContain('urn:li:document:2');
+ });
+ });
+});
diff --git a/datahub-web-react/src/app/document/utils/__tests__/extractMentions.test.ts b/datahub-web-react/src/app/document/utils/__tests__/extractMentions.test.ts
new file mode 100644
index 00000000000000..cd07759c8d7039
--- /dev/null
+++ b/datahub-web-react/src/app/document/utils/__tests__/extractMentions.test.ts
@@ -0,0 +1,156 @@
+import { describe, expect, it } from 'vitest';
+
+import { extractMentions } from '@app/document/utils/extractMentions';
+
+describe('extractMentions', () => {
+ it('should extract document URNs from markdown links', () => {
+ const content = 'Check out [@Document 1](urn:li:document:abc123) for more info.';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:abc123']);
+ expect(result.assetUrns).toEqual([]);
+ });
+
+ it('should extract asset URNs from markdown links', () => {
+ const content = 'See [@Dataset](urn:li:dataset:xyz789) for the data.';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual([]);
+ expect(result.assetUrns).toEqual(['urn:li:dataset:xyz789']);
+ });
+
+ it('should extract multiple document and asset URNs', () => {
+ const content = `
+ Check [@Doc1](urn:li:document:doc1) and [@Doc2](urn:li:document:doc2).
+ Also see [@Dataset1](urn:li:dataset:ds1) and [@Dataset2](urn:li:dataset:ds2).
+ `;
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:doc1', 'urn:li:document:doc2']);
+ expect(result.assetUrns).toEqual(['urn:li:dataset:ds1', 'urn:li:dataset:ds2']);
+ });
+
+ it('should handle mixed document and asset URNs', () => {
+ const content = `
+ [@Document](urn:li:document:123)
+ [@Dataset](urn:li:dataset:456)
+ [@Chart](urn:li:chart:789)
+ [@Another Doc](urn:li:document:abc)
+ `;
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:123', 'urn:li:document:abc']);
+ expect(result.assetUrns).toEqual(['urn:li:dataset:456', 'urn:li:chart:789']);
+ });
+
+ it('should not extract duplicate URNs', () => {
+ const content = `
+ [@Doc](urn:li:document:123)
+ [@Same Doc](urn:li:document:123)
+ [@Dataset](urn:li:dataset:456)
+ [@Same Dataset](urn:li:dataset:456)
+ `;
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:123']);
+ expect(result.assetUrns).toEqual(['urn:li:dataset:456']);
+ });
+
+ it('should handle empty content', () => {
+ const content = '';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual([]);
+ expect(result.assetUrns).toEqual([]);
+ });
+
+ it('should handle content without URNs', () => {
+ const content = 'This is just plain text without any mentions.';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual([]);
+ expect(result.assetUrns).toEqual([]);
+ });
+
+ it('should handle markdown links without @ symbol', () => {
+ const content = `
+ [@Broken Link](not-a-urn)
+ [No @ symbol](urn:li:document:123)
+ Regular text
+ `;
+
+ const result = extractMentions(content);
+
+ // Note: The regex matches any markdown link with URN format, not just those with @
+ expect(result.documentUrns).toEqual(['urn:li:document:123']);
+ expect(result.assetUrns).toEqual([]);
+ });
+
+ it('should handle URNs with special characters', () => {
+ const content = '[@Entity](urn:li:dataFlow:some-flow_123)';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual([]);
+ expect(result.assetUrns).toEqual(['urn:li:dataFlow:some-flow_123']);
+ });
+
+ it('should handle complex markdown with multiple entity types', () => {
+ const content = `
+ # Documentation
+
+ See these resources:
+ - [@User Guide](urn:li:document:guide-123)
+ - [@API Docs](urn:li:document:api-456)
+ - [@Dataset A](urn:li:dataset:dataset-a)
+ - [@Dataset B](urn:li:dataset:dataset-b)
+ - [@Dashboard](urn:li:dashboard:dash-1)
+ - [@ML Model](urn:li:mlModel:model-xyz)
+ `;
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:guide-123', 'urn:li:document:api-456']);
+ expect(result.assetUrns).toEqual([
+ 'urn:li:dataset:dataset-a',
+ 'urn:li:dataset:dataset-b',
+ 'urn:li:dashboard:dash-1',
+ 'urn:li:mlModel:model-xyz',
+ ]);
+ });
+
+ it('should handle URNs in inline code blocks', () => {
+ const content = 'Use `[@Doc](urn:li:document:123)` in your code.';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual(['urn:li:document:123']);
+ });
+
+ it('should handle null or undefined content gracefully', () => {
+ const resultNull = extractMentions(null as any);
+ expect(resultNull.documentUrns).toEqual([]);
+ expect(resultNull.assetUrns).toEqual([]);
+
+ const resultUndefined = extractMentions(undefined as any);
+ expect(resultUndefined.documentUrns).toEqual([]);
+ expect(resultUndefined.assetUrns).toEqual([]);
+ });
+
+ it('should handle URNs with parentheses', () => {
+ const content = '[@Complex Dataset](urn:li:dataset:(urn:li:dataPlatform:kafka,topic-123,PROD))';
+
+ const result = extractMentions(content);
+
+ expect(result.documentUrns).toEqual([]);
+ // The regex now correctly handles nested parentheses in URNs
+ expect(result.assetUrns).toEqual(['urn:li:dataset:(urn:li:dataPlatform:kafka,topic-123,PROD)']);
+ });
+});
diff --git a/datahub-web-react/src/app/document/utils/documentUtils.ts b/datahub-web-react/src/app/document/utils/documentUtils.ts
new file mode 100644
index 00000000000000..742ea42e70ff3c
--- /dev/null
+++ b/datahub-web-react/src/app/document/utils/documentUtils.ts
@@ -0,0 +1,85 @@
+import { DocumentTreeNode } from '@app/document/DocumentTreeContext';
+
+import { Document, DocumentState } from '@types';
+
+/**
+ * Converts a Document to a DocumentTreeNode.
+ *
+ * @param doc - The document to convert
+ * @param hasChildren - Whether this document has children
+ * @returns A DocumentTreeNode representation of the document
+ */
+export function documentToTreeNode(doc: Document, hasChildren: boolean): DocumentTreeNode {
+ return {
+ urn: doc.urn,
+ title: doc.info?.title || 'Untitled',
+ parentUrn: doc.info?.parentDocument?.document?.urn || null,
+ hasChildren,
+ children: undefined, // Not loaded yet
+ };
+}
+
+/**
+ * Sorts documents by creation time in descending order (most recent first).
+ *
+ * @param documents - Array of documents to sort
+ * @returns A new sorted array (does not mutate the original)
+ */
+export function sortDocumentsByCreationTime(documents: Document[]): Document[] {
+ return [...documents].sort((a, b) => {
+ const timeA = a.info?.created?.time || 0;
+ const timeB = b.info?.created?.time || 0;
+ return timeB - timeA; // DESC order
+ });
+}
+
+/**
+ * Extracts related asset URNs from a document.
+ * Handles documents from GraphQL queries where relatedAssets may be null or undefined.
+ *
+ * @param document - Document with info.relatedAssets structure
+ * @returns Array of asset URNs (empty array if none found)
+ */
+export function extractRelatedAssetUrns(
+ document: { info?: { relatedAssets?: Array<{ asset: { urn: string } }> | null } | null } | null,
+): string[] {
+ return document?.info?.relatedAssets?.map((relatedAsset) => relatedAsset.asset.urn) || [];
+}
+
+/**
+ * Merges multiple arrays of URNs and removes duplicates.
+ *
+ * @param urnArrays - Variable number of URN arrays to merge
+ * @returns A new array with unique URNs
+ */
+export function mergeUrns(...urnArrays: (string[] | undefined | null)[]): string[] {
+ return [...new Set(urnArrays.flat().filter((urn): urn is string => Boolean(urn)))];
+}
+
+/**
+ * Creates a default document input for creating a new document.
+ *
+ * @param options - Configuration options for the document
+ * @param options.title - Document title (defaults to 'New Document')
+ * @param options.parentUrn - Optional parent document URN
+ * @param options.relatedAssetUrns - Optional array of related asset URNs
+ * @param options.state - Document state (defaults to Published)
+ * @param options.showInGlobalContext - Whether to show in global context (defaults to true)
+ * @returns CreateDocumentInput object ready for mutation
+ */
+export function createDefaultDocumentInput(options?: {
+ title?: string;
+ parentUrn?: string | null;
+ relatedAssetUrns?: string[];
+ state?: DocumentState;
+ showInGlobalContext?: boolean;
+}) {
+ return {
+ title: options?.title || 'New Document',
+ parentDocument: options?.parentUrn || undefined,
+ relatedAssets: options?.relatedAssetUrns || undefined,
+ contents: { text: '' },
+ state: options?.state || DocumentState.Published,
+ settings: { showInGlobalContext: options?.showInGlobalContext ?? true },
+ };
+}
diff --git a/datahub-web-react/src/app/document/utils/extractMentions.ts b/datahub-web-react/src/app/document/utils/extractMentions.ts
new file mode 100644
index 00000000000000..9b4f596fe3bcfc
--- /dev/null
+++ b/datahub-web-react/src/app/document/utils/extractMentions.ts
@@ -0,0 +1,36 @@
+/**
+ * Extracts @ mentions (URNs) from markdown text.
+ * Searches for markdown link patterns like [@Entity](urn:li:entityType:id)
+ *
+ * @param content - The markdown content to extract mentions from
+ * @returns An object containing arrays of document URNs and asset URNs
+ */
+export function extractMentions(content: string): { documentUrns: string[]; assetUrns: string[] } {
+ if (!content) return { documentUrns: [], assetUrns: [] };
+
+ // Match markdown link syntax: [text](urn:li:entityType:id)
+ // Handle URNs with nested parentheses by matching everything between the markdown link's parens
+ // The pattern matches: [text](urn:li:entityType:...) where ... can include nested parens
+ // We match the URN prefix, then allow nested paren groups or non-paren characters (one or more)
+ const urnPattern = /\[([^\]]+)\]\((urn:li:[a-zA-Z]+:(?:[^)(]+|\([^)]*\))+)\)/g;
+ const matches = Array.from(content.matchAll(urnPattern));
+
+ const documentUrns: string[] = [];
+ const assetUrns: string[] = [];
+
+ matches.forEach((match) => {
+ const urn = match[2]; // URN is in the second capture group (inside parentheses)
+
+ // Check if it's a document URN
+ if (urn.includes(':document:')) {
+ if (!documentUrns.includes(urn)) {
+ documentUrns.push(urn);
+ }
+ } else if (!assetUrns.includes(urn)) {
+ // Everything else is considered an asset
+ assetUrns.push(urn);
+ }
+ });
+
+ return { documentUrns, assetUrns };
+}
diff --git a/datahub-web-react/src/app/entity/Entity.tsx b/datahub-web-react/src/app/entity/Entity.tsx
index e1ab00be3ed9ae..b7a85aad9f4589 100644
--- a/datahub-web-react/src/app/entity/Entity.tsx
+++ b/datahub-web-react/src/app/entity/Entity.tsx
@@ -104,6 +104,10 @@ export enum EntityCapabilityType {
* Assigning an application to a entity
*/
APPLICATIONS,
+ /**
+ * Related context documents for this entity
+ */
+ RELATED_DOCUMENTS,
}
/**
diff --git a/datahub-web-react/src/app/entity/shared/utils.ts b/datahub-web-react/src/app/entity/shared/utils.ts
index 6074051e00ab02..dab5ad0dec0c08 100644
--- a/datahub-web-react/src/app/entity/shared/utils.ts
+++ b/datahub-web-react/src/app/entity/shared/utils.ts
@@ -193,7 +193,8 @@ export function formatEntityType(type: string): string {
return EntityType.Test;
case 'schemafield':
return EntityType.SchemaField;
-
+ case 'document':
+ return EntityType.Document;
// these are const in the java app
case 'dataprocessinstance': // Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME
return EntityType.DataProcessInstance;
diff --git a/datahub-web-react/src/app/entityV2/Entity.tsx b/datahub-web-react/src/app/entityV2/Entity.tsx
index 4f30fab80aa0f3..365c1093cd22a4 100644
--- a/datahub-web-react/src/app/entityV2/Entity.tsx
+++ b/datahub-web-react/src/app/entityV2/Entity.tsx
@@ -101,6 +101,10 @@ export enum EntityCapabilityType {
* Assigning the entity to an application
*/
APPLICATIONS,
+ /**
+ * Related context documents for this entity
+ */
+ RELATED_DOCUMENTS,
}
export interface EntityMenuActions {
diff --git a/datahub-web-react/src/app/entityV2/application/ApplicationEntity.tsx b/datahub-web-react/src/app/entityV2/application/ApplicationEntity.tsx
index 19954f0c61cd1f..87d2a74a557321 100644
--- a/datahub-web-react/src/app/entityV2/application/ApplicationEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/application/ApplicationEntity.tsx
@@ -239,6 +239,7 @@ export class ApplicationEntity implements Entity {
EntityCapabilityType.GLOSSARY_TERMS,
EntityCapabilityType.TAGS,
EntityCapabilityType.DOMAINS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/chart/ChartEntity.tsx b/datahub-web-react/src/app/entityV2/chart/ChartEntity.tsx
index cc46cae6d70182..29ca71b6c30dda 100644
--- a/datahub-web-react/src/app/entityV2/chart/ChartEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/chart/ChartEntity.tsx
@@ -383,6 +383,7 @@ export class ChartEntity implements Entity {
EntityCapabilityType.LINEAGE,
EntityCapabilityType.HEALTH,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/container/ContainerEntity.tsx b/datahub-web-react/src/app/entityV2/container/ContainerEntity.tsx
index 28155113ef07be..b9bd1c8044367a 100644
--- a/datahub-web-react/src/app/entityV2/container/ContainerEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/container/ContainerEntity.tsx
@@ -285,6 +285,7 @@ export class ContainerEntity implements Entity {
EntityCapabilityType.SOFT_DELETE,
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.TEST,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/dashboard/DashboardEntity.tsx b/datahub-web-react/src/app/entityV2/dashboard/DashboardEntity.tsx
index 7d773913f85a98..b6c86e1a9e90a9 100644
--- a/datahub-web-react/src/app/entityV2/dashboard/DashboardEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/dashboard/DashboardEntity.tsx
@@ -385,6 +385,7 @@ export class DashboardEntity implements Entity {
EntityCapabilityType.LINEAGE,
EntityCapabilityType.HEALTH,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/dataFlow/DataFlowEntity.tsx b/datahub-web-react/src/app/entityV2/dataFlow/DataFlowEntity.tsx
index 0c14d1c4c5cc52..c8876296235ad7 100644
--- a/datahub-web-react/src/app/entityV2/dataFlow/DataFlowEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/dataFlow/DataFlowEntity.tsx
@@ -269,6 +269,7 @@ export class DataFlowEntity implements Entity {
EntityCapabilityType.LINEAGE,
EntityCapabilityType.HEALTH,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/dataJob/DataJobEntity.tsx b/datahub-web-react/src/app/entityV2/dataJob/DataJobEntity.tsx
index 6f44e209299f2e..34cf65745a19bc 100644
--- a/datahub-web-react/src/app/entityV2/dataJob/DataJobEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/dataJob/DataJobEntity.tsx
@@ -319,6 +319,7 @@ export class DataJobEntity implements Entity {
EntityCapabilityType.LINEAGE,
EntityCapabilityType.HEALTH,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/dataProduct/DataProductEntity.tsx b/datahub-web-react/src/app/entityV2/dataProduct/DataProductEntity.tsx
index 1510ec3cf6fd15..73c5153a8c2c9f 100644
--- a/datahub-web-react/src/app/entityV2/dataProduct/DataProductEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/dataProduct/DataProductEntity.tsx
@@ -276,6 +276,7 @@ export class DataProductEntity implements Entity {
EntityCapabilityType.TAGS,
EntityCapabilityType.DOMAINS,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/dataset/DatasetEntity.tsx b/datahub-web-react/src/app/entityV2/dataset/DatasetEntity.tsx
index 68843f8bc51722..84969461d1f5d4 100644
--- a/datahub-web-react/src/app/entityV2/dataset/DatasetEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/dataset/DatasetEntity.tsx
@@ -534,6 +534,7 @@ export class DatasetEntity implements Entity {
EntityCapabilityType.LINEAGE,
EntityCapabilityType.HEALTH,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/document/DocumentModal.tsx b/datahub-web-react/src/app/entityV2/document/DocumentModal.tsx
new file mode 100644
index 00000000000000..70a58f997da68c
--- /dev/null
+++ b/datahub-web-react/src/app/entityV2/document/DocumentModal.tsx
@@ -0,0 +1,184 @@
+import { LoadingOutlined } from '@ant-design/icons';
+import { colors } from '@components';
+import React from 'react';
+import styled from 'styled-components';
+
+import { DocumentTreeNode } from '@app/document/DocumentTreeContext';
+import EntityContext from '@app/entity/shared/EntityContext';
+import { DocumentSummaryTab } from '@app/entityV2/document/summary/DocumentSummaryTab';
+import { PageTemplateProvider } from '@app/homeV3/context/PageTemplateContext';
+import { Modal } from '@src/alchemy-components';
+
+import { useGetDocumentQuery } from '@graphql/document.generated';
+import { EntityType, PageTemplateSurfaceType } from '@types';
+
+const StyledModal = styled(Modal)`
+ &&& .ant-modal {
+ max-height: 95vh;
+ top: 2.5vh;
+ }
+
+ &&& .ant-modal-content {
+ max-height: 95vh;
+ height: 95vh;
+ position: relative;
+ display: flex;
+ flex-direction: column;
+ }
+
+ .ant-modal-header {
+ display: none;
+ }
+
+ .ant-modal-body {
+ padding: 8px 0 0 0;
+ flex: 1 1 auto;
+ overflow: hidden;
+ display: flex;
+ flex-direction: column;
+ min-height: 0;
+ }
+
+ .ant-modal-footer {
+ display: none;
+ }
+
+ .ant-modal-close {
+ top: 16px;
+ right: 16px;
+ z-index: 1001;
+ width: 24px;
+ height: 24px;
+ line-height: 1;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ opacity: 0.8;
+ transition: opacity 0.2s ease;
+ padding: 0;
+ border: none;
+ background: transparent;
+
+ &:hover {
+ opacity: 1;
+ }
+
+ .ant-modal-close-x {
+ height: 24px;
+ width: 24px;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ font-size: 16px;
+ color: inherit;
+ line-height: 1;
+ }
+ }
+`;
+
+const ModalContent = styled.div`
+ height: 100%;
+ display: flex;
+ flex-direction: column;
+ overflow: hidden;
+`;
+
+const ContentWrapper = styled.div`
+ flex: 1;
+ overflow-y: auto;
+ padding: 0;
+`;
+
+const LoadingWrapper = styled.div`
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ min-height: 200px;
+ padding: 40px;
+`;
+
+interface DocumentViewModalProps {
+ documentUrn: string;
+ onClose: () => void;
+ onDocumentDeleted?: () => void;
+}
+
+/**
+ * Modal component for viewing a document in-place without navigating away.
+ * Loads document data lazily when the modal opens.
+ */
+export const DocumentModal: React.FC = ({
+ documentUrn: initialDocumentUrn,
+ onClose,
+ onDocumentDeleted,
+}) => {
+ // Use state to allow breadcrumb navigation within the modal
+ const [currentDocumentUrn, setCurrentDocumentUrn] = React.useState(initialDocumentUrn);
+
+ // Update current document URN when initial prop changes
+ React.useEffect(() => {
+ setCurrentDocumentUrn(initialDocumentUrn);
+ }, [initialDocumentUrn]);
+
+ // Lazy load document data when modal opens or when document URN changes
+ const { data, loading, refetch } = useGetDocumentQuery({
+ variables: { urn: currentDocumentUrn, includeParentDocuments: true },
+ skip: !currentDocumentUrn,
+ });
+
+ const document = data?.document;
+
+ const wrappedRefetch = async () => {
+ return refetch();
+ };
+
+ const handleDelete = React.useCallback(
+ (_deletedNode: DocumentTreeNode | null) => {
+ // Delete mutation is handled in DocumentActionsMenu
+ if (onDocumentDeleted) {
+ onDocumentDeleted();
+ }
+ onClose();
+ },
+ [onClose, onDocumentDeleted],
+ );
+
+ return (
+
+
+ {loading || !document ? (
+
+
+
+ ) : (
+ {},
+ refetch: wrappedRefetch,
+ lineage: undefined,
+ }}
+ >
+
+
+
+
+
+
+ )}
+
+
+ );
+};
diff --git a/datahub-web-react/src/app/entityV2/document/DocumentModalNavigationContext.tsx b/datahub-web-react/src/app/entityV2/document/DocumentModalNavigationContext.tsx
new file mode 100644
index 00000000000000..8b7a1dee326414
--- /dev/null
+++ b/datahub-web-react/src/app/entityV2/document/DocumentModalNavigationContext.tsx
@@ -0,0 +1,13 @@
+import { createContext, useContext } from 'react';
+
+// Context for handling breadcrumb navigation within the document modal
+export const DocumentModalNavigationContext = createContext<{
+ navigateToDocument: ((urn: string) => void) | null;
+}>({
+ navigateToDocument: null,
+});
+
+export const useDocumentModalNavigation = () => {
+ const context = useContext(DocumentModalNavigationContext);
+ return context;
+};
diff --git a/datahub-web-react/src/app/entityV2/document/DocumentNativeProfile.tsx b/datahub-web-react/src/app/entityV2/document/DocumentNativeProfile.tsx
index ee3e87ac9dc698..6b9f3a697c8aa0 100644
--- a/datahub-web-react/src/app/entityV2/document/DocumentNativeProfile.tsx
+++ b/datahub-web-react/src/app/entityV2/document/DocumentNativeProfile.tsx
@@ -1,6 +1,6 @@
import { LoadingOutlined } from '@ant-design/icons';
import { BookOpen, ListBullets } from '@phosphor-icons/react';
-import React, { useState } from 'react';
+import React, { useContext, useState } from 'react';
import styled from 'styled-components';
import EntityContext from '@app/entity/shared/EntityContext';
@@ -90,7 +90,7 @@ const sidebarSections = [
*/
export const DocumentNativeProfile: React.FC = ({ urn, document, loading = false, refetch }) => {
const [sidebarClosed, setSidebarClosed] = useState(true); // Start closed by default
- const isCompact = React.useContext(CompactContext);
+ const isCompact = useContext(CompactContext);
if (!document) {
return null;
diff --git a/datahub-web-react/src/app/entityV2/document/summary/AddRelatedEntityDropdown.tsx b/datahub-web-react/src/app/entityV2/document/summary/AddRelatedEntityDropdown.tsx
index 38aa0e80133363..593b2350b0327e 100644
--- a/datahub-web-react/src/app/entityV2/document/summary/AddRelatedEntityDropdown.tsx
+++ b/datahub-web-react/src/app/entityV2/document/summary/AddRelatedEntityDropdown.tsx
@@ -1,4 +1,5 @@
import { Button, Tooltip } from '@components';
+import { message } from 'antd';
import React, { useCallback, useState } from 'react';
import styled from 'styled-components';
@@ -78,6 +79,7 @@ export const AddRelatedEntityDropdown: React.FC =
// Reset to initial state when closing
setSelectedUrns(initialSelectedUrns);
} catch (error) {
+ message.error('Failed to update related entities. An unexpected error occurred!');
console.error('Failed to update related entities:', error);
}
}, [selectedUrns, documentUrn, onConfirm, initialSelectedUrns]);
diff --git a/datahub-web-react/src/app/entityV2/document/summary/DocumentSummaryTab.tsx b/datahub-web-react/src/app/entityV2/document/summary/DocumentSummaryTab.tsx
index d6aab294ccd872..312f8f68aa4d52 100644
--- a/datahub-web-react/src/app/entityV2/document/summary/DocumentSummaryTab.tsx
+++ b/datahub-web-react/src/app/entityV2/document/summary/DocumentSummaryTab.tsx
@@ -3,11 +3,15 @@ import React, { useState } from 'react';
import { useHistory } from 'react-router-dom';
import styled from 'styled-components';
+import { DocumentTreeNode } from '@app/document/DocumentTreeContext';
+import { useDocumentPermissions } from '@app/document/hooks/useDocumentPermissions';
import { useEntityData } from '@app/entity/shared/EntityContext';
import { DocumentChangeHistoryDrawer } from '@app/entityV2/document/changeHistory/DocumentChangeHistoryDrawer';
import { EditableContent } from '@app/entityV2/document/summary/EditableContent';
import { EditableTitle } from '@app/entityV2/document/summary/EditableTitle';
import PropertiesHeader from '@app/entityV2/summary/properties/PropertiesHeader';
+import { DocumentActionsMenu } from '@app/homeV2/layout/sidebar/documents/DocumentActionsMenu';
+import { useModalContext } from '@app/sharedV2/modals/ModalContext';
import { useEntityRegistry } from '@app/useEntityRegistry';
import { Document, EntityType } from '@types';
@@ -17,13 +21,28 @@ const SummaryWrapper = styled.div`
display: flex;
flex-direction: column;
gap: 16px;
- position: relative;
`;
-const HistoryIconButton = styled(Button)`
- position: absolute;
- top: 40px;
- right: 20%;
+const HeaderRow = styled.div`
+ display: flex;
+ align-items: flex-start;
+ gap: 16px;
+ width: 100%;
+`;
+
+const TitleSection = styled.div`
+ flex: 1;
+ min-width: 0; /* Critical: allow flex item to shrink below its content size */
+`;
+
+const TopRightButtonsContainer = styled.div`
+ display: flex;
+ align-items: center;
+ gap: 8px;
+ flex-shrink: 0; /* Buttons take precedence - don't shrink */
+`;
+
+const TopRightButton = styled(Button)`
background: transparent;
border: none;
cursor: pointer;
@@ -33,6 +52,35 @@ const HistoryIconButton = styled(Button)`
align-items: center;
justify-content: center;
color: ${colors.gray[400]};
+
+ &:hover {
+ background-color: ${colors.gray[100]};
+ }
+`;
+
+const ActionsMenuWrapper = styled.div`
+ display: flex;
+ align-items: center;
+
+ /* Style the menu button to match other top right buttons */
+ button {
+ background: transparent;
+ border: none;
+ cursor: pointer;
+ padding: 8px;
+ border-radius: 6px;
+ display: flex;
+ align-items: center;
+ justify-content: center;
+ color: ${colors.gray[400]};
+ min-width: auto;
+ width: auto;
+ height: auto;
+
+ &:hover {
+ background-color: ${colors.gray[100]};
+ }
+ }
`;
const Breadcrumb = styled.div`
@@ -59,51 +107,95 @@ const BreadcrumbSeparator = styled.span`
margin: 0 4px;
`;
-export const DocumentSummaryTab = () => {
+interface DocumentSummaryTabProps {
+ onDelete?: (deletedNode: DocumentTreeNode | null) => void;
+ onMove?: (documentUrn: string) => void;
+}
+
+export const DocumentSummaryTab: React.FC = ({ onDelete, onMove }) => {
const { urn, entityData } = useEntityData();
const document = entityData as Document;
const history = useHistory();
const entityRegistry = useEntityRegistry();
+ const { isInsideModal } = useModalContext();
const [isHistoryDrawerOpen, setIsHistoryDrawerOpen] = useState(false);
+ const { canDelete, canMove } = useDocumentPermissions(urn);
const documentContent = document?.info?.contents?.text || '';
// Get parent documents hierarchy (ordered: direct parent, parent's parent, ...)
const parentDocuments = document?.parentDocuments?.documents || [];
+ // Get the direct parent URN (first in the array)
+ const currentParentUrn = parentDocuments.length > 0 ? parentDocuments[0].urn : null;
const handleParentClick = (parentUrn: string) => {
history.push(entityRegistry.getEntityUrl(EntityType.Document, parentUrn));
};
+ const handleGoToDocument = () => {
+ const url = entityRegistry.getEntityUrl(EntityType.Document, urn);
+ history.push(url);
+ };
+
return (
<>
- {/* History icon button - top right */}
-
- setIsHistoryDrawerOpen(true)}
- aria-label="View change history"
- icon={{ icon: 'Clock', source: 'phosphor', size: '2xl' }}
- />
-
-
- {/* Parent documents breadcrumb - show full hierarchy */}
- {parentDocuments.length > 0 && (
-
- {[...parentDocuments].reverse().map((parent, index) => (
-
- handleParentClick(parent.urn)}>
- {parent.info?.title || 'Untitled'}
-
- {index < parentDocuments.length - 1 && /}
-
- ))}
-
- )}
-
- {/* Simple Notion-style title input - click to edit */}
-
+ {/* Header row with title and buttons - uses flexbox for natural wrapping */}
+
+
+ {/* Parent documents breadcrumb - show full hierarchy */}
+ {parentDocuments.length > 0 && (
+
+ {[...parentDocuments].reverse().map((parent, index) => (
+
+ handleParentClick(parent.urn)}>
+ {parent.info?.title || 'Untitled'}
+
+ {index < parentDocuments.length - 1 && (
+ /
+ )}
+
+ ))}
+
+ )}
+
+ {/* Simple Notion-style title input - click to edit */}
+
+
+
+ {/* Top right buttons - History and Expand (when in modal) */}
+
+
+ setIsHistoryDrawerOpen(true)}
+ aria-label="View change history"
+ icon={{ icon: 'Clock', source: 'phosphor', size: '2xl' }}
+ />
+
+ {isInsideModal && (
+
+
+
+ )}
+
+
+
+
+
{/* Properties list - reuses SummaryTab component */}
diff --git a/datahub-web-react/src/app/entityV2/document/summary/EditableContent.tsx b/datahub-web-react/src/app/entityV2/document/summary/EditableContent.tsx
index f6213a9d799ad9..b70094e6a6cfac 100644
--- a/datahub-web-react/src/app/entityV2/document/summary/EditableContent.tsx
+++ b/datahub-web-react/src/app/entityV2/document/summary/EditableContent.tsx
@@ -30,7 +30,7 @@ const StyledEditor = styled(Editor)<{ $hideToolbar?: boolean }>`
&&& {
.remirror-editor {
padding: 0px 0;
- min-height: 400px;
+ min-height: 460px;
}
.remirror-editor.ProseMirror {
font-size: 15px;
diff --git a/datahub-web-react/src/app/entityV2/document/summary/EditableTitle.tsx b/datahub-web-react/src/app/entityV2/document/summary/EditableTitle.tsx
index 9ee4f969977976..4eefc9b2967539 100644
--- a/datahub-web-react/src/app/entityV2/document/summary/EditableTitle.tsx
+++ b/datahub-web-react/src/app/entityV2/document/summary/EditableTitle.tsx
@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useRef, useState } from 'react';
import styled from 'styled-components';
import { useDocumentPermissions } from '@app/document/hooks/useDocumentPermissions';
@@ -7,6 +7,7 @@ import colors from '@src/alchemy-components/theme/foundations/colors';
const TitleContainer = styled.div`
width: 100%;
+ min-width: 0;
`;
const TitleInput = styled.textarea<{ $editable: boolean }>`
@@ -18,16 +19,21 @@ const TitleInput = styled.textarea<{ $editable: boolean }>`
outline: none;
background: transparent;
width: 100%;
+ min-width: 0;
padding: 6px 8px;
margin: -6px -8px;
cursor: ${(props) => (props.$editable ? 'text' : 'default')};
border-radius: 4px;
resize: none;
- overflow: hidden;
+ overflow-y: auto;
+ overflow-x: hidden;
font-family: inherit;
white-space: pre-wrap;
word-wrap: break-word;
-
+ overflow-wrap: break-word;
+ box-sizing: border-box;
+ min-height: calc(32px * 1.4 + 12px); /* 1 row: font-size * line-height + padding */
+ max-height: calc(32px * 1.4 * 3 + 12px); /* 3 rows: font-size * line-height * 3 + padding */
&:hover {
background-color: transparent;
}
@@ -50,6 +56,7 @@ interface Props {
export const EditableTitle: React.FC = ({ documentUrn, initialTitle }) => {
const [title, setTitle] = useState(initialTitle || '');
const [isSaving, setIsSaving] = useState(false);
+ const textareaRef = useRef(null);
const { canEditTitle } = useDocumentPermissions(documentUrn);
const { updateTitle } = useUpdateDocumentTitleMutation();
@@ -57,12 +64,22 @@ export const EditableTitle: React.FC = ({ documentUrn, initialTitle }) =>
setTitle(initialTitle || '');
}, [initialTitle]);
- // Auto-resize textarea to fit content
- const handleInput = (e: React.FormEvent) => {
- const target = e.currentTarget;
- target.style.height = 'auto';
- target.style.height = `${target.scrollHeight}px`;
- };
+ // Auto-resize textarea up to 3 rows, then scroll
+ useEffect(() => {
+ const textarea = textareaRef.current;
+ if (!textarea) return;
+
+ const { style, scrollHeight } = textarea;
+
+ // Reset height to auto to get the correct scrollHeight
+ style.height = 'auto';
+
+ // Calculate max height for 3 rows (font-size * line-height * 3 + padding)
+ const maxHeight = 32 * 1.4 * 3 + 12; // ~146px
+
+ // Set height to scrollHeight, but cap at maxHeight
+ style.height = `${Math.min(scrollHeight, maxHeight)}px`;
+ }, [title]);
const handleBlur = async () => {
if (title !== initialTitle && !isSaving) {
@@ -85,10 +102,10 @@ export const EditableTitle: React.FC = ({ documentUrn, initialTitle }) =>
return (
setTitle(e.target.value)}
- onInput={handleInput}
onBlur={handleBlur}
onKeyDown={handleKeyDown}
$editable={canEditTitle}
diff --git a/datahub-web-react/src/app/entityV2/document/summary/RelatedAssetsSection.tsx b/datahub-web-react/src/app/entityV2/document/summary/RelatedAssetsSection.tsx
deleted file mode 100644
index 725caadaf52b5c..00000000000000
--- a/datahub-web-react/src/app/entityV2/document/summary/RelatedAssetsSection.tsx
+++ /dev/null
@@ -1,56 +0,0 @@
-import React from 'react';
-import styled from 'styled-components';
-
-import { SectionContainer } from '@app/entityV2/shared/summary/HeaderComponents';
-import { EntityLink } from '@app/homeV2/reference/sections/EntityLink';
-import { useEntityRegistry } from '@app/useEntityRegistry';
-import colors from '@src/alchemy-components/theme/foundations/colors';
-
-import { DocumentRelatedAsset, EntityType } from '@types';
-
-const SectionHeader = styled.h4`
- font-size: 16px;
- font-weight: 600;
- margin: 0;
- color: ${colors.gray[600]};
-`;
-
-const List = styled.div`
- display: flex;
- flex-direction: column;
- gap: 4px;
-`;
-
-const EmptyState = styled.div`
- font-size: 14px;
- font-weight: 400;
- color: ${colors.gray[600]};
- text-align: center;
- padding: 8px;
-`;
-
-interface RelatedAssetsSectionProps {
- relatedAssets?: DocumentRelatedAsset[];
-}
-
-export const RelatedAssetsSection: React.FC = ({ relatedAssets }) => {
- const entityRegistry = useEntityRegistry();
-
- return (
-
- Related
-
- {relatedAssets?.map((relatedAsset) => {
- const { asset } = relatedAsset;
- const genericProperties = entityRegistry.getGenericEntityProperties(
- asset.type as EntityType,
- asset,
- );
-
- return ;
- })}
- {relatedAssets?.length === 0 || !relatedAssets ? No related assets yet : null}
-
-
- );
-};
diff --git a/datahub-web-react/src/app/entityV2/document/summary/RelatedDocumentsSection.tsx b/datahub-web-react/src/app/entityV2/document/summary/RelatedDocumentsSection.tsx
deleted file mode 100644
index 593e8c5443853a..00000000000000
--- a/datahub-web-react/src/app/entityV2/document/summary/RelatedDocumentsSection.tsx
+++ /dev/null
@@ -1,55 +0,0 @@
-import React from 'react';
-import styled from 'styled-components';
-
-import { SectionContainer } from '@app/entityV2/shared/summary/HeaderComponents';
-import { EntityLink } from '@app/homeV2/reference/sections/EntityLink';
-import { useEntityRegistry } from '@app/useEntityRegistry';
-import colors from '@src/alchemy-components/theme/foundations/colors';
-
-import { DocumentRelatedDocument, EntityType } from '@types';
-
-const SectionHeader = styled.h4`
- font-size: 16px;
- font-weight: 600;
- margin: 0;
- color: ${colors.gray[600]};
-`;
-
-const List = styled.div`
- display: flex;
- flex-direction: column;
- gap: 4px;
-`;
-
-interface RelatedDocumentsSectionProps {
- relatedDocuments?: DocumentRelatedDocument[];
-}
-
-export const RelatedDocumentsSection: React.FC = ({ relatedDocuments }) => {
- const entityRegistry = useEntityRegistry();
-
- if (!relatedDocuments || relatedDocuments.length === 0) {
- return null;
- }
-
- return (
-
- Related Context
-
- {relatedDocuments.map((relatedDoc) => {
- const { document } = relatedDoc;
- const genericProperties = entityRegistry.getGenericEntityProperties(EntityType.Document, document);
-
- return (
-
- );
- })}
-
-
- );
-};
diff --git a/datahub-web-react/src/app/entityV2/document/summary/RelatedSection.tsx b/datahub-web-react/src/app/entityV2/document/summary/RelatedSection.tsx
index bba0ee3d072fbd..f8a183ccc09db2 100644
--- a/datahub-web-react/src/app/entityV2/document/summary/RelatedSection.tsx
+++ b/datahub-web-react/src/app/entityV2/document/summary/RelatedSection.tsx
@@ -4,22 +4,33 @@ import styled from 'styled-components';
import { useUserContext } from '@app/context/useUserContext';
import { AddRelatedEntityDropdown } from '@app/entityV2/document/summary/AddRelatedEntityDropdown';
-import { SectionContainer } from '@app/entityV2/shared/summary/HeaderComponents';
import { EntityLink } from '@app/homeV2/reference/sections/EntityLink';
import { useEntityRegistry } from '@app/useEntityRegistry';
import colors from '@src/alchemy-components/theme/foundations/colors';
import { AndFilterInput, DocumentRelatedAsset, DocumentRelatedDocument, EntityType, FilterOperator } from '@types';
+const Section = styled.div`
+ display: flex;
+ flex-direction: column;
+ gap: 4px;
+
+ &:hover {
+ .hover-btn {
+ display: flex;
+ }
+ }
+ padding-top: 20px;
+`;
+
const SectionHeader = styled.div`
display: flex;
justify-content: space-between;
align-items: center;
- margin-bottom: 8px;
`;
const SectionTitle = styled.h4`
- font-size: 16px;
+ font-size: 14px;
font-weight: 600;
margin: 0;
color: ${colors.gray[600]};
@@ -35,8 +46,8 @@ const EmptyState = styled.div`
font-size: 14px;
font-weight: 400;
color: ${colors.gray[1800]};
- text-align: center;
- padding: 8px;
+ text-align: start;
+ padding: 0px;
`;
const EntityItemContainer = styled.div`
@@ -172,7 +183,7 @@ export const RelatedSection: React.FC = ({
};
return (
-
+ Related
{canEdit && (
@@ -224,9 +235,9 @@ export const RelatedSection: React.FC = ({
);
})
) : (
- No related assets or context
+ Add related assets or context
)}
-
+
);
};
diff --git a/datahub-web-react/src/app/entityV2/domain/DomainEntity.tsx b/datahub-web-react/src/app/entityV2/domain/DomainEntity.tsx
index 0d25e27672fde3..3f819e26f85147 100644
--- a/datahub-web-react/src/app/entityV2/domain/DomainEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/domain/DomainEntity.tsx
@@ -237,6 +237,6 @@ export class DomainEntity implements Entity {
supportedCapabilities = () => {
// TODO.. Determine whether SOFT_DELETE should go into here.
- return new Set([EntityCapabilityType.OWNERS]);
+ return new Set([EntityCapabilityType.OWNERS, EntityCapabilityType.RELATED_DOCUMENTS]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/glossaryNode/GlossaryNodeEntity.tsx b/datahub-web-react/src/app/entityV2/glossaryNode/GlossaryNodeEntity.tsx
index 433bffaeed65f9..b6f06b76d89d02 100644
--- a/datahub-web-react/src/app/entityV2/glossaryNode/GlossaryNodeEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/glossaryNode/GlossaryNodeEntity.tsx
@@ -221,6 +221,7 @@ class GlossaryNodeEntity implements Entity {
EntityCapabilityType.DEPRECATION,
EntityCapabilityType.SOFT_DELETE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/glossaryTerm/GlossaryTermEntity.tsx b/datahub-web-react/src/app/entityV2/glossaryTerm/GlossaryTermEntity.tsx
index 7b2590a8d570e0..2d8d61c139b1c7 100644
--- a/datahub-web-react/src/app/entityV2/glossaryTerm/GlossaryTermEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/glossaryTerm/GlossaryTermEntity.tsx
@@ -262,6 +262,7 @@ export class GlossaryTermEntity implements Entity {
EntityCapabilityType.DEPRECATION,
EntityCapabilityType.SOFT_DELETE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
diff --git a/datahub-web-react/src/app/entityV2/mlFeature/MLFeatureEntity.tsx b/datahub-web-react/src/app/entityV2/mlFeature/MLFeatureEntity.tsx
index 977a9605b1c858..53d0f8ec9c66fa 100644
--- a/datahub-web-react/src/app/entityV2/mlFeature/MLFeatureEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/mlFeature/MLFeatureEntity.tsx
@@ -270,6 +270,7 @@ export class MLFeatureEntity implements Entity {
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.LINEAGE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/mlFeatureTable/MLFeatureTableEntity.tsx b/datahub-web-react/src/app/entityV2/mlFeatureTable/MLFeatureTableEntity.tsx
index 2bac0f3b278ffd..9808e19545efe8 100644
--- a/datahub-web-react/src/app/entityV2/mlFeatureTable/MLFeatureTableEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/mlFeatureTable/MLFeatureTableEntity.tsx
@@ -238,6 +238,7 @@ export class MLFeatureTableEntity implements Entity {
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.LINEAGE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/mlModel/MLModelEntity.tsx b/datahub-web-react/src/app/entityV2/mlModel/MLModelEntity.tsx
index c185f7bd3ceb41..77c8bab4dd7d2c 100644
--- a/datahub-web-react/src/app/entityV2/mlModel/MLModelEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/mlModel/MLModelEntity.tsx
@@ -257,6 +257,7 @@ export class MLModelEntity implements Entity {
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.LINEAGE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/mlModelGroup/MLModelGroupEntity.tsx b/datahub-web-react/src/app/entityV2/mlModelGroup/MLModelGroupEntity.tsx
index db65cd973aa66e..5f2c7a09b5a5a8 100644
--- a/datahub-web-react/src/app/entityV2/mlModelGroup/MLModelGroupEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/mlModelGroup/MLModelGroupEntity.tsx
@@ -234,6 +234,7 @@ export class MLModelGroupEntity implements Entity {
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.LINEAGE,
EntityCapabilityType.APPLICATIONS,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/mlPrimaryKey/MLPrimaryKeyEntity.tsx b/datahub-web-react/src/app/entityV2/mlPrimaryKey/MLPrimaryKeyEntity.tsx
index f33a245842e68f..34f7c4cf238a63 100644
--- a/datahub-web-react/src/app/entityV2/mlPrimaryKey/MLPrimaryKeyEntity.tsx
+++ b/datahub-web-react/src/app/entityV2/mlPrimaryKey/MLPrimaryKeyEntity.tsx
@@ -231,6 +231,7 @@ export class MLPrimaryKeyEntity implements Entity {
EntityCapabilityType.SOFT_DELETE,
EntityCapabilityType.DATA_PRODUCTS,
EntityCapabilityType.LINEAGE,
+ EntityCapabilityType.RELATED_DOCUMENTS,
]);
};
}
diff --git a/datahub-web-react/src/app/entityV2/shared/components/links/LinkIcon.tsx b/datahub-web-react/src/app/entityV2/shared/components/links/LinkIcon.tsx
index cdf1be60f0b698..ee4e023af0dd52 100644
--- a/datahub-web-react/src/app/entityV2/shared/components/links/LinkIcon.tsx
+++ b/datahub-web-react/src/app/entityV2/shared/components/links/LinkIcon.tsx
@@ -17,9 +17,15 @@ const Container = styled.div`
interface Props {
url: string;
className?: string;
+ /**
+ * Whether to use primary color. Defaults to true (primary color).
+ * If false, uses gray color with level 600.
+ */
+ usePrimaryColor?: boolean;
+ style?: React.CSSProperties;
}
-export function LinkIcon({ url, className }: Props) {
+export function LinkIcon({ url, className, usePrimaryColor = true, style }: Props) {
const isDocumentationFileUploadV1Enabled = useIsDocumentationFileUploadV1Enabled();
const renderIcon = useCallback(() => {
@@ -29,8 +35,20 @@ export function LinkIcon({ url, className }: Props) {
return ;
}
- return ;
- }, [isDocumentationFileUploadV1Enabled, url]);
-
- return {renderIcon()};
+ return (
+
+ );
+ }, [isDocumentationFileUploadV1Enabled, url, usePrimaryColor]);
+
+ return (
+
+ {renderIcon()}
+
+ );
}
diff --git a/datahub-web-react/src/app/entityV2/shared/tabs/Documentation/DocumentationTab.tsx b/datahub-web-react/src/app/entityV2/shared/tabs/Documentation/DocumentationTab.tsx
index 3717477be00582..0db4b2c5cc273e 100644
--- a/datahub-web-react/src/app/entityV2/shared/tabs/Documentation/DocumentationTab.tsx
+++ b/datahub-web-react/src/app/entityV2/shared/tabs/Documentation/DocumentationTab.tsx
@@ -1,25 +1,24 @@
import { EditOutlined, ExpandAltOutlined, PlusOutlined } from '@ant-design/icons';
-import { Button as AntButton, Divider, Typography } from 'antd';
+import { Button as AntButton, Typography } from 'antd';
import queryString from 'query-string';
import React, { useEffect } from 'react';
import { useLocation } from 'react-router-dom';
import styled from 'styled-components';
import { useEntityData, useRouteToTab } from '@app/entity/shared/EntityContext';
-import { AddLinkModal } from '@app/entityV2/shared/components/styled/AddLinkModal';
import { EmptyTab } from '@app/entityV2/shared/components/styled/EmptyTab';
import TabToolbar from '@app/entityV2/shared/components/styled/TabToolbar';
import { REDESIGN_COLORS } from '@app/entityV2/shared/constants';
import { DescriptionEditor } from '@app/entityV2/shared/tabs/Documentation/components/DescriptionEditor';
import { DescriptionPreviewModal } from '@app/entityV2/shared/tabs/Documentation/components/DescriptionPreviewModal';
-import { LinkList } from '@app/entityV2/shared/tabs/Documentation/components/LinkList';
+import { RelatedSection } from '@app/entityV2/shared/tabs/Documentation/components/RelatedSection';
import { getAssetDescriptionDetails } from '@app/entityV2/shared/tabs/Documentation/utils';
import { EDITED_DESCRIPTIONS_CACHE_NAME } from '@app/entityV2/shared/utils';
import { Button, Editor } from '@src/alchemy-components';
const DocumentationContainer = styled.div`
- margin: 0 32px;
- padding: 40px 0;
+ margin: 0 16px;
+ padding: 32px 0;
max-width: calc(100% - 10px);
`;
@@ -85,7 +84,6 @@ export const DocumentationTab = ({ properties }: { properties?: Props }) => {
>
Edit
- {!hideLinksButton && }