datahub.protobuf.visitors.field.ProtobufExtensionFieldVisitor 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
DataHub integration with Protobuf schemas for metadata
There is a newer version: 0.14.1
package datahub.protobuf.visitors.field;

import com.linkedin.common.GlobalTags;
import com.linkedin.common.GlossaryTermAssociation;
import com.linkedin.common.GlossaryTermAssociationArray;
import com.linkedin.common.GlossaryTerms;
import com.linkedin.common.TagAssociation;
import com.linkedin.common.TagAssociationArray;
import com.linkedin.common.urn.TagUrn;
import com.linkedin.schema.SchemaField;
import com.linkedin.tag.TagProperties;
import com.linkedin.util.Pair;
import datahub.protobuf.model.ProtobufField;
import datahub.protobuf.visitors.ProtobufExtensionUtil;
import datahub.protobuf.visitors.VisitContext;

import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static datahub.protobuf.ProtobufUtils.getFieldOptions;
import static datahub.protobuf.ProtobufUtils.getMessageOptions;

public class ProtobufExtensionFieldVisitor extends SchemaFieldVisitor {

    @Override
    public Stream> visitField(ProtobufField field, VisitContext context) {
        boolean isPrimaryKey = getFieldOptions(field.getFieldProto()).stream().map(Pair::getKey)
                .anyMatch(fieldDesc -> fieldDesc.getName().matches("(?i).*primary_?key"));

        List tags = Stream.concat(
                ProtobufExtensionUtil.extractTagPropertiesFromOptions(
                        getFieldOptions(field.getFieldProto()),
                        context.getGraph().getRegistry()),
                        promotedTags(field, context))
                .distinct().map(tag -> new TagAssociation().setTag(new TagUrn(tag.getName())))
                .sorted(Comparator.comparing(t -> t.getTag().getName()))
                .collect(Collectors.toList());

        List terms =  Stream.concat(
                ProtobufExtensionUtil.extractTermAssociationsFromOptions(
                                getFieldOptions(field.getFieldProto()), context.getGraph().getRegistry()),
                promotedTerms(field, context))
                .distinct()
                .sorted(Comparator.comparing(a -> a.getUrn().getNameEntity()))
                .collect(Collectors.toList());

        return context.streamAllPaths(field).map(path -> Pair.of(
                new SchemaField()
                        .setFieldPath(context.getFieldPath(path))
                        .setNullable(!isPrimaryKey)
                        .setIsPartOfKey(isPrimaryKey)
                        .setDescription(field.comment())
                        .setNativeDataType(field.nativeType())
                        .setType(field.schemaFieldDataType())
                        .setGlobalTags(new GlobalTags().setTags(new TagAssociationArray(tags)))
                        .setGlossaryTerms(new GlossaryTerms()
                                .setTerms(new GlossaryTermAssociationArray(terms))
                                .setAuditStamp(context.getAuditStamp())),
                context.calculateSortOrder(path, field)));
    }

    /**
     * Promote tags from nested message to field.
     * @return tags
     */
    private Stream promotedTags(ProtobufField field, VisitContext context) {
        if (field.isMessage()) {
            return context.getGraph().outgoingEdgesOf(field).stream().flatMap(e ->
                    ProtobufExtensionUtil.extractTagPropertiesFromOptions(getMessageOptions(e.getEdgeTarget().messageProto()),
                            context.getGraph().getRegistry())
            ).distinct();
        } else {
            return Stream.of();
        }
    }

    /**
     * Promote terms from nested message to field.
     * @return terms
     */
    private Stream promotedTerms(ProtobufField field, VisitContext context) {
        if (field.isMessage()) {
            return context.getGraph().outgoingEdgesOf(field).stream().flatMap(e ->
                    ProtobufExtensionUtil.extractTermAssociationsFromOptions(getMessageOptions(e.getEdgeTarget().messageProto()),
                            context.getGraph().getRegistry())
            ).distinct();
        } else {
            return Stream.of();
        }
    }

}