Configure Apache Airflow data lineage tracking with OpenLineage for comprehensive workflow observability

What this solves

Data lineage tracking helps you understand how data flows through your Airflow workflows, identifying dependencies between datasets and transformations. OpenLineage provides standardized lineage collection across different data processing tools, while Marquez offers visualization capabilities. This setup enables data governance, impact analysis, and debugging of complex data pipelines.

Prerequisites

You'll need a working Apache Airflow installation with PostgreSQL backend. If you don't have this setup, follow our Apache Airflow installation guide first.

Ensure your Airflow installation has internet access for downloading OpenLineage packages and connecting to external lineage backends.

Step-by-step configuration

pip install openlineage-airflow openlineage-python

docker run -d \
  --name marquez \
  -p 3000:3000 \
  -p 5432:5432 \
  -e MARQUEZ_CONFIG=/usr/src/app/marquez.yml \
  marquezproject/marquez:latest

transport:
  type: http
  url: http://localhost:3000
  endpoint: /api/v1/lineage
  timeout: 5000
  auth:
    type: api_key
    api_key: your-api-key-here

facets:
  spark_version: true
  spark_logicalPlan: true
  processing_engine: true
  schema: true
  datasource: true
  lifecycle: true
  ownership: true
  columnLineage: true

dataset:
  namespaceResolvers:
    - type: hostname
      hosts:
        - localhost
        - 127.0.0.1

[openlineage]
openlineage_config_path = /opt/airflow/openlineage.yml
extracts = true
log_level = INFO
namespace = production_airflow
transport = http
url = http://localhost:3000

[core]
dags_folder = /opt/airflow/dags
executor = LocalExecutor
sql_alchemy_conn = postgresql+psycopg2://airflow:password@localhost:5432/airflow

pip install apache-airflow-providers-openlineage

OPENLINEAGE_CONFIG=/opt/airflow/openlineage.yml
OPENLINEAGE_NAMESPACE=production_airflow
OPENLINEAGE_URL=http://localhost:3000
AIRFLOW__OPENLINEAGE__CONFIG_PATH=/opt/airflow/openlineage.yml
AIRFLOW__OPENLINEAGE__TRANSPORT_TYPE=http

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.postgres_operator import PostgresOperator
from airflow.operators.python_operator import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from openlineage.airflow.extractors import OperatorLineage

default_args = {
    'owner': 'data-team',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'data_lineage_example',
    default_args=default_args,
    description='DAG demonstrating data lineage tracking',
    schedule_interval=timedelta(days=1),
    catchup=False,
    tags=['lineage', 'example']
)

# SQL transformation with lineage
create_staging_table = PostgresOperator(
    task_id='create_staging_table',
    postgres_conn_id='postgres_default',
    sql="""
    CREATE TABLE IF NOT EXISTS staging.customer_metrics AS
    SELECT 
        customer_id,
        COUNT(*) as order_count,
        SUM(total_amount) as total_spent,
        AVG(total_amount) as avg_order_value,
        MAX(order_date) as last_order_date
    FROM raw.orders 
    WHERE order_date >= '{{ ds }}'
    GROUP BY customer_id;
    """,
    dag=dag
)

def process_customer_data(**context):
    """
    Python function with explicit lineage metadata
    """
    hook = PostgresHook(postgres_conn_id='postgres_default')
    
    # Input dataset information for lineage
    input_datasets = [
        {'namespace': 'postgresql://localhost:5432', 'name': 'staging.customer_metrics'}
    ]
    
    # Output dataset information
    output_datasets = [
        {'namespace': 'postgresql://localhost:5432', 'name': 'analytics.customer_segments'}
    ]
    
    # Perform data processing
    sql = """
    INSERT INTO analytics.customer_segments
    SELECT 
        customer_id,
        CASE 
            WHEN total_spent > 1000 THEN 'high_value'
            WHEN total_spent > 500 THEN 'medium_value'
            ELSE 'low_value'
        END as segment,
        order_count,
        total_spent,
        avg_order_value,
        last_order_date,
        '{{ ds }}' as processing_date
    FROM staging.customer_metrics;
    """
    
    hook.run(sql)
    
    return {
        'input_datasets': input_datasets,
        'output_datasets': output_datasets,
        'transformation': 'customer_segmentation'
    }

process_segments = PythonOperator(
    task_id='process_customer_segments',
    python_callable=process_customer_data,
    dag=dag
)

# Data quality check with lineage
quality_check = PostgresOperator(
    task_id='data_quality_check',
    postgres_conn_id='postgres_default',
    sql="""
    INSERT INTO data_quality.check_results
    SELECT 
        'customer_segments' as table_name,
        COUNT(*) as record_count,
        COUNT(CASE WHEN segment IS NULL THEN 1 END) as null_segments,
        '{{ ds }}' as check_date
    FROM analytics.customer_segments
    WHERE processing_date = '{{ ds }}';
    """,
    dag=dag
)

create_staging_table >> process_segments >> quality_check

from typing import List, Optional, Union
from openlineage.airflow.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.facets import (
    SchemaDatasetFacet,
    SchemaField,
    DataSourceDatasetFacet,
    LifecycleStateChangeDatasetFacet,
    LifecycleStateChange
)
from openlineage.client.run import Dataset
from airflow.models import BaseOperator

class CustomPostgresExtractor(BaseExtractor):
    """
    Custom extractor for PostgreSQL operations with enhanced lineage
    """
    
    @classmethod
    def get_operator_classnames(cls) -> List[str]:
        return ['PostgresOperator', 'PostgresInsertOperator']
    
    def extract(self) -> Optional[OperatorLineage]:
        """
        Extract lineage information from PostgreSQL operators
        """
        if not hasattr(self.operator, 'sql'):
            return None
            
        sql = self.operator.sql
        if not sql:
            return None
            
        # Parse SQL to identify input/output tables
        inputs = self._parse_input_tables(sql)
        outputs = self._parse_output_tables(sql)
        
        input_datasets = []
        output_datasets = []
        
        # Create input datasets
        for table in inputs:
            dataset = Dataset(
                namespace=f"postgresql://{self._get_connection_host()}",
                name=table,
                facets={
                    'dataSource': DataSourceDatasetFacet(
                        name='postgresql',
                        uri=f"postgresql://{self._get_connection_host()}"
                    ),
                    'schema': self._get_table_schema(table)
                }
            )
            input_datasets.append(dataset)
            
        # Create output datasets
        for table in outputs:
            dataset = Dataset(
                namespace=f"postgresql://{self._get_connection_host()}",
                name=table,
                facets={
                    'dataSource': DataSourceDatasetFacet(
                        name='postgresql',
                        uri=f"postgresql://{self._get_connection_host()}"
                    ),
                    'lifecycleStateChange': LifecycleStateChangeDatasetFacet(
                        lifecycleStateChange=LifecycleStateChange.CREATE
                    ),
                    'schema': self._get_table_schema(table)
                }
            )
            output_datasets.append(dataset)
            
        return OperatorLineage(
            inputs=input_datasets,
            outputs=output_datasets
        )
    
    def _parse_input_tables(self, sql: str) -> List[str]:
        """Parse SQL to extract input table names"""
        import re
        # Simple regex to find tables in FROM and JOIN clauses
        pattern = r'(?:FROM|JOIN)\s+([\w\.]+)'
        matches = re.findall(pattern, sql, re.IGNORECASE)
        return list(set(matches))
    
    def _parse_output_tables(self, sql: str) -> List[str]:
        """Parse SQL to extract output table names"""
        import re
        # Simple regex to find tables in INSERT INTO and CREATE TABLE
        pattern = r'(?:INSERT\s+INTO|CREATE\s+TABLE(?:\s+IF\s+NOT\s+EXISTS)?)\s+([\w\.]+)'
        matches = re.findall(pattern, sql, re.IGNORECASE)
        return list(set(matches))
    
    def _get_connection_host(self) -> str:
        """Get database host from connection"""
        from airflow.hooks.base import BaseHook
        try:
            conn = BaseHook.get_connection(self.operator.postgres_conn_id)
            return f"{conn.host}:{conn.port or 5432}"
        except:
            return "localhost:5432"
    
    def _get_table_schema(self, table_name: str) -> Optional[SchemaDatasetFacet]:
        """Get table schema information"""
        try:
            from airflow.providers.postgres.hooks.postgres import PostgresHook
            hook = PostgresHook(postgres_conn_id=self.operator.postgres_conn_id)
            
            # Get column information
            sql = """
            SELECT column_name, data_type, is_nullable
            FROM information_schema.columns 
            WHERE table_name = %s
            ORDER BY ordinal_position
            """
            
            table_only = table_name.split('.')[-1]  # Remove schema prefix if present
            result = hook.get_records(sql, parameters=[table_only])
            
            if not result:
                return None
                
            fields = []
            for row in result:
                field = SchemaField(
                    name=row[0],
                    type=row[1],
                    description=f"Column {row[0]} ({'nullable' if row[2] == 'YES' else 'not null'})"
                )
                fields.append(field)
                
            return SchemaDatasetFacet(fields=fields)
        except Exception:
            return None

server:
  applicationConnectors:
    - type: http
      port: 3000
      bindHost: 0.0.0.0
  adminConnectors:
    - type: http
      port: 3001
      bindHost: 0.0.0.0

database:
  driverClass: org.postgresql.Driver
  url: jdbc:postgresql://localhost:5432/marquez
  user: marquez
  password: marquez_password
  maxWaitForConnection: 1s
  validationQuery: SELECT 1
  validationQueryTimeout: 3s
  minSize: 8
  maxSize: 32
  checkConnectionWhileIdle: false
  evictionInterval: 10s
  minIdleTime: 1m

logging:
  level: INFO
  loggers:
    marquez: INFO
    org.eclipse.jetty: WARN
  appenders:
    - type: console
      threshold: ALL
      timeZone: UTC
      target: stdout
      logFormat: "%d{ISO8601} [%thread] %-5level %logger{36} - %msg%n"

migrateOnStartup: true

lineage:
  writeEnabled: true
  readEnabled: true

from openlineage.airflow import DAGLineageExtractor
from openlineage.client.run import RunEvent, RunState, Job, Run
from openlineage.client.facets import JobFacet, RunFacet
from typing import List, Dict, Any
from datetime import datetime

class CustomLineageConfig:
    """
    Configuration class for custom lineage tracking
    """
    
    def __init__(self, namespace: str = "airflow"):
        self.namespace = namespace
        
    def create_job_facets(self, 
                         task_id: str, 
                         dag_id: str, 
                         owner: str = None,
                         tags: List[str] = None) -> Dict[str, JobFacet]:
        """
        Create job facets with metadata
        """
        facets = {}
        
        if owner:
            facets['ownership'] = JobFacet(
                _producer="custom_lineage",
                _schemaURL="custom_schema",
                properties={'owner': owner, 'dag_id': dag_id}
            )
            
        if tags:
            facets['tags'] = JobFacet(
                _producer="custom_lineage",
                _schemaURL="custom_schema",
                properties={'tags': tags}
            )
            
        return facets
    
    def create_run_facets(self, 
                         execution_date: datetime,
                         task_instance: Any = None) -> Dict[str, RunFacet]:
        """
        Create run facets with execution metadata
        """
        facets = {}
        
        facets['processing'] = RunFacet(
            _producer="custom_lineage",
            _schemaURL="custom_schema",
            properties={
                'execution_date': execution_date.isoformat(),
                'processing_engine': 'airflow'
            }
        )
        
        if task_instance:
            facets['performance'] = RunFacet(
                _producer="custom_lineage",
                _schemaURL="custom_schema",
                properties={
                    'start_time': task_instance.start_date.isoformat() if task_instance.start_date else None,
                    'end_time': task_instance.end_date.isoformat() if task_instance.end_date else None,
                    'duration_seconds': (task_instance.end_date - task_instance.start_date).total_seconds() if task_instance.end_date and task_instance.start_date else None
                }
            )
            
        return facets
    
    def create_lineage_event(self,
                           job_name: str,
                           run_id: str,
                           run_state: RunState,
                           inputs: List[Dict] = None,
                           outputs: List[Dict] = None,
                           job_facets: Dict[str, JobFacet] = None,
                           run_facets: Dict[str, RunFacet] = None) -> RunEvent:
        """
        Create a complete lineage event
        """
        job = Job(
            namespace=self.namespace,
            name=job_name,
            facets=job_facets or {}
        )
        
        run = Run(
            runId=run_id,
            facets=run_facets or {}
        )
        
        event = RunEvent(
            eventType=run_state,
            eventTime=datetime.utcnow().isoformat(),
            run=run,
            job=job,
            inputs=inputs or [],
            outputs=outputs or [],
            producer="custom_airflow_lineage",
            schemaURL="https://openlineage.io/spec/1.0.0/OpenLineage.json"
        )
        
        return event

# Global configuration instance
lineage_config = CustomLineageConfig(namespace="production_airflow")

sudo systemctl stop airflow-webserver
sudo systemctl stop airflow-scheduler
sudo systemctl start airflow-scheduler
sudo systemctl start airflow-webserver
sudo systemctl status airflow-webserver airflow-scheduler

Configure advanced lineage features

from openlineage.client.facets import ColumnLineageDatasetFacet, Fields, InputField
from openlineage.client.run import Dataset
from typing import List, Dict

def create_column_lineage_facet(
    input_fields: List[Dict[str, str]],
    output_fields: List[Dict[str, str]],
    transformations: Dict[str, List[str]]
) -> ColumnLineageDatasetFacet:
    """
    Create column-level lineage facet
    
    Args:
        input_fields: List of {"name": "field_name", "namespace": "dataset_namespace"}
        output_fields: List of {"name": "field_name"}
        transformations: Dict mapping output fields to input fields
    """
    
    fields = {}
    
    for output_field in output_fields:
        field_name = output_field["name"]
        
        if field_name in transformations:
            input