creating-openlineage-extractors

OpenLineage エクストラクターの作成

このスキルでは、組み込みサポートのない Airflow オペレーターからリネージを取得するカスタム OpenLineage エクストラクターの作成方法を説明します。

各アプローチの使い分け

シナリオ	アプローチ
自分が所有または保守しているオペレーター	OpenLineage Methods（推奨）
変更できないサードパーティ製オペレーター	Custom Extractor
カラムレベルのリネージが必要	OpenLineage Methods または Custom Extractor
複雑な抽出ロジック	OpenLineage Methods または Custom Extractor
シンプルなテーブルレベルのリネージ	Inlets/Outlets（最も簡単で、優先度は最低）

アプローチ 1：OpenLineage Methods（推奨）

from airflow.models import BaseOperator

class MyCustomOperator(BaseOperator):
    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

    def execute(self, context):
        self._rows_processed = self._process_data()
        return self._rows_processed

    def get_openlineage_facets_on_start(self):
        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):
        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
                        )
                    },
                )
            ],
        )

アプローチ 2：カスタムエクストラクター

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

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

    def _execute_extraction(self) -> OperatorLineage | None:
        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:
        return None

エクストラクターの登録

設定ファイル：

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

環境変数：

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