Airflow 動手玩：（三）動手寫 Operator

在上一篇動手寫 DAG 中，我們用了 Operator 來定義不同的 Task，像是執行 Bash 的 BashOperator，或是執行 Python 的 PythonOperator，不過如果剛好有我們想要的功能是 Airflow 沒有定義的，該怎麼辦呢？還好 Airflow 讓我們可以自定義 Operator，讓 Airflow 更有擴展性！

實作 Operator

設置 Operator 路徑

與 DAG 類似，自定義的 Operator 也會需要一個指定的資料夾存放，這個資料夾預設是在 {AIRFLOW_HOME}/plugins，資料夾的位置也可以在 Airflow 的設定檔裡調整，在第五篇《Airflow 設定檔簡介》裡會再詳細的跟大家介紹。

> mkdir ~/airflow/plugins
> touch ~/airflow/plugins/my_operator.py

目標

這次要實作的 Operator 可以傳入使用者的名字，並且在執行時紀錄跟使用者打招呼的時間。

Constructor

自定義的 Operator 會繼承 BaseOperator，所以我們只需要在 Constructor 中定義要傳入的參數，並實作 execute 函式就能完成 Operator 的基礎功能。
在 Constructor 中，我們可以使用 Airflow 定義好的裝飾器 @apply_defaults ，@apply_defaults 可以幫我們的 Operator 載入 Default Arguments，並存放在 kwargs['params']，讓我們可以在 constructor 中調用。

class GreetOperator(BaseOperator):

    @apply_defaults
    def __init__(
            self,
            name: str,
            *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.name = name

Execute

execute 函式帶有 context 變數，context 帶有一些好用的功能，像是 ti 或 task_instance 代表 TaskInstance，透過 TaskInstance 我們可以跟 XCom 互動，上一篇我們展示了 BashOperator 透過 xcom_push 參數，把 bash 執行的結果存入 XCom。

在 Operator 中，執行結束回傳的值會被記錄在 XCom 的 return_value 裡，如果我們想要存更多東西在 XCom 裡，可以怎麼做呢？我們可以透過 TaskInstance 的 xcom_push 功能，存入其他的值，像是下面的例子，我們除了回傳 message 存入 return_value 外，我們另外還存入了現在的 timestamp 在 time 裡面。

    def execute(self, context):
        message = "Hello {}".format(self.name)
        context['ti'].xcom_push(key='time', value=datetime.now().timestamp())
        print(message)
        return message

定義 Plugin

最後定義 AirflowPlugin，讓我們待會再可以在 DAG 中引入我們的 Operator。

class AirflowTestPlugin(AirflowPlugin):
    name = 'greet_operator'
    operators = [GreetOperator]

完整的 Operator Code

from datetime import datetime

from airflow.plugins_manager import AirflowPlugin
from airflow.models.baseoperator import BaseOperator
from airflow.utils.decorators import apply_defaults


class GreetOperator(BaseOperator):

    @apply_defaults
    def __init__(
            self,
            name: str,
            *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.name = name

    def execute(self, context):
        message = "Hello {}".format(self.name)
        context['ti'].xcom_push(key='time', value=datetime.now().timestamp())
        print(message)
        return message


class AirflowTestPlugin(AirflowPlugin):
    name = 'greet_operator'
    operators = [GreetOperator]

再次實作 DAG

目標

這次我們要實作的 DAG 會使用剛剛定義的 Operator 來打招呼，並且把在我們 Operator 中存入的 time 印出來。

完整 DAG Code

由於上一篇講解過如何實作 DAG，所以這邊就直接給完整的 Code，大家可以先在 {AIRFLOW_HOME}/dags 底下創立另一個 DAG 檔案。

> touch ~/airflow/dags/my_operator_dag.py

from datetime import datetime, timedelta


from airflow import DAG
from airflow.operators.python_operator import PythonOperator


from airflow.operators.greet_operator import GreetOperator
# These args will get passed on to each operator
# You can override them on a per-task basis during operator initialization
default_args = {
    'owner': 'someone',
    'depends_on_past': False,
    'start_date': datetime(2020, 2, 25),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(seconds=10),
}


dag = DAG(
    dag_id='my_operator_dag',
    description='my operator dag',
    default_args=default_args,
    schedule_interval='*/1 * * * *'
)


greet = GreetOperator(
    task_id='greet',
    name='someone',
    dag=dag
)


def get_timestamp(**context):
    print(f"greeting time: {context['task_instance'].xcom_pull(task_ids='greet', key='time')}")


show_time = PythonOperator(
    task_id='show_time',
    python_callable=get_timestamp,
    provide_context=True,
    dag=dag
)


greet >> show_time

查詢 XCom

最後大家可以再次開始 Webserver 以及 Scheduler，執行剛剛的 DAG 後，我們可以在 Tree View 裡點擊任一個 greet 方框，之後點選 Task Instance Details。

接下來再點選 XCom，我們就可以在介面上看到 GreetOperator 存入的值。

總結

這次我們從頭實作了自定義的 Operator，透過自定義的 Operator，我們也可以讓 Airflow 有更多擴展，像是在 Airflow 裡面有預設的 redshift_to_s3_operator，但有些公司用的 Warehouse 並不是 Redshift，可能是 Snowflake，這時就可以自定義一個 snowflake_to_s3_operator。
到這邊我們介紹了 Airflow 的基礎功能，接下來我們要瞭解 Airflow 背後是如何運作的，以便我們未來在 Prod 環境部署 Airflow 時，可以更好的做 Troubleshooting。