Change to Lineage as a base abstract class

shnapz · shnapz · commit d054d3ec6fcf · 2025-11-10T20:28:08.000-05:00
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageRegistrar.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageRegistrar.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.lineage;
+
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.metrics.Lineage;
+import org.apache.beam.sdk.options.PipelineOptions;
+
+public interface LineageRegistrar {
+
+  @Nullable
+  Lineage fromOptions(PipelineOptions options, Lineage.Type type);
+}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java
@@ -1,40 +1,53 @@
-package org.apache.beam.sdk.lineage;
-
-import com.facebook.presto.hadoop.$internal.org.apache.avro.reflect.Nullable;
-
-public interface LineageReporter {
-  /**
-   * Adds lineage information using pre-formatted FQN segments.
-   *
-   * @param rollupSegments FQN segments already escaped per Dataplex format
-   */
-  void add(Iterable<String> rollupSegments);
-
-  /**
-   * Adds lineage with system, optional subtype, and hierarchical segments.
-   *
-   * @param system The data system identifier (e.g., "bigquery", "kafka")
-   * @param subtype Optional subtype (e.g., "table", "topic"), may be null
-   * @param segments Hierarchical path segments
-   * @param lastSegmentSep Separator for the last segment, may be null
-   */
-  void add(
-      String system,
-      @Nullable String subtype,
-      Iterable<String> segments,
-      @Nullable String lastSegmentSep);
-
-  /**
-   * Add a FQN (fully-qualified name) to Lineage.
-   */
-  default void add(String system, Iterable<String> segments, @Nullable String sep) {
-    add(system, null, segments, sep);
-  }
-
-  /**
-   * Add a FQN (fully-qualified name) to Lineage.
-   */
-  default void add(String system, Iterable<String> segments) {
-    add(system, segments, null);
-  }
-}
+///*
+// * Licensed to the Apache Software Foundation (ASF) under one
+// * or more contributor license agreements.  See the NOTICE file
+// * distributed with this work for additional information
+// * regarding copyright ownership.  The ASF licenses this file
+// * to you under the Apache License, Version 2.0 (the
+// * "License"); you may not use this file except in compliance
+// * with the License.  You may obtain a copy of the License at
+// *
+// *     http://www.apache.org/licenses/LICENSE-2.0
+// *
+// * Unless required by applicable law or agreed to in writing, software
+// * distributed under the License is distributed on an "AS IS" BASIS,
+// * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// * See the License for the specific language governing permissions and
+// * limitations under the License.
+// */
+//package org.apache.beam.sdk.lineage;
+//
+//import com.facebook.presto.hadoop.$internal.org.apache.avro.reflect.Nullable;
+//
+//public interface LineageReporter {
+//  /**
+//   * Adds lineage information using pre-formatted FQN segments.
+//   *
+//   * @param rollupSegments FQN segments already escaped per Dataplex format
+//   */
+//  void add(Iterable<String> rollupSegments);
+//
+//  /**
+//   * Adds lineage with system, optional subtype, and hierarchical segments.
+//   *
+//   * @param system The data system identifier (e.g., "bigquery", "kafka")
+//   * @param subtype Optional subtype (e.g., "table", "topic"), may be null
+//   * @param segments Hierarchical path segments
+//   * @param lastSegmentSep Separator for the last segment, may be null
+//   */
+//  void add(
+//      String system,
+//      @Nullable String subtype,
+//      Iterable<String> segments,
+//      @Nullable String lastSegmentSep);
+//
+//  /** Add a FQN (fully-qualified name) to Lineage. */
+//  default void add(String system, Iterable<String> segments, @Nullable String sep) {
+//    add(system, null, segments, sep);
+//  }
+//
+//  /** Add a FQN (fully-qualified name) to Lineage. */
+//  default void add(String system, Iterable<String> segments) {
+//    add(system, segments, null);
+//  }
+//}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Lineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Lineage.java
@@ -28,16 +28,14 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Pattern;
 import org.apache.beam.sdk.annotations.Internal;
-import org.apache.beam.sdk.lineage.LineageReporter;
-import org.apache.beam.sdk.lineage.LineageReporterRegistrar;
+import org.apache.beam.sdk.lineage.LineageRegistrar;
 import org.apache.beam.sdk.metrics.Metrics.MetricsFlag;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.util.common.ReflectHelpers;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -47,31 +45,19 @@
 /**
  * Standard collection of metrics used to record source and sinks information for lineage tracking.
  */
-public class Lineage {
-  // Namespace for lineage metrics; used to filter queries in Lineage.query() and in MetricsLineageReporter
+public abstract class Lineage {
   public static final String LINEAGE_NAMESPACE = "lineage";
   private static final Logger LOG = LoggerFactory.getLogger(Lineage.class);
-  private static final AtomicReference<LineageReporter> SOURCES = new AtomicReference<>();
-  private static final AtomicReference<LineageReporter> SINKS = new AtomicReference<>();
+  private static final AtomicReference<Lineage> SOURCES = new AtomicReference<>();
+  private static final AtomicReference<Lineage> SINKS = new AtomicReference<>();
 
   private static final AtomicReference<KV<Long, Integer>> LINEAGE_REVISION =
       new AtomicReference<>();
 
   // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot.
   private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]");
 
-  private final Metric metric;
-
-  private Lineage(Type type) {
-    if (MetricsFlag.lineageRollupEnabled()) {
-      this.metric =
-          Metrics.boundedTrie(
-              LINEAGE_NAMESPACE,
-              type == Type.SOURCE ? Type.SOURCEV2.toString() : Type.SINKV2.toString());
-    } else {
-      this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString());
-    }
-  }
+  protected Lineage() {}
 
   @Internal
   public static void initialize(PipelineOptions options) {
@@ -82,70 +68,70 @@ public static void initialize(PipelineOptions options) {
     while (true) {
       KV<Long, Integer> currentRevision = LINEAGE_REVISION.get();
 
-      // Skip re-initialization if same options and revision hasn't changed
       if (currentRevision != null
           && currentRevision.getKey().equals(optionsId)
           && currentRevision.getValue() >= nextRevision) {
-        LOG.debug("Lineage already initialized with options ID {} revision {}, skipping",
-            optionsId, currentRevision.getValue());
+        LOG.debug(
+            "Lineage already initialized with options ID {} revision {}, skipping",
+            optionsId,
+            currentRevision.getValue());
         return;
       }
 
       if (LINEAGE_REVISION.compareAndSet(currentRevision, KV.of(optionsId, nextRevision))) {
-        LineageReporter sources = createReporter(options, Type.SOURCE);
-        LineageReporter sinks = createReporter(options, Type.SINK);
+        Lineage sources = createLineage(options, Type.SOURCE);
+        Lineage sinks = createLineage(options, Type.SINK);
 
         SOURCES.set(sources);
         SINKS.set(sinks);
 
         if (currentRevision == null) {
           LOG.info("Lineage initialized with options ID {} revision {}", optionsId, nextRevision);
         } else {
-          LOG.info("Lineage re-initialized from options ID {} to {} (revision {} -> {})",
-              currentRevision.getKey(), optionsId,
-              currentRevision.getValue(), nextRevision);
+          LOG.info(
+              "Lineage re-initialized from options ID {} to {} (revision {} -> {})",
+              currentRevision.getKey(),
+              optionsId,
+              currentRevision.getValue(),
+              nextRevision);
         }
         return;
       }
     }
   }
 
-  /// //// NEW METHOD
-  private static LineageReporter createReporter(PipelineOptions options, Type type) {
-    Set<LineageReporterRegistrar> registrars = Sets.newTreeSet(
-        ReflectHelpers.ObjectsClassComparator.INSTANCE);
-    registrars.addAll(Lists.newArrayList(
-        ServiceLoader.load(LineageReporterRegistrar.class,
-            ReflectHelpers.findClassLoader())));
+  private static Lineage createLineage(PipelineOptions options, Type type) {
+    Set<LineageRegistrar> registrars =
+        Sets.newTreeSet(ReflectHelpers.ObjectsClassComparator.INSTANCE);
+    registrars.addAll(
+        Lists.newArrayList(
+            ServiceLoader.load(LineageRegistrar.class, ReflectHelpers.findClassLoader())));
 
-    for (LineageReporterRegistrar registrar : registrars) {
-      LineageReporter reporter = registrar.fromOptions(options, type);
+    for (LineageRegistrar registrar : registrars) {
+      Lineage reporter = registrar.fromOptions(options, type);
       if (reporter != null) {
-        LOG.info("Using {} for lineage type {}",
-            reporter.getClass().getName(), type);
+        LOG.info("Using {} for lineage type {}", reporter.getClass().getName(), type);
         return reporter;
       }
     }
 
     LOG.debug("Using default Metrics-based lineage for type {}", type);
-    return new MetricsLineageReporter(type);
+    return new MetricsLineage(type);
   }
 
-  /**
-   * Get {@link LineageReporter} representing sources and optionally side inputs.
-   */
-  public static LineageReporter getSources() {
-    LineageReporter sources = SOURCES.get();
+  /** Get {@link Lineage} representing sources and optionally side inputs. */
+  public static Lineage getSources() {
+    Lineage sources = SOURCES.get();
     if (sources == null) {
       initialize(PipelineOptionsFactory.create());
       sources = SOURCES.get();
     }
     return sources;
   }
 
-  /** {@link LineageReporter} representing sinks. */
-  public static LineageReporter getSinks() {
-    LineageReporter sinks = SINKS.get();
+  /** {@link Lineage} representing sinks. */
+  public static Lineage getSinks() {
+    Lineage sinks = SINKS.get();
     if (sinks == null) {
       initialize(PipelineOptionsFactory.create());
       sinks = SINKS.get();
@@ -228,14 +214,7 @@ public void add(String system, Iterable<String> segments) {
    *     which is already escaped.
    *     <p>In particular, this means they will often have trailing delimiters.
    */
-  public void add(Iterable<String> rollupSegments) {
-    ImmutableList<String> segments = ImmutableList.copyOf(rollupSegments);
-    if (MetricsFlag.lineageRollupEnabled()) {
-      ((BoundedTrie) this.metric).add(segments);
-    } else {
-      ((StringSet) this.metric).add(String.join("", segments));
-    }
-  }
+  public abstract void add(Iterable<String> rollupSegments);
 
   /**
    * Query {@link BoundedTrie} metrics from {@link MetricResults}.
@@ -245,9 +224,8 @@ public void add(Iterable<String> rollupSegments) {
    * @param truncatedMarker the marker to use to represent truncated FQNs.
    * @return A flat representation of all FQNs. If the FQN was truncated then it has a trailing
    *     truncatedMarker.
-   *
-   * <p>NOTE: When using a custom LineageReporter plugin, this method
-   * will return empty results since lineage is not stored in Metrics.
+   *     <p>NOTE: When using a custom LineageReporter plugin, this method will return empty results
+   *     since lineage is not stored in Metrics.
    */
   public static Set<String> query(MetricResults results, Type type, String truncatedMarker) {
     MetricQueryResults lineageQueryResults = getLineageQueryResults(results, type);
@@ -276,9 +254,8 @@ public static Set<String> query(MetricResults results, Type type, String truncat
    * @param results FQNs from the result
    * @param type sources or sinks
    * @return A flat representation of all FQNs. If the FQN was truncated then it has a trailing '*'.
-   *
-   * <p>NOTE: When using a custom LineageReporter plugin, this method
-   * will return empty results since lineage is not stored in Metrics.
+   *     <p>NOTE: When using a custom LineageReporter plugin, this method will return empty results
+   *     since lineage is not stored in Metrics.
    */
   public static Set<String> query(MetricResults results, Type type) {
     if (MetricsFlag.lineageRollupEnabled()) {
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineage.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.metrics;
+
+import org.apache.beam.sdk.lineage.LineageReporter;
+import org.apache.beam.sdk.metrics.Metrics.MetricsFlag;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+
+public class MetricsLineage extends Lineage {
+
+  private final Metric metric;
+
+  public MetricsLineage(final Lineage.Type type) {
+    if (MetricsFlag.lineageRollupEnabled()) {
+      this.metric =
+          Metrics.boundedTrie(
+              Lineage.LINEAGE_NAMESPACE,
+              type == Lineage.Type.SOURCE
+                  ? Lineage.Type.SOURCEV2.toString()
+                  : Lineage.Type.SINKV2.toString());
+    } else {
+      this.metric = Metrics.stringSet(Lineage.LINEAGE_NAMESPACE, type.toString());
+    }
+  }
+
+  @Override
+  public void add(final Iterable<String> rollupSegments) {
+    ImmutableList<String> segments = ImmutableList.copyOf(rollupSegments);
+    if (MetricsFlag.lineageRollupEnabled()) {
+      ((BoundedTrie) this.metric).add(segments);
+    } else {
+      ((StringSet) this.metric).add(String.join("", segments));
+    }
+  }
+
+  @Override
+  public void add(
+      final String system,
+      final String subtype,
+      final Iterable<String> segments,
+      final String lastSegmentSep) {
+    add(Lineage.getFQNParts(system, subtype, segments, lastSegmentSep));
+  }
+}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java
