From 55c516c3e01d31a12e8c475a2b59564655a7e15d Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 10 Nov 2025 17:20:37 -0500 Subject: [PATCH 01/28] LineageReporter plugin initial implementation --- .../beam/sdk/lineage/LineageReporter.java | 40 +++++++ .../sdk/lineage/LineageReporterRegistrar.java | 12 ++ .../org/apache/beam/sdk/metrics/Lineage.java | 113 ++++++++++++++++-- .../sdk/metrics/MetricsLineageReporter.java | 36 ++++++ 4 files changed, 192 insertions(+), 9 deletions(-) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java 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 new file mode 100644 index 000000000000..a50a72eca9a9 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java @@ -0,0 +1,40 @@ +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 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 segments, + @Nullable String lastSegmentSep); + + /** + * Add a FQN (fully-qualified name) to Lineage. + */ + default void add(String system, Iterable segments, @Nullable String sep) { + add(system, null, segments, sep); + } + + /** + * Add a FQN (fully-qualified name) to Lineage. + */ + default void add(String system, Iterable 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 new file mode 100644 index 000000000000..beba19d4f243 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java @@ -0,0 +1,12 @@ +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 LineageReporterRegistrar { + + @Nullable + LineageReporter fromOptions(PipelineOptions options, Lineage.Type type); + +} 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 index 1e0124fc518b..260b2cd5c181 100644 --- 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 @@ -17,27 +17,46 @@ */ package org.apache.beam.sdk.metrics; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; + import java.util.ArrayList; import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.ServiceLoader; import java.util.Set; +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.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; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * 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 static final String LINEAGE_NAMESPACE = "lineage"; - private static final Lineage SOURCES = new Lineage(Type.SOURCE); - private static final Lineage SINKS = new Lineage(Type.SINK); + private static final Logger LOG = LoggerFactory.getLogger(Lineage.class); + private static final AtomicReference SOURCES = new AtomicReference<>(); + private static final AtomicReference SINKS = new AtomicReference<>(); + + private static final AtomicReference> LINEAGE_REVISION = + new AtomicReference<>(); + // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot. private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]"); @@ -54,14 +73,84 @@ private Lineage(Type type) { } } - /** {@link Lineage} representing sources and optionally side inputs. */ - public static Lineage getSources() { - return SOURCES; + @Internal + public static void initialize(PipelineOptions options) { + checkNotNull(options, "options cannot be null"); + long optionsId = options.getOptionsId(); + int nextRevision = options.revision(); + + while (true) { + KV 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()); + return; + } + + if (LINEAGE_REVISION.compareAndSet(currentRevision, KV.of(optionsId, nextRevision))) { + LineageReporter sources = createReporter(options, Type.SOURCE); + LineageReporter sinks = createReporter(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); + } + return; + } + } } - /** {@link Lineage} representing sinks. */ - public static Lineage getSinks() { - return SINKS; + /// //// NEW METHOD + private static LineageReporter createReporter(PipelineOptions options, Type type) { + Set registrars = Sets.newTreeSet( + ReflectHelpers.ObjectsClassComparator.INSTANCE); + registrars.addAll(Lists.newArrayList( + ServiceLoader.load(LineageReporterRegistrar.class, + ReflectHelpers.findClassLoader()))); + + for (LineageReporterRegistrar registrar : registrars) { + LineageReporter reporter = registrar.fromOptions(options, type); + if (reporter != null) { + 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); + } + + /** + * Get {@link LineageReporter} representing sources and optionally side inputs. + */ + public static LineageReporter getSources() { + LineageReporter 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(); + if (sinks == null) { + initialize(PipelineOptionsFactory.create()); + sinks = SINKS.get(); + } + return sinks; } @VisibleForTesting @@ -156,6 +245,9 @@ public void add(Iterable 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. + * + *

NOTE: When using a custom LineageReporter plugin, this method + * will return empty results since lineage is not stored in Metrics. */ public static Set query(MetricResults results, Type type, String truncatedMarker) { MetricQueryResults lineageQueryResults = getLineageQueryResults(results, type); @@ -184,6 +276,9 @@ public static Set 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 '*'. + * + *

NOTE: When using a custom LineageReporter plugin, this method + * will return empty results since lineage is not stored in Metrics. */ public static Set query(MetricResults results, Type type) { if (MetricsFlag.lineageRollupEnabled()) { 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 new file mode 100644 index 000000000000..73f57188a3d2 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java @@ -0,0 +1,36 @@ +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 MetricsLineageReporter implements LineageReporter { + + private final Metric metric; + + public MetricsLineageReporter(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 rollupSegments) { + ImmutableList 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 segments, + final String lastSegmentSep) { + add(Lineage.getFQNParts(system, subtype, segments, lastSegmentSep)); + } +} From a2bcd2ae84d8c88794cffdfe3511555e74b93a6b Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 10 Nov 2025 20:28:08 -0500 Subject: [PATCH 02/28] Change to Lineage as a base abstract class --- .../beam/sdk/lineage/LineageRegistrar.java | 28 ++++++ .../beam/sdk/lineage/LineageReporter.java | 93 +++++++++-------- .../sdk/lineage/LineageReporterRegistrar.java | 12 --- .../org/apache/beam/sdk/metrics/Lineage.java | 99 +++++++------------ .../beam/sdk/metrics/MetricsLineage.java | 59 +++++++++++ .../sdk/metrics/MetricsLineageReporter.java | 36 ------- 6 files changed, 178 insertions(+), 149 deletions(-) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageRegistrar.java delete mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineage.java delete mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java 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 new file mode 100644 index 000000000000..a716e1f529d7 --- /dev/null +++ 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 index a50a72eca9a9..45c6605308d9 100644 --- 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 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 segments, - @Nullable String lastSegmentSep); - - /** - * Add a FQN (fully-qualified name) to Lineage. - */ - default void add(String system, Iterable segments, @Nullable String sep) { - add(system, null, segments, sep); - } - - /** - * Add a FQN (fully-qualified name) to Lineage. - */ - default void add(String system, Iterable 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 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 segments, +// @Nullable String lastSegmentSep); +// +// /** Add a FQN (fully-qualified name) to Lineage. */ +// default void add(String system, Iterable segments, @Nullable String sep) { +// add(system, null, segments, sep); +// } +// +// /** Add a FQN (fully-qualified name) to Lineage. */ +// default void add(String system, Iterable 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 deleted file mode 100644 index beba19d4f243..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporterRegistrar.java +++ /dev/null @@ -1,12 +0,0 @@ -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 LineageReporterRegistrar { - - @Nullable - LineageReporter fromOptions(PipelineOptions options, Lineage.Type type); - -} 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 index 260b2cd5c181..5ef0819eb44f 100644 --- 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,8 +28,7 @@ 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; @@ -37,7 +36,6 @@ 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,12 +45,11 @@ /** * 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 SOURCES = new AtomicReference<>(); - private static final AtomicReference SINKS = new AtomicReference<>(); + private static final AtomicReference SOURCES = new AtomicReference<>(); + private static final AtomicReference SINKS = new AtomicReference<>(); private static final AtomicReference> LINEAGE_REVISION = new AtomicReference<>(); @@ -60,18 +57,7 @@ public class Lineage { // 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,18 +68,19 @@ public static void initialize(PipelineOptions options) { while (true) { KV 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); @@ -101,41 +88,40 @@ public static void initialize(PipelineOptions options) { 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 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 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(); @@ -143,9 +129,9 @@ public static LineageReporter getSources() { 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 segments) { * which is already escaped. *

In particular, this means they will often have trailing delimiters. */ - public void add(Iterable rollupSegments) { - ImmutableList 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 rollupSegments); /** * Query {@link BoundedTrie} metrics from {@link MetricResults}. @@ -245,9 +224,8 @@ public void add(Iterable 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. - * - *

NOTE: When using a custom LineageReporter plugin, this method - * will return empty results since lineage is not stored in Metrics. + *

NOTE: When using a custom LineageReporter plugin, this method will return empty results + * since lineage is not stored in Metrics. */ public static Set query(MetricResults results, Type type, String truncatedMarker) { MetricQueryResults lineageQueryResults = getLineageQueryResults(results, type); @@ -276,9 +254,8 @@ public static Set 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 '*'. - * - *

NOTE: When using a custom LineageReporter plugin, this method - * will return empty results since lineage is not stored in Metrics. + *

NOTE: When using a custom LineageReporter plugin, this method will return empty results + * since lineage is not stored in Metrics. */ public static Set 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 new file mode 100644 index 000000000000..9076cfc525eb --- /dev/null +++ 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 rollupSegments) { + ImmutableList 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 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 deleted file mode 100644 index 73f57188a3d2..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineageReporter.java +++ /dev/null @@ -1,36 +0,0 @@ -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 MetricsLineageReporter implements LineageReporter { - - private final Metric metric; - - public MetricsLineageReporter(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 rollupSegments) { - ImmutableList 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 segments, - final String lastSegmentSep) { - add(Lineage.getFQNParts(system, subtype, segments, lastSegmentSep)); - } -} From f14a5ecf79e5cec2734cddf717af7c670ff847a2 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 10 Nov 2025 21:14:50 -0500 Subject: [PATCH 03/28] Add tests --- .../beam/sdk/lineage/LineageRegistrar.java | 2 +- .../beam/sdk/lineage/LineageReporter.java | 53 --------------- .../org/apache/beam/sdk/metrics/Lineage.java | 29 ++++---- .../beam/sdk/metrics/MetricsLineage.java | 18 ++--- .../sdk/lineage/LineageRegistrarTest.java | 67 +++++++++++++++++++ .../apache/beam/sdk/lineage/TestLineage.java | 41 ++++++++++++ .../sdk/lineage/TestLineageRegistrar.java | 38 +++++++++++ 7 files changed, 170 insertions(+), 78 deletions(-) delete mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java 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 index a716e1f529d7..278564701f0c 100644 --- 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 @@ -24,5 +24,5 @@ public interface LineageRegistrar { @Nullable - Lineage fromOptions(PipelineOptions options, Lineage.Type type); + Lineage fromOptions(PipelineOptions options, Lineage.LineageDirection direction); } 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 deleted file mode 100644 index 45c6605308d9..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageReporter.java +++ /dev/null @@ -1,53 +0,0 @@ -///* -// * 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 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 segments, -// @Nullable String lastSegmentSep); -// -// /** Add a FQN (fully-qualified name) to Lineage. */ -// default void add(String system, Iterable segments, @Nullable String sep) { -// add(system, null, segments, sep); -// } -// -// /** Add a FQN (fully-qualified name) to Lineage. */ -// default void add(String system, Iterable segments) { -// add(system, segments, null); -// } -//} 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 index 5ef0819eb44f..cafbb9fa6d39 100644 --- 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 @@ -57,6 +57,11 @@ public abstract class Lineage { // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot. private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]"); + public enum LineageDirection { + SOURCE, + SINK + } + protected Lineage() {} @Internal @@ -79,8 +84,8 @@ public static void initialize(PipelineOptions options) { } if (LINEAGE_REVISION.compareAndSet(currentRevision, KV.of(optionsId, nextRevision))) { - Lineage sources = createLineage(options, Type.SOURCE); - Lineage sinks = createLineage(options, Type.SINK); + Lineage sources = createLineage(options, LineageDirection.SOURCE); + Lineage sinks = createLineage(options, LineageDirection.SINK); SOURCES.set(sources); SINKS.set(sinks); @@ -100,7 +105,7 @@ public static void initialize(PipelineOptions options) { } } - private static Lineage createLineage(PipelineOptions options, Type type) { + private static Lineage createLineage(PipelineOptions options, LineageDirection direction) { Set registrars = Sets.newTreeSet(ReflectHelpers.ObjectsClassComparator.INSTANCE); registrars.addAll( @@ -108,18 +113,18 @@ private static Lineage createLineage(PipelineOptions options, Type type) { ServiceLoader.load(LineageRegistrar.class, ReflectHelpers.findClassLoader()))); for (LineageRegistrar registrar : registrars) { - Lineage reporter = registrar.fromOptions(options, type); + Lineage reporter = registrar.fromOptions(options, direction); if (reporter != null) { - LOG.info("Using {} for lineage type {}", reporter.getClass().getName(), type); + LOG.info("Using {} for lineage direction {}", reporter.getClass().getName(), direction); return reporter; } } - LOG.debug("Using default Metrics-based lineage for type {}", type); - return new MetricsLineage(type); + LOG.debug("Using default Metrics-based lineage for direction {}", direction); + return new MetricsLineage(direction); } - /** Get {@link Lineage} representing sources and optionally side inputs. */ + /** {@link Lineage} representing sources and optionally side inputs. */ public static Lineage getSources() { Lineage sources = SOURCES.get(); if (sources == null) { @@ -224,8 +229,8 @@ public void add(String system, Iterable segments) { * @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. - *

NOTE: When using a custom LineageReporter plugin, this method will return empty results - * since lineage is not stored in Metrics. + *

NOTE: When using a custom Lineage plugin, this method will return empty results since + * lineage is not stored in Metrics. */ public static Set query(MetricResults results, Type type, String truncatedMarker) { MetricQueryResults lineageQueryResults = getLineageQueryResults(results, type); @@ -254,8 +259,8 @@ public static Set 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 '*'. - *

NOTE: When using a custom LineageReporter plugin, this method will return empty results - * since lineage is not stored in Metrics. + *

NOTE: When using a custom Lineage plugin, this method will return empty results since + * lineage is not stored in Metrics. */ public static Set 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 index 9076cfc525eb..836abe4c4cc3 100644 --- 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 @@ -17,7 +17,6 @@ */ 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; @@ -25,12 +24,16 @@ public class MetricsLineage extends Lineage { private final Metric metric; - public MetricsLineage(final Lineage.Type type) { + public MetricsLineage(final Lineage.LineageDirection direction) { + // Derive Metrics-specific Type from LineageDirection + Lineage.Type type = + (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK; + if (MetricsFlag.lineageRollupEnabled()) { this.metric = Metrics.boundedTrie( Lineage.LINEAGE_NAMESPACE, - type == Lineage.Type.SOURCE + direction == Lineage.LineageDirection.SOURCE ? Lineage.Type.SOURCEV2.toString() : Lineage.Type.SINKV2.toString()); } else { @@ -47,13 +50,4 @@ public void add(final Iterable rollupSegments) { ((StringSet) this.metric).add(String.join("", segments)); } } - - @Override - public void add( - final String system, - final String subtype, - final Iterable segments, - final String lastSegmentSep) { - add(Lineage.getFQNParts(system, subtype, segments, lastSegmentSep)); - } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java new file mode 100644 index 000000000000..20370a2ac161 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -0,0 +1,67 @@ +/* + * 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 static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.notNullValue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; + +import java.util.ServiceLoader; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link LineageRegistrar} ServiceLoader discovery. */ +@RunWith(JUnit4.class) +public class LineageRegistrarTest { + + @Test + public void testServiceLoaderDiscovery() { + // Load all LineageRegistrar implementations via ServiceLoader + for (LineageRegistrar registrar : + Lists.newArrayList(ServiceLoader.load(LineageRegistrar.class).iterator())) { + + // Check if we found the TestLineageRegistrar + if (registrar instanceof TestLineageRegistrar) { + + // Test with SOURCE direction + Lineage sourceLineage = + registrar.fromOptions(PipelineOptionsFactory.create(), Lineage.LineageDirection.SOURCE); + assertThat(sourceLineage, notNullValue()); + assertThat(sourceLineage, instanceOf(TestLineage.class)); + assertEquals(Lineage.LineageDirection.SOURCE, ((TestLineage) sourceLineage).getDirection()); + + // Test with SINK direction + Lineage sinkLineage = + registrar.fromOptions(PipelineOptionsFactory.create(), Lineage.LineageDirection.SINK); + assertThat(sinkLineage, notNullValue()); + assertThat(sinkLineage, instanceOf(TestLineage.class)); + assertEquals(Lineage.LineageDirection.SINK, ((TestLineage) sinkLineage).getDirection()); + + return; + } + } + + fail("Expected to find " + TestLineageRegistrar.class); + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java new file mode 100644 index 000000000000..7ccd73a1e2d0 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -0,0 +1,41 @@ +/* + * 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 org.apache.beam.sdk.metrics.Lineage; + +/** + * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery. + */ +public class TestLineage extends Lineage { + + private final LineageDirection direction; + + public TestLineage(LineageDirection direction) { + this.direction = direction; + } + + @Override + public void add(Iterable rollupSegments) { + // Test implementation - no-op for discovery testing + } + + public LineageDirection getDirection() { + return direction; + } +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java new file mode 100644 index 000000000000..3e06bfca4c82 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java @@ -0,0 +1,38 @@ +/* + * 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.google.auto.service.AutoService; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.PipelineOptions; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A test {@link LineageRegistrar} for ServiceLoader discovery testing. This registrar always + * returns a TestLineage instance. + */ +@AutoService(LineageRegistrar.class) +public class TestLineageRegistrar implements LineageRegistrar { + + @Override + public @Nullable Lineage fromOptions( + PipelineOptions options, Lineage.LineageDirection direction) { + // For testing, always return a TestLineage instance + return new TestLineage(direction); + } +} From 2280a2133155712f41c0fb1fcd2666cd491a1207 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 10:17:54 -0500 Subject: [PATCH 04/28] Add a TestPipeline integration test --- .../sdk/lineage/LineageRegistrarTest.java | 212 +++++++++++++++++- .../apache/beam/sdk/lineage/TestLineage.java | 35 ++- .../beam/sdk/lineage/TestLineageOptions.java | 32 +++ .../sdk/lineage/TestLineageRegistrar.java | 15 +- 4 files changed, 283 insertions(+), 11 deletions(-) create mode 100644 sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 20370a2ac161..66959dee8c0e 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -18,23 +18,46 @@ package org.apache.beam.sdk.lineage; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; +import java.util.Arrays; +import java.util.List; import java.util.ServiceLoader; +import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +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.junit.Before; +import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Tests for {@link LineageRegistrar} ServiceLoader discovery. */ +/** + * Tests for {@link LineageRegistrar} ServiceLoader discovery and DirectRunner integration. + */ @RunWith(JUnit4.class) public class LineageRegistrarTest { + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + @Before + public void setUp() { + // Clear any recorded lineage from previous tests + TestLineage.clearRecorded(); + } + @Test public void testServiceLoaderDiscovery() { // Load all LineageRegistrar implementations via ServiceLoader @@ -44,16 +67,18 @@ public void testServiceLoaderDiscovery() { // Check if we found the TestLineageRegistrar if (registrar instanceof TestLineageRegistrar) { + // Create options with test lineage enabled + TestLineageOptions options = PipelineOptionsFactory.create().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + // Test with SOURCE direction - Lineage sourceLineage = - registrar.fromOptions(PipelineOptionsFactory.create(), Lineage.LineageDirection.SOURCE); + Lineage sourceLineage = registrar.fromOptions(options, Lineage.LineageDirection.SOURCE); assertThat(sourceLineage, notNullValue()); assertThat(sourceLineage, instanceOf(TestLineage.class)); assertEquals(Lineage.LineageDirection.SOURCE, ((TestLineage) sourceLineage).getDirection()); // Test with SINK direction - Lineage sinkLineage = - registrar.fromOptions(PipelineOptionsFactory.create(), Lineage.LineageDirection.SINK); + Lineage sinkLineage = registrar.fromOptions(options, Lineage.LineageDirection.SINK); assertThat(sinkLineage, notNullValue()); assertThat(sinkLineage, instanceOf(TestLineage.class)); assertEquals(Lineage.LineageDirection.SINK, ((TestLineage) sinkLineage).getDirection()); @@ -64,4 +89,181 @@ public void testServiceLoaderDiscovery() { fail("Expected to find " + TestLineageRegistrar.class); } + + @Test + public void testLineageIntegrationWithSimpleFQN() { + // Enable test lineage plugin + TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + Lineage.initialize(pipeline.getOptions()); + + // Run pipeline that records lineage + pipeline + .apply(Create.of("a", "b", "c")) + .apply(ParDo.of(new RecordSourceLineageDoFn("testsystem", Arrays.asList("db", "table")))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // Verify lineage was recorded + List sources = TestLineage.getRecordedSources(); + assertThat(sources, hasItem("testsystem:db.table")); + } + + @Test + public void testLineageIntegrationWithSubtype() { + // Enable test lineage plugin + TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + Lineage.initialize(pipeline.getOptions()); + + // Run pipeline that records lineage with subtype + pipeline + .apply(Create.of(1, 2, 3)) + .apply( + ParDo.of( + new RecordSourceLineageWithSubtypeDoFn( + "spanner", "table", Arrays.asList("project", "instance", "database", "table")))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // Verify lineage was recorded with subtype + List sources = TestLineage.getRecordedSources(); + assertThat(sources, hasItem("spanner:table:project.instance.database.table")); + } + + @Test + public void testLineageIntegrationWithLastSegmentSeparator() { + // Enable test lineage plugin + TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + Lineage.initialize(pipeline.getOptions()); + + // Run pipeline that records lineage with custom separator + pipeline + .apply(Create.of("x", "y", "z")) + .apply( + ParDo.of( + new RecordSourceLineageWithSeparatorDoFn( + "gcs", Arrays.asList("bucket", "path/to/file.txt"), "/"))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // Verify lineage was recorded with separator + List sources = TestLineage.getRecordedSources(); + assertThat(sources, hasItem("gcs:bucket.path/to/file.txt")); + } + + @Test + public void testLineageIntegrationWithBothSourcesAndSinks() { + // Enable test lineage plugin + TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + Lineage.initialize(pipeline.getOptions()); + + // Run pipeline that records both source and sink lineage + pipeline + .apply(Create.of("data1", "data2")) + .apply(ParDo.of(new RecordBothSourceAndSinkLineageDoFn())); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // Verify both source and sink lineage were recorded + List sources = TestLineage.getRecordedSources(); + List sinks = TestLineage.getRecordedSinks(); + + assertThat(sources, hasItem("input-system:input-db.input-table")); + assertThat(sinks, hasItem("output-system:output-db.output-table")); + } + + @Test + public void testLineageIntegrationWithMultipleElements() { + // Enable test lineage plugin + TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + Lineage.initialize(pipeline.getOptions()); + + // Run pipeline with multiple elements to test thread safety + pipeline + .apply(Create.of(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)) + .apply(ParDo.of(new RecordSourceLineageDoFn("system", Arrays.asList("resource")))); + + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // Verify lineage was recorded for all elements (may have duplicates) + List sources = TestLineage.getRecordedSources(); + assertThat(sources, hasSize(10)); // One per element + assertThat(sources, hasItem("system:resource")); + } + + // Helper DoFn classes for recording lineage + + /** DoFn that records source lineage with simple FQN. */ + private static class RecordSourceLineageDoFn extends DoFn { + private final String system; + private final List segments; + + RecordSourceLineageDoFn(String system, List segments) { + this.system = system; + this.segments = segments; + } + + @ProcessElement + public void processElement(ProcessContext c) { + Lineage.getSources().add(system, segments); + c.output(c.element()); + } + } + + /** DoFn that records source lineage with subtype. */ + private static class RecordSourceLineageWithSubtypeDoFn extends DoFn { + private final String system; + private final String subtype; + private final List segments; + + RecordSourceLineageWithSubtypeDoFn(String system, String subtype, List segments) { + this.system = system; + this.subtype = subtype; + this.segments = segments; + } + + @ProcessElement + public void processElement(ProcessContext c) { + Lineage.getSources().add(system, subtype, segments, null); + c.output(c.element()); + } + } + + /** DoFn that records source lineage with custom last segment separator. */ + private static class RecordSourceLineageWithSeparatorDoFn extends DoFn { + private final String system; + private final List segments; + private final String separator; + + RecordSourceLineageWithSeparatorDoFn(String system, List segments, String separator) { + this.system = system; + this.segments = segments; + this.separator = separator; + } + + @ProcessElement + public void processElement(ProcessContext c) { + Lineage.getSources().add(system, segments, separator); + c.output(c.element()); + } + } + + /** DoFn that records both source and sink lineage. */ + private static class RecordBothSourceAndSinkLineageDoFn extends DoFn { + @ProcessElement + public void processElement(ProcessContext c) { + Lineage.getSources().add("input-system", ImmutableList.of("input-db", "input-table")); + Lineage.getSinks().add("output-system", ImmutableList.of("output-db", "output-table")); + c.output(c.element()); + } + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java index 7ccd73a1e2d0..c0b3e67470c2 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -17,13 +17,25 @@ */ package org.apache.beam.sdk.lineage; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** - * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery. + * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery + * and integration testing with DirectRunner. + * + *

This implementation records all lineage FQNs in thread-safe static storage for test + * assertions. */ public class TestLineage extends Lineage { + // Thread-safe storage for recorded lineage, keyed by direction + private static final ConcurrentHashMap> RECORDED_LINEAGE = + new ConcurrentHashMap<>(); + private final LineageDirection direction; public TestLineage(LineageDirection direction) { @@ -32,10 +44,29 @@ public TestLineage(LineageDirection direction) { @Override public void add(Iterable rollupSegments) { - // Test implementation - no-op for discovery testing + // Record the FQN for test assertions + String fqn = String.join("", rollupSegments); + RECORDED_LINEAGE.computeIfAbsent(direction, k -> new CopyOnWriteArrayList<>()).add(fqn); } public LineageDirection getDirection() { return direction; } + + /** Returns all recorded source lineage FQNs. */ + public static List getRecordedSources() { + return ImmutableList.copyOf( + RECORDED_LINEAGE.getOrDefault(LineageDirection.SOURCE, ImmutableList.of())); + } + + /** Returns all recorded sink lineage FQNs. */ + public static List getRecordedSinks() { + return ImmutableList.copyOf( + RECORDED_LINEAGE.getOrDefault(LineageDirection.SINK, ImmutableList.of())); + } + + /** Clears all recorded lineage. Should be called in @Before to ensure test isolation. */ + public static void clearRecorded() { + RECORDED_LINEAGE.clear(); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java new file mode 100644 index 000000000000..e3437a55bb3c --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java @@ -0,0 +1,32 @@ +/* + * 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 org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptions; + +/** PipelineOptions for configuring the test lineage plugin. */ +public interface TestLineageOptions extends PipelineOptions { + + @Description("Enable test lineage plugin for integration testing") + @Default.Boolean(false) + Boolean getEnableTestLineage(); + + void setEnableTestLineage(Boolean value); +} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java index 3e06bfca4c82..23598d4420cf 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java @@ -23,8 +23,10 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * A test {@link LineageRegistrar} for ServiceLoader discovery testing. This registrar always - * returns a TestLineage instance. + * A test {@link LineageRegistrar} for ServiceLoader discovery testing. + * + *

This registrar only activates when {@link TestLineageOptions#getEnableTestLineage()} is true, + * ensuring it doesn't interfere with other tests in the suite. */ @AutoService(LineageRegistrar.class) public class TestLineageRegistrar implements LineageRegistrar { @@ -32,7 +34,12 @@ public class TestLineageRegistrar implements LineageRegistrar { @Override public @Nullable Lineage fromOptions( PipelineOptions options, Lineage.LineageDirection direction) { - // For testing, always return a TestLineage instance - return new TestLineage(direction); + // Only activate if explicitly enabled via TestLineageOptions + TestLineageOptions testOptions = options.as(TestLineageOptions.class); + if (testOptions.getEnableTestLineage()) { + return new TestLineage(direction); + } + // Return null to use default MetricsLineage + return null; } } From f96c578105b48dff7f920df90b020f4bd2e1b9c0 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 14:23:24 -0500 Subject: [PATCH 05/28] Initialize Lineage from FileSystem --- .../main/java/org/apache/beam/sdk/io/FileSystems.java | 1 + .../main/java/org/apache/beam/sdk/metrics/Lineage.java | 6 +++--- .../apache/beam/sdk/lineage/LineageRegistrarTest.java | 10 +++++----- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileSystems.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileSystems.java index 155df53c6c2e..6133ca9fdb39 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileSystems.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileSystems.java @@ -578,6 +578,7 @@ public static void setDefaultPipelineOptions(PipelineOptions options) { // entry to set other PipelineOption determined flags Metrics.setDefaultPipelineOptions(options); + Lineage.setDefaultPipelineOptions(options); while (true) { KV revision = FILESYSTEM_REVISION.get(); 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 index cafbb9fa6d39..0b01579f367b 100644 --- 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 @@ -65,7 +65,7 @@ public enum LineageDirection { protected Lineage() {} @Internal - public static void initialize(PipelineOptions options) { + public static void setDefaultPipelineOptions(PipelineOptions options) { checkNotNull(options, "options cannot be null"); long optionsId = options.getOptionsId(); int nextRevision = options.revision(); @@ -128,7 +128,7 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d public static Lineage getSources() { Lineage sources = SOURCES.get(); if (sources == null) { - initialize(PipelineOptionsFactory.create()); + setDefaultPipelineOptions(PipelineOptionsFactory.create()); sources = SOURCES.get(); } return sources; @@ -138,7 +138,7 @@ public static Lineage getSources() { public static Lineage getSinks() { Lineage sinks = SINKS.get(); if (sinks == null) { - initialize(PipelineOptionsFactory.create()); + setDefaultPipelineOptions(PipelineOptionsFactory.create()); sinks = SINKS.get(); } return sinks; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 66959dee8c0e..c03a487e0ba9 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -95,7 +95,7 @@ public void testLineageIntegrationWithSimpleFQN() { // Enable test lineage plugin TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); options.setEnableTestLineage(true); - Lineage.initialize(pipeline.getOptions()); + Lineage.setDefaultPipelineOptions(pipeline.getOptions()); // Run pipeline that records lineage pipeline @@ -115,7 +115,7 @@ public void testLineageIntegrationWithSubtype() { // Enable test lineage plugin TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); options.setEnableTestLineage(true); - Lineage.initialize(pipeline.getOptions()); + Lineage.setDefaultPipelineOptions(pipeline.getOptions()); // Run pipeline that records lineage with subtype pipeline @@ -138,7 +138,7 @@ public void testLineageIntegrationWithLastSegmentSeparator() { // Enable test lineage plugin TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); options.setEnableTestLineage(true); - Lineage.initialize(pipeline.getOptions()); + Lineage.setDefaultPipelineOptions(pipeline.getOptions()); // Run pipeline that records lineage with custom separator pipeline @@ -161,7 +161,7 @@ public void testLineageIntegrationWithBothSourcesAndSinks() { // Enable test lineage plugin TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); options.setEnableTestLineage(true); - Lineage.initialize(pipeline.getOptions()); + Lineage.setDefaultPipelineOptions(pipeline.getOptions()); // Run pipeline that records both source and sink lineage pipeline @@ -184,7 +184,7 @@ public void testLineageIntegrationWithMultipleElements() { // Enable test lineage plugin TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); options.setEnableTestLineage(true); - Lineage.initialize(pipeline.getOptions()); + Lineage.setDefaultPipelineOptions(pipeline.getOptions()); // Run pipeline with multiple elements to test thread safety pipeline From 9631d07d05f54f944f4d794d91beb8f3e811f940 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 15:50:54 -0500 Subject: [PATCH 06/28] Fix formatting --- .../apache/beam/sdk/lineage/LineageRegistrarTest.java | 9 ++++----- .../java/org/apache/beam/sdk/lineage/TestLineage.java | 4 ++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index c03a487e0ba9..2de4f9a4260f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -18,7 +18,6 @@ package org.apache.beam.sdk.lineage; import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.hasItem; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.instanceOf; @@ -44,9 +43,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** - * Tests for {@link LineageRegistrar} ServiceLoader discovery and DirectRunner integration. - */ +/** Tests for {@link LineageRegistrar} ServiceLoader discovery and DirectRunner integration. */ @RunWith(JUnit4.class) public class LineageRegistrarTest { @@ -123,7 +120,9 @@ public void testLineageIntegrationWithSubtype() { .apply( ParDo.of( new RecordSourceLineageWithSubtypeDoFn( - "spanner", "table", Arrays.asList("project", "instance", "database", "table")))); + "spanner", + "table", + Arrays.asList("project", "instance", "database", "table")))); PipelineResult result = pipeline.run(); result.waitUntilFinish(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java index c0b3e67470c2..89997661915f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -24,8 +24,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** - * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery - * and integration testing with DirectRunner. + * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery and + * integration testing with DirectRunner. * *

This implementation records all lineage FQNs in thread-safe static storage for test * assertions. From 63539e3ab5a73e8531ba0d703eb5c20cf5ceaff1 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 16:21:51 -0500 Subject: [PATCH 07/28] Fix Flaky JmsIOTest --- .../src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index b3233f866172..04fb6f7cad63 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java @@ -661,6 +661,9 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { final int delay = 10; return connectorClass == JmsConnectionFactory.class ? (JmsTextMessage message) -> { + if (message == null) { + return null; + } final JmsAcknowledgeCallback originalCallback = message.getAcknowledgeCallback(); JmsAcknowledgeCallback jmsAcknowledgeCallbackMock = Mockito.mock(JmsAcknowledgeCallback.class); @@ -680,6 +683,9 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { return message; } : (ActiveMQMessage message) -> { + if (message == null) { + return null; + } final Callback originalCallback = message.getAcknowledgeCallback(); message.setAcknowledgeCallback( () -> { From e60bd8c4a70877753d40e83b9192d48a3d61800e Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 16:32:23 -0500 Subject: [PATCH 08/28] Adding to change log and fixing style error --- CHANGES.md | 1 + .../apache/beam/sdk/lineage/package-info.java | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java diff --git a/CHANGES.md b/CHANGES.md index bdcbd3451c7b..cc1ec48ba188 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -74,6 +74,7 @@ compatible. Both coders can decode encoded bytes from the other coder ([#38139](https://github.com/apache/beam/issues/38139)). * (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). +* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). ## Breaking Changes diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java new file mode 100644 index 000000000000..d65e9fc966d1 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java @@ -0,0 +1,32 @@ +/* + * 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. + */ +/** + * Lineage tracking support for Apache Beam pipelines. + * + *

This package provides a plugin mechanism to support different lineage implementations through + * the {@link org.apache.beam.sdk.lineage.LineageRegistrar} interface. Lineage implementations can + * be registered and discovered at runtime to track data lineage information during pipeline + * execution. + * + *

For lineage capabilities, see {@link org.apache.beam.sdk.metrics.Lineage}. + */ +@DefaultAnnotation(NonNull.class) +package org.apache.beam.sdk.lineage; + +import edu.umd.cs.findbugs.annotations.DefaultAnnotation; +import org.checkerframework.checker.nullness.qual.NonNull; \ No newline at end of file From ec67d259c32a2070abe26fc63a04d94b0ddd8936 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 11 Nov 2025 17:18:02 -0500 Subject: [PATCH 09/28] fix build --- .../src/main/java/org/apache/beam/sdk/lineage/package-info.java | 2 +- .../jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java index d65e9fc966d1..30fbde839023 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java @@ -29,4 +29,4 @@ package org.apache.beam.sdk.lineage; import edu.umd.cs.findbugs.annotations.DefaultAnnotation; -import org.checkerframework.checker.nullness.qual.NonNull; \ No newline at end of file +import org.checkerframework.checker.nullness.qual.NonNull; diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java index 266d04342d1f..212a27b1b2a7 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java @@ -215,7 +215,7 @@ public void testPublishingThenReadingAll() throws IOException, JMSException { int unackRecords = countRemain(QUEUE); assertTrue( String.format("Too many unacknowledged messages: %d", unackRecords), - unackRecords < OPTIONS.getNumberOfRecords() * 0.003); + unackRecords < OPTIONS.getNumberOfRecords() * 0.005); // acknowledged records int ackRecords = OPTIONS.getNumberOfRecords() - unackRecords; From f1a6e0613e97b8f415e22df3229e12de3a391e95 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Wed, 12 Nov 2025 16:05:28 -0500 Subject: [PATCH 10/28] fix tests --- .../sdk/lineage/LineageRegistrarTest.java | 50 ++++++++++--------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 2de4f9a4260f..86b8ad075f89 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -31,6 +31,7 @@ import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.NeedsRunner; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; @@ -38,8 +39,8 @@ 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.junit.Before; -import org.junit.Rule; import org.junit.Test; +import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -47,14 +48,22 @@ @RunWith(JUnit4.class) public class LineageRegistrarTest { - @Rule public final transient TestPipeline pipeline = TestPipeline.create(); - @Before public void setUp() { // Clear any recorded lineage from previous tests TestLineage.clearRecorded(); } + /** Helper to create a TestPipeline with test lineage enabled. */ + private TestPipeline createTestPipelineWithLineage() { + TestLineageOptions options = PipelineOptionsFactory.create().as(TestLineageOptions.class); + options.setEnableTestLineage(true); + TestPipeline pipeline = TestPipeline.fromOptions(options); + // Disable enforcement since we're not using @Rule + pipeline.enableAbandonedNodeEnforcement(false); + return pipeline; + } + @Test public void testServiceLoaderDiscovery() { // Load all LineageRegistrar implementations via ServiceLoader @@ -88,11 +97,10 @@ public void testServiceLoaderDiscovery() { } @Test + @Category(NeedsRunner.class) public void testLineageIntegrationWithSimpleFQN() { - // Enable test lineage plugin - TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - Lineage.setDefaultPipelineOptions(pipeline.getOptions()); + // Create pipeline with test lineage enabled - Lineage will be initialized during pipeline.run() + TestPipeline pipeline = createTestPipelineWithLineage(); // Run pipeline that records lineage pipeline @@ -108,11 +116,10 @@ public void testLineageIntegrationWithSimpleFQN() { } @Test + @Category(NeedsRunner.class) public void testLineageIntegrationWithSubtype() { - // Enable test lineage plugin - TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - Lineage.setDefaultPipelineOptions(pipeline.getOptions()); + // Create pipeline with test lineage enabled - Lineage will be initialized during pipeline.run() + TestPipeline pipeline = createTestPipelineWithLineage(); // Run pipeline that records lineage with subtype pipeline @@ -133,11 +140,10 @@ public void testLineageIntegrationWithSubtype() { } @Test + @Category(NeedsRunner.class) public void testLineageIntegrationWithLastSegmentSeparator() { - // Enable test lineage plugin - TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - Lineage.setDefaultPipelineOptions(pipeline.getOptions()); + // Create pipeline with test lineage enabled - Lineage will be initialized during pipeline.run() + TestPipeline pipeline = createTestPipelineWithLineage(); // Run pipeline that records lineage with custom separator pipeline @@ -156,11 +162,10 @@ public void testLineageIntegrationWithLastSegmentSeparator() { } @Test + @Category(NeedsRunner.class) public void testLineageIntegrationWithBothSourcesAndSinks() { - // Enable test lineage plugin - TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - Lineage.setDefaultPipelineOptions(pipeline.getOptions()); + // Create pipeline with test lineage enabled - Lineage will be initialized during pipeline.run() + TestPipeline pipeline = createTestPipelineWithLineage(); // Run pipeline that records both source and sink lineage pipeline @@ -179,11 +184,10 @@ public void testLineageIntegrationWithBothSourcesAndSinks() { } @Test + @Category(NeedsRunner.class) public void testLineageIntegrationWithMultipleElements() { - // Enable test lineage plugin - TestLineageOptions options = pipeline.getOptions().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - Lineage.setDefaultPipelineOptions(pipeline.getOptions()); + // Create pipeline with test lineage enabled - Lineage will be initialized during pipeline.run() + TestPipeline pipeline = createTestPipelineWithLineage(); // Run pipeline with multiple elements to test thread safety pipeline From b9cbb8aeaa53ba660efcfeeaaf11e492583827af Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Thu, 13 Nov 2025 17:05:13 -0500 Subject: [PATCH 11/28] Improve test logging --- .../sdk/lineage/LineageRegistrarTest.java | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 86b8ad075f89..96dd445d753a 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -39,8 +39,11 @@ 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.junit.Before; +import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; +import org.junit.rules.TestWatcher; +import org.junit.runner.Description; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -48,6 +51,35 @@ @RunWith(JUnit4.class) public class LineageRegistrarTest { + /** + * TestWatcher that logs detailed lineage diagnostics only when tests fail. + * This keeps successful test output clean while providing deep debugging for failures. + */ + @Rule + public TestWatcher lineageDebugLogger = new TestWatcher() { + @Override + protected void failed(Throwable e, Description description) { + System.err.println("=== Lineage Test Failure Diagnostics ==="); + System.err.println("Test: " + description.getMethodName()); + System.err.println("Error: " + e.getMessage()); + + List sources = TestLineage.getRecordedSources(); + List sinks = TestLineage.getRecordedSinks(); + + System.err.println("\nRecorded Sources (" + sources.size() + "):"); + for (int i = 0; i < sources.size(); i++) { + System.err.println(" [" + i + "] \"" + sources.get(i) + "\""); + } + + System.err.println("\nRecorded Sinks (" + sinks.size() + "):"); + for (int i = 0; i < sinks.size(); i++) { + System.err.println(" [" + i + "] \"" + sinks.get(i) + "\""); + } + + System.err.println("========================================"); + } + }; + @Before public void setUp() { // Clear any recorded lineage from previous tests From ea14cba7b609172255e424f1265b6cf812495031 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 30 Dec 2025 12:46:11 -0500 Subject: [PATCH 12/28] fix test --- .../org/apache/beam/sdk/lineage/LineageRegistrarTest.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 96dd445d753a..17bd95bfd352 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -190,7 +190,7 @@ public void testLineageIntegrationWithLastSegmentSeparator() { // Verify lineage was recorded with separator List sources = TestLineage.getRecordedSources(); - assertThat(sources, hasItem("gcs:bucket.path/to/file.txt")); + assertThat(sources, hasItem("gcs:bucket.`path/to/file.txt`")); } @Test @@ -249,6 +249,7 @@ private static class RecordSourceLineageDoFn extends DoFn { @ProcessElement public void processElement(ProcessContext c) { + // !!! Lineage Caller !!! Lineage.getSources().add(system, segments); c.output(c.element()); } @@ -268,6 +269,7 @@ private static class RecordSourceLineageWithSubtypeDoFn extends DoFn { @ProcessElement public void processElement(ProcessContext c) { + // !!! Lineage Caller !!! Lineage.getSources().add("input-system", ImmutableList.of("input-db", "input-table")); + // !!! Lineage Caller !!! Lineage.getSinks().add("output-system", ImmutableList.of("output-db", "output-table")); c.output(c.element()); } From 3f68c768e6c8e0724d48b6c312d6546995ab5663 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 30 Dec 2025 17:19:18 -0500 Subject: [PATCH 13/28] fix formatting --- .../sdk/lineage/LineageRegistrarTest.java | 51 ++++++++++--------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 17bd95bfd352..efccd6d20370 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -52,33 +52,34 @@ public class LineageRegistrarTest { /** - * TestWatcher that logs detailed lineage diagnostics only when tests fail. - * This keeps successful test output clean while providing deep debugging for failures. + * TestWatcher that logs detailed lineage diagnostics only when tests fail. This keeps successful + * test output clean while providing deep debugging for failures. */ @Rule - public TestWatcher lineageDebugLogger = new TestWatcher() { - @Override - protected void failed(Throwable e, Description description) { - System.err.println("=== Lineage Test Failure Diagnostics ==="); - System.err.println("Test: " + description.getMethodName()); - System.err.println("Error: " + e.getMessage()); - - List sources = TestLineage.getRecordedSources(); - List sinks = TestLineage.getRecordedSinks(); - - System.err.println("\nRecorded Sources (" + sources.size() + "):"); - for (int i = 0; i < sources.size(); i++) { - System.err.println(" [" + i + "] \"" + sources.get(i) + "\""); - } - - System.err.println("\nRecorded Sinks (" + sinks.size() + "):"); - for (int i = 0; i < sinks.size(); i++) { - System.err.println(" [" + i + "] \"" + sinks.get(i) + "\""); - } - - System.err.println("========================================"); - } - }; + public TestWatcher lineageDebugLogger = + new TestWatcher() { + @Override + protected void failed(Throwable e, Description description) { + System.err.println("=== Lineage Test Failure Diagnostics ==="); + System.err.println("Test: " + description.getMethodName()); + System.err.println("Error: " + e.getMessage()); + + List sources = TestLineage.getRecordedSources(); + List sinks = TestLineage.getRecordedSinks(); + + System.err.println("\nRecorded Sources (" + sources.size() + "):"); + for (int i = 0; i < sources.size(); i++) { + System.err.println(" [" + i + "] \"" + sources.get(i) + "\""); + } + + System.err.println("\nRecorded Sinks (" + sinks.size() + "):"); + for (int i = 0; i < sinks.size(); i++) { + System.err.println(" [" + i + "] \"" + sinks.get(i) + "\""); + } + + System.err.println("========================================"); + } + }; @Before public void setUp() { From 7c9a8bceb0c6c3a8df352bf9093d6fb2ec71823e Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 30 Dec 2025 18:34:27 -0500 Subject: [PATCH 14/28] Revert unnecessary fixes in tests --- .../src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java | 2 +- .../src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java | 6 ------ 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java index 212a27b1b2a7..266d04342d1f 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java @@ -215,7 +215,7 @@ public void testPublishingThenReadingAll() throws IOException, JMSException { int unackRecords = countRemain(QUEUE); assertTrue( String.format("Too many unacknowledged messages: %d", unackRecords), - unackRecords < OPTIONS.getNumberOfRecords() * 0.005); + unackRecords < OPTIONS.getNumberOfRecords() * 0.003); // acknowledged records int ackRecords = OPTIONS.getNumberOfRecords() - unackRecords; diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java index 04fb6f7cad63..b3233f866172 100644 --- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java @@ -661,9 +661,6 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { final int delay = 10; return connectorClass == JmsConnectionFactory.class ? (JmsTextMessage message) -> { - if (message == null) { - return null; - } final JmsAcknowledgeCallback originalCallback = message.getAcknowledgeCallback(); JmsAcknowledgeCallback jmsAcknowledgeCallbackMock = Mockito.mock(JmsAcknowledgeCallback.class); @@ -683,9 +680,6 @@ private JmsIO.UnboundedJmsReader setupReaderForTest() throws JMSException { return message; } : (ActiveMQMessage message) -> { - if (message == null) { - return null; - } final Callback originalCallback = message.getAcknowledgeCallback(); message.setAcknowledgeCallback( () -> { From be86d8fd31296aef795ab21881a0210cbf55c8e9 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Fri, 6 Mar 2026 12:07:32 -0500 Subject: [PATCH 15/28] Extract separate LineageBase interface --- .../beam/sdk/lineage/LineageRegistrar.java | 9 +++- .../apache/beam/sdk/lineage/package-info.java | 4 -- .../org/apache/beam/sdk/metrics/Lineage.java | 22 +++++++--- .../apache/beam/sdk/metrics/LineageBase.java | 42 +++++++++++++++++++ .../beam/sdk/metrics/MetricsLineage.java | 8 +++- .../sdk/lineage/LineageRegistrarTest.java | 5 ++- .../apache/beam/sdk/lineage/TestLineage.java | 19 +++++---- .../sdk/lineage/TestLineageRegistrar.java | 3 +- 8 files changed, 88 insertions(+), 24 deletions(-) create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java 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 index 278564701f0c..ae7b08f3874e 100644 --- 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 @@ -19,10 +19,17 @@ import javax.annotation.Nullable; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptions; +/** + * Interface for discovering and creating lineage plugin implementations. + * + *

Plugins should return {@link LineageBase} implementations that will be wrapped in {@link + * Lineage} facade instances for end users. + */ public interface LineageRegistrar { @Nullable - Lineage fromOptions(PipelineOptions options, Lineage.LineageDirection direction); + LineageBase fromOptions(PipelineOptions options, Lineage.LineageDirection direction); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java index 30fbde839023..1e203fe90fef 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java @@ -25,8 +25,4 @@ * *

For lineage capabilities, see {@link org.apache.beam.sdk.metrics.Lineage}. */ -@DefaultAnnotation(NonNull.class) package org.apache.beam.sdk.lineage; - -import edu.umd.cs.findbugs.annotations.DefaultAnnotation; -import org.checkerframework.checker.nullness.qual.NonNull; 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 index 0b01579f367b..e2e180831e7c 100644 --- 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 @@ -44,8 +44,12 @@ /** * Standard collection of metrics used to record source and sinks information for lineage tracking. + * + *

This is a facade class that provides utility methods and delegates actual lineage recording to + * {@link LineageBase} implementations. Plugins should implement {@link LineageBase} and register + * via {@link org.apache.beam.sdk.lineage.LineageRegistrar}. */ -public abstract class Lineage { +public final class Lineage { public static final String LINEAGE_NAMESPACE = "lineage"; private static final Logger LOG = LoggerFactory.getLogger(Lineage.class); private static final AtomicReference SOURCES = new AtomicReference<>(); @@ -57,12 +61,16 @@ public abstract class Lineage { // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot. private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]"); + private final LineageBase delegate; + public enum LineageDirection { SOURCE, SINK } - protected Lineage() {} + private Lineage(LineageBase delegate) { + this.delegate = checkNotNull(delegate, "delegate cannot be null"); + } @Internal public static void setDefaultPipelineOptions(PipelineOptions options) { @@ -113,15 +121,15 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d ServiceLoader.load(LineageRegistrar.class, ReflectHelpers.findClassLoader()))); for (LineageRegistrar registrar : registrars) { - Lineage reporter = registrar.fromOptions(options, direction); + LineageBase reporter = registrar.fromOptions(options, direction); if (reporter != null) { LOG.info("Using {} for lineage direction {}", reporter.getClass().getName(), direction); - return reporter; + return new Lineage(reporter); } } LOG.debug("Using default Metrics-based lineage for direction {}", direction); - return new MetricsLineage(direction); + return new Lineage(new MetricsLineage(direction)); } /** {@link Lineage} representing sources and optionally side inputs. */ @@ -219,7 +227,9 @@ public void add(String system, Iterable segments) { * which is already escaped. *

In particular, this means they will often have trailing delimiters. */ - public abstract void add(Iterable rollupSegments); + public void add(Iterable rollupSegments) { + delegate.add(rollupSegments); + } /** * Query {@link BoundedTrie} metrics from {@link MetricResults}. diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java new file mode 100644 index 000000000000..2ef1d1948df3 --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java @@ -0,0 +1,42 @@ +/* + * 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.annotations.Internal; + +/** + * Plugin interface for lineage implementations. + * + *

This is the core contract that lineage plugins must implement. Plugins should implement this + * interface and register via {@link org.apache.beam.sdk.lineage.LineageRegistrar}. + * + *

End users should use the {@link Lineage} facade class instead of implementing this interface + * directly. + */ +@Internal +public interface LineageBase { + /** + * Adds the given FQN as lineage. + * + * @param rollupSegments should be an iterable of strings whose concatenation is a valid Dataplex FQN + * which is already escaped. + *

In particular, this means they will often have trailing delimiters. + */ + void add(Iterable rollupSegments); +} 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 index 836abe4c4cc3..3ddb666a6483 100644 --- 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 @@ -20,7 +20,13 @@ 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 { +/** + * Default lineage implementation that stores lineage information in Beam metrics. + * + *

This implementation uses either {@link BoundedTrie} or {@link StringSet} metrics depending on + * the {@link MetricsFlag#lineageRollupEnabled()} flag. + */ +public class MetricsLineage implements LineageBase { private final Metric metric; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index efccd6d20370..8c47911f576b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -30,6 +30,7 @@ import java.util.ServiceLoader; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.NeedsRunner; import org.apache.beam.sdk.testing.TestPipeline; @@ -111,13 +112,13 @@ public void testServiceLoaderDiscovery() { options.setEnableTestLineage(true); // Test with SOURCE direction - Lineage sourceLineage = registrar.fromOptions(options, Lineage.LineageDirection.SOURCE); + LineageBase sourceLineage = registrar.fromOptions(options, Lineage.LineageDirection.SOURCE); assertThat(sourceLineage, notNullValue()); assertThat(sourceLineage, instanceOf(TestLineage.class)); assertEquals(Lineage.LineageDirection.SOURCE, ((TestLineage) sourceLineage).getDirection()); // Test with SINK direction - Lineage sinkLineage = registrar.fromOptions(options, Lineage.LineageDirection.SINK); + LineageBase sinkLineage = registrar.fromOptions(options, Lineage.LineageDirection.SINK); assertThat(sinkLineage, notNullValue()); assertThat(sinkLineage, instanceOf(TestLineage.class)); assertEquals(Lineage.LineageDirection.SINK, ((TestLineage) sinkLineage).getDirection()); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java index 89997661915f..742b67f43dfb 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -21,24 +21,25 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** - * A test implementation of {@link Lineage} for testing LineageRegistrar ServiceLoader discovery and - * integration testing with DirectRunner. + * A test implementation of {@link LineageBase} for testing LineageRegistrar ServiceLoader discovery + * and integration testing with DirectRunner. * *

This implementation records all lineage FQNs in thread-safe static storage for test * assertions. */ -public class TestLineage extends Lineage { +public class TestLineage implements LineageBase { // Thread-safe storage for recorded lineage, keyed by direction - private static final ConcurrentHashMap> RECORDED_LINEAGE = + private static final ConcurrentHashMap> RECORDED_LINEAGE = new ConcurrentHashMap<>(); - private final LineageDirection direction; + private final Lineage.LineageDirection direction; - public TestLineage(LineageDirection direction) { + public TestLineage(Lineage.LineageDirection direction) { this.direction = direction; } @@ -49,20 +50,20 @@ public void add(Iterable rollupSegments) { RECORDED_LINEAGE.computeIfAbsent(direction, k -> new CopyOnWriteArrayList<>()).add(fqn); } - public LineageDirection getDirection() { + public Lineage.LineageDirection getDirection() { return direction; } /** Returns all recorded source lineage FQNs. */ public static List getRecordedSources() { return ImmutableList.copyOf( - RECORDED_LINEAGE.getOrDefault(LineageDirection.SOURCE, ImmutableList.of())); + RECORDED_LINEAGE.getOrDefault(Lineage.LineageDirection.SOURCE, ImmutableList.of())); } /** Returns all recorded sink lineage FQNs. */ public static List getRecordedSinks() { return ImmutableList.copyOf( - RECORDED_LINEAGE.getOrDefault(LineageDirection.SINK, ImmutableList.of())); + RECORDED_LINEAGE.getOrDefault(Lineage.LineageDirection.SINK, ImmutableList.of())); } /** Clears all recorded lineage. Should be called in @Before to ensure test isolation. */ diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java index 23598d4420cf..866a95a891db 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java @@ -19,6 +19,7 @@ import com.google.auto.service.AutoService; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptions; import org.checkerframework.checker.nullness.qual.Nullable; @@ -32,7 +33,7 @@ public class TestLineageRegistrar implements LineageRegistrar { @Override - public @Nullable Lineage fromOptions( + public @Nullable LineageBase fromOptions( PipelineOptions options, Lineage.LineageDirection direction) { // Only activate if explicitly enabled via TestLineageOptions TestLineageOptions testOptions = options.as(TestLineageOptions.class); From b0ed836802d9ef16a86f19760546796746f66329 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Fri, 6 Mar 2026 12:22:23 -0500 Subject: [PATCH 16/28] small fixes --- .../apache/beam/sdk/{metrics => lineage}/LineageBase.java | 8 ++++---- .../org/apache/beam/sdk/lineage/LineageRegistrar.java | 8 +------- .../main/java/org/apache/beam/sdk/metrics/Lineage.java | 5 +---- .../java/org/apache/beam/sdk/metrics/MetricsLineage.java | 1 + .../org/apache/beam/sdk/lineage/LineageRegistrarTest.java | 1 - .../java/org/apache/beam/sdk/lineage/TestLineage.java | 1 - .../org/apache/beam/sdk/lineage/TestLineageRegistrar.java | 1 - 7 files changed, 7 insertions(+), 18 deletions(-) rename sdks/java/core/src/main/java/org/apache/beam/sdk/{metrics => lineage}/LineageBase.java (86%) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java similarity index 86% rename from sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java rename to sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java index 2ef1d1948df3..56bcbbeb9e3b 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/LineageBase.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java @@ -15,7 +15,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.sdk.metrics; +package org.apache.beam.sdk.lineage; import org.apache.beam.sdk.annotations.Internal; @@ -23,10 +23,10 @@ * Plugin interface for lineage implementations. * *

This is the core contract that lineage plugins must implement. Plugins should implement this - * interface and register via {@link org.apache.beam.sdk.lineage.LineageRegistrar}. + * interface and register via {@link LineageRegistrar}. * - *

End users should use the {@link Lineage} facade class instead of implementing this interface - * directly. + *

End users should use the {@link org.apache.beam.sdk.metrics.Lineage} facade class instead of + * implementing this interface directly. */ @Internal public interface LineageBase { 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 index ae7b08f3874e..53aa665c50e3 100644 --- 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 @@ -19,15 +19,9 @@ import javax.annotation.Nullable; import org.apache.beam.sdk.metrics.Lineage; -import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptions; -/** - * Interface for discovering and creating lineage plugin implementations. - * - *

Plugins should return {@link LineageBase} implementations that will be wrapped in {@link - * Lineage} facade instances for end users. - */ +/** Interface for discovering and creating lineage plugin implementations. */ public interface LineageRegistrar { @Nullable 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 index e2e180831e7c..b1c9ff3ec292 100644 --- 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,6 +28,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.lineage.LineageBase; import org.apache.beam.sdk.lineage.LineageRegistrar; import org.apache.beam.sdk.metrics.Metrics.MetricsFlag; import org.apache.beam.sdk.options.PipelineOptions; @@ -44,10 +45,6 @@ /** * Standard collection of metrics used to record source and sinks information for lineage tracking. - * - *

This is a facade class that provides utility methods and delegates actual lineage recording to - * {@link LineageBase} implementations. Plugins should implement {@link LineageBase} and register - * via {@link org.apache.beam.sdk.lineage.LineageRegistrar}. */ public final class Lineage { public static final String LINEAGE_NAMESPACE = "lineage"; 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 index 3ddb666a6483..a22ddc6bd0c9 100644 --- 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 @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.metrics; +import org.apache.beam.sdk.lineage.LineageBase; import org.apache.beam.sdk.metrics.Metrics.MetricsFlag; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java index 8c47911f576b..e1f1a8f4d4a5 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java @@ -30,7 +30,6 @@ import java.util.ServiceLoader; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.Lineage; -import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.NeedsRunner; import org.apache.beam.sdk.testing.TestPipeline; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java index 742b67f43dfb..1e9bbf65149c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -21,7 +21,6 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import org.apache.beam.sdk.metrics.Lineage; -import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java index 866a95a891db..ff691b216542 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java @@ -19,7 +19,6 @@ import com.google.auto.service.AutoService; import org.apache.beam.sdk.metrics.Lineage; -import org.apache.beam.sdk.metrics.LineageBase; import org.apache.beam.sdk.options.PipelineOptions; import org.checkerframework.checker.nullness.qual.Nullable; From 95b71bc2f1ff4def95ef9b8739c55e289129ab15 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 30 Mar 2026 13:01:15 -0400 Subject: [PATCH 17/28] Address PR review --- .../apache/beam/sdk/lineage/LineageBase.java | 8 ++- .../beam/sdk/lineage/LineageOptions.java} | 32 +++++----- .../beam/sdk/lineage/LineageRegistrar.java | 29 --------- .../apache/beam/sdk/lineage/package-info.java | 6 +- .../metrics/BoundedTrieMetricsLineage.java | 44 ++++++++++++++ .../org/apache/beam/sdk/metrics/Lineage.java | 40 +++++++------ .../beam/sdk/metrics/MetricsLineage.java | 60 ------------------- .../sdk/metrics/StringSetMetricsLineage.java} | 29 +++++---- ...istrarTest.java => LineagePluginTest.java} | 58 +++++++----------- .../apache/beam/sdk/lineage/TestLineage.java | 9 ++- 10 files changed, 137 insertions(+), 178 deletions(-) rename sdks/java/core/src/{test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java => main/java/org/apache/beam/sdk/lineage/LineageOptions.java} (53%) delete mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageRegistrar.java create mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java delete mode 100644 sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineage.java rename sdks/java/core/src/{test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java => main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java} (52%) rename sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/{LineageRegistrarTest.java => LineagePluginTest.java} (84%) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java index 56bcbbeb9e3b..c11ff9411a9a 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageBase.java @@ -22,8 +22,12 @@ /** * Plugin interface for lineage implementations. * - *

This is the core contract that lineage plugins must implement. Plugins should implement this - * interface and register via {@link LineageRegistrar}. + *

This is the core contract that lineage plugins must implement. Custom implementations are + * selected via the {@code --lineageType} pipeline option (see {@link LineageOptions}). + * + *

Implementations must provide a public constructor accepting ({@link + * org.apache.beam.sdk.options.PipelineOptions}, {@link + * org.apache.beam.sdk.metrics.Lineage.LineageDirection}). * *

End users should use the {@link org.apache.beam.sdk.metrics.Lineage} facade class instead of * implementing this interface directly. diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageOptions.java similarity index 53% rename from sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java rename to sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageOptions.java index ff691b216542..274874ac0c5f 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageRegistrar.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageOptions.java @@ -17,29 +17,25 @@ */ package org.apache.beam.sdk.lineage; -import com.google.auto.service.AutoService; -import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PipelineOptions; import org.checkerframework.checker.nullness.qual.Nullable; /** - * A test {@link LineageRegistrar} for ServiceLoader discovery testing. + * Pipeline options for selecting a custom {@link LineageBase} implementation. * - *

This registrar only activates when {@link TestLineageOptions#getEnableTestLineage()} is true, - * ensuring it doesn't interfere with other tests in the suite. + *

When not set, the default Metrics-based lineage is used. Can be set from the command line: + * {@code --lineageType=com.example.MyLineage} */ -@AutoService(LineageRegistrar.class) -public class TestLineageRegistrar implements LineageRegistrar { +public interface LineageOptions extends PipelineOptions { - @Override - public @Nullable LineageBase fromOptions( - PipelineOptions options, Lineage.LineageDirection direction) { - // Only activate if explicitly enabled via TestLineageOptions - TestLineageOptions testOptions = options.as(TestLineageOptions.class); - if (testOptions.getEnableTestLineage()) { - return new TestLineage(direction); - } - // Return null to use default MetricsLineage - return null; - } + @Description( + "The fully qualified class name of the LineageBase implementation to use for recording " + + "lineage. The class must implement LineageBase and have a public constructor accepting " + + "(PipelineOptions, Lineage.LineageDirection). " + + "If not specified, the default Metrics-based lineage is used.") + @Nullable + Class getLineageType(); + + void setLineageType(@Nullable Class lineageClass); } 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 deleted file mode 100644 index 53aa665c50e3..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/LineageRegistrar.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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; - -/** Interface for discovering and creating lineage plugin implementations. */ -public interface LineageRegistrar { - - @Nullable - LineageBase fromOptions(PipelineOptions options, Lineage.LineageDirection direction); -} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java index 1e203fe90fef..4fa2dcbb6b71 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/lineage/package-info.java @@ -19,9 +19,9 @@ * Lineage tracking support for Apache Beam pipelines. * *

This package provides a plugin mechanism to support different lineage implementations through - * the {@link org.apache.beam.sdk.lineage.LineageRegistrar} interface. Lineage implementations can - * be registered and discovered at runtime to track data lineage information during pipeline - * execution. + * the {@link org.apache.beam.sdk.lineage.LineageBase} interface. Lineage implementations can be + * selected via the {@code --lineageType} pipeline option to track data lineage information during + * pipeline execution. * *

For lineage capabilities, see {@link org.apache.beam.sdk.metrics.Lineage}. */ diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java new file mode 100644 index 000000000000..78dff258532a --- /dev/null +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java @@ -0,0 +1,44 @@ +/* + * 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.LineageBase; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; + +/** + * Lineage implementation that stores lineage information in {@link BoundedTrie} metrics. + * + *

Used when {@link Metrics.MetricsFlag#lineageRollupEnabled()} is true. + */ +class BoundedTrieMetricsLineage implements LineageBase { + + private final BoundedTrie metric; + + BoundedTrieMetricsLineage(Lineage.LineageDirection direction) { + Lineage.Type type = + (direction == Lineage.LineageDirection.SOURCE) + ? Lineage.Type.SOURCEV2 + : Lineage.Type.SINKV2; + this.metric = Metrics.boundedTrie(Lineage.LINEAGE_NAMESPACE, type.toString()); + } + + @Override + public void add(Iterable rollupSegments) { + metric.add(ImmutableList.copyOf(rollupSegments)); + } +} 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 index b1c9ff3ec292..05d34847a746 100644 --- 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 @@ -23,22 +23,18 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; -import java.util.ServiceLoader; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.lineage.LineageBase; -import org.apache.beam.sdk.lineage.LineageRegistrar; +import org.apache.beam.sdk.lineage.LineageOptions; 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.Lists; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -111,22 +107,32 @@ public static void setDefaultPipelineOptions(PipelineOptions options) { } private static Lineage createLineage(PipelineOptions options, LineageDirection direction) { - Set registrars = - Sets.newTreeSet(ReflectHelpers.ObjectsClassComparator.INSTANCE); - registrars.addAll( - Lists.newArrayList( - ServiceLoader.load(LineageRegistrar.class, ReflectHelpers.findClassLoader()))); - - for (LineageRegistrar registrar : registrars) { - LineageBase reporter = registrar.fromOptions(options, direction); - if (reporter != null) { - LOG.info("Using {} for lineage direction {}", reporter.getClass().getName(), direction); - return new Lineage(reporter); + Class lineageClass = options.as(LineageOptions.class).getLineageType(); + + if (lineageClass != null) { + try { + LineageBase lineage = + lineageClass + .getDeclaredConstructor(PipelineOptions.class, LineageDirection.class) + .newInstance(options, direction); + LOG.info("Using {} for lineage direction {}", lineageClass.getName(), direction); + return new Lineage(lineage); + } catch (ReflectiveOperationException e) { + throw new IllegalArgumentException( + "Failed to instantiate lineage implementation: " + + lineageClass.getName() + + ". The class must have a public constructor accepting " + + "(PipelineOptions, Lineage.LineageDirection).", + e); } } LOG.debug("Using default Metrics-based lineage for direction {}", direction); - return new Lineage(new MetricsLineage(direction)); + LineageBase defaultLineage = + MetricsFlag.lineageRollupEnabled() + ? new BoundedTrieMetricsLineage(direction) + : new StringSetMetricsLineage(direction); + return new Lineage(defaultLineage); } /** {@link Lineage} representing sources and optionally side inputs. */ 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 deleted file mode 100644 index a22ddc6bd0c9..000000000000 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsLineage.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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.LineageBase; -import org.apache.beam.sdk.metrics.Metrics.MetricsFlag; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; - -/** - * Default lineage implementation that stores lineage information in Beam metrics. - * - *

This implementation uses either {@link BoundedTrie} or {@link StringSet} metrics depending on - * the {@link MetricsFlag#lineageRollupEnabled()} flag. - */ -public class MetricsLineage implements LineageBase { - - private final Metric metric; - - public MetricsLineage(final Lineage.LineageDirection direction) { - // Derive Metrics-specific Type from LineageDirection - Lineage.Type type = - (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK; - - if (MetricsFlag.lineageRollupEnabled()) { - this.metric = - Metrics.boundedTrie( - Lineage.LINEAGE_NAMESPACE, - direction == Lineage.LineageDirection.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 rollupSegments) { - ImmutableList segments = ImmutableList.copyOf(rollupSegments); - if (MetricsFlag.lineageRollupEnabled()) { - ((BoundedTrie) this.metric).add(segments); - } else { - ((StringSet) this.metric).add(String.join("", segments)); - } - } -} diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java similarity index 52% rename from sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java rename to sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java index e3437a55bb3c..b12a49182a8c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineageOptions.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java @@ -15,18 +15,27 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.beam.sdk.lineage; +package org.apache.beam.sdk.metrics; -import org.apache.beam.sdk.options.Default; -import org.apache.beam.sdk.options.Description; -import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.lineage.LineageBase; -/** PipelineOptions for configuring the test lineage plugin. */ -public interface TestLineageOptions extends PipelineOptions { +/** + * Lineage implementation that stores lineage information in {@link StringSet} metrics. + * + *

Used when {@link Metrics.MetricsFlag#lineageRollupEnabled()} is false. + */ +class StringSetMetricsLineage implements LineageBase { + + private final StringSet metric; - @Description("Enable test lineage plugin for integration testing") - @Default.Boolean(false) - Boolean getEnableTestLineage(); + StringSetMetricsLineage(Lineage.LineageDirection direction) { + Lineage.Type type = + (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK; + this.metric = Metrics.stringSet(Lineage.LINEAGE_NAMESPACE, type.toString()); + } - void setEnableTestLineage(Boolean value); + @Override + public void add(Iterable rollupSegments) { + metric.add(String.join("", rollupSegments)); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java similarity index 84% rename from sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java rename to sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java index e1f1a8f4d4a5..75c329ae6842 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineageRegistrarTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java @@ -23,11 +23,9 @@ import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.notNullValue; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.fail; import java.util.Arrays; import java.util.List; -import java.util.ServiceLoader; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -37,7 +35,6 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.ParDo; 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.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -47,9 +44,9 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Tests for {@link LineageRegistrar} ServiceLoader discovery and DirectRunner integration. */ +/** Tests for {@link LineageBase} pipeline option selection and DirectRunner integration. */ @RunWith(JUnit4.class) -public class LineageRegistrarTest { +public class LineagePluginTest { /** * TestWatcher that logs detailed lineage diagnostics only when tests fail. This keeps successful @@ -87,10 +84,10 @@ public void setUp() { TestLineage.clearRecorded(); } - /** Helper to create a TestPipeline with test lineage enabled. */ + /** Helper to create a TestPipeline with test lineage configured. */ private TestPipeline createTestPipelineWithLineage() { - TestLineageOptions options = PipelineOptionsFactory.create().as(TestLineageOptions.class); - options.setEnableTestLineage(true); + LineageOptions options = PipelineOptionsFactory.create().as(LineageOptions.class); + options.setLineageType(TestLineage.class); TestPipeline pipeline = TestPipeline.fromOptions(options); // Disable enforcement since we're not using @Rule pipeline.enableAbandonedNodeEnforcement(false); @@ -98,35 +95,22 @@ private TestPipeline createTestPipelineWithLineage() { } @Test - public void testServiceLoaderDiscovery() { - // Load all LineageRegistrar implementations via ServiceLoader - for (LineageRegistrar registrar : - Lists.newArrayList(ServiceLoader.load(LineageRegistrar.class).iterator())) { - - // Check if we found the TestLineageRegistrar - if (registrar instanceof TestLineageRegistrar) { - - // Create options with test lineage enabled - TestLineageOptions options = PipelineOptionsFactory.create().as(TestLineageOptions.class); - options.setEnableTestLineage(true); - - // Test with SOURCE direction - LineageBase sourceLineage = registrar.fromOptions(options, Lineage.LineageDirection.SOURCE); - assertThat(sourceLineage, notNullValue()); - assertThat(sourceLineage, instanceOf(TestLineage.class)); - assertEquals(Lineage.LineageDirection.SOURCE, ((TestLineage) sourceLineage).getDirection()); - - // Test with SINK direction - LineageBase sinkLineage = registrar.fromOptions(options, Lineage.LineageDirection.SINK); - assertThat(sinkLineage, notNullValue()); - assertThat(sinkLineage, instanceOf(TestLineage.class)); - assertEquals(Lineage.LineageDirection.SINK, ((TestLineage) sinkLineage).getDirection()); - - return; - } - } - - fail("Expected to find " + TestLineageRegistrar.class); + public void testExplicitLineageSelection() { + // Instantiate TestLineage directly and verify behavior + LineageOptions options = PipelineOptionsFactory.create().as(LineageOptions.class); + options.setLineageType(TestLineage.class); + + // Test with SOURCE direction + TestLineage sourceLineage = new TestLineage(options, Lineage.LineageDirection.SOURCE); + assertThat(sourceLineage, notNullValue()); + assertThat(sourceLineage, instanceOf(LineageBase.class)); + assertEquals(Lineage.LineageDirection.SOURCE, sourceLineage.getDirection()); + + // Test with SINK direction + TestLineage sinkLineage = new TestLineage(options, Lineage.LineageDirection.SINK); + assertThat(sinkLineage, notNullValue()); + assertThat(sinkLineage, instanceOf(LineageBase.class)); + assertEquals(Lineage.LineageDirection.SINK, sinkLineage.getDirection()); } @Test diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java index 1e9bbf65149c..5344c190fdab 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/TestLineage.java @@ -21,11 +21,11 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** - * A test implementation of {@link LineageBase} for testing LineageRegistrar ServiceLoader discovery - * and integration testing with DirectRunner. + * A test implementation of {@link LineageBase} for integration testing with DirectRunner. * *

This implementation records all lineage FQNs in thread-safe static storage for test * assertions. @@ -38,6 +38,11 @@ public class TestLineage implements LineageBase { private final Lineage.LineageDirection direction; + /** Constructor used by reflection via {@code --lineageType} pipeline option. */ + public TestLineage(PipelineOptions options, Lineage.LineageDirection direction) { + this(direction); + } + public TestLineage(Lineage.LineageDirection direction) { this.direction = direction; } From b7c1dd17a14903694e42f5d1e04e35a55cc73cf2 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 30 Mar 2026 14:12:26 -0400 Subject: [PATCH 18/28] Added LineageOptions registration --- PR_description.md | 165 ++++++++++++++++++ .../org/apache/beam/sdk/metrics/Lineage.java | 2 +- .../DefaultPipelineOptionsRegistrar.java | 2 + 3 files changed, 168 insertions(+), 1 deletion(-) create mode 100644 PR_description.md diff --git a/PR_description.md b/PR_description.md new file mode 100644 index 000000000000..42ebbb52e728 --- /dev/null +++ b/PR_description.md @@ -0,0 +1,165 @@ +Addresses #36790: "[Feature Request]: Make lineage tracking pluggable" + +## Changes + +- **Created** `org.apache.beam.sdk.lineage.LineageBase` - Plugin interface with single `add()` method +- **Refactored** `org.apache.beam.sdk.metrics.Lineage` - Hardcoded metrics → Delegation to `LineageBase` plugins +- **Created** `org.apache.beam.sdk.lineage.LineageRegistrar` - Plugin discovery interface (ServiceLoader) +- **Extracted** `org.apache.beam.sdk.metrics.MetricsLineage` - Default metrics-based implementation (implements `LineageBase`) +- **Added** Plugin initialization in `FileSystems.setDefaultPipelineOptions()` + +## Architecture + +**Before (master)**: `Lineage` was a concrete class hardcoded to use Beam metrics: + +```java +public class Lineage { + private static final Lineage SOURCES = new Lineage(Type.SOURCE); + private static final Lineage SINKS = new Lineage(Type.SINK); + private final Metric metric; // Hardcoded to Beam metrics + + private Lineage(Type type) { + this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString()); + } + + public void add(Iterable segments) { + ((StringSet) metric).add(String.join("", segments)); // Always metrics + } +} +``` + +**After (this PR)**: Clean separation via composition pattern: + +```java +// Plugin contract (simple interface) +public interface LineageBase { + void add(Iterable rollupSegments); +} + +// Public API (final facade delegating to plugin) +public final class Lineage { + private final LineageBase delegate; // Plugin implementation + + private Lineage(LineageBase delegate) { + this.delegate = delegate; + } + + // Delegates to plugin + public void add(Iterable segments) { + delegate.add(segments); + } + + // Convenience overloads + public void add(String system, Iterable segments) { ... } + public void add(String system, String subtype, ...) { ... } + + // Static utilities (unchanged) + public static Lineage getSources() { ... } + public static Lineage getSinks() { ... } + public static String wrapSegment(String value) { ... } + public static Set query(MetricResults results, Type type) { ... } +} + +// Default implementation (backward compatible) +public class MetricsLineage implements LineageBase { + private final Metric metric; + + @Override + public void add(Iterable segments) { + ((BoundedTrie) metric).add(segments); + } +} +``` + +**Plugin Selection**: ServiceLoader discovery, first match wins, fallback to `MetricsLineage`. + +**Backward Compatibility**: ✅ All existing code works unchanged (24+ call sites, static utilities, enums). + +## Why Pluggable Lineage? + +### 1. Runner Fragmentation + +Metrics-based lineage is **scattered across runners** with inconsistent support: +- **Dataflow**: Real-time metrics export to Cloud Monitoring +- **Flink**: Batch-only aggregation (no streaming support yet) +- **Spark/Direct**: Varying levels of support + +**Impact**: Multi-runner organizations must consolidate lineage from different metrics backends, each with different APIs and formats. + +**Plugin Solution**: Single implementation works consistently across all runners. + +### 2. Enterprise Integration + +Organizations with existing lineage infrastructure need: +- **Direct API integration** (Atlan, Collibra, Marquez, DataHub, OpenLineage) +- **Custom metadata enrichment** not in metrics subsystem + +**Example**: Flyte workflow executing a Beam pipeline needs to tag lineage with Flyte execution ID and cost allocation. This context exists in the orchestrator, not in Beam workers' metrics. + +### 3. Standard Formats + +OpenLineage is the industry standard. Plugin enables direct emission vs. export metrics → parse → transform → send. + +## Initialization + +`Lineage.setDefaultPipelineOptions(options)` is called from `FileSystems.setDefaultPipelineOptions()` (same pattern as `Metrics`). + +**Rationale**: `FileSystems.setDefaultPipelineOptions()` is called at 48+ locations covering all execution scenarios (pipeline construction, worker startup, deserialization). + +**Known Limitation**: Follows existing `FileSystems` pattern despite known issues ([#18430](https://github.com/apache/beam/issues/18430)). Architectural improvements would address all subsystems together. + +## Thread Safety + +Uses `AtomicReference` with `compareAndSet` loop (same pattern as `FileSystems`/`Metrics`): +- `AtomicReference>` tracks PipelineOptions identity +- `AtomicReference` for SOURCES/SINKS instances + +## Example: OpenLineage Plugin + +_For demonstration only (OpenLineage integration out of scope)_ + +```java +// 1. Plugin options +public interface OpenLineageOptions extends PipelineOptions { + @Description("OpenLineage endpoint URL") + String getOpenLineageUrl(); + void setOpenLineageUrl(String url); + + @Description("Enable OpenLineage plugin") + @Default.Boolean(false) + Boolean getEnableOpenLineage(); + void setEnableOpenLineage(Boolean enable); +} + +// 2. Implement LineageBase +class OpenLineageReporter implements LineageBase { + private final String endpoint; + private final Lineage.LineageDirection direction; + + @Override + public void add(Iterable rollupSegments) { + String fqn = String.join("", rollupSegments); + // POST to OpenLineage API with workflow context + sendToOpenLineage(endpoint, direction, fqn); + } +} + +// 3. Register via ServiceLoader +@AutoService(LineageRegistrar.class) +public class OpenLineageRegistrar implements LineageRegistrar { + @Override + public LineageBase fromOptions(PipelineOptions options, Lineage.LineageDirection direction) { + OpenLineageOptions opts = options.as(OpenLineageOptions.class); + if (opts.getEnableOpenLineage()) { + return new OpenLineageReporter(opts.getOpenLineageUrl(), direction); + } + return null; // Fall back to MetricsLineage + } +} + +// 4. Usage +PipelineOptions options = PipelineOptionsFactory.create(); +options.as(OpenLineageOptions.class).setEnableOpenLineage(true); +options.as(OpenLineageOptions.class).setOpenLineageUrl("https://lineage-api.example.com"); +Pipeline p = Pipeline.create(options); +``` \ No newline at end of file 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 index 05d34847a746..5e98cf0f2060 100644 --- 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 @@ -127,7 +127,7 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d } } - LOG.debug("Using default Metrics-based lineage for direction {}", direction); + LOG.info("Using default Metrics-based lineage for direction {}", direction); LineageBase defaultLineage = MetricsFlag.lineageRollupEnabled() ? new BoundedTrieMetricsLineage(direction) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/DefaultPipelineOptionsRegistrar.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/DefaultPipelineOptionsRegistrar.java index 7fc8a829482d..5689e7f0bf5b 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/DefaultPipelineOptionsRegistrar.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/DefaultPipelineOptionsRegistrar.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.options; import com.google.auto.service.AutoService; +import org.apache.beam.sdk.lineage.LineageOptions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** @@ -35,6 +36,7 @@ public Iterable> getPipelineOptions() { .add(ExperimentalOptions.class) .add(SdkHarnessOptions.class) .add(PortablePipelineOptions.class) + .add(LineageOptions.class) .build(); } } From 85990af8e2d938d72441345e076e1e0d75fb98e0 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 30 Mar 2026 14:15:04 -0400 Subject: [PATCH 19/28] remove unused files --- PR_description.md | 165 ---------------------------------------------- 1 file changed, 165 deletions(-) delete mode 100644 PR_description.md diff --git a/PR_description.md b/PR_description.md deleted file mode 100644 index 42ebbb52e728..000000000000 --- a/PR_description.md +++ /dev/null @@ -1,165 +0,0 @@ -Addresses #36790: "[Feature Request]: Make lineage tracking pluggable" - -## Changes - -- **Created** `org.apache.beam.sdk.lineage.LineageBase` - Plugin interface with single `add()` method -- **Refactored** `org.apache.beam.sdk.metrics.Lineage` - Hardcoded metrics → Delegation to `LineageBase` plugins -- **Created** `org.apache.beam.sdk.lineage.LineageRegistrar` - Plugin discovery interface (ServiceLoader) -- **Extracted** `org.apache.beam.sdk.metrics.MetricsLineage` - Default metrics-based implementation (implements `LineageBase`) -- **Added** Plugin initialization in `FileSystems.setDefaultPipelineOptions()` - -## Architecture - -**Before (master)**: `Lineage` was a concrete class hardcoded to use Beam metrics: - -```java -public class Lineage { - private static final Lineage SOURCES = new Lineage(Type.SOURCE); - private static final Lineage SINKS = new Lineage(Type.SINK); - private final Metric metric; // Hardcoded to Beam metrics - - private Lineage(Type type) { - this.metric = Metrics.stringSet(LINEAGE_NAMESPACE, type.toString()); - } - - public void add(Iterable segments) { - ((StringSet) metric).add(String.join("", segments)); // Always metrics - } -} -``` - -**After (this PR)**: Clean separation via composition pattern: - -```java -// Plugin contract (simple interface) -public interface LineageBase { - void add(Iterable rollupSegments); -} - -// Public API (final facade delegating to plugin) -public final class Lineage { - private final LineageBase delegate; // Plugin implementation - - private Lineage(LineageBase delegate) { - this.delegate = delegate; - } - - // Delegates to plugin - public void add(Iterable segments) { - delegate.add(segments); - } - - // Convenience overloads - public void add(String system, Iterable segments) { ... } - public void add(String system, String subtype, ...) { ... } - - // Static utilities (unchanged) - public static Lineage getSources() { ... } - public static Lineage getSinks() { ... } - public static String wrapSegment(String value) { ... } - public static Set query(MetricResults results, Type type) { ... } -} - -// Default implementation (backward compatible) -public class MetricsLineage implements LineageBase { - private final Metric metric; - - @Override - public void add(Iterable segments) { - ((BoundedTrie) metric).add(segments); - } -} -``` - -**Plugin Selection**: ServiceLoader discovery, first match wins, fallback to `MetricsLineage`. - -**Backward Compatibility**: ✅ All existing code works unchanged (24+ call sites, static utilities, enums). - -## Why Pluggable Lineage? - -### 1. Runner Fragmentation - -Metrics-based lineage is **scattered across runners** with inconsistent support: -- **Dataflow**: Real-time metrics export to Cloud Monitoring -- **Flink**: Batch-only aggregation (no streaming support yet) -- **Spark/Direct**: Varying levels of support - -**Impact**: Multi-runner organizations must consolidate lineage from different metrics backends, each with different APIs and formats. - -**Plugin Solution**: Single implementation works consistently across all runners. - -### 2. Enterprise Integration - -Organizations with existing lineage infrastructure need: -- **Direct API integration** (Atlan, Collibra, Marquez, DataHub, OpenLineage) -- **Custom metadata enrichment** not in metrics subsystem - -**Example**: Flyte workflow executing a Beam pipeline needs to tag lineage with Flyte execution ID and cost allocation. This context exists in the orchestrator, not in Beam workers' metrics. - -### 3. Standard Formats - -OpenLineage is the industry standard. Plugin enables direct emission vs. export metrics → parse → transform → send. - -## Initialization - -`Lineage.setDefaultPipelineOptions(options)` is called from `FileSystems.setDefaultPipelineOptions()` (same pattern as `Metrics`). - -**Rationale**: `FileSystems.setDefaultPipelineOptions()` is called at 48+ locations covering all execution scenarios (pipeline construction, worker startup, deserialization). - -**Known Limitation**: Follows existing `FileSystems` pattern despite known issues ([#18430](https://github.com/apache/beam/issues/18430)). Architectural improvements would address all subsystems together. - -## Thread Safety - -Uses `AtomicReference` with `compareAndSet` loop (same pattern as `FileSystems`/`Metrics`): -- `AtomicReference>` tracks PipelineOptions identity -- `AtomicReference` for SOURCES/SINKS instances - -## Example: OpenLineage Plugin - -_For demonstration only (OpenLineage integration out of scope)_ - -```java -// 1. Plugin options -public interface OpenLineageOptions extends PipelineOptions { - @Description("OpenLineage endpoint URL") - String getOpenLineageUrl(); - void setOpenLineageUrl(String url); - - @Description("Enable OpenLineage plugin") - @Default.Boolean(false) - Boolean getEnableOpenLineage(); - void setEnableOpenLineage(Boolean enable); -} - -// 2. Implement LineageBase -class OpenLineageReporter implements LineageBase { - private final String endpoint; - private final Lineage.LineageDirection direction; - - @Override - public void add(Iterable rollupSegments) { - String fqn = String.join("", rollupSegments); - // POST to OpenLineage API with workflow context - sendToOpenLineage(endpoint, direction, fqn); - } -} - -// 3. Register via ServiceLoader -@AutoService(LineageRegistrar.class) -public class OpenLineageRegistrar implements LineageRegistrar { - @Override - public LineageBase fromOptions(PipelineOptions options, Lineage.LineageDirection direction) { - OpenLineageOptions opts = options.as(OpenLineageOptions.class); - if (opts.getEnableOpenLineage()) { - return new OpenLineageReporter(opts.getOpenLineageUrl(), direction); - } - return null; // Fall back to MetricsLineage - } -} - -// 4. Usage -PipelineOptions options = PipelineOptionsFactory.create(); -options.as(OpenLineageOptions.class).setEnableOpenLineage(true); -options.as(OpenLineageOptions.class).setOpenLineageUrl("https://lineage-api.example.com"); -Pipeline p = Pipeline.create(options); -``` \ No newline at end of file From 77915f4d2ecc1e5f15816e70dce2f63ce98c24fd Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Mon, 30 Mar 2026 16:21:57 -0400 Subject: [PATCH 20/28] fix failing test with moquito --- CHANGES.md | 4 ++++ .../src/main/java/org/apache/beam/sdk/metrics/Lineage.java | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index cc1ec48ba188..602fe7e9337d 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -86,7 +86,11 @@ ## Bugfixes +<<<<<<< HEAD * Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). +======= +* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +>>>>>>> 7a91283a2b (fix failing test with moquito) ## Security Fixes 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 index 5e98cf0f2060..639ab1f31820 100644 --- 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 @@ -42,7 +42,7 @@ /** * Standard collection of metrics used to record source and sinks information for lineage tracking. */ -public final class Lineage { +public class Lineage { public static final String LINEAGE_NAMESPACE = "lineage"; private static final Logger LOG = LoggerFactory.getLogger(Lineage.class); private static final AtomicReference SOURCES = new AtomicReference<>(); From 6469f7f6d30ad131a476b48d8f05afc30df5e567 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 31 Mar 2026 14:58:44 -0400 Subject: [PATCH 21/28] Address review comments --- .../org/apache/beam/sdk/metrics/Lineage.java | 51 +++++-------------- .../beam/sdk/lineage/LineagePluginTest.java | 20 +++++--- 2 files changed, 25 insertions(+), 46 deletions(-) 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 index 639ab1f31820..a727658141b4 100644 --- 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 @@ -23,6 +23,7 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Objects; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; @@ -32,7 +33,6 @@ 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.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.checkerframework.checker.nullness.qual.Nullable; @@ -48,8 +48,8 @@ public class Lineage { private static final AtomicReference SOURCES = new AtomicReference<>(); private static final AtomicReference SINKS = new AtomicReference<>(); - private static final AtomicReference> LINEAGE_REVISION = - new AtomicReference<>(); + private static final AtomicReference<@Nullable Class> + CURRENT_LINEAGE_TYPE = new AtomicReference<>(); // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot. private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]"); @@ -68,41 +68,16 @@ private Lineage(LineageBase delegate) { @Internal public static void setDefaultPipelineOptions(PipelineOptions options) { checkNotNull(options, "options cannot be null"); - long optionsId = options.getOptionsId(); - int nextRevision = options.revision(); - - while (true) { - KV currentRevision = LINEAGE_REVISION.get(); - - if (currentRevision != null - && currentRevision.getKey().equals(optionsId) - && currentRevision.getValue() >= nextRevision) { - LOG.debug( - "Lineage already initialized with options ID {} revision {}, skipping", - optionsId, - currentRevision.getValue()); - return; - } - - if (LINEAGE_REVISION.compareAndSet(currentRevision, KV.of(optionsId, nextRevision))) { - Lineage sources = createLineage(options, LineageDirection.SOURCE); - Lineage sinks = createLineage(options, LineageDirection.SINK); + Class requestedType = options.as(LineageOptions.class).getLineageType(); - 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); - } - return; - } + Class currentType = CURRENT_LINEAGE_TYPE.get(); + if (Objects.equals(currentType, requestedType) && SOURCES.get() != null) { + return; + } + if (CURRENT_LINEAGE_TYPE.compareAndSet(currentType, requestedType)) { + SOURCES.set(createLineage(options, LineageDirection.SOURCE)); + SINKS.set(createLineage(options, LineageDirection.SINK)); + LOG.debug("Lineage initialized with type {}", requestedType); } } @@ -127,7 +102,7 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d } } - LOG.info("Using default Metrics-based lineage for direction {}", direction); + LOG.debug("Using default Metrics-based lineage for direction {}", direction); LineageBase defaultLineage = MetricsFlag.lineageRollupEnabled() ? new BoundedTrieMetricsLineage(direction) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java index 75c329ae6842..ff50dc1ffbe2 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/lineage/LineagePluginTest.java @@ -43,11 +43,15 @@ import org.junit.runner.Description; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Tests for {@link LineageBase} pipeline option selection and DirectRunner integration. */ @RunWith(JUnit4.class) public class LineagePluginTest { + private static final Logger LOG = LoggerFactory.getLogger(LineagePluginTest.class); + /** * TestWatcher that logs detailed lineage diagnostics only when tests fail. This keeps successful * test output clean while providing deep debugging for failures. @@ -57,24 +61,24 @@ public class LineagePluginTest { new TestWatcher() { @Override protected void failed(Throwable e, Description description) { - System.err.println("=== Lineage Test Failure Diagnostics ==="); - System.err.println("Test: " + description.getMethodName()); - System.err.println("Error: " + e.getMessage()); + LOG.error("=== Lineage Test Failure Diagnostics ==="); + LOG.error("Test: {}", description.getMethodName()); + LOG.error("Error:", e); List sources = TestLineage.getRecordedSources(); List sinks = TestLineage.getRecordedSinks(); - System.err.println("\nRecorded Sources (" + sources.size() + "):"); + LOG.error("Recorded Sources ({}):", sources.size()); for (int i = 0; i < sources.size(); i++) { - System.err.println(" [" + i + "] \"" + sources.get(i) + "\""); + LOG.error(" [{}] \"{}\"", i, sources.get(i)); } - System.err.println("\nRecorded Sinks (" + sinks.size() + "):"); + LOG.error("Recorded Sinks ({}):", sinks.size()); for (int i = 0; i < sinks.size(); i++) { - System.err.println(" [" + i + "] \"" + sinks.get(i) + "\""); + LOG.error(" [{}] \"{}\"", i, sinks.get(i)); } - System.err.println("========================================"); + LOG.error("========================================"); } }; From dfae036998efc69736e8b7185e5b9c65d9dced04 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Thu, 2 Apr 2026 17:46:55 -0400 Subject: [PATCH 22/28] Simplify Lineage thread sync during init --- .../java/org/apache/beam/sdk/metrics/Lineage.java | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) 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 index a727658141b4..859cd9cd330e 100644 --- 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 @@ -23,7 +23,6 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; -import java.util.Objects; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; @@ -48,8 +47,7 @@ public class Lineage { private static final AtomicReference SOURCES = new AtomicReference<>(); private static final AtomicReference SINKS = new AtomicReference<>(); - private static final AtomicReference<@Nullable Class> - CURRENT_LINEAGE_TYPE = new AtomicReference<>(); + private static final Object INIT_LOCK = new Object(); // Reserved characters are backtick, colon, whitespace (space, \t, \n) and dot. private static final Pattern RESERVED_CHARS = Pattern.compile("[:\\s.`]"); @@ -68,16 +66,11 @@ private Lineage(LineageBase delegate) { @Internal public static void setDefaultPipelineOptions(PipelineOptions options) { checkNotNull(options, "options cannot be null"); - Class requestedType = options.as(LineageOptions.class).getLineageType(); - - Class currentType = CURRENT_LINEAGE_TYPE.get(); - if (Objects.equals(currentType, requestedType) && SOURCES.get() != null) { - return; - } - if (CURRENT_LINEAGE_TYPE.compareAndSet(currentType, requestedType)) { + synchronized (INIT_LOCK) { SOURCES.set(createLineage(options, LineageDirection.SOURCE)); SINKS.set(createLineage(options, LineageDirection.SINK)); - LOG.debug("Lineage initialized with type {}", requestedType); + LOG.debug( + "Lineage initialized with type {}", options.as(LineageOptions.class).getLineageType()); } } From 12b41afc175cc8027b1b16d78198a921bf2157db Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Thu, 2 Apr 2026 17:53:09 -0400 Subject: [PATCH 23/28] target 2.73.0 --- CHANGES.md | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 602fe7e9337d..bdd5ea3d1f39 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -74,7 +74,6 @@ compatible. Both coders can decode encoded bytes from the other coder ([#38139](https://github.com/apache/beam/issues/38139)). * (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). -* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). ## Breaking Changes @@ -86,11 +85,7 @@ ## Bugfixes -<<<<<<< HEAD * Fixed BigQueryEnrichmentHandler batch mode dropping earlier requests when multiple requests share the same enrichment key (Python) ([#38035](https://github.com/apache/beam/issues/38035)). -======= -* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ->>>>>>> 7a91283a2b (fix failing test with moquito) ## Security Fixes @@ -111,6 +106,7 @@ ## New Features / Improvements +* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). * Added `ADKAgentModelHandler` for running Google Agent Development Kit (ADK) agents (Python) ([#37917](https://github.com/apache/beam/issues/37917)). * (Python) Added exception chaining to preserve error context in CloudSQLEnrichmentHandler, processes utilities, and core transforms ([#37422](https://github.com/apache/beam/issues/37422)). * (Python) Added a pipeline option `--experiments=pip_no_build_isolation` to disable build isolation when installing dependencies in the runtime environment ([#37331](https://github.com/apache/beam/issues/37331)). From bf61f18b82e04705c8994592843970f4f2c2d69e Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 21 Apr 2026 18:24:36 -0400 Subject: [PATCH 24/28] improve setDefaultPipelineOptions concurrency --- .../org/apache/beam/sdk/metrics/Lineage.java | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) 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 index 859cd9cd330e..2c851f622083 100644 --- 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 @@ -24,14 +24,12 @@ import java.util.Iterator; import java.util.List; import java.util.Set; -import java.util.concurrent.atomic.AtomicReference; import java.util.regex.Pattern; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.lineage.LineageBase; import org.apache.beam.sdk.lineage.LineageOptions; 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.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.checkerframework.checker.nullness.qual.Nullable; @@ -44,8 +42,10 @@ public class Lineage { public static final String LINEAGE_NAMESPACE = "lineage"; private static final Logger LOG = LoggerFactory.getLogger(Lineage.class); - private static final AtomicReference SOURCES = new AtomicReference<>(); - private static final AtomicReference SINKS = new AtomicReference<>(); + + private static volatile @Nullable Lineage SOURCES; + private static volatile @Nullable Lineage SINKS; + private static volatile @Nullable Class CURRENT_LINEAGE_TYPE; private static final Object INIT_LOCK = new Object(); @@ -66,14 +66,31 @@ private Lineage(LineageBase delegate) { @Internal public static void setDefaultPipelineOptions(PipelineOptions options) { checkNotNull(options, "options cannot be null"); + Class requestedType = options.as(LineageOptions.class).getLineageType(); + + if (canSkipInit(requestedType)) { + return; + } synchronized (INIT_LOCK) { - SOURCES.set(createLineage(options, LineageDirection.SOURCE)); - SINKS.set(createLineage(options, LineageDirection.SINK)); - LOG.debug( - "Lineage initialized with type {}", options.as(LineageOptions.class).getLineageType()); + if (canSkipInit(requestedType)) { + return; + } + SOURCES = createLineage(options, LineageDirection.SOURCE); + SINKS = createLineage(options, LineageDirection.SINK); + CURRENT_LINEAGE_TYPE = requestedType; + LOG.debug("Lineage initialized with type {}", requestedType); } } + private static boolean canSkipInit(@Nullable Class requestedType) { + if (SOURCES == null) { + return false; + } + // When no type is requested, preserve whatever is already initialized. + // When a type is requested, only re-init if it differs from the active type. + return requestedType == null || requestedType.equals(CURRENT_LINEAGE_TYPE); + } + private static Lineage createLineage(PipelineOptions options, LineageDirection direction) { Class lineageClass = options.as(LineageOptions.class).getLineageType(); @@ -105,22 +122,16 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d /** {@link Lineage} representing sources and optionally side inputs. */ public static Lineage getSources() { - Lineage sources = SOURCES.get(); - if (sources == null) { - setDefaultPipelineOptions(PipelineOptionsFactory.create()); - sources = SOURCES.get(); - } - return sources; + return checkNotNull( + SOURCES, + "Lineage not initialized. FileSystems.setDefaultPipelineOptions must be called first."); } /** {@link Lineage} representing sinks. */ public static Lineage getSinks() { - Lineage sinks = SINKS.get(); - if (sinks == null) { - setDefaultPipelineOptions(PipelineOptionsFactory.create()); - sinks = SINKS.get(); - } - return sinks; + return checkNotNull( + SINKS, + "Lineage not initialized. FileSystems.setDefaultPipelineOptions must be called first."); } @VisibleForTesting From 050b66a7a6d090f0f0ec542f186e0bd4faf857b8 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Tue, 21 Apr 2026 19:31:36 -0400 Subject: [PATCH 25/28] fix spotless --- .../org/apache/beam/sdk/metrics/Lineage.java | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) 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 index 2c851f622083..55d268cdddd6 100644 --- 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 @@ -43,9 +43,9 @@ public class Lineage { public static final String LINEAGE_NAMESPACE = "lineage"; private static final Logger LOG = LoggerFactory.getLogger(Lineage.class); - private static volatile @Nullable Lineage SOURCES; - private static volatile @Nullable Lineage SINKS; - private static volatile @Nullable Class CURRENT_LINEAGE_TYPE; + private static volatile @Nullable Lineage sources; + private static volatile @Nullable Lineage sinks; + private static volatile @Nullable Class currentLineageType; private static final Object INIT_LOCK = new Object(); @@ -75,20 +75,20 @@ public static void setDefaultPipelineOptions(PipelineOptions options) { if (canSkipInit(requestedType)) { return; } - SOURCES = createLineage(options, LineageDirection.SOURCE); - SINKS = createLineage(options, LineageDirection.SINK); - CURRENT_LINEAGE_TYPE = requestedType; + sources = createLineage(options, LineageDirection.SOURCE); + sinks = createLineage(options, LineageDirection.SINK); + currentLineageType = requestedType; LOG.debug("Lineage initialized with type {}", requestedType); } } private static boolean canSkipInit(@Nullable Class requestedType) { - if (SOURCES == null) { + if (sources == null) { return false; } // When no type is requested, preserve whatever is already initialized. // When a type is requested, only re-init if it differs from the active type. - return requestedType == null || requestedType.equals(CURRENT_LINEAGE_TYPE); + return requestedType == null || requestedType.equals(currentLineageType); } private static Lineage createLineage(PipelineOptions options, LineageDirection direction) { @@ -123,14 +123,14 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d /** {@link Lineage} representing sources and optionally side inputs. */ public static Lineage getSources() { return checkNotNull( - SOURCES, + sources, "Lineage not initialized. FileSystems.setDefaultPipelineOptions must be called first."); } /** {@link Lineage} representing sinks. */ public static Lineage getSinks() { return checkNotNull( - SINKS, + sinks, "Lineage not initialized. FileSystems.setDefaultPipelineOptions must be called first."); } From 18b5d291756bda62eaa2e2ec23d4bd30a2d5423a Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Wed, 22 Apr 2026 15:54:16 -0400 Subject: [PATCH 26/28] final fixes --- CHANGES.md | 2 +- .../apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java | 4 +++- .../src/main/java/org/apache/beam/sdk/metrics/Lineage.java | 4 ++-- .../org/apache/beam/sdk/metrics/StringSetMetricsLineage.java | 4 +++- 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index bdd5ea3d1f39..cc1ec48ba188 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -74,6 +74,7 @@ compatible. Both coders can decode encoded bytes from the other coder ([#38139](https://github.com/apache/beam/issues/38139)). * (Python) Added type alias for with_exception_handling to be used for typehints. ([#38173](https://github.com/apache/beam/issues/38173)). +* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). ## Breaking Changes @@ -106,7 +107,6 @@ ## New Features / Improvements -* Added plugin mechanism to support different Lineage implementations (Java) ([#36790](https://github.com/apache/beam/issues/36790)). * Added `ADKAgentModelHandler` for running Google Agent Development Kit (ADK) agents (Python) ([#37917](https://github.com/apache/beam/issues/37917)). * (Python) Added exception chaining to preserve error context in CloudSQLEnrichmentHandler, processes utilities, and core transforms ([#37422](https://github.com/apache/beam/issues/37422)). * (Python) Added a pipeline option `--experiments=pip_no_build_isolation` to disable build isolation when installing dependencies in the runtime environment ([#37331](https://github.com/apache/beam/issues/37331)). diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java index 78dff258532a..a2f321ba3df3 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.metrics; import org.apache.beam.sdk.lineage.LineageBase; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** @@ -29,7 +30,8 @@ class BoundedTrieMetricsLineage implements LineageBase { private final BoundedTrie metric; - BoundedTrieMetricsLineage(Lineage.LineageDirection direction) { + @SuppressWarnings("unused") + public BoundedTrieMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCEV2 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 index 55d268cdddd6..1a193ec006e0 100644 --- 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 @@ -115,8 +115,8 @@ private static Lineage createLineage(PipelineOptions options, LineageDirection d LOG.debug("Using default Metrics-based lineage for direction {}", direction); LineageBase defaultLineage = MetricsFlag.lineageRollupEnabled() - ? new BoundedTrieMetricsLineage(direction) - : new StringSetMetricsLineage(direction); + ? new BoundedTrieMetricsLineage(options, direction) + : new StringSetMetricsLineage(options, direction); return new Lineage(defaultLineage); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java index b12a49182a8c..a12bffb1bfec 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.metrics; import org.apache.beam.sdk.lineage.LineageBase; +import org.apache.beam.sdk.options.PipelineOptions; /** * Lineage implementation that stores lineage information in {@link StringSet} metrics. @@ -28,7 +29,8 @@ class StringSetMetricsLineage implements LineageBase { private final StringSet metric; - StringSetMetricsLineage(Lineage.LineageDirection direction) { + @SuppressWarnings("unused") + public StringSetMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK; this.metric = Metrics.stringSet(Lineage.LINEAGE_NAMESPACE, type.toString()); From d2c58ccb626cfa67d26386505ae02b00cc101822 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Thu, 23 Apr 2026 10:50:28 -0400 Subject: [PATCH 27/28] address review comments --- .../org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java | 1 - .../org/apache/beam/sdk/metrics/StringSetMetricsLineage.java | 1 - 2 files changed, 2 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java index a2f321ba3df3..32e8ee9f1e59 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java @@ -30,7 +30,6 @@ class BoundedTrieMetricsLineage implements LineageBase { private final BoundedTrie metric; - @SuppressWarnings("unused") public BoundedTrieMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java index a12bffb1bfec..1679ca715bfa 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java @@ -29,7 +29,6 @@ class StringSetMetricsLineage implements LineageBase { private final StringSet metric; - @SuppressWarnings("unused") public StringSetMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK; From 38d40a3669a7d581f83d61b44a98eb99adab2c77 Mon Sep 17 00:00:00 2001 From: Andrew Kabas Date: Thu, 23 Apr 2026 12:03:55 -0400 Subject: [PATCH 28/28] put suppressions back --- .../org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java | 1 + .../org/apache/beam/sdk/metrics/StringSetMetricsLineage.java | 1 + 2 files changed, 2 insertions(+) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java index 32e8ee9f1e59..a2f321ba3df3 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/BoundedTrieMetricsLineage.java @@ -30,6 +30,7 @@ class BoundedTrieMetricsLineage implements LineageBase { private final BoundedTrie metric; + @SuppressWarnings("unused") public BoundedTrieMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java index 1679ca715bfa..a12bffb1bfec 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/StringSetMetricsLineage.java @@ -29,6 +29,7 @@ class StringSetMetricsLineage implements LineageBase { private final StringSet metric; + @SuppressWarnings("unused") public StringSetMetricsLineage(PipelineOptions options, Lineage.LineageDirection direction) { Lineage.Type type = (direction == Lineage.LineageDirection.SOURCE) ? Lineage.Type.SOURCE : Lineage.Type.SINK;