creating-openlineage-extractors

创建 OpenLineage 提取器

本指南将引导你创建自定义 OpenLineage 提取器，以捕获尚未内置支持的 Airflow 算子的数据血缘。

**参考：**有关最新模式和受支持算子／Hook 列表，请参阅 OpenLineage provider 开发者指南。

何时使用哪种方法

场景	方法
你负责开发或维护的算子	OpenLineage 方法（推荐，最简单）
无法修改的第三方算子	自定义提取器
需要列级数据血缘	OpenLineage 方法或自定义提取器
复杂的提取逻辑	OpenLineage 方法或自定义提取器
简单的表级数据血缘	Inlets/Outlets（最简单，但优先级最低）

**重要：**只要可行，就始终优先使用 OpenLineage 方法，而非自定义提取器。提取器更难编写，在算子变更后也更容易与其行为不一致，而且更难调试。

在 Astro 上

Astro 内置 OpenLineage 集成，无需额外配置传输。数据血缘事件会自动收集，并显示在 Astro UI 的 Lineage 标签页中。部署到 Astro 项目的自定义提取器会自动加载，因此你只需在 airflow.cfg 中或通过环境变量注册，然后进行部署。

---

两种方法

1. OpenLineage 方法（推荐）

当你可以直接向自定义算子添加方法时使用。这是处理由你负责的算子的首选方案。

2. 自定义提取器

当你需要从无法修改的第三方或 provider 算子中获取数据血缘时使用。

---

方法 1：OpenLineage 方法（推荐）

当你负责维护算子时，直接添加 OpenLineage 方法：

from airflow.models import BaseOperator


class MyCustomOperator(BaseOperator):
    """Custom operator with built-in OpenLineage support."""

    def __init__(self, source_table: str, target_table: str, **kwargs):
        super().__init__(**kwargs)
        self.source_table = source_table
        self.target_table = target_table
        self._rows_processed = 0  # Set during execution

    def execute(self, context):
        # Do the actual work
        self._rows_processed = self._process_data()
        return self._rows_processed

    def get_openlineage_facets_on_start(self):
        """Called when task starts. Return known inputs/outputs."""
        # Import locally to avoid circular imports
        from openlineage.client.event_v2 import Dataset
        from airflow.providers.openlineage.extractors import OperatorLineage

        return OperatorLineage(
            inputs=[Dataset(namespace="postgres://db", name=self.source_table)],
            outputs=[Dataset(namespace="postgres://db", name=self.target_table)],
        )

    def get_openlineage_facets_on_complete(self, task_instance):
        """Called after success. Add runtime metadata."""
        from openlineage.client.event_v2 import Dataset
        from openlineage.client.facet_v2 import output_statistics_output_dataset
        from airflow.providers.openlineage.extractors import OperatorLineage

        return OperatorLineage(
            inputs=[Dataset(namespace="postgres://db", name=self.source_table)],
            outputs=[
                Dataset(
                    namespace="postgres://db",
                    name=self.target_table,
                    facets={
                        "outputStatistics": output_statistics_output_dataset.OutputStatisticsOutputDatasetFacet(
                            rowCount=self._rows_processed
                        )
                    },
                )
            ],
        )

    def get_openlineage_facets_on_failure(self, task_instance):
        """Called after failure. Optional - for partial lineage."""
        return None

OpenLineage 方法参考

方法	调用时机	必需
get_openlineage_facets_on_start()	任务进入 RUNNING 状态时	否
get_openlineage_facets_on_complete(ti)	任务成功时	否
get_openlineage_facets_on_failure(ti)	任务失败时	否

只实现你需要的方法即可。未实现的方法会回退到 Hook 级数据血缘或 inlets/outlets。

---

方法 2：自定义提取器

仅在你无法修改算子时（例如第三方或 provider 算子）使用此方法。

基本结构

from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset


class MyOperatorExtractor(BaseExtractor):
    """Extract lineage from MyCustomOperator."""

    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        """Return operator class names this extractor handles."""
        return ["MyCustomOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        """Called BEFORE operator executes. Use for known inputs/outputs."""
        # Access operator properties via self.operator
        source_table = self.operator.source_table
        target_table = self.operator.target_table

        return OperatorLineage(
            inputs=[
                Dataset(
                    namespace="postgres://mydb:5432",
                    name=f"public.{source_table}",
                )
            ],
            outputs=[
                Dataset(
                    namespace="postgres://mydb:5432",
                    name=f"public.{target_table}",
                )
            ],
        )

    def extract_on_complete(self, task_instance) -> OperatorLineage | None:
        """Called AFTER operator executes. Use for runtime-determined lineage."""
        # Access properties set during execution
        # Useful for operators that determine outputs at runtime
        return None

OperatorLineage 结构

from airflow.providers.openlineage.extractors.base import OperatorLineage
from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import sql_job

lineage = OperatorLineage(
    inputs=[Dataset(namespace="...", name="...")],      # Input datasets
    outputs=[Dataset(namespace="...", name="...")],     # Output datasets
    run_facets={"sql": sql_job.SQLJobFacet(query="SELECT...")},  # Run metadata
    job_facets={},                                      # Job metadata
)

提取方法

方法	调用时机	用途
_execute_extraction()	算子运行前	静态／已知的数据血缘
extract_on_complete(task_instance)	成功后	运行时确定的数据血缘
extract_on_failure(task_instance)	失败后	错误发生时的部分数据血缘

注册提取器

选项 1：配置文件（airflow.cfg）

[openlineage]
extractors = mypackage.extractors.MyOperatorExtractor;mypackage.extractors.AnotherExtractor

选项 2：环境变量

AIRFLOW__OPENLINEAGE__EXTRACTORS='mypackage.extractors.MyOperatorExtractor;mypackage.extractors.AnotherExtractor'

**重要：**该路径必须可从 Airflow worker 导入。请将提取器放在 DAG 文件夹或已安装的软件包中。

---

常见模式

SQL 算子提取器

from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import sql_job


class MySqlOperatorExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["MySqlOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        sql = self.operator.sql
        conn_id = self.operator.conn_id

        # Parse SQL to find tables (simplified example)
        # In practice, use a SQL parser like sqlglot
        inputs, outputs = self._parse_sql(sql)

        namespace = f"postgres://{conn_id}"

        return OperatorLineage(
            inputs=[Dataset(namespace=namespace, name=t) for t in inputs],
            outputs=[Dataset(namespace=namespace, name=t) for t in outputs],
            job_facets={
                "sql": sql_job.SQLJobFacet(query=sql)
            },
        )

    def _parse_sql(self, sql: str) -> tuple[list[str], list[str]]:
        """Parse SQL to extract table names. Use sqlglot for real parsing."""
        # Simplified example - use proper SQL parser in production
        inputs = []
        outputs = []
        # ... parsing logic ...
        return inputs, outputs

文件传输提取器

from airflow.providers.openlineage.extractors.base import BaseExtractor, OperatorLineage
from openlineage.client.event_v2 import Dataset


class S3ToSnowflakeExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["S3ToSnowflakeOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        s3_bucket = self.operator.s3_bucket
        s3_key = self.operator.s3_key
        table = self.operator.table
        schema = self.operator.schema

        return OperatorLineage(
            inputs=[
                Dataset(
                    namespace=f"s3://{s3_bucket}",
                    name=s3_key,
                )
            ],
            outputs=[
                Dataset(
                    namespace="snowflake://myaccount.snowflakecomputing.com",
                    name=f"{schema}.{table}",
                )
            ],
        )

从执行中获取动态数据血缘

from openlineage.client.event_v2 import Dataset


class DynamicOutputExtractor(BaseExtractor):
    @classmethod
    def get_operator_classnames(cls) -> list[str]:
        return ["DynamicOutputOperator"]

    def _execute_extraction(self) -> OperatorLineage | None:
        # Only inputs known before execution
        return OperatorLineage(
            inputs=[Dataset(namespace="...", name=self.operator.source)],
        )

    def extract_on_complete(self, task_instance) -> OperatorLineage | None:
        # Outputs determined during execution
        # Access via operator properties set in execute()
        outputs = self.operator.created_tables  # Set during execute()

        return OperatorLineage(
            inputs=[Dataset(namespace="...", name=self.operator.source)],
            outputs=[Dataset(namespace="...", name=t) for t in outputs],
        )

---

常见问题

1. 循环导入

**问题：**在顶层导入 Airflow 模块会导致循环导入。

# ❌ BAD - can cause circular import issues
from airflow.models import TaskInstance
from openlineage.client.event_v2 import Dataset

class MyExtractor(BaseExtractor):
    ...

# ✅ GOOD - import inside methods
class MyExtractor(BaseExtractor):
    def _execute_extraction(self):
        from openlineage.client.event_v2 import Dataset
        # ...

2. 导入路径错误

**问题：**提取器路径与实际模块位置不匹配。

# ❌ Wrong - path doesn't exist
AIRFLOW__OPENLINEAGE__EXTRACTORS='extractors.MyExtractor'

# ✅ Correct - full importable path
AIRFLOW__OPENLINEAGE__EXTRACTORS='dags.extractors.my_extractor.MyExtractor'

3. 未处理 None

**问题：**算子属性为 None 时，提取会失败。

# ✅ Handle optional properties
def _execute_extraction(self) -> OperatorLineage | None:
    if not self.operator.source_table:
        return None  # Skip extraction

    return OperatorLineage(...)

---

测试提取器

单元测试

import pytest
from unittest.mock import MagicMock
from mypackage.extractors import MyOperatorExtractor


def test_extractor():
    # Mock the operator
    operator = MagicMock()
    operator.source_table = "input_table"
    operator.target_table = "output_table"

    # Create extractor
    extractor = MyOperatorExtractor(operator)

    # Test extraction
    lineage = extractor._execute_extraction()

    assert len(lineage.inputs) == 1
    assert lineage.inputs[0].name == "input_table"
    assert len(lineage.outputs) == 1
    assert lineage.outputs[0].name == "output_table"

---

优先级规则

OpenLineage 按以下顺序检查数据血缘：

1. 自定义提取器（优先级最高）
2. 算子上的 OpenLineage 方法
3. Hook 级数据血缘（来自 HookLineageCollector）
4. Inlets/Outlets（优先级最低）

如果存在自定义提取器，它会覆盖内置提取逻辑以及 inlets/outlets。

---

相关技能

• annotating-task-lineage：用于通过输入端和输出端追踪简单的表级血缘
• tracing-upstream-lineage：调查数据来源
• tracing-downstream-lineage：调查数据依赖关系