Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,7 @@ public void setFields(Table table, Fields fields, RelationIncludes relationInclu
? EntityUtil.getLatestUsage(daoCollection.usageDAO(), table.getId())
: table.getUsageSummary());
}
if (fields.contains(COLUMN_FIELD) && fields.contains(FIELD_TAGS)) {
if (fields.contains(COLUMN_FIELD)) {
populateEntityFieldTags(entityType, table.getColumns(), table.getFullyQualifiedName(), true);
}
table.setJoins(fields.contains("joins") ? getJoins(table) : table.getJoins());
Expand Down Expand Up @@ -232,8 +232,8 @@ public void setFieldsInBulk(Fields fields, List<Table> entities) {
fetchAndSetFields(entities, fields);
setInheritedFields(entities, fields);

// Column tags come from tag_usage, not table JSON — fetched via fetchAndSetColumnTags when tags
// requested
// Column tags come from tag_usage, not table JSON. Fetch them when columns or tags are
// requested.
entities.forEach(table -> clearFieldsInternal(table, fields));
}

Expand Down Expand Up @@ -283,15 +283,19 @@ private void fetchAndSetCustomMetrics(List<Table> tables, Fields fields) {
}

private void fetchAndSetColumnTags(List<Table> tables, Fields fields) {
if (!fields.contains(FIELD_TAGS) || tables == null || tables.isEmpty()) {
if ((!fields.contains(FIELD_TAGS) && !fields.contains(COLUMN_FIELD))
|| tables == null
|| tables.isEmpty()) {
return;
}
List<String> entityFQNs = tables.stream().map(Table::getFullyQualifiedName).toList();
Map<String, List<TagLabel>> tagsMap = batchFetchTags(entityFQNs);
for (Table table : tables) {
table.setTags(
addDerivedTagsGracefully(
tagsMap.getOrDefault(table.getFullyQualifiedName(), Collections.emptyList())));
if (fields.contains(FIELD_TAGS)) {
List<String> entityFQNs = tables.stream().map(Table::getFullyQualifiedName).toList();
Map<String, List<TagLabel>> tagsMap = batchFetchTags(entityFQNs);
for (Table table : tables) {
table.setTags(
addDerivedTagsGracefully(
tagsMap.getOrDefault(table.getFullyQualifiedName(), Collections.emptyList())));
}
}

if (fields.contains(COLUMN_FIELD)) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,147 @@
package org.openmetadata.service.jdbi3;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.mockito.ArgumentMatchers.anyList;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.openmetadata.schema.api.services.CreateDatabaseService;
import org.openmetadata.schema.entity.data.Table;
import org.openmetadata.schema.type.Column;
import org.openmetadata.schema.type.ColumnDataType;
import org.openmetadata.schema.type.EntityReference;
import org.openmetadata.schema.type.TagLabel;
import org.openmetadata.schema.type.TagLabel.LabelType;
import org.openmetadata.schema.type.TagLabel.State;
import org.openmetadata.schema.type.TagLabel.TagSource;
import org.openmetadata.schema.type.Include;
import org.openmetadata.service.Entity;
import org.openmetadata.service.util.EntityUtil.Fields;
import org.openmetadata.service.util.EntityUtil.RelationIncludes;
import org.openmetadata.service.util.FullyQualifiedName;

class TableRepositoryColumnTagsTest {
private CollectionDAO collectionDAO;
private CollectionDAO.TagUsageDAO tagUsageDAO;
private TableRepository repository;

@BeforeEach
void setUp() {
Entity.cleanup();

collectionDAO = mock(CollectionDAO.class);
tagUsageDAO = mock(CollectionDAO.TagUsageDAO.class);

when(collectionDAO.tableDAO()).thenReturn(mock(CollectionDAO.TableDAO.class));
when(collectionDAO.tagUsageDAO()).thenReturn(tagUsageDAO);
Entity.setCollectionDAO(collectionDAO);

repository = new TableRepository();
}

@AfterEach
void tearDown() {
Entity.cleanup();
}

@Test
void setFields_hydratesColumnTagsWhenOnlyColumnsAreRequested() {
Table table = tableWithColumn("service.database.schema.users", "email");
TagLabel piiTag =
new TagLabel()
.withTagFQN("PII.Sensitive")
.withSource(TagSource.CLASSIFICATION)
.withLabelType(LabelType.MANUAL)
.withState(State.CONFIRMED);
String columnFqn = table.getColumns().getFirst().getFullyQualifiedName();

when(tagUsageDAO.getTagsByPrefix(table.getFullyQualifiedName(), ".%", true))
.thenReturn(Map.of(FullyQualifiedName.buildHash(columnFqn), List.of(piiTag)));

repository.setFields(
table,
new Fields(repository.getAllowedFields(), "columns"),
RelationIncludes.fromInclude(Include.NON_DELETED));

assertNull(table.getTags());
assertColumnTags(table, piiTag);
}

@Test
void setFieldsInBulk_hydratesColumnTagsWhenOnlyColumnsAreRequested() {
Table table = tableWithColumn("service.database.schema.users", "email");
TagLabel piiTag =
new TagLabel()
.withTagFQN("PII.Sensitive")
.withSource(TagSource.CLASSIFICATION)
.withLabelType(LabelType.MANUAL)
.withState(State.CONFIRMED);
String columnFqn = table.getColumns().getFirst().getFullyQualifiedName();

when(tagUsageDAO.getTagsInternalBatch(anyList()))
.thenReturn(List.of(tagUsage(columnFqn, piiTag)));

repository.setFieldsInBulk(
new Fields(repository.getAllowedFields(), "columns"), List.of(table));

assertNull(table.getTags());
assertColumnTags(table, piiTag);
}
Comment on lines +77 to +96
Copy link

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

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

In setFieldsInBulk_hydratesColumnTagsWhenOnlyColumnsAreRequested, assertNull(table.getTags()) doesn’t prove bulk table-level tags were not fetched, since clearFieldsInternal will null out table.tags whenever fields doesn’t include tags (even if they were loaded). To fully regression-test the “keep top-level tags gated behind fields=tags” behavior, add a Mockito verification that the tag DAO is only queried for the column FQN(s) (and not also for the table FQN) when Fields is just columns.

Copilot uses AI. Check for mistakes.

private static Table tableWithColumn(String tableFqn, String columnName) {
String schemaFqn = FullyQualifiedName.getParentFQN(tableFqn);
String databaseFqn = FullyQualifiedName.getParentFQN(schemaFqn);
String serviceFqn = FullyQualifiedName.getParentFQN(databaseFqn);
Column column =
new Column()
.withName(columnName)
.withFullyQualifiedName(FullyQualifiedName.add(tableFqn, columnName))
.withDataType(ColumnDataType.STRING);

return new Table()
.withId(UUID.randomUUID())
.withName(FullyQualifiedName.getShortName(tableFqn))
.withFullyQualifiedName(tableFqn)
.withDatabaseSchema(entityReference(Entity.DATABASE_SCHEMA, schemaFqn))
.withDatabase(entityReference(Entity.DATABASE, databaseFqn))
.withService(entityReference(Entity.DATABASE_SERVICE, serviceFqn))
.withServiceType(CreateDatabaseService.DatabaseServiceType.Mysql)
.withColumns(List.of(column));
}

private static EntityReference entityReference(String type, String fqn) {
return new EntityReference()
.withId(UUID.randomUUID())
.withType(type)
.withName(FullyQualifiedName.getShortName(fqn))
.withFullyQualifiedName(fqn);
}

private static CollectionDAO.TagUsageDAO.TagLabelWithFQNHash tagUsage(
String targetFqn, TagLabel tagLabel) {
CollectionDAO.TagUsageDAO.TagLabelWithFQNHash usage =
new CollectionDAO.TagUsageDAO.TagLabelWithFQNHash();
usage.setTargetFQNHash(FullyQualifiedName.buildHash(targetFqn));
usage.setSource(tagLabel.getSource().ordinal());
usage.setTagFQN(tagLabel.getTagFQN());
usage.setLabelType(tagLabel.getLabelType().ordinal());
usage.setState(tagLabel.getState().ordinal());
return usage;
}

private static void assertColumnTags(Table table, TagLabel expectedTag) {
List<TagLabel> tags = table.getColumns().getFirst().getTags();
assertEquals(1, tags.size());
assertEquals(expectedTag.getTagFQN(), tags.getFirst().getTagFQN());
assertEquals(expectedTag.getSource(), tags.getFirst().getSource());
assertEquals(expectedTag.getLabelType(), tags.getFirst().getLabelType());
assertEquals(expectedTag.getState(), tags.getFirst().getState());
}
}
Loading