diff --git a/CHANGES.md b/CHANGES.md index 1e09ade6c389..66a218d043e5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,6 +69,7 @@ ## New Features / Improvements * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). +* Added an OpenLineage extension (`sdks/java/extensions/openlineage`) that emits OpenLineage events describing the datasets a pipeline reads and writes, via a wrapper runner (`--runner=OpenLineageRunner`) and a pluggable lineage backend (`--lineageType`) (Java) ([#39427](https://github.com/apache/beam/issues/39427)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/build.gradle.kts b/build.gradle.kts index ba2f7aaed67e..69dc97ad0432 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -299,6 +299,7 @@ tasks.register("javaPreCommit") { dependsOn(":sdks:java:extensions:join-library:build") dependsOn(":sdks:java:extensions:kryo:build") dependsOn(":sdks:java:extensions:ml:build") + dependsOn(":sdks:java:extensions:openlineage:build") dependsOn(":sdks:java:extensions:protobuf:build") dependsOn(":sdks:java:extensions:python:build") dependsOn(":sdks:java:extensions:sbe:build") diff --git a/sdks/java/extensions/openlineage/build.gradle b/sdks/java/extensions/openlineage/build.gradle new file mode 100644 index 000000000000..8c11e58bad23 --- /dev/null +++ b/sdks/java/extensions/openlineage/build.gradle @@ -0,0 +1,54 @@ +/* + * 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. + */ + +plugins { id 'org.apache.beam.module' } +applyJavaNature( + automaticModuleName: 'org.apache.beam.sdk.extensions.openlineage', + // matches sdks/java/io/iceberg, which this module inspects at test time + requireJavaVersion: JavaVersion.VERSION_11, +) + +description = "Apache Beam :: SDKs :: Java :: Extensions :: OpenLineage" +ext.summary = "Emit OpenLineage events for Beam pipelines." + +def openlineage_version = "1.51.0" + +dependencies { + implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation "io.openlineage:openlineage-java:$openlineage_version" + implementation library.java.jackson_annotations + implementation library.java.jackson_core + implementation library.java.joda_time + implementation library.java.slf4j_api + implementation library.java.vendored_guava_32_1_2_jre + // IO modules are compile-only: transform-specific lineage visitors activate at runtime only + // when the corresponding IO is already on the user's classpath. + compileOnly project(path: ":sdks:java:io:google-cloud-platform") + compileOnly project(path: ":sdks:java:io:iceberg") + // keep in sync with iceberg_version in sdks/java/io/iceberg/build.gradle + compileOnly "org.apache.iceberg:iceberg-api:1.10.0" + testImplementation library.java.junit + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation project(path: ":sdks:java:io:google-cloud-platform") + testImplementation project(path: ":sdks:java:io:iceberg") + testImplementation project(path: ":sdks:java:managed") + testImplementation "org.apache.iceberg:iceberg-api:1.10.0" + testImplementation "org.apache.iceberg:iceberg-core:1.10.0" + testImplementation library.java.hadoop_common + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfig.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfig.java new file mode 100644 index 000000000000..6371106ec515 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfig.java @@ -0,0 +1,90 @@ +/* + * 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.extensions.openlineage; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.openlineage.client.OpenLineageConfig; +import io.openlineage.client.job.JobConfig; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * OpenLineage configuration for Beam pipelines, extending the generic {@link OpenLineageConfig} + * from openlineage-java with Beam-specific entries — the same pattern as {@code + * SparkOpenLineageConfig} and {@code FlinkOpenLineageConfig} in the official integrations. + * + *

Deserialized from {@code openlineage.yml} or {@code OPENLINEAGE__} environment variables and + * merged with values from {@link OpenLineagePipelineOptions}. + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public class BeamOpenLineageConfig extends OpenLineageConfig { + + /** Default interval between RUNNING events, matching the Flink integration. */ + static final int DEFAULT_TRACKING_INTERVAL_SECONDS = 60; + + @JsonProperty("trackingIntervalInSeconds") + private @Nullable Integer trackingIntervalInSeconds; + + public @Nullable Integer getTrackingIntervalInSeconds() { + return trackingIntervalInSeconds; + } + + public void setTrackingIntervalInSeconds(@Nullable Integer trackingIntervalInSeconds) { + this.trackingIntervalInSeconds = trackingIntervalInSeconds; + } + + @Override + public BeamOpenLineageConfig mergeWithNonNull(BeamOpenLineageConfig other) { + BeamOpenLineageConfig merged = new BeamOpenLineageConfig(); + merged.setTransportConfig(mergePropertyWith(getTransportConfig(), other.getTransportConfig())); + merged.setFacetsConfig(mergePropertyWith(getFacetsConfig(), other.getFacetsConfig())); + merged.setDatasetConfig(mergePropertyWith(getDatasetConfig(), other.getDatasetConfig())); + merged.setCircuitBreaker(mergePropertyWith(getCircuitBreaker(), other.getCircuitBreaker())); + merged.setMetricsConfig(mergePropertyWith(getMetricsConfig(), other.getMetricsConfig())); + merged.setRunConfig(mergePropertyWith(getRunConfig(), other.getRunConfig())); + JobConfig mergedJobConfig = mergeJobConfig(getJobConfig(), other.getJobConfig()); + if (mergedJobConfig != null) { + merged.setJobConfig(mergedJobConfig); + } + merged.setTrackingIntervalInSeconds( + mergePropertyWith(trackingIntervalInSeconds, other.getTrackingIntervalInSeconds())); + return merged; + } + + /** + * Merges job configs field by field. {@link JobConfig#mergeWithNonNull} assumes a non-null owners + * map on both sides and throws otherwise, so it cannot be used directly on configs assembled from + * pipeline options. + */ + private static @Nullable JobConfig mergeJobConfig( + @Nullable JobConfig base, @Nullable JobConfig overrides) { + if (base == null) { + return overrides; + } + if (overrides == null) { + return base; + } + JobConfig merged = new JobConfig(); + merged.setNamespace( + overrides.getNamespace() != null ? overrides.getNamespace() : base.getNamespace()); + merged.setName(overrides.getName() != null ? overrides.getName() : base.getName()); + merged.setOwners(overrides.getOwners() != null ? overrides.getOwners() : base.getOwners()); + merged.setTags(overrides.getTags() != null ? overrides.getTags() : base.getTags()); + return merged; + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParser.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParser.java new file mode 100644 index 000000000000..7875502332dd --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParser.java @@ -0,0 +1,83 @@ +/* + * 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.extensions.openlineage; + +import com.fasterxml.jackson.core.type.TypeReference; +import io.openlineage.client.DefaultConfigPathProvider; +import io.openlineage.client.OpenLineageClientUtils; +import io.openlineage.client.job.JobConfig; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Resolves the effective {@link BeamOpenLineageConfig} for a pipeline, mirroring the Spark + * integration's {@code ArgumentParser}: pipeline options take precedence over {@code + * openlineage.yml}, which takes precedence over {@code OPENLINEAGE__} environment variables. + */ +class BeamOpenLineageConfigParser { + + private static final Logger LOG = LoggerFactory.getLogger(BeamOpenLineageConfigParser.class); + private static final TypeReference TYPE_REFERENCE = + new TypeReference() {}; + + private BeamOpenLineageConfigParser() {} + + static BeamOpenLineageConfig parse(OpenLineagePipelineOptions options) { + BeamOpenLineageConfig fromEnv = loadFromEnvVars(); + BeamOpenLineageConfig fromFile = loadFromFile(); + BeamOpenLineageConfig fromOptions = fromOptions(options); + return fromEnv.mergeWith(fromFile).mergeWith(fromOptions); + } + + private static BeamOpenLineageConfig loadFromFile() { + try { + return OpenLineageClientUtils.loadOpenLineageConfigYaml( + new DefaultConfigPathProvider(), TYPE_REFERENCE); + } catch (RuntimeException e) { + LOG.debug("No openlineage.yml configuration found", e); + return new BeamOpenLineageConfig(); + } + } + + private static BeamOpenLineageConfig loadFromEnvVars() { + try { + return OpenLineageClientUtils.loadOpenLineageConfigFromEnvVars(TYPE_REFERENCE); + } catch (RuntimeException e) { + LOG.debug("No OPENLINEAGE__ environment configuration found", e); + return new BeamOpenLineageConfig(); + } + } + + private static BeamOpenLineageConfig fromOptions(OpenLineagePipelineOptions options) { + BeamOpenLineageConfig config = new BeamOpenLineageConfig(); + String namespace = options.getOpenLineageNamespace(); + String jobName = options.getOpenLineageJobName(); + if (namespace != null || jobName != null) { + JobConfig jobConfig = new JobConfig(); + if (namespace != null) { + jobConfig.setNamespace(namespace); + } + if (jobName != null) { + jobConfig.setName(jobName); + } + config.setJobConfig(jobConfig); + } + config.setTrackingIntervalInSeconds(options.getOpenLineageTrackingIntervalInSeconds()); + return config; + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqns.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqns.java new file mode 100644 index 000000000000..cf03b78f8ee5 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqns.java @@ -0,0 +1,191 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; + +/** + * Converts the Dataplex-style FQN strings reported through {@link + * org.apache.beam.sdk.metrics.Lineage} by Beam IO connectors (e.g. {@code + * pubsub:topic:project.topic}, {@code kafka:`host:9092`.topic}, {@code bigquery:p.d.t}) into {@link + * DatasetIdentifier}s that follow the OpenLineage dataset naming conventions. + */ +class DataplexFqns { + + /** Beam JDBC subprotocols whose OpenLineage spec scheme differs. */ + private static final Map SCHEME_ALIASES = new HashMap<>(); + + static { + SCHEME_ALIASES.put("postgresql", "postgres"); + SCHEME_ALIASES.put("sqlserver", "mssql"); + } + + private DataplexFqns() {} + + /** Parses a flat Beam lineage FQN into an OpenLineage dataset identity. */ + static DatasetIdentifier toDatasetIdentifier(String fqn) { + ParsedFqn parsed = ParsedFqn.parse(fqn); + List seg = parsed.segments; + switch (parsed.system) { + case "pubsub": + { + String subtype = parsed.subtype == null ? "topic" : parsed.subtype; + return new DatasetIdentifier(subtype + ":" + String.join(":", seg), "pubsub"); + } + case "kafka": + { + // Beam reports comma-joined bootstrap servers; the spec namespace is one host:port, + // matching the Spark integration's KafkaBootstrapServerResolver (first server). + String bootstrap = Splitter.on(',').splitToList(first(seg)).get(0); + return new DatasetIdentifier( + seg.size() > 1 ? seg.get(1) : "unknown", "kafka://" + bootstrap); + } + case "bigquery": + return new DatasetIdentifier(String.join(".", seg), "bigquery"); + case "gcs": + return new DatasetIdentifier(objectKey(seg, 1), "gs://" + first(seg)); + case "s3": + return new DatasetIdentifier(objectKey(seg, 1), "s3://" + first(seg)); + case "abs": + { + // Beam segments: account, container, blob. + String account = first(seg); + String container = seg.size() > 1 ? seg.get(1) : "unknown"; + return new DatasetIdentifier( + objectKey(seg, 2), "wasbs://" + container + "@" + account + ".blob.core.windows.net"); + } + case "hdfs": + // The path keeps its leading slash, matching openlineage-java's FilesystemDatasetUtils + // so the same physical dataset gets one identity on every code path. + return new DatasetIdentifier( + seg.size() > 1 ? String.join("/", seg.subList(1, seg.size())) : "/", + "hdfs://" + first(seg)); + case "spanner": + { + // Beam segments: project, instanceConfig, instance, database, table. + String project = first(seg); + String instance = seg.size() > 2 ? seg.get(2) : "unknown"; + String database = seg.size() > 3 ? seg.get(3) : "unknown"; + String name = seg.size() > 4 ? database + "." + seg.get(4) : database; + return new DatasetIdentifier(name, "spanner://" + project + ":" + instance); + } + default: + { + String scheme = SCHEME_ALIASES.getOrDefault(parsed.system, parsed.system); + // JDBC-style FQNs carry host:port as the first segment. + if (seg.size() >= 2 && seg.get(0).contains(":")) { + return new DatasetIdentifier( + String.join(".", seg.subList(1, seg.size())), scheme + "://" + seg.get(0)); + } + return new DatasetIdentifier(String.join(".", seg), scheme); + } + } + } + + /** Object-store names in the spec are the bare object key or path, without a leading slash. */ + private static String objectKey(List seg, int fromIndex) { + if (seg.size() <= fromIndex) { + return "/"; + } + String key = String.join("/", seg.subList(fromIndex, seg.size())); + return key.startsWith("/") ? key.substring(1) : key; + } + + private static String first(List seg) { + return seg.isEmpty() ? "unknown" : seg.get(0); + } + + /** + * Grammar of a Beam lineage FQN: {@code system ':' (subtype ':')? segment ('.' segment)*}. A + * segment containing reserved characters is wrapped in backticks with literal backticks doubled + * (see {@link org.apache.beam.sdk.metrics.Lineage#wrapSegment}). + */ + static final class ParsedFqn { + final String system; + final @org.checkerframework.checker.nullness.qual.Nullable String subtype; + final List segments; + + private ParsedFqn( + String system, + @org.checkerframework.checker.nullness.qual.Nullable String subtype, + List segments) { + this.system = system; + this.subtype = subtype; + this.segments = segments; + } + + static ParsedFqn parse(String fqn) { + // BoundedTrie-based lineage marks truncated FQNs with a trailing '*'. + if (fqn.endsWith("*")) { + fqn = fqn.substring(0, fqn.length() - 1); + } + int firstColon = fqn.indexOf(':'); + if (firstColon < 0) { + return new ParsedFqn(fqn, null, new ArrayList<>()); + } + String system = fqn.substring(0, firstColon); + String rest = fqn.substring(firstColon + 1); + String subtype = null; + // Only Pub/Sub FQNs carry a subtype qualifier (topic/subscription). + if ("pubsub".equals(system)) { + int nextColon = rest.indexOf(':'); + if (nextColon > 0) { + String candidate = rest.substring(0, nextColon); + if (!candidate.contains(".") && !candidate.contains("`")) { + subtype = candidate; + rest = rest.substring(nextColon + 1); + } + } + } + return new ParsedFqn(system, subtype, splitSegments(rest)); + } + + /** Splits on unescaped '.', honoring backtick quoting with doubled-backtick escapes. */ + private static List splitSegments(String s) { + List segments = new ArrayList<>(); + StringBuilder current = new StringBuilder(); + boolean inQuotes = false; + for (int i = 0; i < s.length(); i++) { + char c = s.charAt(i); + if (c == '`') { + if (inQuotes && i + 1 < s.length() && s.charAt(i + 1) == '`') { + current.append('`'); + i++; + } else { + inQuotes = !inQuotes; + } + } else if (c == '.' && !inQuotes) { + segments.add(current.toString()); + current.setLength(0); + } else { + current.append(c); + } + } + if (current.length() > 0) { + segments.add(current.toString()); + } + return segments; + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/EventEmitter.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/EventEmitter.java new file mode 100644 index 000000000000..f73d5a7c3c6f --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/EventEmitter.java @@ -0,0 +1,111 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.OpenLineage; +import io.openlineage.client.OpenLineageClient; +import io.openlineage.client.circuitBreaker.CircuitBreaker; +import io.openlineage.client.circuitBreaker.CircuitBreakerFactory; +import io.openlineage.client.transports.ConsoleConfig; +import io.openlineage.client.transports.FacetsConfig; +import io.openlineage.client.transports.TransportConfig; +import io.openlineage.client.transports.TransportFactory; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Sends OpenLineage events through the configured transport, mirroring the {@code EventEmitter} + * classes of the Spark and Flink integrations: construction and emission never propagate an + * exception to the pipeline — an invalid transport configuration degrades to a non-emitting client + * with an error logged, and emission is wrapped in the configured circuit breaker. + */ +class EventEmitter { + + private static final Logger LOG = LoggerFactory.getLogger(EventEmitter.class); + + private final @Nullable OpenLineageClient client; + private final @Nullable CircuitBreaker circuitBreaker; + + EventEmitter(BeamOpenLineageConfig config) { + this.client = buildClient(config); + this.circuitBreaker = buildCircuitBreaker(config); + } + + private static @Nullable OpenLineageClient buildClient(BeamOpenLineageConfig config) { + try { + TransportConfig transportConfig = + config.getTransportConfig() == null ? new ConsoleConfig() : config.getTransportConfig(); + FacetsConfig facetsConfig = + config.getFacetsConfig() == null ? new FacetsConfig() : config.getFacetsConfig(); + return OpenLineageClient.builder() + .transport(new TransportFactory(transportConfig).build()) + .disableFacets(facetsConfig.getEffectiveDisabledFacets()) + .build(); + } catch (RuntimeException e) { + LOG.error( + "Invalid OpenLineage transport configuration; lineage events will NOT be emitted", e); + return null; + } + } + + private static @Nullable CircuitBreaker buildCircuitBreaker(BeamOpenLineageConfig config) { + try { + return config.getCircuitBreaker() == null + ? null + : new CircuitBreakerFactory(config.getCircuitBreaker()).build(); + } catch (RuntimeException e) { + LOG.warn("Invalid OpenLineage circuit breaker configuration; continuing without one", e); + return null; + } + } + + /** Emits the event; failures are logged and swallowed so lineage never breaks the job. */ + void emit(OpenLineage.RunEvent event) { + final OpenLineageClient localClient = client; + if (localClient == null) { + return; + } + try { + if (circuitBreaker != null) { + circuitBreaker.run( + () -> { + localClient.emit(event); + return Boolean.TRUE; + }); + } else { + localClient.emit(event); + } + LOG.debug("Emitted OpenLineage event: {} run {}", event.getEventType(), event.getRun()); + } catch (RuntimeException e) { + LOG.warn("Failed to emit OpenLineage event, swallowing exception", e); + } + } + + /** Closes the underlying transport (flushing buffered events); never throws. */ + void close() { + if (client == null) { + return; + } + try { + client.close(); + } catch (Exception e) { + LOG.debug("Failed to close OpenLineage client", e); + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/IcebergLineageVisitor.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/IcebergLineageVisitor.java new file mode 100644 index 000000000000..3112ff661900 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/IcebergLineageVisitor.java @@ -0,0 +1,220 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import io.openlineage.client.utils.filesystem.FilesystemDatasetUtils; +import java.lang.reflect.Method; +import java.net.URI; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.iceberg.catalog.Catalog; +import org.apache.iceberg.catalog.TableIdentifier; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Extracts dataset identities from {@code IcebergIO.ReadRows} and {@code IcebergIO.WriteRows}, + * following the Spark integration's {@code IcebergHandler}: the primary identity is the table's + * physical location (via {@link FilesystemDatasetUtils}, so object-store naming matches the + * OpenLineage spec) and the catalog table identity is attached as a {@code TABLE} symlink. When the + * catalog cannot be reached at submission time, falls back to a catalog-based {@code iceberg://} + * identity so the dataset is still reported. + * + *

IcebergIO's configuration getters are package-private, so they are read reflectively; this + * visitor is only instantiated when IcebergIO is on the classpath (see {@link VisitorFactory}). + */ +class IcebergLineageVisitor extends PipelineLineageVisitor { + + private static final Logger LOG = LoggerFactory.getLogger(IcebergLineageVisitor.class); + private static final String ICEBERG_WRITE_CLASS = + "org.apache.beam.sdk.io.iceberg.IcebergIO$WriteRows"; + private static final String ICEBERG_READ_CLASS = + "org.apache.beam.sdk.io.iceberg.IcebergIO$ReadRows"; + + @Override + boolean isDefinedAt(PTransform transform) { + return extendsClass(transform, ICEBERG_WRITE_CLASS) + || extendsClass(transform, ICEBERG_READ_CLASS); + } + + @Override + List applyInputs(PTransform transform) { + if (!extendsClass(transform, ICEBERG_READ_CLASS)) { + return Collections.emptyList(); + } + return extract(transform); + } + + @Override + List applyOutputs(PTransform transform) { + if (!extendsClass(transform, ICEBERG_WRITE_CLASS)) { + return Collections.emptyList(); + } + return extract(transform); + } + + private List extract(PTransform transform) { + try { + Object catalogConfig = invokeDeclared(transform, "getCatalogConfig"); + Object tid = invokeDeclared(transform, "getTableIdentifier"); + if (tid == null) { + // Managed.write(ICEBERG) wraps even a constant table name in dynamic destinations; + // recover it when the destination template has no per-record placeholders. + tid = + staticTableFromDynamicDestinations(invokeDeclared(transform, "getDynamicDestinations")); + } + if (catalogConfig == null || tid == null) { + // Truly dynamic destinations: the table is not known until runtime. + LOG.info("IcebergIO transform without a static table identifier; skipping"); + return Collections.emptyList(); + } + TableIdentifier tableId = + tid instanceof TableIdentifier + ? (TableIdentifier) tid + : TableIdentifier.parse(tid.toString()); + String catalogName = (String) invokeDeclared(catalogConfig, "getCatalogName"); + @SuppressWarnings("unchecked") + Map properties = + (Map) invokeDeclared(catalogConfig, "getCatalogProperties"); + String uri = properties == null ? null : properties.get("uri"); + String warehouse = properties == null ? null : properties.get("warehouse"); + + String symlinkNamespace = + uri != null ? uri : (warehouse != null ? warehouse : "iceberg://" + catalogName); + String location = loadTableLocation(catalogConfig, tableId); + DatasetIdentifier identifier = null; + if (location != null) { + try { + identifier = FilesystemDatasetUtils.fromLocation(URI.create(location)); + } catch (RuntimeException e) { + // Fall through to the catalog-based identity instead of dropping the dataset. + LOG.debug("Unparseable Iceberg table location '{}'", location, e); + } + } + if (identifier == null) { + identifier = + new DatasetIdentifier(tableId.toString(), "iceberg://" + authority(uri, catalogName)); + } + return Collections.singletonList( + identifier.withSymlink( + tableId.toString(), symlinkNamespace, DatasetIdentifier.SymlinkType.TABLE)); + } catch (ReflectiveOperationException | RuntimeException e) { + LOG.warn("Unable to extract lineage from IcebergIO transform", e); + return Collections.emptyList(); + } + } + + /** + * Recovers the table identifier from IcebergIO's {@code PortableIcebergDestinations} (used by + * {@code Managed.write(ICEBERG)}) when the destination template is a constant — i.e. its {@code + * RowStringInterpolator} has no fields to replace. Returns null for genuinely per-record + * destinations, which are only knowable at runtime. + */ + private static @Nullable String staticTableFromDynamicDestinations( + @Nullable Object dynamicDestinations) { + if (dynamicDestinations == null + || !dynamicDestinations + .getClass() + .getName() + .equals("org.apache.beam.sdk.io.iceberg.PortableIcebergDestinations")) { + return null; + } + try { + Object interpolator = readField(dynamicDestinations, "interpolator"); + if (interpolator == null) { + return null; + } + Object fieldsToReplace = readField(interpolator, "fieldsToReplace"); + if (!(fieldsToReplace instanceof java.util.Collection) + || !((java.util.Collection) fieldsToReplace).isEmpty()) { + return null; + } + Object template = readField(interpolator, "template"); + return template == null ? null : template.toString(); + } catch (ReflectiveOperationException | RuntimeException e) { + LOG.debug("Could not recover static table from dynamic destinations", e); + return null; + } + } + + /** Reads a (possibly private) field found anywhere in the class hierarchy. */ + private static @Nullable Object readField(Object target, String fieldName) + throws ReflectiveOperationException { + Class cls = target.getClass(); + while (cls != null) { + try { + java.lang.reflect.Field field = cls.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(target); + } catch (NoSuchFieldException e) { + cls = cls.getSuperclass(); + } + } + return null; + } + + /** Loads the table's physical location from the catalog; null when unreachable. */ + private static @Nullable String loadTableLocation(Object catalogConfig, TableIdentifier tableId) { + try { + Object catalog = invokeDeclared(catalogConfig, "catalog"); + if (catalog instanceof Catalog) { + return ((Catalog) catalog).loadTable(tableId).location(); + } + } catch (ReflectiveOperationException | RuntimeException | NoClassDefFoundError e) { + LOG.info( + "Could not resolve Iceberg table location for {} (catalog unreachable at submit " + + "time?): {}", + tableId, + e.toString()); + } + return null; + } + + private static String authority(@Nullable String uri, @Nullable String catalogName) { + if (uri != null) { + try { + String authority = URI.create(uri).getAuthority(); + if (authority != null) { + return authority; + } + } catch (RuntimeException e) { + // fall through to the catalog name + } + } + return catalogName != null ? catalogName : "unknown"; + } + + private static @Nullable Object invokeDeclared(Object target, String methodName) + throws ReflectiveOperationException { + Class cls = target.getClass(); + while (cls != null) { + try { + Method method = cls.getDeclaredMethod(methodName); + method.setAccessible(true); + return method.invoke(target); + } catch (NoSuchMethodException e) { + cls = cls.getSuperclass(); + } + } + return null; + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/LineageProvider.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/LineageProvider.java new file mode 100644 index 000000000000..ca0e9ee0760f --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/LineageProvider.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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.util.Collections; +import java.util.List; + +/** + * Extension point for {@link org.apache.beam.sdk.transforms.PTransform}s to expose the datasets + * they read or write, mirroring the Flink integration's {@code + * io.openlineage.flink.api.LineageProvider}. Transforms implementing this interface are picked up + * automatically during pipeline-graph traversal, without a dedicated visitor. + */ +public interface LineageProvider { + + /** Datasets this transform reads, in OpenLineage naming-convention form. */ + default List getInputDatasets() { + return Collections.emptyList(); + } + + /** Datasets this transform writes, in OpenLineage naming-convention form. */ + default List getOutputDatasets() { + return Collections.emptyList(); + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageContext.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageContext.java new file mode 100644 index 000000000000..467fccfc675f --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageContext.java @@ -0,0 +1,423 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.OpenLineage; +import io.openlineage.client.utils.DatasetIdentifier; +import io.openlineage.client.utils.UUIDUtils; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.ApplicationNameOptions; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Holds the OpenLineage state of one pipeline execution in this JVM — run identity, discovered + * datasets, and event lifecycle — mirroring the {@code OpenLineageContext} classes of the Spark and + * Flink integrations. One pipeline execution is tracked per JVM (driver or worker); every JVM of + * the same execution resolves the identical run id so events merge into one run downstream. When a + * JVM is reused for a new execution after the previous one finished (e.g. a Flink session cluster), + * a fresh context replaces the finished one. + * + *

Construction never throws: invalid run ids fall back to the deterministic id and transport + * failures degrade to a non-emitting client, so lineage can never fail the pipeline. + */ +class OpenLineageContext { + + private static final Logger LOG = LoggerFactory.getLogger(OpenLineageContext.class); + private static final String DEFAULT_JOB_NAMESPACE = "beam_jobs"; + + private static @Nullable OpenLineageContext instance; + private static @Nullable BeamOpenLineageConfig configOverride; + + /** + * Returns the per-JVM context, creating it from the given options on first use. If the current + * context already finished its run, a fresh context is created so a reused worker JVM tracks the + * new execution instead of staying silent. + */ + static synchronized OpenLineageContext getOrCreate(PipelineOptions options) { + OpenLineageContext local = instance; + if (local == null || local.finished.get()) { + local = new OpenLineageContext(options); + instance = local; + } + return local; + } + + /** + * Re-resolves the context at submission time, after {@link OpenLineageRunner} has minted the run + * id. The {@code LineageBase} plugin may already have been constructed while the runner was being + * resolved (before the id existed); as long as nothing has been emitted yet, the context is + * rebuilt so driver and workers agree on one run id. + */ + static synchronized OpenLineageContext refreshForRunner(PipelineOptions options) { + OpenLineageContext local = instance; + if (local == null || local.finished.get() || !local.started.get()) { + local = new OpenLineageContext(options); + instance = local; + } + return local; + } + + /** Test seam, mirroring the Spark listener's overrideDefaultFactoryForTests. */ + @VisibleForTesting + static synchronized void overrideConfigForTests(@Nullable BeamOpenLineageConfig config) { + configOverride = config; + } + + @VisibleForTesting + static synchronized void resetForTests() { + instance = null; + configOverride = null; + } + + private final OpenLineage openLineage = new OpenLineage(Versions.OPEN_LINEAGE_PRODUCER_URI); + private final BeamOpenLineageConfig config; + private final EventEmitter emitter; + private final String jobNamespace; + private final String jobName; + private final UUID runId; + private final OpenLineage.@Nullable ParentRunFacet parentRunFacet; + private final boolean disabled; + + private final Map inputs = new ConcurrentHashMap<>(); + private final Map outputs = new ConcurrentHashMap<>(); + private final AtomicBoolean started = new AtomicBoolean(false); + private final AtomicBoolean finished = new AtomicBoolean(false); + // Guarded by the synchronized emit path; guarantees START precedes every other event. + private boolean startEmitted; + private volatile boolean streaming; + + private OpenLineageContext(PipelineOptions options) { + OpenLineagePipelineOptions olOptions = options.as(OpenLineagePipelineOptions.class); + BeamOpenLineageConfig resolved = BeamOpenLineageConfigParser.parse(olOptions); + if (configOverride != null) { + resolved = resolved.mergeWith(configOverride); + } + this.config = resolved; + this.disabled = + olOptions.isOpenLineageDisabled() + || Boolean.parseBoolean(System.getenv("OPENLINEAGE_DISABLED")); + this.emitter = new EventEmitter(resolved); + this.jobNamespace = resolveJobNamespace(resolved); + this.jobName = resolveJobName(resolved, options); + this.runId = resolveRunId(olOptions, options); + this.parentRunFacet = buildParentRunFacet(openLineage, olOptions); + this.streaming = options.as(org.apache.beam.sdk.options.StreamingOptions.class).isStreaming(); + } + + private static String resolveJobNamespace(BeamOpenLineageConfig config) { + if (config.getJobConfig() != null && config.getJobConfig().getNamespace() != null) { + return config.getJobConfig().getNamespace(); + } + return DEFAULT_JOB_NAMESPACE; + } + + private static String resolveJobName(BeamOpenLineageConfig config, PipelineOptions options) { + if (config.getJobConfig() != null && config.getJobConfig().getName() != null) { + return config.getJobConfig().getName(); + } + String appName = options.as(ApplicationNameOptions.class).getAppName(); + return appName != null ? appName : options.getJobName(); + } + + /** + * Resolves the run id. An explicit id (minted by {@link OpenLineageRunner} at submission and + * propagated through serialized options, like the Spark integration's driver-minted application + * run id) wins. Otherwise — or when the explicit id is malformed — a deterministic UUID is + * derived from the job identity with the same helper the Airflow provider uses, so independent + * worker JVMs still agree. Never throws. + */ + private static UUID resolveRunId(OpenLineagePipelineOptions olOptions, PipelineOptions options) { + String explicit = olOptions.getOpenLineageRunId(); + if (explicit != null) { + try { + return UUID.fromString(explicit); + } catch (IllegalArgumentException e) { + LOG.warn( + "Malformed openLineageRunId '{}'; falling back to a deterministic run id", explicit); + } + } + byte[] identity = + (options.getJobName() + "/" + options.getOptionsId()).getBytes(StandardCharsets.UTF_8); + return UUIDUtils.generateStaticUUID(Instant.EPOCH, identity); + } + + /** Builds the parent run facet when fully configured; never throws on malformed input. */ + private static OpenLineage.@Nullable ParentRunFacet buildParentRunFacet( + OpenLineage openLineage, OpenLineagePipelineOptions options) { + String parentRunId = options.getOpenLineageParentRunId(); + String parentJobName = options.getOpenLineageParentJobName(); + String parentJobNamespace = options.getOpenLineageParentJobNamespace(); + if (parentRunId == null || parentJobName == null || parentJobNamespace == null) { + return null; + } + try { + UUID parentUuid = UUID.fromString(parentRunId); + return openLineage + .newParentRunFacetBuilder() + .run(openLineage.newParentRunFacetRun(parentUuid)) + .job(openLineage.newParentRunFacetJob(parentJobNamespace, parentJobName)) + .root( + openLineage.newParentRunFacetRoot( + openLineage.newRootRun(parentUuid), + openLineage.newRootJob(parentJobNamespace, parentJobName))) + .build(); + } catch (IllegalArgumentException e) { + LOG.warn("Malformed openLineageParentRunId '{}'; omitting the parent run facet", parentRunId); + return null; + } + } + + BeamOpenLineageConfig getConfig() { + return config; + } + + UUID getRunId() { + return runId; + } + + boolean isDisabled() { + return disabled; + } + + void setStreaming(boolean streaming) { + this.streaming = this.streaming || streaming; + } + + /** Emits START once per run in this JVM. */ + void onJobSubmitted() { + if (disabled) { + return; + } + if (started.compareAndSet(false, true)) { + emit(OpenLineage.RunEvent.EventType.START, null); + } + } + + /** Registers a dataset without emitting; returns true when it is new to this run. */ + boolean registerDataset(LineageDirection direction, DatasetIdentifier dataset) { + if (disabled) { + return false; + } + Map target = direction == LineageDirection.INPUT ? inputs : outputs; + String key = dataset.getNamespace() + "\n" + dataset.getName(); + boolean added = target.putIfAbsent(key, dataset) == null; + if (added) { + LOG.info( + "OpenLineage discovered {} dataset {} {}", + direction, + dataset.getNamespace(), + dataset.getName()); + } + return added; + } + + /** + * Registers a dataset discovered at runtime (worker-side plugin path) and emits a RUNNING event + * when it is new, so long-running streaming jobs surface datasets as they appear. + */ + void onDatasetDiscovered(LineageDirection direction, DatasetIdentifier dataset) { + if (registerDataset(direction, dataset) && started.get() && !finished.get()) { + emit(OpenLineage.RunEvent.EventType.RUNNING, null); + } + } + + /** Emits the periodic RUNNING event, driven by {@link OpenLineageJobTracker}. */ + void onTrackingTick() { + if (disabled || finished.get()) { + return; + } + onJobSubmitted(); + emit(OpenLineage.RunEvent.EventType.RUNNING, null); + } + + /** Pulls IO-reported lineage out of the metrics-based lineage store, if the runner has it. */ + void sweepLineageMetrics(PipelineResult result) { + if (disabled) { + return; + } + try { + Set sources = Lineage.query(result.metrics(), Lineage.Type.SOURCE); + for (String fqn : sources) { + registerDataset(LineageDirection.INPUT, DataplexFqns.toDatasetIdentifier(fqn)); + } + Set sinks = Lineage.query(result.metrics(), Lineage.Type.SINK); + for (String fqn : sinks) { + registerDataset(LineageDirection.OUTPUT, DataplexFqns.toDatasetIdentifier(fqn)); + } + } catch (RuntimeException e) { + LOG.debug("Lineage metrics sweep failed", e); + } + } + + /** Emits the terminal event exactly once. */ + void onJobFinished(OpenLineage.RunEvent.EventType eventType, @Nullable Throwable error) { + if (disabled) { + return; + } + if (finished.compareAndSet(false, true)) { + started.set(true); + emit(eventType, error); + emitter.close(); + } + } + + /** + * Builds and sends one event. Synchronized so that concurrent worker threads cannot reorder + * events on the transport: START is always emitted (exactly once) before any other event, and + * nothing is emitted after the terminal event. + */ + private synchronized void emit( + OpenLineage.RunEvent.EventType eventType, @Nullable Throwable error) { + if (finished.get() && eventType == OpenLineage.RunEvent.EventType.RUNNING) { + return; + } + try { + if (eventType == OpenLineage.RunEvent.EventType.START) { + if (startEmitted) { + return; + } + startEmitted = true; + } else if (!startEmitted) { + startEmitted = true; + emitter.emit(buildEvent(OpenLineage.RunEvent.EventType.START, null)); + } + emitter.emit(buildEvent(eventType, error)); + } catch (RuntimeException e) { + LOG.warn("Failed to build OpenLineage {} event", eventType, e); + } + } + + private OpenLineage.RunEvent buildEvent( + OpenLineage.RunEvent.EventType eventType, @Nullable Throwable error) { + OpenLineage.RunFacetsBuilder runFacets = + openLineage + .newRunFacetsBuilder() + .processing_engine( + openLineage + .newProcessingEngineRunFacetBuilder() + .name("beam") + .version(Versions.getVersion()) + .openlineageAdapterVersion(Versions.getVersion()) + .build()); + if (parentRunFacet != null) { + runFacets.parent(parentRunFacet); + } + if (error != null) { + String message = error.getMessage() == null ? error.toString() : error.getMessage(); + runFacets.errorMessage( + openLineage.newErrorMessageRunFacet(message, "JAVA", stackTraceOf(error))); + } + return openLineage + .newRunEventBuilder() + .eventType(eventType) + .eventTime(ZonedDateTime.now(ZoneOffset.UTC)) + .run(openLineage.newRunBuilder().runId(runId).facets(runFacets.build()).build()) + .job( + openLineage + .newJobBuilder() + .namespace(jobNamespace) + .name(jobName) + .facets( + openLineage + .newJobFacetsBuilder() + .jobType( + openLineage + .newJobTypeJobFacetBuilder() + .processingType(streaming ? "STREAMING" : "BATCH") + .integration("BEAM") + .jobType("JOB") + .build()) + .build()) + .build()) + .inputs(buildInputs()) + .outputs(buildOutputs()) + .build(); + } + + private List buildInputs() { + List result = new ArrayList<>(); + for (DatasetIdentifier di : inputs.values()) { + OpenLineage.InputDatasetBuilder builder = + openLineage.newInputDatasetBuilder().namespace(di.getNamespace()).name(di.getName()); + OpenLineage.DatasetFacets facets = datasetFacets(di); + if (facets != null) { + builder.facets(facets); + } + result.add(builder.build()); + } + return result; + } + + private List buildOutputs() { + List result = new ArrayList<>(); + for (DatasetIdentifier di : outputs.values()) { + OpenLineage.OutputDatasetBuilder builder = + openLineage.newOutputDatasetBuilder().namespace(di.getNamespace()).name(di.getName()); + OpenLineage.DatasetFacets facets = datasetFacets(di); + if (facets != null) { + builder.facets(facets); + } + result.add(builder.build()); + } + return result; + } + + private OpenLineage.@Nullable DatasetFacets datasetFacets(DatasetIdentifier di) { + if (di.getSymlinks().isEmpty()) { + return null; + } + List identifiers = new ArrayList<>(); + for (DatasetIdentifier.Symlink symlink : di.getSymlinks()) { + identifiers.add( + openLineage.newSymlinksDatasetFacetIdentifiers( + symlink.getNamespace(), symlink.getName(), symlink.getType().toString())); + } + return openLineage + .newDatasetFacetsBuilder() + .symlinks(openLineage.newSymlinksDatasetFacetBuilder().identifiers(identifiers).build()) + .build(); + } + + private static String stackTraceOf(Throwable error) { + java.io.StringWriter writer = new java.io.StringWriter(); + error.printStackTrace(new java.io.PrintWriter(writer)); + return writer.toString(); + } + + /** Direction of a dataset relative to the pipeline. */ + enum LineageDirection { + INPUT, + OUTPUT + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageJobTracker.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageJobTracker.java new file mode 100644 index 000000000000..64201e0d1018 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageJobTracker.java @@ -0,0 +1,108 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.OpenLineage; +import org.apache.beam.sdk.PipelineResult; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Watches a running pipeline and emits periodic RUNNING events, mirroring the Flink integration's + * {@code OpenLineageContinousJobTracker}: a daemon thread wakes up every {@code + * trackingIntervalInSeconds} (default 60), sweeps any lineage reported through Beam metrics by the + * IO connectors, and emits a RUNNING event. When it observes a terminal state it emits the terminal + * event (DONE and UPDATED map to COMPLETE, CANCELLED to ABORT, other terminal states to FAIL) and + * stops. The deterministic terminal trigger is {@link OpenLineagePipelineResult}; this tracker + * covers callers that never invoke waitUntilFinish (e.g. detached submission), in which case a + * runner whose {@code getState()} is unsupported may only ever produce RUNNING heartbeats. + */ +class OpenLineageJobTracker { + + private static final Logger LOG = LoggerFactory.getLogger(OpenLineageJobTracker.class); + + private final OpenLineageContext context; + private final PipelineResult result; + private final long intervalMillis; + private @Nullable Thread trackingThread; + + OpenLineageJobTracker(OpenLineageContext context, PipelineResult result, int intervalSeconds) { + this.context = context; + this.result = result; + this.intervalMillis = intervalSeconds * 1000L; + } + + /** Starts the tracking thread; returns immediately. */ + void startTracking() { + Thread thread = + new Thread( + () -> { + LOG.info("Starting OpenLineage job tracking every {} ms", intervalMillis); + while (true) { + context.sweepLineageMetrics(result); + PipelineResult.State state = currentState(); + if (state != null && state.isTerminal()) { + context.onJobFinished(terminalEventType(state), null); + return; + } + context.onTrackingTick(); + try { + Thread.sleep(intervalMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + }, + "beam-openlineage-job-tracker"); + thread.setDaemon(true); + thread.start(); + this.trackingThread = thread; + } + + void stopTracking() { + Thread thread = trackingThread; + if (thread != null) { + thread.interrupt(); + } + } + + private PipelineResult.@Nullable State currentState() { + try { + return result.getState(); + } catch (RuntimeException e) { + LOG.debug("Unable to poll pipeline state", e); + return null; + } + } + + static OpenLineage.RunEvent.EventType terminalEventType(PipelineResult.State state) { + switch (state) { + case DONE: + case UPDATED: + // UPDATED means the job was replaced by a newer version (e.g. Dataflow --update); the + // run itself ended successfully. + return OpenLineage.RunEvent.EventType.COMPLETE; + case CANCELLED: + return OpenLineage.RunEvent.EventType.ABORT; + default: + return OpenLineage.RunEvent.EventType.FAIL; + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineage.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineage.java new file mode 100644 index 000000000000..101e74c6b4d0 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineage.java @@ -0,0 +1,99 @@ +/* + * 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.extensions.openlineage; + +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.lineage.LineageBase; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link LineageBase} plugin that forwards every source/sink FQN reported by Beam IO connectors + * to OpenLineage, live, from whichever JVM the IO executes in (e.g. a Flink TaskManager). Activated + * with {@code --lineageType=org.apache.beam.sdk.extensions.openlineage.OpenLineageLineage}. + * + *

Each FQN is also teed back into the default metrics-based lineage store (both the StringSet + * and BoundedTrie forms, so {@link Lineage#query} keeps working whether or not the {@code + * enable_lineage_rollup} experiment is active) and runner-native lineage consumers keep working + * unchanged. + * + *

This class is instantiated reflectively by the SDK during worker initialization, so it must + * never throw: any failure degrades to metrics-only lineage with a warning logged. + * + *

All JVMs of the same pipeline execution resolve the same run id (see {@link + * OpenLineageContext}), so events from parallel workers merge into a single run downstream. + */ +public class OpenLineageLineage implements LineageBase { + + private static final Logger LOG = LoggerFactory.getLogger(OpenLineageLineage.class); + + private final PipelineOptions options; + private final Lineage.LineageDirection direction; + private final String stringSetMetricName; + private final String boundedTrieMetricName; + private volatile boolean emissionBroken; + + public OpenLineageLineage(PipelineOptions options, Lineage.LineageDirection direction) { + this.options = options; + this.direction = direction; + boolean isSource = direction == Lineage.LineageDirection.SOURCE; + this.stringSetMetricName = + isSource ? Lineage.Type.SOURCE.toString() : Lineage.Type.SINK.toString(); + this.boundedTrieMetricName = + isSource ? Lineage.Type.SOURCEV2.toString() : Lineage.Type.SINKV2.toString(); + } + + @Override + public void add(Iterable rollupSegments) { + List parts = new ArrayList<>(); + rollupSegments.forEach(parts::add); + String fqn = String.join("", parts); + + // Tee into the default metrics-based lineage so metrics consumers keep working. Both forms + // are populated because Lineage.query reads StringSet or BoundedTrie depending on the + // enable_lineage_rollup experiment. + Metrics.stringSet(Lineage.LINEAGE_NAMESPACE, stringSetMetricName).add(fqn); + Metrics.boundedTrie(Lineage.LINEAGE_NAMESPACE, boundedTrieMetricName) + .add(ImmutableList.copyOf(parts)); + + if (emissionBroken) { + return; + } + try { + // The context is resolved per call (cheap after initialization) rather than captured at + // construction, because at submission time OpenLineageRunner may replace the context after + // minting the run id. + OpenLineageContext context = OpenLineageContext.getOrCreate(options); + context.onJobSubmitted(); + context.onDatasetDiscovered( + direction == Lineage.LineageDirection.SOURCE + ? OpenLineageContext.LineageDirection.INPUT + : OpenLineageContext.LineageDirection.OUTPUT, + DataplexFqns.toDatasetIdentifier(fqn)); + } catch (RuntimeException | NoClassDefFoundError e) { + // Lineage must never fail the pipeline; keep the metrics tee working and stop trying. + emissionBroken = true; + LOG.warn("OpenLineage emission disabled after failure; metrics-based lineage continues", e); + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineOptions.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineOptions.java new file mode 100644 index 000000000000..4dc1aba9030c --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineOptions.java @@ -0,0 +1,124 @@ +/* + * 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.extensions.openlineage; + +import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptions; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Pipeline options controlling OpenLineage emission. + * + *

These options mirror the engine-native configuration of the official OpenLineage integrations + * (the {@code spark.openlineage.*} Spark conf keys and the {@code openlineage.*} Flink conf keys). + * Transport, facet, dataset and circuit-breaker configuration is NOT exposed here; it is read the + * standard OpenLineage way — from an {@code openlineage.yml} file (located via the {@code + * OPENLINEAGE_CONFIG} environment variable, the working directory, or {@code ~/.openlineage/}) or + * from {@code OPENLINEAGE__}-prefixed environment variables. Values set via these pipeline options + * take precedence over the file, which takes precedence over environment variables, matching the + * Spark integration's precedence rules. + * + *

Because these are {@link PipelineOptions}, they are serialized with the pipeline and are + * visible to every worker JVM, which lets worker-side emission agree on the run identity. + */ +public interface OpenLineagePipelineOptions extends PipelineOptions { + + @Description( + "Disables OpenLineage emission entirely. The OPENLINEAGE_DISABLED environment variable is " + + "also honored, matching other OpenLineage integrations.") + @Default.Boolean(false) + boolean isOpenLineageDisabled(); + + void setOpenLineageDisabled(boolean disabled); + + @Description( + "OpenLineage job namespace. Overrides the job.namespace from openlineage.yml. " + + "Defaults to \"beam_jobs\" when not configured anywhere.") + @Nullable + String getOpenLineageNamespace(); + + void setOpenLineageNamespace(@Nullable String namespace); + + @Description( + "OpenLineage job name. Overrides the job.name from openlineage.yml. Defaults to the " + + "Beam application name when not configured anywhere.") + @Nullable + String getOpenLineageJobName(); + + void setOpenLineageJobName(@Nullable String jobName); + + @Description( + "Run UUID for this pipeline execution. OpenLineageRunner mints a fresh UUIDv7 at " + + "submission and stores it here so that all worker JVMs share it. When absent, a " + + "deterministic UUID is derived from the job identity so independent JVMs still " + + "agree.") + @Nullable + String getOpenLineageRunId(); + + void setOpenLineageRunId(@Nullable String runId); + + @Description( + "Interval in seconds between RUNNING lifecycle events while the pipeline executes, " + + "matching the Flink integration's openlineage.trackingIntervalInSeconds. Defaults " + + "to 60. Set openLineageDisableTracking to turn periodic events off.") + @Nullable + Integer getOpenLineageTrackingIntervalInSeconds(); + + void setOpenLineageTrackingIntervalInSeconds(@Nullable Integer seconds); + + @Description( + "Disables the periodic RUNNING event tracker, matching the Flink integration's " + + "openlineage.flink.disableCheckpointTracking.") + @Default.Boolean(false) + boolean isOpenLineageDisableTracking(); + + void setOpenLineageDisableTracking(boolean disabled); + + @Description( + "The runner to delegate to when using --runner=OpenLineageRunner. Accepts a simple class " + + "name registered via PipelineRunnerRegistrar (e.g. FlinkRunner, DirectRunner) or a " + + "fully qualified class name.") + @Default.String("DirectRunner") + String getOpenLineageDelegateRunner(); + + void setOpenLineageDelegateRunner(String delegateRunner); + + @Description( + "Parent job name for the parent run facet, matching spark.openlineage.parentJobName. " + + "The facet is attached only when parent run id, job name and namespace are all set.") + @Nullable + String getOpenLineageParentJobName(); + + void setOpenLineageParentJobName(@Nullable String parentJobName); + + @Description( + "Parent job namespace for the parent run facet, matching " + + "spark.openlineage.parentJobNamespace.") + @Nullable + String getOpenLineageParentJobNamespace(); + + void setOpenLineageParentJobNamespace(@Nullable String parentJobNamespace); + + @Description( + "Parent run id (UUID) for the parent run facet, matching spark.openlineage.parentRunId.") + @Nullable + String getOpenLineageParentRunId(); + + void setOpenLineageParentRunId(@Nullable String parentRunId); +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineResult.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineResult.java new file mode 100644 index 000000000000..17d064a8b6a9 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineagePipelineResult.java @@ -0,0 +1,86 @@ +/* + * 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.extensions.openlineage; + +import java.io.IOException; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.metrics.MetricResults; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; + +/** + * Delegating {@link PipelineResult} returned by {@link OpenLineageRunner} that emits the terminal + * OpenLineage event deterministically when the caller observes job completion through {@link + * #waitUntilFinish()} or {@link #cancel()} — instead of relying solely on the periodic {@link + * OpenLineageJobTracker} poll, which could miss a fast batch job if the JVM exits before the next + * tick. + */ +class OpenLineagePipelineResult implements PipelineResult { + + private final PipelineResult delegate; + private final OpenLineageContext context; + + OpenLineagePipelineResult(PipelineResult delegate, OpenLineageContext context) { + this.delegate = delegate; + this.context = context; + } + + @Override + public State getState() { + return delegate.getState(); + } + + @Override + public State cancel() throws IOException { + try { + State state = delegate.cancel(); + afterTerminal(state == null ? State.CANCELLED : state); + return state; + } catch (IOException | RuntimeException e) { + afterTerminal(State.CANCELLED); + throw e; + } + } + + @Override + public State waitUntilFinish(Duration duration) { + State state = delegate.waitUntilFinish(duration); + afterTerminal(state); + return state; + } + + @Override + public State waitUntilFinish() { + State state = delegate.waitUntilFinish(); + afterTerminal(state); + return state; + } + + @Override + public MetricResults metrics() { + return delegate.metrics(); + } + + private void afterTerminal(@Nullable State state) { + if (state == null || !state.isTerminal()) { + return; + } + context.sweepLineageMetrics(delegate); + context.onJobFinished(OpenLineageJobTracker.terminalEventType(state), null); + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRegistrars.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRegistrars.java new file mode 100644 index 000000000000..1c41aaea85e4 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRegistrars.java @@ -0,0 +1,49 @@ +/* + * 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.extensions.openlineage; + +import com.google.auto.service.AutoService; +import java.util.Collections; +import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsRegistrar; +import org.apache.beam.sdk.runners.PipelineRunnerRegistrar; + +/** Registers {@link OpenLineageRunner} and {@link OpenLineagePipelineOptions} with the SDK. */ +public class OpenLineageRegistrars { + + private OpenLineageRegistrars() {} + + /** Makes {@code --runner=OpenLineageRunner} resolvable from the command line. */ + @AutoService(PipelineRunnerRegistrar.class) + public static class Runner implements PipelineRunnerRegistrar { + @Override + public Iterable>> getPipelineRunners() { + return Collections.singletonList(OpenLineageRunner.class); + } + } + + /** Makes the {@code --openLineage*} options parseable from the command line. */ + @AutoService(PipelineOptionsRegistrar.class) + public static class Options implements PipelineOptionsRegistrar { + @Override + public Iterable> getPipelineOptions() { + return Collections.singletonList(OpenLineagePipelineOptions.class); + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunner.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunner.java new file mode 100644 index 000000000000..301e7b72674d --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunner.java @@ -0,0 +1,154 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.OpenLineage; +import io.openlineage.client.utils.UUIDUtils; +import java.util.ServiceLoader; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.runners.PipelineRunnerRegistrar; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link PipelineRunner} that adds OpenLineage emission to any Beam runner. + * + *

Usage is configuration-only: + * + * 

{@code
+ * --runner=OpenLineageRunner
+ * --openLineageDelegateRunner=FlinkRunner
+ * }
+ * + *

At submission it mints a UUIDv7 run id (like the Spark integration's driver-minted application + * run id) and stores it in the serialized options so worker-side emission shares the run; extracts + * datasets from the pipeline graph; emits START; delegates to the real runner; and starts a {@link + * OpenLineageJobTracker} that emits periodic RUNNING events. The terminal COMPLETE/ABORT/FAIL event + * is emitted deterministically when the caller observes completion through the returned {@link + * OpenLineagePipelineResult}. Transport configuration is read the standard OpenLineage way (see + * {@link OpenLineagePipelineOptions}). + * + *

Note: the delegate runner defaults to {@code DirectRunner}, matching Beam's own default when + * no {@code --runner} is set. + */ +public class OpenLineageRunner extends PipelineRunner { + + private static final Logger LOG = LoggerFactory.getLogger(OpenLineageRunner.class); + + private final PipelineOptions options; + + private OpenLineageRunner(PipelineOptions options) { + this.options = options; + } + + /** Constructs the runner from options; called by Beam's runner factory. */ + public static OpenLineageRunner fromOptions(PipelineOptions options) { + return new OpenLineageRunner(options); + } + + @Override + public PipelineResult run(Pipeline pipeline) { + OpenLineagePipelineOptions olOptions = options.as(OpenLineagePipelineOptions.class); + if (olOptions.isOpenLineageDisabled() + || Boolean.parseBoolean(System.getenv("OPENLINEAGE_DISABLED"))) { + LOG.info("OpenLineage is disabled; delegating without lineage emission"); + return resolveDelegate(olOptions).run(pipeline); + } + + if (olOptions.getOpenLineageRunId() == null) { + olOptions.setOpenLineageRunId(UUIDUtils.generateNewUUID().toString()); + } + // Re-resolve the context AFTER minting the run id: the LineageBase plugin may already have + // created a context (without the id) while this runner was being constructed. + OpenLineageContext context = OpenLineageContext.refreshForRunner(options); + try { + PipelineGraphExtractor.extract(pipeline, context); + } catch (RuntimeException e) { + LOG.warn("OpenLineage graph extraction failed; continuing without static datasets", e); + } + context.onJobSubmitted(); + + PipelineRunner delegate = resolveDelegate(olOptions); + LOG.info("OpenLineageRunner delegating to {}", delegate.getClass().getName()); + PipelineResult result; + try { + result = delegate.run(pipeline); + } catch (RuntimeException e) { + context.onJobFinished(OpenLineage.RunEvent.EventType.FAIL, e); + throw e; + } + + if (olOptions.isOpenLineageDisableTracking()) { + LOG.info("OpenLineage tracking disabled; no periodic RUNNING events will be emitted"); + } else { + Integer configured = context.getConfig().getTrackingIntervalInSeconds(); + int intervalSeconds = + configured == null || configured <= 0 + ? BeamOpenLineageConfig.DEFAULT_TRACKING_INTERVAL_SECONDS + : configured; + if (configured != null && configured <= 0) { + LOG.warn( + "Ignoring non-positive trackingIntervalInSeconds {}; using default {}", + configured, + BeamOpenLineageConfig.DEFAULT_TRACKING_INTERVAL_SECONDS); + } + new OpenLineageJobTracker(context, result, intervalSeconds).startTracking(); + } + return new OpenLineagePipelineResult(result, context); + } + + private PipelineRunner resolveDelegate( + OpenLineagePipelineOptions olOptions) { + String name = olOptions.getOpenLineageDelegateRunner(); + Class> runnerClass = resolveRunnerClass(name); + if (runnerClass == null || runnerClass == OpenLineageRunner.class) { + throw new IllegalArgumentException( + "Cannot resolve --openLineageDelegateRunner=" + name + " to a pipeline runner"); + } + // Point the options at the real runner and let Beam's own factory construct it. Workers + // see the delegate runner in their serialized options, exactly as if it had been used + // directly. + options.setRunner(runnerClass); + return PipelineRunner.fromOptions(options); + } + + private static @Nullable Class> resolveRunnerClass(String name) { + try { + Class cls = Class.forName(name); + if (PipelineRunner.class.isAssignableFrom(cls)) { + @SuppressWarnings("unchecked") + Class> runnerClass = (Class>) cls; + return runnerClass; + } + return null; + } catch (ClassNotFoundException e) { + for (PipelineRunnerRegistrar registrar : ServiceLoader.load(PipelineRunnerRegistrar.class)) { + for (Class> cls : registrar.getPipelineRunners()) { + if (cls.getSimpleName().equals(name)) { + return cls; + } + } + } + return null; + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractor.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractor.java new file mode 100644 index 000000000000..4650a926ed37 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractor.java @@ -0,0 +1,75 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PValue; + +/** + * Walks the pipeline graph at submission time, feeding every transform through the {@link + * VisitorFactory} and registering discovered datasets on the {@link OpenLineageContext}. Also + * detects whether the pipeline is streaming (any unbounded {@link PCollection}), which drives the + * {@code jobType} facet's processing type — the same heuristic the Flink integration's {@code + * JobTypeUtils} applies to unbounded sources. + */ +class PipelineGraphExtractor { + + private PipelineGraphExtractor() {} + + static void extract(Pipeline pipeline, OpenLineageContext context) { + VisitorFactory visitorFactory = new VisitorFactory(); + pipeline.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + inspect(node); + return CompositeBehavior.ENTER_TRANSFORM; + } + + @Override + public void visitPrimitiveTransform(TransformHierarchy.Node node) { + inspect(node); + } + + @Override + public void visitValue(PValue value, TransformHierarchy.Node producer) { + if (value instanceof PCollection + && ((PCollection) value).isBounded() == PCollection.IsBounded.UNBOUNDED) { + context.setStreaming(true); + } + } + + private void inspect(TransformHierarchy.Node node) { + PTransform transform = node.getTransform(); + if (transform == null) { + return; + } + for (DatasetIdentifier dataset : visitorFactory.extractInputs(transform)) { + context.registerDataset(OpenLineageContext.LineageDirection.INPUT, dataset); + } + for (DatasetIdentifier dataset : visitorFactory.extractOutputs(transform)) { + context.registerDataset(OpenLineageContext.LineageDirection.OUTPUT, dataset); + } + } + }); + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineLineageVisitor.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineLineageVisitor.java new file mode 100644 index 000000000000..e9b099cfa0cd --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PipelineLineageVisitor.java @@ -0,0 +1,56 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.transforms.PTransform; + +/** + * Extracts dataset identities from a specific IO transform in the pipeline graph, mirroring the + * Flink integration's per-connector {@code Visitor} classes. Implementations must be safe to probe + * with arbitrary transforms and must never throw from {@link #isDefinedAt}. + */ +abstract class PipelineLineageVisitor { + + /** Returns true when this visitor understands the given transform. */ + abstract boolean isDefinedAt(PTransform transform); + + /** Datasets the transform reads. */ + List applyInputs(PTransform transform) { + return Collections.emptyList(); + } + + /** Datasets the transform writes. */ + List applyOutputs(PTransform transform) { + return Collections.emptyList(); + } + + /** True when the transform's class hierarchy contains the given class name. */ + static boolean extendsClass(Object object, String className) { + Class cls = object.getClass(); + while (cls != null) { + if (cls.getName().equals(className)) { + return true; + } + cls = cls.getSuperclass(); + } + return false; + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PubsubLineageVisitor.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PubsubLineageVisitor.java new file mode 100644 index 000000000000..842ee0f8c5ad --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/PubsubLineageVisitor.java @@ -0,0 +1,130 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient; +import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource; +import org.apache.beam.sdk.options.ValueProvider; +import org.apache.beam.sdk.transforms.PTransform; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Extracts Pub/Sub topics and subscriptions from {@code PubsubIO} transforms. Reads use the public + * getters of the expanded {@link PubsubUnboundedSource}; writes reflect the package-private topic + * getter of {@code PubsubIO.Write}. Naming follows the OpenLineage spec: namespace {@code pubsub}, + * name {@code topic:{project}:{topic}} or {@code subscription:{project}:{subscription}}. + */ +class PubsubLineageVisitor extends PipelineLineageVisitor { + + private static final Logger LOG = LoggerFactory.getLogger(PubsubLineageVisitor.class); + private static final String PUBSUB_WRITE_CLASS = + "org.apache.beam.sdk.io.gcp.pubsub.PubsubIO$Write"; + + @Override + boolean isDefinedAt(PTransform transform) { + return transform instanceof PubsubUnboundedSource + || extendsClass(transform, PUBSUB_WRITE_CLASS); + } + + @Override + List applyInputs(PTransform transform) { + if (!(transform instanceof PubsubUnboundedSource)) { + return Collections.emptyList(); + } + PubsubUnboundedSource source = (PubsubUnboundedSource) transform; + List result = new ArrayList<>(); + // Each value is guarded independently: a runtime-only ValueProvider (templated pipelines) + // throws from get(), and must not prevent extraction of the other value. + PubsubClient.TopicPath topic = accessibleValue(source.getTopicProvider()); + if (topic != null) { + addPath(result, "topic", topic.getPath()); + } + PubsubClient.SubscriptionPath subscription = accessibleValue(source.getSubscriptionProvider()); + if (subscription != null) { + addPath(result, "subscription", subscription.getPath()); + } + return result; + } + + @Override + List applyOutputs(PTransform transform) { + if (!extendsClass(transform, PUBSUB_WRITE_CLASS)) { + return Collections.emptyList(); + } + try { + Object provider = invokeDeclared(transform, "getTopicProvider"); + if (provider instanceof ValueProvider && ((ValueProvider) provider).isAccessible()) { + Object pubsubTopic = ((ValueProvider) provider).get(); + if (pubsubTopic != null) { + Object path = invokeDeclared(pubsubTopic, "asPath"); + if (path != null) { + List result = new ArrayList<>(); + addPath(result, "topic", path.toString()); + return result; + } + } + } + } catch (ReflectiveOperationException | RuntimeException e) { + LOG.warn("Unable to extract topic from PubsubIO.Write", e); + } + return Collections.emptyList(); + } + + private static @Nullable T accessibleValue(@Nullable ValueProvider provider) { + try { + if (provider != null && provider.isAccessible()) { + return provider.get(); + } + } catch (RuntimeException e) { + LOG.debug("Pubsub value not available at graph time", e); + } + return null; + } + + /** Converts "projects/{project}/topics/{topic}" into the OpenLineage spec name. */ + private static void addPath(List result, String kind, String path) { + String[] parts = path.split("/", -1); + if (parts.length >= 4) { + result.add(new DatasetIdentifier(kind + ":" + parts[1] + ":" + parts[3], "pubsub")); + } else { + result.add(new DatasetIdentifier(kind + ":" + path, "pubsub")); + } + } + + private static @Nullable Object invokeDeclared(Object target, String methodName) + throws ReflectiveOperationException { + Class cls = target.getClass(); + while (cls != null) { + try { + Method method = cls.getDeclaredMethod(methodName); + method.setAccessible(true); + return method.invoke(target); + } catch (NoSuchMethodException e) { + cls = cls.getSuperclass(); + } + } + return null; + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/Versions.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/Versions.java new file mode 100644 index 000000000000..55a2364297ad --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/Versions.java @@ -0,0 +1,37 @@ +/* + * 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.extensions.openlineage; + +import java.net.URI; +import org.apache.beam.sdk.util.ReleaseInfo; + +/** Identifies this integration as the producer of emitted OpenLineage events. */ +class Versions { + + private Versions() {} + + static final URI OPEN_LINEAGE_PRODUCER_URI = + URI.create( + "https://github.com/apache/beam/tree/v" + + getVersion() + + "/sdks/java/extensions/openlineage"); + + static String getVersion() { + return ReleaseInfo.getReleaseInfo().getVersion(); + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/VisitorFactory.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/VisitorFactory.java new file mode 100644 index 000000000000..3dc677717e5c --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/VisitorFactory.java @@ -0,0 +1,103 @@ +/* + * 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.extensions.openlineage; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.transforms.PTransform; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Assembles the list of {@link PipelineLineageVisitor}s, guarding each IO-specific visitor with a + * classpath check so this extension works without the corresponding IO module — the same pattern as + * the Flink integration's {@code VisitorFactoryImpl} and {@code ClassUtils}. + */ +class VisitorFactory { + + private static final Logger LOG = LoggerFactory.getLogger(VisitorFactory.class); + + private final List visitors = new ArrayList<>(); + + VisitorFactory() { + if (hasClass("org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource")) { + visitors.add(new PubsubLineageVisitor()); + } + if (hasClass("org.apache.beam.sdk.io.iceberg.IcebergIO") + && hasClass("org.apache.iceberg.catalog.TableIdentifier")) { + visitors.add(new IcebergLineageVisitor()); + } + visitors.add(new LineageProviderVisitor()); + } + + List extractInputs(PTransform transform) { + List result = new ArrayList<>(); + for (PipelineLineageVisitor visitor : visitors) { + try { + if (visitor.isDefinedAt(transform)) { + result.addAll(visitor.applyInputs(transform)); + } + } catch (RuntimeException | NoClassDefFoundError e) { + LOG.warn("Lineage visitor {} failed", visitor.getClass().getSimpleName(), e); + } + } + return result; + } + + List extractOutputs(PTransform transform) { + List result = new ArrayList<>(); + for (PipelineLineageVisitor visitor : visitors) { + try { + if (visitor.isDefinedAt(transform)) { + result.addAll(visitor.applyOutputs(transform)); + } + } catch (RuntimeException | NoClassDefFoundError e) { + LOG.warn("Lineage visitor {} failed", visitor.getClass().getSimpleName(), e); + } + } + return result; + } + + private static boolean hasClass(String className) { + try { + Class.forName(className, false, VisitorFactory.class.getClassLoader()); + return true; + } catch (ClassNotFoundException | NoClassDefFoundError e) { + return false; + } + } + + /** Picks up transforms that implement the public {@link LineageProvider} extension point. */ + private static class LineageProviderVisitor extends PipelineLineageVisitor { + @Override + boolean isDefinedAt(PTransform transform) { + return transform instanceof LineageProvider; + } + + @Override + List applyInputs(PTransform transform) { + return ((LineageProvider) transform).getInputDatasets(); + } + + @Override + List applyOutputs(PTransform transform) { + return ((LineageProvider) transform).getOutputDatasets(); + } + } +} diff --git a/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/package-info.java b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/package-info.java new file mode 100644 index 000000000000..45bd54bba824 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/main/java/org/apache/beam/sdk/extensions/openlineage/package-info.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. + */ + +/** + * Emits OpenLineage events describing the datasets a Beam + * pipeline reads and writes, so OpenLineage-compatible backends (such as Marquez) can build a + * lineage graph for Beam jobs on any runner. + */ +@DefaultAnnotation(NonNull.class) +package org.apache.beam.sdk.extensions.openlineage; + +import edu.umd.cs.findbugs.annotations.DefaultAnnotation; +import org.checkerframework.checker.nullness.qual.NonNull; diff --git a/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParserTest.java b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParserTest.java new file mode 100644 index 000000000000..09f4f5b90453 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/BeamOpenLineageConfigParserTest.java @@ -0,0 +1,81 @@ +/* + * 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.extensions.openlineage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; + +import io.openlineage.client.job.JobConfig; +import io.openlineage.client.transports.ConsoleConfig; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link BeamOpenLineageConfigParser} and {@link BeamOpenLineageConfig} merging. */ +@RunWith(JUnit4.class) +public class BeamOpenLineageConfigParserTest { + + @Test + public void testOptionsProvideJobIdentity() { + OpenLineagePipelineOptions options = + PipelineOptionsFactory.as(OpenLineagePipelineOptions.class); + options.setOpenLineageNamespace("my_namespace"); + options.setOpenLineageJobName("my_job"); + options.setOpenLineageTrackingIntervalInSeconds(30); + + BeamOpenLineageConfig config = BeamOpenLineageConfigParser.parse(options); + assertNotNull(config.getJobConfig()); + assertEquals("my_namespace", config.getJobConfig().getNamespace()); + assertEquals("my_job", config.getJobConfig().getName()); + assertEquals(Integer.valueOf(30), config.getTrackingIntervalInSeconds()); + } + + @Test + public void testOptionsWinOverBaseConfigOnMerge() { + BeamOpenLineageConfig base = new BeamOpenLineageConfig(); + JobConfig baseJob = new JobConfig(); + baseJob.setNamespace("from_file"); + baseJob.setName("file_job"); + base.setJobConfig(baseJob); + base.setTransportConfig(new ConsoleConfig()); + base.setTrackingIntervalInSeconds(120); + + BeamOpenLineageConfig overrides = new BeamOpenLineageConfig(); + JobConfig overrideJob = new JobConfig(); + overrideJob.setNamespace("from_options"); + overrides.setJobConfig(overrideJob); + + BeamOpenLineageConfig merged = base.mergeWith(overrides); + // The override's non-null values win; everything else survives from the base. + assertNotNull(merged.getJobConfig()); + assertEquals("from_options", merged.getJobConfig().getNamespace()); + assertEquals("file_job", merged.getJobConfig().getName()); + assertEquals(Integer.valueOf(120), merged.getTrackingIntervalInSeconds()); + assertNotNull(merged.getTransportConfig()); + } + + @Test + public void testEmptyOptionsYieldEmptyJobConfig() { + OpenLineagePipelineOptions options = + PipelineOptionsFactory.as(OpenLineagePipelineOptions.class); + BeamOpenLineageConfig config = BeamOpenLineageConfigParser.parse(options); + assertNull(config.getTrackingIntervalInSeconds()); + } +} diff --git a/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqnsTest.java b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqnsTest.java new file mode 100644 index 000000000000..779e0eeb1521 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/DataplexFqnsTest.java @@ -0,0 +1,119 @@ +/* + * 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.extensions.openlineage; + +import static org.junit.Assert.assertEquals; + +import io.openlineage.client.utils.DatasetIdentifier; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for {@link DataplexFqns}: every mapping must follow the OpenLineage dataset naming + * conventions (https://openlineage.io/docs/spec/naming/). + */ +@RunWith(JUnit4.class) +public class DataplexFqnsTest { + + private static void assertMapsTo(String fqn, String expectedNamespace, String expectedName) { + DatasetIdentifier identifier = DataplexFqns.toDatasetIdentifier(fqn); + assertEquals(expectedNamespace, identifier.getNamespace()); + assertEquals(expectedName, identifier.getName()); + } + + @Test + public void testPubsubTopic() { + assertMapsTo("pubsub:topic:acme-prod.orders-events", "pubsub", "topic:acme-prod:orders-events"); + } + + @Test + public void testPubsubSubscription() { + assertMapsTo( + "pubsub:subscription:acme-prod.orders-sub", "pubsub", "subscription:acme-prod:orders-sub"); + } + + @Test + public void testKafkaWithEscapedBootstrapServers() { + assertMapsTo( + "kafka:`broker-1:9092,broker-2:9092`.payments", "kafka://broker-1:9092", "payments"); + } + + @Test + public void testBigQuery() { + assertMapsTo("bigquery:acme-prod.sales.orders", "bigquery", "acme-prod.sales.orders"); + } + + @Test + public void testGcsWithEscapedObjectKey() { + assertMapsTo( + "gcs:acme-lakehouse.`warehouse/sales.db/orders`", + "gs://acme-lakehouse", + "warehouse/sales.db/orders"); + } + + @Test + public void testS3() { + assertMapsTo("s3:my-bucket.`data/file.parquet`", "s3://my-bucket", "data/file.parquet"); + } + + @Test + public void testAzureBlobStorage() { + assertMapsTo( + "abs:myaccount.mycontainer.path/to/blob", + "wasbs://mycontainer@myaccount.blob.core.windows.net", + "path/to/blob"); + } + + @Test + public void testHdfs() { + assertMapsTo("hdfs:`namenode:8020`./user/data", "hdfs://namenode:8020", "/user/data"); + } + + @Test + public void testSpanner() { + assertMapsTo( + "spanner:acme-prod.regional-us.instance1.salesdb.orders", + "spanner://acme-prod:instance1", + "salesdb.orders"); + } + + @Test + public void testJdbcPostgresAlias() { + assertMapsTo( + "postgresql:`db-host:5432`.mydb.public.users", + "postgres://db-host:5432", + "mydb.public.users"); + } + + @Test + public void testUnknownSystemFallsBackToGenericMapping() { + // Systems without an OpenLineage naming-spec entry pass through conservatively. + assertMapsTo("bigtable:acme-prod.instance1.events", "bigtable", "acme-prod.instance1.events"); + } + + @Test + public void testTruncatedFqnMarkerIsStripped() { + assertMapsTo("bigquery:acme-prod.sales.orders*", "bigquery", "acme-prod.sales.orders"); + } + + @Test + public void testDoubledBacktickEscapesLiteralBacktick() { + assertMapsTo("kafka:`h:9092`.`odd``topic`", "kafka://h:9092", "odd`topic"); + } +} diff --git a/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineageTest.java b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineageTest.java new file mode 100644 index 000000000000..811e8e5c3a99 --- /dev/null +++ b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageLineageTest.java @@ -0,0 +1,140 @@ +/* + * 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.extensions.openlineage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import io.openlineage.client.OpenLineage; +import io.openlineage.client.OpenLineageClientUtils; +import io.openlineage.client.transports.FileConfig; +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.lineage.LineageOptions; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End-to-end test for the {@link OpenLineageLineage} plugin: activated via {@code lineageType}, it + * must capture lineage live from worker code and tee it back into the metrics store. + */ +@RunWith(JUnit4.class) +public class OpenLineageLineageTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private File eventsFile; + + @Before + public void setUp() { + eventsFile = new File(temporaryFolder.getRoot(), "plugin-events.jsonl"); + BeamOpenLineageConfig config = new BeamOpenLineageConfig(); + config.setTransportConfig(new FileConfig(eventsFile.getAbsolutePath())); + OpenLineageContext.resetForTests(); + OpenLineageContext.overrideConfigForTests(config); + } + + @After + public void tearDown() { + OpenLineageContext.resetForTests(); + } + + /** Mirrors the runtime Lineage calls IO connectors make. */ + private static class ReportingFn extends DoFn { + private transient boolean reported; + + @ProcessElement + public void processElement(ProcessContext context) { + if (!reported) { + reported = true; + Lineage.getSources().add("kafka", Arrays.asList("broker-1:9092,broker-2:9092", "payments")); + } + context.output(context.element()); + } + } + + @Test + public void testPluginCapturesLineageLiveAndTeesToMetrics() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(LineageOptions.class).setLineageType(OpenLineageLineage.class); + options.as(OpenLineagePipelineOptions.class).setOpenLineageJobName("plugin_job"); + + // Pin the per-JVM context to this test's options before running: in a shared test JVM an + // earlier test may have installed the plugin bound to its own options (Beam's Lineage + // plugin registration is static and cannot be reset); production worker JVMs initialize + // exactly once per job. + OpenLineageContext.getOrCreate(options); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1, 2, 3)).apply(ParDo.of(new ReportingFn())); + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + // The plugin emitted events live (START plus at least one RUNNING with the dataset). + List events = awaitDatasetEvent(); + assertEquals(OpenLineage.RunEvent.EventType.START, events.get(0).getEventType()); + OpenLineage.RunEvent last = events.get(events.size() - 1); + assertEquals("kafka://broker-1:9092", last.getInputs().get(0).getNamespace()); + assertEquals("payments", last.getInputs().get(0).getName()); + assertEquals("plugin_job", last.getJob().getName()); + + // The tee keeps the metrics-based lineage store populated for other consumers. + Set sources = Lineage.query(result.metrics(), Lineage.Type.SOURCE); + assertFalse(sources.isEmpty()); + assertTrue( + sources.toString(), sources.contains("kafka:`broker-1:9092,broker-2:9092`.payments")); + } + + private List awaitDatasetEvent() throws Exception { + long deadline = System.currentTimeMillis() + 30_000; + List events = null; + while (System.currentTimeMillis() < deadline) { + if (eventsFile.exists()) { + events = + Files.readAllLines(eventsFile.toPath(), StandardCharsets.UTF_8).stream() + .filter(line -> !line.isEmpty()) + .map(OpenLineageClientUtils::runEventFromJson) + .collect(Collectors.toList()); + if (!events.isEmpty() && !events.get(events.size() - 1).getInputs().isEmpty()) { + return events; + } + } + Thread.sleep(250); + } + throw new AssertionError("No event with datasets observed; got " + events); + } +} diff --git a/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunnerTest.java b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunnerTest.java new file mode 100644 index 000000000000..ccad7c55577e --- /dev/null +++ b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/OpenLineageRunnerTest.java @@ -0,0 +1,346 @@ +/* + * 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.extensions.openlineage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import io.openlineage.client.OpenLineage; +import io.openlineage.client.OpenLineageClientUtils; +import io.openlineage.client.transports.FileConfig; +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.metrics.Lineage; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * End-to-end test for {@link OpenLineageRunner}: runs a pipeline on the direct runner through the + * wrapper with a file transport and asserts the emitted event sequence, mirroring the event-file + * assertions of the Spark and Flink integration tests. + */ +@RunWith(JUnit4.class) +public class OpenLineageRunnerTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private File eventsFile; + + @Before + public void setUp() throws Exception { + eventsFile = new File(temporaryFolder.getRoot(), "events.jsonl"); + BeamOpenLineageConfig config = new BeamOpenLineageConfig(); + config.setTransportConfig(new FileConfig(eventsFile.getAbsolutePath())); + OpenLineageContext.resetForTests(); + OpenLineageContext.overrideConfigForTests(config); + } + + @After + public void tearDown() { + OpenLineageContext.resetForTests(); + } + + private static class ThrowingFn extends DoFn { + @ProcessElement + public void processElement() { + throw new IllegalStateException("boom"); + } + } + + /** Mirrors the runtime Lineage calls IO connectors make (e.g. PubsubIO.java). */ + private static class ReportingFn extends DoFn { + private transient boolean reported; + + @ProcessElement + public void processElement(ProcessContext context) { + if (!reported) { + reported = true; + Lineage.getSources() + .add("pubsub", "topic", Arrays.asList("acme-prod", "orders-events"), null); + Lineage.getSinks().add("bigquery", Arrays.asList("acme-prod", "sales", "orders")); + } + context.output(context.element()); + } + } + + @Test + public void testStartAndCompleteEventsWithSweptDatasets() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + options.as(OpenLineagePipelineOptions.class).setOpenLineageTrackingIntervalInSeconds(1); + options.as(OpenLineagePipelineOptions.class).setOpenLineageNamespace("test_namespace"); + options.as(OpenLineagePipelineOptions.class).setOpenLineageJobName("test_job"); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1, 2, 3)).apply(ParDo.of(new ReportingFn())); + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List events = awaitTerminalEvent(); + assertEquals(OpenLineage.RunEvent.EventType.START, events.get(0).getEventType()); + OpenLineage.RunEvent terminal = events.get(events.size() - 1); + assertEquals(OpenLineage.RunEvent.EventType.COMPLETE, terminal.getEventType()); + + // All events belong to one run, identified by the driver-minted UUID. + assertEquals(1, events.stream().map(e -> e.getRun().getRunId()).distinct().count()); + // Job identity comes from the pipeline options. + assertEquals("test_namespace", terminal.getJob().getNamespace()); + assertEquals("test_job", terminal.getJob().getName()); + // The jobType facet marks this integration. + assertEquals("BEAM", terminal.getJob().getFacets().getJobType().getIntegration()); + // Runtime lineage swept from metrics lands on the terminal event. + assertEquals("topic:acme-prod:orders-events", terminal.getInputs().get(0).getName()); + assertEquals("pubsub", terminal.getInputs().get(0).getNamespace()); + assertEquals("acme-prod.sales.orders", terminal.getOutputs().get(0).getName()); + assertEquals("bigquery", terminal.getOutputs().get(0).getNamespace()); + } + + @Test + public void testMalformedRunIdFallsBackWithoutFailing() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + options.as(OpenLineagePipelineOptions.class).setOpenLineageRunId("not-a-uuid"); + options + .as(org.apache.beam.sdk.lineage.LineageOptions.class) + .setLineageType(OpenLineageLineage.class); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1, 2, 3)).apply(ParDo.of(new ReportingFn())); + // Must not throw anywhere — not at submission and not in worker initialization. + pipeline.run().waitUntilFinish(); + + List events = awaitTerminalEvent(); + // A valid (deterministic fallback) UUID was used instead. + java.util.UUID.fromString(events.get(0).getRun().getRunId().toString()); + } + + @Test + public void testMintedRunIdSharedAcrossRunnerAndPluginWithoutTrackerTiming() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + // Interval so large the tracker cannot fire: the terminal event must come from the + // returned PipelineResult, deterministically. + options.as(OpenLineagePipelineOptions.class).setOpenLineageTrackingIntervalInSeconds(3600); + options + .as(org.apache.beam.sdk.lineage.LineageOptions.class) + .setLineageType(OpenLineageLineage.class); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1, 2, 3)).apply(ParDo.of(new ReportingFn())); + PipelineResult result = pipeline.run(); + result.waitUntilFinish(); + + List events = awaitTerminalEvent(); + assertEquals( + OpenLineage.RunEvent.EventType.COMPLETE, events.get(events.size() - 1).getEventType()); + // Every event — the runner's START/COMPLETE and the plugin's live RUNNING — must carry the + // run id minted at submission. + String minted = + pipeline.getOptions().as(OpenLineagePipelineOptions.class).getOpenLineageRunId(); + assertEquals(1, events.stream().map(e -> e.getRun().getRunId().toString()).distinct().count()); + assertEquals(minted, events.get(0).getRun().getRunId().toString()); + // The plugin's live capture reached the events. + assertTrue( + events.stream() + .flatMap(e -> e.getInputs().stream()) + .anyMatch(d -> d.getName().equals("topic:acme-prod:orders-events"))); + } + + @Test + public void testParentRunFacetAttachedWhenFullyConfigured() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + OpenLineagePipelineOptions olOptions = options.as(OpenLineagePipelineOptions.class); + olOptions.setOpenLineageParentRunId("11111111-2222-3333-4444-555555555555"); + olOptions.setOpenLineageParentJobName("parent_dag.parent_task"); + olOptions.setOpenLineageParentJobNamespace("airflow_namespace"); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1)); + pipeline.run().waitUntilFinish(); + + List events = awaitTerminalEvent(); + OpenLineage.ParentRunFacet parent = events.get(0).getRun().getFacets().getParent(); + assertEquals("11111111-2222-3333-4444-555555555555", parent.getRun().getRunId().toString()); + assertEquals("parent_dag.parent_task", parent.getJob().getName()); + assertEquals("airflow_namespace", parent.getJob().getNamespace()); + } + + @Test + public void testManagedIcebergWriteEmitsDatasetWithSymlink() throws Exception { + String warehouse = "file://" + temporaryFolder.newFolder("e2e-warehouse").getAbsolutePath(); + org.apache.iceberg.catalog.TableIdentifier tableId = + org.apache.iceberg.catalog.TableIdentifier.parse("demo.orders"); + try (org.apache.iceberg.hadoop.HadoopCatalog catalog = + new org.apache.iceberg.hadoop.HadoopCatalog( + new org.apache.hadoop.conf.Configuration(), warehouse)) { + catalog.createTable( + tableId, + new org.apache.iceberg.Schema( + org.apache.iceberg.types.Types.NestedField.required( + 1, "id", org.apache.iceberg.types.Types.LongType.get()), + org.apache.iceberg.types.Types.NestedField.required( + 2, "name", org.apache.iceberg.types.Types.StringType.get()))); + } + java.util.Map config = new java.util.HashMap<>(); + config.put("table", "demo.orders"); + config.put("catalog_name", "local"); + java.util.Map catalogProps = new java.util.HashMap<>(); + catalogProps.put("type", "hadoop"); + catalogProps.put("warehouse", warehouse); + config.put("catalog_properties", catalogProps); + + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + org.apache.beam.sdk.schemas.Schema beamSchema = + org.apache.beam.sdk.schemas.Schema.builder() + .addInt64Field("id") + .addStringField("name") + .build(); + + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + Create.of( + org.apache.beam.sdk.values.Row.withSchema(beamSchema) + .addValues(1L, "laptop") + .build()) + .withRowSchema(beamSchema)) + .apply( + org.apache.beam.sdk.managed.Managed.write(org.apache.beam.sdk.managed.Managed.ICEBERG) + .withConfig(config)); + pipeline.run().waitUntilFinish(); + + List events = awaitTerminalEvent(); + // The graph-extracted Iceberg dataset must be on the START event already. + OpenLineage.OutputDataset output = events.get(0).getOutputs().get(0); + assertEquals("file", output.getNamespace()); + assertTrue(output.getName().endsWith("/e2e-warehouse/demo/orders")); + OpenLineage.SymlinksDatasetFacetIdentifiers symlink = + output.getFacets().getSymlinks().getIdentifiers().get(0); + assertEquals("demo.orders", symlink.getName()); + assertEquals("TABLE", symlink.getType()); + } + + @Test + public void testStreamingPipelineEmitsRunningHeartbeatsAndAbortOnCancel() throws Exception { + PipelineOptions options = PipelineOptionsFactory.fromArgs("--blockOnRun=false").create(); + options.setRunner(OpenLineageRunner.class); + options.as(OpenLineagePipelineOptions.class).setOpenLineageTrackingIntervalInSeconds(1); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(org.apache.beam.sdk.io.GenerateSequence.from(0)); + PipelineResult result = pipeline.run(); + + // The tracker must emit periodic RUNNING events while the job runs... + awaitEventCount(OpenLineage.RunEvent.EventType.RUNNING, 2); + // ...and CANCELLED must map to ABORT, per the Flink integration's terminal mapping. + result.cancel(); + awaitEventCount(OpenLineage.RunEvent.EventType.ABORT, 1); + } + + @Test + public void testFailingPipelineEmitsFailWithErrorMessageFacet() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1)).apply(ParDo.of(new ThrowingFn())); + try { + pipeline.run().waitUntilFinish(); + throw new AssertionError("pipeline should have failed"); + } catch (RuntimeException expected) { + // expected + } + + List events = awaitEventCount(OpenLineage.RunEvent.EventType.FAIL, 1); + OpenLineage.RunEvent fail = events.get(events.size() - 1); + String message = fail.getRun().getFacets().getErrorMessage().getMessage(); + assertTrue("errorMessage was: " + message, message.contains("boom")); + assertEquals("JAVA", fail.getRun().getFacets().getErrorMessage().getProgrammingLanguage()); + } + + private List awaitEventCount( + OpenLineage.RunEvent.EventType type, int minCount) throws Exception { + long deadline = System.currentTimeMillis() + 30_000; + List events = readEvents(); + while (System.currentTimeMillis() < deadline) { + events = readEvents(); + if (events.stream().filter(e -> e.getEventType() == type).count() >= minCount) { + return events; + } + Thread.sleep(250); + } + throw new AssertionError("Expected " + minCount + " " + type + " events; got " + events); + } + + @Test + public void testDisabledOptionSuppressesAllEvents() throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(OpenLineageRunner.class); + options.as(OpenLineagePipelineOptions.class).setOpenLineageDisabled(true); + + Pipeline pipeline = Pipeline.create(options); + pipeline.apply(Create.of(1)); + pipeline.run().waitUntilFinish(); + Thread.sleep(2000); + + assertTrue(!eventsFile.exists() || readEvents().isEmpty()); + } + + private List awaitTerminalEvent() throws Exception { + long deadline = System.currentTimeMillis() + 30_000; + while (System.currentTimeMillis() < deadline) { + List events = readEvents(); + if (!events.isEmpty() + && events.get(events.size() - 1).getEventType() + == OpenLineage.RunEvent.EventType.COMPLETE) { + return events; + } + Thread.sleep(250); + } + throw new AssertionError("No COMPLETE event observed in " + readEvents()); + } + + private List readEvents() throws Exception { + if (!eventsFile.exists()) { + return new ArrayList<>(); + } + return Files.readAllLines(eventsFile.toPath(), StandardCharsets.UTF_8).stream() + .filter(line -> !line.isEmpty()) + .map(OpenLineageClientUtils::runEventFromJson) + .collect(Collectors.toList()); + } +} diff --git a/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractorTest.java b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractorTest.java new file mode 100644 index 000000000000..4024afcf5dda --- /dev/null +++ b/sdks/java/extensions/openlineage/src/test/java/org/apache/beam/sdk/extensions/openlineage/PipelineGraphExtractorTest.java @@ -0,0 +1,275 @@ +/* + * 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.extensions.openlineage; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import io.openlineage.client.utils.DatasetIdentifier; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO; +import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergIO; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.hadoop.HadoopCatalog; +import org.apache.iceberg.types.Types; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link VisitorFactory} and the per-IO {@link PipelineLineageVisitor}s. */ +@RunWith(JUnit4.class) +public class PipelineGraphExtractorTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private static final Schema BEAM_SCHEMA = + Schema.builder().addInt64Field("id").addStringField("name").build(); + + @Test + public void testPubsubReadTopicExtraction() { + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline.apply(PubsubIO.readStrings().fromTopic("projects/acme-prod/topics/orders-events")); + + List inputs = extract(pipeline, true); + assertEquals(1, inputs.size()); + assertEquals("pubsub", inputs.get(0).getNamespace()); + assertEquals("topic:acme-prod:orders-events", inputs.get(0).getName()); + } + + @Test + public void testPubsubReadSubscriptionExtraction() { + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline.apply( + PubsubIO.readStrings().fromSubscription("projects/acme-prod/subscriptions/orders-sub")); + + List inputs = extract(pipeline, true); + assertEquals(1, inputs.size()); + assertEquals("pubsub", inputs.get(0).getNamespace()); + assertEquals("subscription:acme-prod:orders-sub", inputs.get(0).getName()); + } + + @Test + public void testIcebergWriteWithReachableCatalogUsesTableLocation() throws Exception { + String warehouse = "file://" + temporaryFolder.newFolder("warehouse").getAbsolutePath(); + TableIdentifier tableId = TableIdentifier.parse("demo.orders"); + try (HadoopCatalog catalog = new HadoopCatalog(new Configuration(), warehouse)) { + catalog.createTable( + tableId, + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.required(2, "name", Types.StringType.get()))); + } + + List outputs = + extract(icebergWritePipeline(warehouse, "hadoop", null), false); + assertEquals(1, outputs.size()); + DatasetIdentifier identifier = outputs.get(0); + // Primary identity is the physical table location, per the Spark integration's + // IcebergHandler; local file locations map to the spec's "file" namespace. + assertEquals("file", identifier.getNamespace()); + assertTrue(identifier.getName().endsWith("/warehouse/demo/orders")); + // The catalog table identity rides in a TABLE symlink. + assertEquals(1, identifier.getSymlinks().size()); + assertEquals("demo.orders", identifier.getSymlinks().get(0).getName()); + assertEquals(DatasetIdentifier.SymlinkType.TABLE, identifier.getSymlinks().get(0).getType()); + } + + @Test + public void testIcebergWriteWithUnreachableRestCatalogFallsBack() { + // Nothing listens on this port; extraction must degrade gracefully, not fail submission. + List outputs = + extract( + icebergWritePipeline( + "gs://acme-lakehouse/warehouse", "rest", "http://127.0.0.1:1/catalog"), + false); + assertEquals(1, outputs.size()); + DatasetIdentifier identifier = outputs.get(0); + assertEquals("iceberg://127.0.0.1:1", identifier.getNamespace()); + assertEquals("demo.orders", identifier.getName()); + assertEquals(1, identifier.getSymlinks().size()); + assertEquals("http://127.0.0.1:1/catalog", identifier.getSymlinks().get(0).getNamespace()); + } + + @Test + public void testManagedIcebergWriteExtraction() throws Exception { + String warehouse = "file://" + temporaryFolder.newFolder("managed-warehouse").getAbsolutePath(); + TableIdentifier tableId = TableIdentifier.parse("demo.orders"); + try (HadoopCatalog catalog = new HadoopCatalog(new Configuration(), warehouse)) { + catalog.createTable( + tableId, + new org.apache.iceberg.Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.required(2, "name", Types.StringType.get()))); + } + Map config = new HashMap<>(); + config.put("table", "demo.orders"); + config.put("catalog_name", "local"); + Map catalogProps = new HashMap<>(); + catalogProps.put("type", "hadoop"); + catalogProps.put("warehouse", warehouse); + config.put("catalog_properties", catalogProps); + + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline + .apply( + org.apache.beam.sdk.transforms.Create.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "a").build()) + .withRowSchema(BEAM_SCHEMA)) + .apply( + org.apache.beam.sdk.managed.Managed.write(org.apache.beam.sdk.managed.Managed.ICEBERG) + .withConfig(config)); + + // Managed wraps the constant table name in dynamic destinations; the visitor must still + // recover it and resolve the physical location. + List outputs = extract(pipeline, false); + assertEquals(1, outputs.size()); + assertEquals("file", outputs.get(0).getNamespace()); + assertTrue(outputs.get(0).getName().endsWith("/managed-warehouse/demo/orders")); + assertEquals("demo.orders", outputs.get(0).getSymlinks().get(0).getName()); + } + + @Test + public void testManagedIcebergWriteWithPerRecordTemplateIsSkipped() { + Map config = new HashMap<>(); + config.put("table", "demo.orders_{name}"); // per-record destination, unknowable at submit + config.put("catalog_name", "local"); + Map catalogProps = new HashMap<>(); + catalogProps.put("type", "hadoop"); + catalogProps.put("warehouse", "file:///tmp/never-used"); + config.put("catalog_properties", catalogProps); + + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline + .apply( + org.apache.beam.sdk.transforms.Create.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "a").build()) + .withRowSchema(BEAM_SCHEMA)) + .apply( + org.apache.beam.sdk.managed.Managed.write(org.apache.beam.sdk.managed.Managed.ICEBERG) + .withConfig(config)); + + assertEquals(0, extract(pipeline, false).size()); + } + + @Test + public void testPubsubWriteTopicExtraction() { + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline + .apply(org.apache.beam.sdk.transforms.Create.of("a")) + .apply(PubsubIO.writeStrings().to("projects/acme-prod/topics/enriched-orders")); + + List outputs = extract(pipeline, false); + assertEquals(1, outputs.size()); + assertEquals("pubsub", outputs.get(0).getNamespace()); + assertEquals("topic:acme-prod:enriched-orders", outputs.get(0).getName()); + } + + @Test + public void testLineageProviderTransformIsPickedUp() { + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline.apply(new ProviderTransform()); + + List inputs = extract(pipeline, true); + assertEquals(1, inputs.size()); + assertEquals("kafka://broker:9092", inputs.get(0).getNamespace()); + assertEquals("events", inputs.get(0).getName()); + } + + private static Pipeline icebergWritePipeline(String warehouse, String type, String uri) { + Map catalogProps = new HashMap<>(); + catalogProps.put("type", type); + catalogProps.put("warehouse", warehouse); + if (uri != null) { + catalogProps.put("uri", uri); + } + Pipeline pipeline = Pipeline.create(PipelineOptionsFactory.create()); + pipeline + .apply( + org.apache.beam.sdk.transforms.Create.of( + Row.withSchema(BEAM_SCHEMA).addValues(1L, "a").build()) + .withRowSchema(BEAM_SCHEMA)) + .apply( + IcebergIO.writeRows( + IcebergCatalogConfig.builder() + .setCatalogName("local") + .setCatalogProperties(catalogProps) + .build()) + .to(TableIdentifier.parse("demo.orders"))); + return pipeline; + } + + /** Runs the visitor factory over the pipeline graph and collects one direction. */ + private static List extract(Pipeline pipeline, boolean inputs) { + VisitorFactory visitorFactory = new VisitorFactory(); + List result = new java.util.ArrayList<>(); + pipeline.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform( + org.apache.beam.sdk.runners.TransformHierarchy.Node node) { + inspect(node.getTransform()); + return CompositeBehavior.ENTER_TRANSFORM; + } + + @Override + public void visitPrimitiveTransform( + org.apache.beam.sdk.runners.TransformHierarchy.Node node) { + inspect(node.getTransform()); + } + + private void inspect(PTransform transform) { + if (transform == null) { + return; + } + result.addAll( + inputs + ? visitorFactory.extractInputs(transform) + : visitorFactory.extractOutputs(transform)); + } + }); + return result; + } + + /** A user transform exposing its datasets through the public extension point. */ + private static class ProviderTransform + extends PTransform> + implements LineageProvider { + @Override + public List getInputDatasets() { + return java.util.Collections.singletonList( + new DatasetIdentifier("events", "kafka://broker:9092")); + } + + @Override + public PCollection expand(org.apache.beam.sdk.values.PBegin input) { + return input.apply(org.apache.beam.sdk.transforms.Create.of("a")); + } + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index d9dbbf9021ef..ca15756076ff 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -197,6 +197,7 @@ include(":sdks:java:extensions:jackson") include(":sdks:java:extensions:join-library") include(":sdks:java:extensions:kafka-factories") include(":sdks:java:extensions:ml") +include(":sdks:java:extensions:openlineage") include(":sdks:java:extensions:ordered") include(":sdks:java:extensions:protobuf") include(":sdks:java:extensions:python")