datahub.protobuf.visitors.dataset.DatasetVisitor Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of datahub-protobuf Show documentation
Show all versions of datahub-protobuf Show documentation
DataHub integration with Protobuf schemas for metadata
package datahub.protobuf.visitors.dataset;
import com.linkedin.common.Deprecation;
import com.linkedin.common.GlobalTags;
import com.linkedin.common.GlossaryTermAssociation;
import com.linkedin.common.GlossaryTermAssociationArray;
import com.linkedin.common.GlossaryTerms;
import com.linkedin.common.InstitutionalMemory;
import com.linkedin.common.InstitutionalMemoryMetadata;
import com.linkedin.common.InstitutionalMemoryMetadataArray;
import com.linkedin.common.Owner;
import com.linkedin.common.OwnerArray;
import com.linkedin.common.Ownership;
import com.linkedin.common.TagAssociation;
import com.linkedin.common.TagAssociationArray;
import com.linkedin.common.UrnArray;
import com.linkedin.common.urn.DatasetUrn;
import com.linkedin.data.DataMap;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.data.template.StringMap;
import com.linkedin.dataset.DatasetProperties;
import com.linkedin.domain.Domains;
import com.linkedin.events.metadata.ChangeType;
import datahub.event.MetadataChangeProposalWrapper;
import datahub.protobuf.model.ProtobufGraph;
import datahub.protobuf.visitors.ProtobufModelVisitor;
import datahub.protobuf.visitors.VisitContext;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.AllArgsConstructor;
import lombok.Builder;
@Builder
@AllArgsConstructor
public class DatasetVisitor
implements ProtobufModelVisitor> {
@Builder.Default
private final List>
institutionalMemoryMetadataVisitors = List.of();
@Builder.Default
private final List> datasetPropertyVisitors = List.of();
@Builder.Default
private final List> tagAssociationVisitors = List.of();
@Builder.Default
private final List> termAssociationVisitors =
List.of();
@Builder.Default private final List> ownershipVisitors = List.of();
@Builder.Default
private final List> domainVisitors = List.of();
@Builder.Default private final String protocBase64 = "";
@Builder.Default
private final ProtobufModelVisitor descriptionVisitor = new DescriptionVisitor();
@Builder.Default
private final ProtobufModelVisitor deprecationVisitor = new DeprecationVisitor();
@Builder.Default private final boolean enableProtocCustomProperty = false;
@Override
public Stream> visitGraph(
VisitContext context) {
final String datasetUrn = context.getDatasetUrn().toString();
final ProtobufGraph g = context.getGraph();
return Stream.of(
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new DatasetProperties()
.setName(
context
.getDatasetUrn()
.getDatasetNameEntity()
.substring(
context.getDatasetUrn().getDatasetNameEntity().lastIndexOf(".")
+ 1))
.setQualifiedName(context.getDatasetUrn().getDatasetNameEntity())
.setDescription(
g.accept(context, List.of(descriptionVisitor))
.collect(Collectors.joining("\n")))
.setCustomProperties(
new StringMap(
Stream.concat(
enableProtocCustomProperty
? Stream.of(Map.entry("protoc", protocBase64))
: Stream.empty(),
g.accept(context, datasetPropertyVisitors)
.flatMap(
props ->
props.getCustomProperties().entrySet().stream()))
.collect(
Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))),
"datasetProperties"),
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new InstitutionalMemory()
.setElements(
new InstitutionalMemoryMetadataArray(
g.accept(context, institutionalMemoryMetadataVisitors)
.map(inst -> inst.setCreateStamp(context.getAuditStamp()))
.collect(
Collectors.toMap(
InstitutionalMemoryMetadata::getUrl,
Function.identity(),
(a1, a2) -> a1,
LinkedHashMap::new))
.values())),
"institutionalMemory"),
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new GlobalTags()
.setTags(
new TagAssociationArray(
g.accept(context, tagAssociationVisitors)
.collect(Collectors.toList()))),
"globalTags"),
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new GlossaryTerms()
.setTerms(
new GlossaryTermAssociationArray(
g.accept(context, termAssociationVisitors)
.collect(Collectors.toList())))
.setAuditStamp(context.getAuditStamp()),
"glossaryTerms"),
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new Ownership()
.setOwners(
new OwnerArray(
g.accept(context, ownershipVisitors).collect(Collectors.toList())))
.setLastModified(context.getAuditStamp()),
"ownership"),
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
new Domains(
new DataMap(
Map.of(
"domains",
new UrnArray(
g.accept(context, domainVisitors).collect(Collectors.toList()))
.data()))),
"domains"),
g.accept(context, List.of(deprecationVisitor))
.findFirst()
.map(
dep ->
new MetadataChangeProposalWrapper<>(
DatasetUrn.ENTITY_TYPE,
datasetUrn,
ChangeType.UPSERT,
dep,
"deprecation"))
.orElse(null))
.filter(Objects::nonNull);
}
}