[FLINK-34466] create KafkaDatasetFacet

Signed-off-by: Pawel Leszczynski <[email protected]>

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/DefaultKafkaDatasetFacet.java

@@ -0,0 +1,63 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import org.apache.flink.connector.kafka.source.KafkaPropertiesUtil;
+
+import java.util.Objects;
+import java.util.Properties;
+
+/** Default implementation of {@link KafkaDatasetFacet}. */
+public class DefaultKafkaDatasetFacet implements KafkaDatasetFacet {
+
+    public static final String KAFKA_FACET_NAME = "kafka";
+
+    private Properties properties;
+
+    private final KafkaDatasetIdentifier topicIdentifier;
+
+    public DefaultKafkaDatasetFacet(KafkaDatasetIdentifier topicIdentifier, Properties properties) {
+        this(topicIdentifier);
+
+        this.properties = new Properties();
+        KafkaPropertiesUtil.copyProperties(properties, this.properties);
+    }
+
+    public DefaultKafkaDatasetFacet(KafkaDatasetIdentifier topicIdentifier) {
+        this.topicIdentifier = topicIdentifier;
+    }
+
+    public void setProperties(Properties properties) {
+        this.properties = new Properties();
+        KafkaPropertiesUtil.copyProperties(properties, this.properties);
+    }
+
+    public Properties getProperties() {
+        return properties;
+    }
+
+    public KafkaDatasetIdentifier getTopicIdentifier() {
+        return topicIdentifier;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        DefaultKafkaDatasetFacet that = (DefaultKafkaDatasetFacet) o;
+        return Objects.equals(properties, that.properties)
+                && Objects.equals(topicIdentifier, that.topicIdentifier);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(properties, topicIdentifier);
+    }
+
+    @Override
+    public String name() {
+        return KAFKA_FACET_NAME;
+    }
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/DefaultKafkaDatasetIdentifier.java

@@ -0,0 +1,56 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import javax.annotation.Nullable;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.regex.Pattern;
+
+/** Default implementation of {@link KafkaDatasetIdentifier}. */
+public class DefaultKafkaDatasetIdentifier implements KafkaDatasetIdentifier {
+
+    @Nullable private final List<String> topics;
+    @Nullable private final Pattern topicPattern;
+
+    public DefaultKafkaDatasetIdentifier(List<String> fixedTopics, Pattern topicPattern) {
+        this.topics = fixedTopics;
+        this.topicPattern = topicPattern;
+    }
+
+    public static DefaultKafkaDatasetIdentifier ofPattern(Pattern pattern) {
+        return new DefaultKafkaDatasetIdentifier(Collections.emptyList(), pattern);
+    }
+
+    public static DefaultKafkaDatasetIdentifier ofTopics(List<String> fixedTopics) {
+        return new DefaultKafkaDatasetIdentifier(fixedTopics, null);
+    }
+
+    @Nullable
+    public List<String> getTopics() {
+        return topics;
+    }
+
+    @Nullable
+    public Pattern getTopicPattern() {
+        return topicPattern;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        DefaultKafkaDatasetIdentifier that = (DefaultKafkaDatasetIdentifier) o;
+        return Objects.equals(topics, that.topics)
+                && Objects.equals(topicPattern, that.topicPattern);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(topics, topicPattern);
+    }
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/DefaultTypeDatasetFacet.java

@@ -0,0 +1,42 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+
+import java.util.Objects;
+
+/** Default implementation of {@link KafkaDatasetFacet}. */
+public class DefaultTypeDatasetFacet implements TypeDatasetFacet {
+
+    public static final String TYPE_FACET_NAME = "type";
+
+    private final TypeInformation typeInformation;
+
+    public DefaultTypeDatasetFacet(TypeInformation typeInformation) {
+        this.typeInformation = typeInformation;
+    }
+
+    public TypeInformation getTypeInformation() {
+        return typeInformation;
+    }
+
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        DefaultTypeDatasetFacet that = (DefaultTypeDatasetFacet) o;
+        return Objects.equals(typeInformation, that.typeInformation);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(typeInformation);
+    }
+
+    @Override
+    public String name() {
+        return TYPE_FACET_NAME;
+    }
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/KafkaDatasetFacet.java

@@ -0,0 +1,14 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import org.apache.flink.streaming.api.lineage.LineageDatasetFacet;
+
+import java.util.Properties;
+
+/** Facet definition to contain all Kafka specific information on Kafka sources and sinks. */
+public interface KafkaDatasetFacet extends LineageDatasetFacet {
+    Properties getProperties();
+
+    KafkaDatasetIdentifier getTopicIdentifier();
+
+    void setProperties(Properties properties);
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/KafkaDatasetFacetProvider.java

@@ -0,0 +1,15 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import java.util.Optional;
+
+/** Contains method to extract {@link KafkaDatasetFacet}. */
+public interface KafkaDatasetFacetProvider {
+
+    /**
+     * Returns a Kafka dataset facet or `Optional.empty` in case an implementing class is not able
+     * to identify a dataset.
+     *
+     * @return
+     */
+    Optional<KafkaDatasetFacet> getKafkaDatasetFacet();
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/KafkaDatasetIdentifier.java

@@ -0,0 +1,28 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import javax.annotation.Nullable;
+
+import java.util.List;
+import java.util.regex.Pattern;
+
+/** Kafka dataset identifier which can contain either a list of topics or a topic pattern. */
+public interface KafkaDatasetIdentifier {
+    @Nullable
+    List<String> getTopics();
+
+    @Nullable
+    Pattern getTopicPattern();
+
+    /**
+     * Assigns lineage dataset's name which is topic pattern if it is present or comma separated
+     * list of topics.
+     *
+     * @return
+     */
+    default String toLineageName() {
+        if (getTopicPattern() != null) {
+            return getTopicPattern().toString();
+        }
+        return String.join(",", getTopics());
+    }
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/KafkaDatasetIdentifierProvider.java

@@ -0,0 +1,15 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import java.util.Optional;
+
+/** Contains method which allows extracting topic identifier. */
+public interface KafkaDatasetIdentifierProvider {
+
+    /**
+     * Gets Kafka dataset identifier or empty in case a class implementing is not able to extract
+     * dataset identifier.
+     *
+     * @return
+     */
+    Optional<DefaultKafkaDatasetIdentifier> getDatasetIdentifier();
+}

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/LineageUtil.java

@@ -20,89 +20,43 @@
 package org.apache.flink.connector.kafka.lineage;
 
 import org.apache.flink.api.connector.source.Boundedness;
-import org.apache.flink.connector.kafka.lineage.facets.KafkaTopicListFacet;
-import org.apache.flink.connector.kafka.lineage.facets.KafkaTopicPatternFacet;
 import org.apache.flink.streaming.api.lineage.LineageDataset;
 import org.apache.flink.streaming.api.lineage.LineageDatasetFacet;
 import org.apache.flink.streaming.api.lineage.LineageVertex;
 import org.apache.flink.streaming.api.lineage.SourceLineageVertex;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Properties;
 import java.util.stream.Collectors;
 
-/** Utility class with useful methods for managing dataset facets. */
+/** Utility class with useful methods for managing lineage objects. */
 public class LineageUtil {
 
     private static final String KAFKA_DATASET_PREFIX = "kafka://";
     private static final String COMMA = ",";
     private static final String SEMICOLON = ";";
 
-    /**
-     * Loads facet from any object implementing @link{DatasetFacetProvider} interface.
-     *
-     * @param object
-     * @return
-     */
-    public static Collection<LineageDatasetFacet> facetsFrom(Object object) {
-        return Optional.of(object)
-                .filter(LineageFacetProvider.class::isInstance)
-                .map(LineageFacetProvider.class::cast)
-                .map(LineageFacetProvider::getDatasetFacets)
-                .orElse(Collections.emptyList());
+    public static LineageDataset datasetOf(String namespace, KafkaDatasetFacet kafkaDatasetFacet) {
+        return datasetOf(namespace, kafkaDatasetFacet, Collections.emptyList());
     }
 
-    /**
-     * Creates dataset from a list of facets. Uses {@link KafkaTopicListFacet} to extract dataset
-     * name from. Dataset per each element of topic list is created
-     *
-     * @param facets
-     * @return
-     */
-    public static Collection<LineageDataset> datasetsFrom(
-            String namespace, Collection<LineageDatasetFacet> facets) {
-        // Check if topic list facet is available -> if so explode the list of facets
-        Optional<KafkaTopicListFacet> topicList =
-                facets.stream()
-                        .filter(KafkaTopicListFacet.class::isInstance)
-                        .map(KafkaTopicListFacet.class::cast)
-                        .findAny();
-
-        List<LineageDataset> datasets = new ArrayList<>();
-
-        // Explode list of other facets
-        if (topicList.isPresent()) {
-            List<LineageDatasetFacet> facetsWithoutTopicList =
-                    facets.stream().filter(f -> !f.equals(topicList)).collect(Collectors.toList());
-
-            datasets.addAll(
-                    topicList.get().topics.stream()
-                            .map(t -> datasetOf(namespace, t, facetsWithoutTopicList))
-                            .collect(Collectors.toList()));
-        }
-
-        // Check if topic pattern is present
-        // If so topic pattern will be used as a dataset name
-        datasets.addAll(
-                facets.stream()
-                        .filter(KafkaTopicPatternFacet.class::isInstance)
-                        .map(KafkaTopicPatternFacet.class::cast)
-                        .map(f -> datasetOf(namespace, f.pattern.toString(), facets))
-                        .collect(Collectors.toList()));
-        return datasets;
+    public static LineageDataset datasetOf(
+            String namespace, KafkaDatasetFacet kafkaDatasetFacet, TypeDatasetFacet typeFacet) {
+        return datasetOf(namespace, kafkaDatasetFacet, Collections.singletonList(typeFacet));
     }
 
     private static LineageDataset datasetOf(
-            String namespace, String name, Collection<LineageDatasetFacet> facets) {
+            String namespace,
+            KafkaDatasetFacet kafkaDatasetFacet,
+            List<LineageDatasetFacet> facets) {
         return new LineageDataset() {
             @Override
             public String name() {
-                return name;
+                return kafkaDatasetFacet.getTopicIdentifier().toLineageName();
             }
 
             @Override
@@ -112,16 +66,24 @@ public String namespace() {
 
             @Override
             public Map<String, LineageDatasetFacet> facets() {
-                return facets.stream()
-                        .distinct()
-                        .collect(Collectors.toMap(LineageDatasetFacet::name, item -> item));
+                Map facetMap = new HashMap<String, LineageDatasetFacet>();
+                facetMap.put(DefaultKafkaDatasetFacet.KAFKA_FACET_NAME, kafkaDatasetFacet);
+                facetMap.putAll(
+                        facets.stream()
+                                .collect(
+                                        Collectors.toMap(LineageDatasetFacet::name, item -> item)));
+                return facetMap;
             }
         };
     }
 
-    public static String datasetNamespaceOf(Properties properties) {
+    public static String namespaceOf(Properties properties) {
         String bootstrapServers = properties.getProperty("bootstrap.servers");
 
+        if (bootstrapServers == null) {
+            return KAFKA_DATASET_PREFIX;
+        }
+
         if (bootstrapServers.contains(COMMA)) {
             bootstrapServers = bootstrapServers.split(COMMA)[0];
         } else if (bootstrapServers.contains(SEMICOLON)) {

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/TypeDatasetFacet.java

@@ -0,0 +1,9 @@
+package org.apache.flink.connector.kafka.lineage;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.streaming.api.lineage.LineageDatasetFacet;
+
+/** Facet definition to contain type information of source and sink. */
+public interface TypeDatasetFacet extends LineageDatasetFacet {
+    TypeInformation getTypeInformation();
+}