Skip to content

Commit feae18e

Browse files
rahul-madaanrahul-madaan
committed
Harden lifecycle and worker-safety guarantees
Fixes issues found in an end-to-end review of the extension: - Context construction can no longer throw: malformed run ids fall back to the deterministic id, malformed parent run ids omit the facet, and invalid transport or circuit-breaker configuration degrades to a non-emitting client with an error logged. The LineageBase plugin is instantiated reflectively during worker initialization, so a configuration mistake previously crash-looped the job. - The plugin no longer freezes the run id before OpenLineageRunner has minted it: the context is re-resolved at submission and the plugin looks the context up per call, so driver and worker events always share one run id. - The terminal event is now emitted deterministically through the returned PipelineResult wrapper on waitUntilFinish()/cancel(), instead of relying on the periodic tracker poll racing JVM exit. UPDATED now maps to COMPLETE rather than FAIL. - A finished context is replaced on the next execution so reused worker JVMs (e.g. Flink session clusters) track the new run. - PubsubIO extraction guards each ValueProvider access independently, so a runtime-only topic no longer discards an accessible subscription in templated pipelines. - Iceberg extraction falls back to the catalog identity when a table location cannot be parsed as a URI instead of dropping the dataset. - The metrics tee now populates both the StringSet and BoundedTrie forms so Lineage.query keeps working under enable_lineage_rollup. - Azure datasets use the wasbs blob endpoint and HDFS paths keep their leading slash, matching openlineage-java's canonical naming. - The tracker is not started when lineage is disabled, non-positive tracking intervals fall back to the default, the client is closed after the terminal event, and the dataset dedup key can no longer collide across namespace/name boundaries. Addresses #39427.
1 parent 4ed44f9 commit feae18e

14 files changed

Lines changed: 396 additions & 105 deletions

File tree

‎sdks/java/extensions/openlineage/build.gradle‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,15 +33,16 @@ dependencies {
3333
implementation "io.openlineage:openlineage-java:$openlineage_version"
3434
implementation library.java.jackson_annotations
3535
implementation library.java.jackson_core
36+
implementation library.java.joda_time
3637
implementation library.java.slf4j_api
3738
implementation library.java.vendored_guava_32_1_2_jre
3839
// IO modules are compile-only: transform-specific lineage visitors activate at runtime only
3940
// when the corresponding IO is already on the user's classpath.
4041
compileOnly project(path: ":sdks:java:io:google-cloud-platform")
4142
compileOnly project(path: ":sdks:java:io:iceberg")
43+
// keep in sync with iceberg_version in sdks/java/io/iceberg/build.gradle
4244
compileOnly "org.apache.iceberg:iceberg-api:1.10.0"
4345
testImplementation library.java.junit
44-
testImplementation library.java.hamcrest
4546
testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
4647
testImplementation project(path: ":sdks:java:io:google-cloud-platform")
4748
testImplementation project(path: ":sdks:java:io:iceberg")

‎sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqns.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,10 +73,14 @@ static DatasetIdentifier toDatasetIdentifier(String fqn) {
7373
String account = first(seg);
7474
String container = seg.size() > 1 ? seg.get(1) : "unknown";
7575
return new DatasetIdentifier(
76-
objectKey(seg, 2), "wasbs://" + container + "@" + account + ".dfs.core.windows.net");
76+
objectKey(seg, 2), "wasbs://" + container + "@" + account + ".blob.core.windows.net");
7777
}
7878
case "hdfs":
79-
return new DatasetIdentifier(objectKey(seg, 1), "hdfs://" + first(seg));
79+
// The path keeps its leading slash, matching openlineage-java's FilesystemDatasetUtils
80+
// so the same physical dataset gets one identity on every code path.
81+
return new DatasetIdentifier(
82+
seg.size() > 1 ? String.join("/", seg.subList(1, seg.size())) : "/",
83+
"hdfs://" + first(seg));
8084
case "spanner":
8185
{
8286
// Beam segments: project, instanceConfig, instance, database, table.

‎sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/EventEmitter.java‎

Lines changed: 52 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -31,47 +31,81 @@
3131

3232
/**
3333
* Sends OpenLineage events through the configured transport, mirroring the {@code EventEmitter}
34-
* classes of the Spark and Flink integrations: emission is wrapped in the configured circuit
35-
* breaker and never propagates an exception to the pipeline.
34+
* classes of the Spark and Flink integrations: construction and emission never propagate an
35+
* exception to the pipeline — an invalid transport configuration degrades to a non-emitting client
36+
* with an error logged, and emission is wrapped in the configured circuit breaker.
3637
*/
3738
class EventEmitter {
3839

3940
private static final Logger LOG = LoggerFactory.getLogger(EventEmitter.class);
4041

41-
private final OpenLineageClient client;
42+
private final @Nullable OpenLineageClient client;
4243
private final @Nullable CircuitBreaker circuitBreaker;
4344

4445
EventEmitter(BeamOpenLineageConfig config) {
45-
TransportConfig transportConfig =
46-
config.getTransportConfig() == null ? new ConsoleConfig() : config.getTransportConfig();
47-
FacetsConfig facetsConfig =
48-
config.getFacetsConfig() == null ? new FacetsConfig() : config.getFacetsConfig();
49-
this.client =
50-
OpenLineageClient.builder()
51-
.transport(new TransportFactory(transportConfig).build())
52-
.disableFacets(facetsConfig.getEffectiveDisabledFacets())
53-
.build();
54-
this.circuitBreaker =
55-
config.getCircuitBreaker() == null
56-
? null
57-
: new CircuitBreakerFactory(config.getCircuitBreaker()).build();
46+
this.client = buildClient(config);
47+
this.circuitBreaker = buildCircuitBreaker(config);
48+
}
49+
50+
private static @Nullable OpenLineageClient buildClient(BeamOpenLineageConfig config) {
51+
try {
52+
TransportConfig transportConfig =
53+
config.getTransportConfig() == null ? new ConsoleConfig() : config.getTransportConfig();
54+
FacetsConfig facetsConfig =
55+
config.getFacetsConfig() == null ? new FacetsConfig() : config.getFacetsConfig();
56+
return OpenLineageClient.builder()
57+
.transport(new TransportFactory(transportConfig).build())
58+
.disableFacets(facetsConfig.getEffectiveDisabledFacets())
59+
.build();
60+
} catch (RuntimeException e) {
61+
LOG.error(
62+
"Invalid OpenLineage transport configuration; lineage events will NOT be emitted", e);
63+
return null;
64+
}
65+
}
66+
67+
private static @Nullable CircuitBreaker buildCircuitBreaker(BeamOpenLineageConfig config) {
68+
try {
69+
return config.getCircuitBreaker() == null
70+
? null
71+
: new CircuitBreakerFactory(config.getCircuitBreaker()).build();
72+
} catch (RuntimeException e) {
73+
LOG.warn("Invalid OpenLineage circuit breaker configuration; continuing without one", e);
74+
return null;
75+
}
5876
}
5977

6078
/** Emits the event; failures are logged and swallowed so lineage never breaks the job. */
6179
void emit(OpenLineage.RunEvent event) {
80+
final OpenLineageClient localClient = client;
81+
if (localClient == null) {
82+
return;
83+
}
6284
try {
6385
if (circuitBreaker != null) {
6486
circuitBreaker.run(
6587
() -> {
66-
client.emit(event);
88+
localClient.emit(event);
6789
return Boolean.TRUE;
6890
});
6991
} else {
70-
client.emit(event);
92+
localClient.emit(event);
7193
}
7294
LOG.debug("Emitted OpenLineage event: {} run {}", event.getEventType(), event.getRun());
7395
} catch (RuntimeException e) {
7496
LOG.warn("Failed to emit OpenLineage event, swallowing exception", e);
7597
}
7698
}
99+
100+
/** Closes the underlying transport (flushing buffered events); never throws. */
101+
void close() {
102+
if (client == null) {
103+
return;
104+
}
105+
try {
106+
client.close();
107+
} catch (Exception e) {
108+
LOG.debug("Failed to close OpenLineage client", e);
109+
}
110+
}
77111
}

‎sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/IcebergLineageVisitor.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -101,10 +101,16 @@ private List<DatasetIdentifier> extract(PTransform<?, ?> transform) {
101101
String symlinkNamespace =
102102
uri != null ? uri : (warehouse != null ? warehouse : "iceberg://" + catalogName);
103103
String location = loadTableLocation(catalogConfig, tableId);
104-
DatasetIdentifier identifier;
104+
DatasetIdentifier identifier = null;
105105
if (location != null) {
106-
identifier = FilesystemDatasetUtils.fromLocation(URI.create(location));
107-
} else {
106+
try {
107+
identifier = FilesystemDatasetUtils.fromLocation(URI.create(location));
108+
} catch (RuntimeException e) {
109+
// Fall through to the catalog-based identity instead of dropping the dataset.
110+
LOG.debug("Unparseable Iceberg table location '{}'", location, e);
111+
}
112+
}
113+
if (identifier == null) {
108114
identifier =
109115
new DatasetIdentifier(tableId.toString(), "iceberg://" + authority(uri, catalogName));
110116
}

‎sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageContext.java‎

Lines changed: 86 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,12 @@
2727
import java.util.ArrayList;
2828
import java.util.List;
2929
import java.util.Map;
30+
import java.util.Set;
3031
import java.util.UUID;
3132
import java.util.concurrent.ConcurrentHashMap;
3233
import java.util.concurrent.atomic.AtomicBoolean;
34+
import org.apache.beam.sdk.PipelineResult;
35+
import org.apache.beam.sdk.metrics.Lineage;
3336
import org.apache.beam.sdk.options.ApplicationNameOptions;
3437
import org.apache.beam.sdk.options.PipelineOptions;
3538
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
@@ -40,8 +43,13 @@
4043
/**
4144
* Holds the OpenLineage state of one pipeline execution in this JVM — run identity, discovered
4245
* datasets, and event lifecycle — mirroring the {@code OpenLineageContext} classes of the Spark and
43-
* Flink integrations. One instance exists per JVM (driver or worker); every JVM of the same
44-
* pipeline execution resolves the identical run id so events merge into one run downstream.
46+
* Flink integrations. One pipeline execution is tracked per JVM (driver or worker); every JVM of
47+
* the same execution resolves the identical run id so events merge into one run downstream. When a
48+
* JVM is reused for a new execution after the previous one finished (e.g. a Flink session cluster),
49+
* a fresh context replaces the finished one.
50+
*
51+
* <p>Construction never throws: invalid run ids fall back to the deterministic id and transport
52+
* failures degrade to a non-emitting client, so lineage can never fail the pipeline.
4553
*/
4654
class OpenLineageContext {
4755

@@ -51,12 +59,33 @@ class OpenLineageContext {
5159
private static @Nullable OpenLineageContext instance;
5260
private static @Nullable BeamOpenLineageConfig configOverride;
5361

54-
/** Returns the per-JVM context, creating it from the given options on first use. */
62+
/**
63+
* Returns the per-JVM context, creating it from the given options on first use. If the current
64+
* context already finished its run, a fresh context is created so a reused worker JVM tracks the
65+
* new execution instead of staying silent.
66+
*/
5567
static synchronized OpenLineageContext getOrCreate(PipelineOptions options) {
56-
if (instance == null) {
57-
instance = new OpenLineageContext(options);
68+
OpenLineageContext local = instance;
69+
if (local == null || local.finished.get()) {
70+
local = new OpenLineageContext(options);
71+
instance = local;
5872
}
59-
return instance;
73+
return local;
74+
}
75+
76+
/**
77+
* Re-resolves the context at submission time, after {@link OpenLineageRunner} has minted the run
78+
* id. The {@code LineageBase} plugin may already have been constructed while the runner was being
79+
* resolved (before the id existed); as long as nothing has been emitted yet, the context is
80+
* rebuilt so driver and workers agree on one run id.
81+
*/
82+
static synchronized OpenLineageContext refreshForRunner(PipelineOptions options) {
83+
OpenLineageContext local = instance;
84+
if (local == null || local.finished.get() || !local.started.get()) {
85+
local = new OpenLineageContext(options);
86+
instance = local;
87+
}
88+
return local;
6089
}
6190

6291
/** Test seam, mirroring the Spark listener's overrideDefaultFactoryForTests. */
@@ -122,19 +151,26 @@ private static String resolveJobName(BeamOpenLineageConfig config, PipelineOptio
122151
/**
123152
* Resolves the run id. An explicit id (minted by {@link OpenLineageRunner} at submission and
124153
* propagated through serialized options, like the Spark integration's driver-minted application
125-
* run id) wins. Otherwise a deterministic UUID is derived from the job identity with the same
126-
* helper the Airflow provider uses, so independent worker JVMs agree.
154+
* run id) wins. Otherwise — or when the explicit id is malformed — a deterministic UUID is
155+
* derived from the job identity with the same helper the Airflow provider uses, so independent
156+
* worker JVMs still agree. Never throws.
127157
*/
128158
private static UUID resolveRunId(OpenLineagePipelineOptions olOptions, PipelineOptions options) {
129159
String explicit = olOptions.getOpenLineageRunId();
130160
if (explicit != null) {
131-
return UUID.fromString(explicit);
161+
try {
162+
return UUID.fromString(explicit);
163+
} catch (IllegalArgumentException e) {
164+
LOG.warn(
165+
"Malformed openLineageRunId '{}'; falling back to a deterministic run id", explicit);
166+
}
132167
}
133168
byte[] identity =
134169
(options.getJobName() + "/" + options.getOptionsId()).getBytes(StandardCharsets.UTF_8);
135170
return UUIDUtils.generateStaticUUID(Instant.EPOCH, identity);
136171
}
137172

173+
/** Builds the parent run facet when fully configured; never throws on malformed input. */
138174
private static OpenLineage.@Nullable ParentRunFacet buildParentRunFacet(
139175
OpenLineage openLineage, OpenLineagePipelineOptions options) {
140176
String parentRunId = options.getOpenLineageParentRunId();
@@ -143,15 +179,21 @@ private static UUID resolveRunId(OpenLineagePipelineOptions olOptions, PipelineO
143179
if (parentRunId == null || parentJobName == null || parentJobNamespace == null) {
144180
return null;
145181
}
146-
return openLineage
147-
.newParentRunFacetBuilder()
148-
.run(openLineage.newParentRunFacetRun(UUID.fromString(parentRunId)))
149-
.job(openLineage.newParentRunFacetJob(parentJobNamespace, parentJobName))
150-
.root(
151-
openLineage.newParentRunFacetRoot(
152-
openLineage.newRootRun(UUID.fromString(parentRunId)),
153-
openLineage.newRootJob(parentJobNamespace, parentJobName)))
154-
.build();
182+
try {
183+
UUID parentUuid = UUID.fromString(parentRunId);
184+
return openLineage
185+
.newParentRunFacetBuilder()
186+
.run(openLineage.newParentRunFacetRun(parentUuid))
187+
.job(openLineage.newParentRunFacetJob(parentJobNamespace, parentJobName))
188+
.root(
189+
openLineage.newParentRunFacetRoot(
190+
openLineage.newRootRun(parentUuid),
191+
openLineage.newRootJob(parentJobNamespace, parentJobName)))
192+
.build();
193+
} catch (IllegalArgumentException e) {
194+
LOG.warn("Malformed openLineageParentRunId '{}'; omitting the parent run facet", parentRunId);
195+
return null;
196+
}
155197
}
156198

157199
BeamOpenLineageConfig getConfig() {
@@ -186,7 +228,8 @@ boolean registerDataset(LineageDirection direction, DatasetIdentifier dataset) {
186228
return false;
187229
}
188230
Map<String, DatasetIdentifier> target = direction == LineageDirection.INPUT ? inputs : outputs;
189-
boolean added = target.putIfAbsent(dataset.getNamespace() + dataset.getName(), dataset) == null;
231+
String key = dataset.getNamespace() + "\n" + dataset.getName();
232+
boolean added = target.putIfAbsent(key, dataset) == null;
190233
if (added) {
191234
LOG.info(
192235
"OpenLineage discovered {} dataset {} {}",
@@ -216,6 +259,25 @@ void onTrackingTick() {
216259
emit(OpenLineage.RunEvent.EventType.RUNNING, null);
217260
}
218261

262+
/** Pulls IO-reported lineage out of the metrics-based lineage store, if the runner has it. */
263+
void sweepLineageMetrics(PipelineResult result) {
264+
if (disabled) {
265+
return;
266+
}
267+
try {
268+
Set<String> sources = Lineage.query(result.metrics(), Lineage.Type.SOURCE);
269+
for (String fqn : sources) {
270+
registerDataset(LineageDirection.INPUT, DataplexFqns.toDatasetIdentifier(fqn));
271+
}
272+
Set<String> sinks = Lineage.query(result.metrics(), Lineage.Type.SINK);
273+
for (String fqn : sinks) {
274+
registerDataset(LineageDirection.OUTPUT, DataplexFqns.toDatasetIdentifier(fqn));
275+
}
276+
} catch (RuntimeException e) {
277+
LOG.debug("Lineage metrics sweep failed", e);
278+
}
279+
}
280+
219281
/** Emits the terminal event exactly once. */
220282
void onJobFinished(OpenLineage.RunEvent.EventType eventType, @Nullable Throwable error) {
221283
if (disabled) {
@@ -224,10 +286,15 @@ void onJobFinished(OpenLineage.RunEvent.EventType eventType, @Nullable Throwable
224286
if (finished.compareAndSet(false, true)) {
225287
started.set(true);
226288
emit(eventType, error);
289+
emitter.close();
227290
}
228291
}
229292

230293
private void emit(OpenLineage.RunEvent.EventType eventType, @Nullable Throwable error) {
294+
// Best-effort guard against a RUNNING event racing past the terminal event.
295+
if (finished.get() && eventType == OpenLineage.RunEvent.EventType.RUNNING) {
296+
return;
297+
}
231298
try {
232299
emitter.emit(buildEvent(eventType, error));
233300
} catch (RuntimeException e) {

0 commit comments

Comments
 (0)