ETL Tool Integration (Apache Spark, Flink)

Lineage from streaming ETL frameworks like Apache Spark and Flink is not being captured, or the integration is complex. Understanding how to instrument these frameworks for DataHub is important for complete lineage.

You Might Be Asking:

• How do I capture Spark job lineage?
• Can DataHub track Flink stream processing?
• How do I instrument custom ETL code?

1. Instrument Spark jobs with DataHub:

Example for Apache Spark with Hive metastore. This can be adapted for Spark with other data sources like Delta Lake, Iceberg, or cloud data warehouses.

from pyspark.sql import SparkSession
from datahub.api.entities.dataproduct import DataProduct
from datahub.emitter.mce_builder import make_dataset_urn
from datahub.emitter.rest_emitter import DatahubRestEmitter

# Create Spark session with DataHub listener
spark = SparkSession.builder \
    .appName("DataHub Instrumented Job") \
    .config("spark.extraListeners", "io.datahubproject.spark.DatahubSparkListener") \
    .config("spark.datahub.rest.server", "http://localhost:8080") \
    .config("spark.datahub.rest.token", "your-token") \
    .getOrCreate()

# Your Spark transformations
df_source = spark.read.table("source_db.source_table")
df_transformed = df_source.filter(df_source.amount > 0) \
                          .groupBy("customer_id") \
                          .agg({"amount": "sum"})
df_transformed.write.saveAsTable("target_db.target_table")

# Lineage is automatically captured:
# source_db.source_table → target_db.target_table

2. Manual Spark lineage emission:

Example for Apache Spark with Hive. This can be adapted for other Spark use cases.

from datahub.emitter.rest_emitter import DatahubRestEmitter
from datahub.metadata.schema_classes import UpstreamLineageClass, UpstreamClass

# After Spark job completes
emitter = DatahubRestEmitter("http://localhost:8080")

# Define lineage
source_urn = make_dataset_urn("hive", "source_db.source_table")
target_urn = make_dataset_urn("hive", "target_db.target_table")

upstream_lineage = UpstreamLineageClass(
    upstreams=[
        UpstreamClass(
            dataset=source_urn,
            type="TRANSFORMED"
        )
    ]
)

# Emit lineage
emitter.emit_mcp(
    entity_urn=target_urn,
    aspect_name="upstreamLineage",
    aspect=upstream_lineage
)

# Emit job execution metadata
from datahub.metadata.schema_classes import DataProcessInstancePropertiesClass

job_instance = DataProcessInstancePropertiesClass(
    name="spark_etl_job",
    type="SPARK",
    created=int(time.time() * 1000),
    externalUrl="https://spark-ui:4040/job/123"
)

3. Instrument Flink jobs:

Example for Apache Flink with Kafka. This can be adapted for Flink with other sources and sinks.

// Flink DataHub integration
import io.datahubproject.flink.DatahubLineageListener;

// Configure Flink environment
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

// Add DataHub listener
env.registerJobListener(new DatahubLineageListener(
    "http://localhost:8080",
    "your-token"
));

// Your Flink pipeline
DataStream source = env.addSource(new FlinkKafkaConsumer<>(...));
DataStream transformed = source.map(...).filter(...);
transformed.addSink(new FlinkKafkaProducer<>(...));

// Execute - lineage automatically captured
env.execute("Flink ETL Job");

4. Create custom lineage for ETL frameworks:

class ETLJobLineageEmitter:
    """
    Generic ETL job lineage emitter
    """
    def __init__(self, datahub_url, token):
        self.emitter = DatahubRestEmitter(datahub_url, token=token)
        self.job_id = None
        self.sources = []
        self.targets = []

    def start_job(self, job_name):
        """Track job start"""
        self.job_id = f"{job_name}_{int(time.time())}"
        self.start_time = time.time()

    def add_source(self, platform, dataset_name):
        """Add source dataset"""
        urn = make_dataset_urn(platform, dataset_name)
        self.sources.append(urn)

    def add_target(self, platform, dataset_name):
        """Add target dataset"""
        urn = make_dataset_urn(platform, dataset_name)
        self.targets.append(urn)

    def emit_lineage(self):
        """Emit lineage for all targets"""
        for target_urn in self.targets:
            upstream_lineage = UpstreamLineageClass(
                upstreams=[
                    UpstreamClass(dataset=src, type="TRANSFORMED")
                    for src in self.sources
                ]
            )

            self.emitter.emit_mcp(
                entity_urn=target_urn,
                aspect_name="upstreamLineage",
                aspect=upstream_lineage
            )

    def complete_job(self, status="SUCCESS"):
        """Track job completion"""
        duration = time.time() - self.start_time
        # Emit job metadata
        print(f"Job {self.job_id} completed in {duration}s with status {status}")
        self.emit_lineage()

# Usage in ETL code
lineage = ETLJobLineageEmitter("http://localhost:8080", "token")
lineage.start_job("customer_aggregation")

# Read from sources
lineage.add_source("kafka", "raw_events")
lineage.add_source("postgres", "customers")

# Process data
# ... your ETL logic ...

# Write to targets
lineage.add_target("snowflake", "analytics.customer_summary")

lineage.complete_job("SUCCESS")

5. Capture Spark streaming lineage:

Example for Spark Structured Streaming with Kafka. This can be adapted for other streaming sources and sinks.

from pyspark.sql import SparkSession
from pyspark.sql.streaming import StreamingQuery

# Structured Streaming with DataHub
spark = SparkSession.builder \
    .config("spark.sql.streaming.schemaInference", "true") \
    .config("spark.datahub.rest.server", "http://localhost:8080") \
    .getOrCreate()

# Read stream
df_stream = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "input_topic") \
    .load()

# Transform
df_processed = df_stream.selectExpr("CAST(value AS STRING) as json") \
    .select(from_json("json", schema).alias("data")) \
    .select("data.*")

# Write stream
query = df_processed.writeStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("topic", "output_topic") \
    .option("checkpointLocation", "/tmp/checkpoint") \
    .start()

# Lineage captured:
# kafka.input_topic → kafka.output_topic

6. Best practices for ETL lineage:

# 1. Emit lineage at job boundaries
#    - Start of job: record sources
#    - End of job: record targets and emit lineage

# 2. Include transformation metadata
#    - Transformation type (filter, aggregation, join)
#    - Business logic description
#    - Code reference

# 3. Track intermediate datasets
#    - Temp tables
#    - Cached data
#    - Checkpoint locations

# 4. Handle failures gracefully
#    - Emit partial lineage even on failure
#    - Include error context

# 5. Performance considerations
#    - Batch lineage emissions
#    - Async emission to avoid blocking ETL

ETL lineage instrumentation should be lightweight to avoid impacting job performance. Consider using async emission. For complex transformations, manual lineage may be more accurate than automated extraction. Test lineage accuracy with sample jobs.

• Spark Integration
• Lineage Emission