Airflow 是 Airbnb 公司開源的任務調度系統, 通過使用 Python 開發 DAG, 非常方便的調度計算任務. 介紹一下在 Airflow 提供的 Operator 不滿足需求的場景下, 如何自己開發 Operator.
0x00 DAG 的最基本執行單元: Operator
在 Airflow 的一個 DAG 中, 最基本的執行單元是 Operator. 例如如下示例 DAG 中, 使用的都是 BashOperator, 執行一個 bash 腳本.
"""
Code that goes along with the Airflow tutorial located at:
https://github.com/airbnb/airflow/blob/master/airflow/example_dags/tutorial.py
"""
from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime(2015, 6, 1),
'email': ['airflow@airflow.com'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5)
}
dag = DAG('tutorial', default_args=default_args)
# t1, t2 and t3 are examples of tasks created by instantiating operators
t1 = BashOperator(
task_id='print_date',
bash_command='date',
dag=dag)
t2 = BashOperator(
task_id='sleep',
bash_command='sleep 5',
retries=3,
dag=dag)
templated_command = """
{% for i in range(5) %}
echo "{{ ds }}"
echo "{{ macros.ds_add(ds, 7)}}"
echo "{{ params.my_param }}"
{% endfor %}
"""
t3 = BashOperator(
task_id='templated',
bash_command=templated_command,
params={'my_param': 'Parameter I passed in'},
dag=dag)
t2.set_upstream(t1)
t3.set_upstream(t1)
Airflow 實現了很多 Operator(參見 Airflow 源代碼), 涵蓋了常用的功能, 例如執行 Hive 查詢, 執行 bash 腳本等. 有幾種特殊的 Operator:
-
XXXSensor用作其他外界條件的 sensor, 實現也很簡單, 在 Operator 的execute方法中進行 long poll, 直到poke方法返回 True 則完成執行.
# BaseSensorOperator 部分源碼
def poke(self, context):
'''
Function that the sensors defined while deriving this class should
override.
'''
raise AirflowException('Override me.')
def execute(self, context):
started_at = datetime.now()
while not self.poke(context):
sleep(self.poke_interval)
if (datetime.now() - started_at).seconds > self.timeout:
raise AirflowSensorTimeout('Snap. Time is OUT.')
logging.info("Success criteria met. Exiting.")
-
PythonOperator用來執行 Python 函數, 這也是使用 Python 代碼來定義 DAG 的好處 -
BranchPythonOperator用來支持分支, 通過函數返回要執行的分支
Airflow Operator 相關 class 繼承關系如下:
.
└── BaseOperator
├── BaseSensorOperator
│ └── ...Sensor
├── PythonOperator
│ ├── BranchPythonOperator
│ └── ShortCircuitOperator
└── ...Operator
0x01 Operator 開發
如果官方的 Operator 都不滿足需求, 那么我們就要來開發一個 Operator. 開發 Operator 也很簡單, 直接繼承 BaseOperator并實現 execute 方法即可.
from airflow.models import BaseOperator
class DemoOperator(BaseOperator):
def __init__(*args, **kwargs):
super(DemoOperator, self).__init__(*args, **kwargs)
def execute(self, context):
print "hello"
除了 execute 方法必須實現外, 還有一個 hook 方法:
-
pre_execute: 在execute方法前調用, 實現點兒準備邏輯 -
post_execute: 在execute方法完成后調用, cleanup 一下 -
on_kill: 在 task 被 kill 的時候執行.
Operator 獲取模板變量
Aiflow 是支持 Templating with Jinja 功能的, 具體來說就是 Operator 中支持模板語言 Jinja, 寫一些 for 循環, 或者通過 {{param}} 語法替換一些變量等(例如 {{ds}} 被替換成執行任務的日期)
# 官方示例的 jinja 語句
{% for i in range(5) %}
echo "{{ ds }}"
echo "{{ macros.ds_add(ds, 7) }}"
echo "{{ params.my_param }}"
{% endfor %}
那么, 自己開發的 Operator 中如何使用這個功能呢?
其實也很簡單, 在自己的 Operator 中添加屬性 template_fields = (attributes_to_be_rendered_with_jinja). 在任務被執行前, Airflow 會自動渲染 template_fields 中的屬性再執行任務.
# 節選自 Airflow 中 BashOperator 源碼
class BashOperator(BaseOperator):
# 這里定義需要被渲染的屬性名稱
template_fields = ('bash_command', 'env')
Operator 部署
開發的 Operator 代碼作為一個 Python 的 Package, 使用 distutil 打包安裝到 Airflow 對應的服務器上即可.
0x02 Operator 跟其他系統交互
Airflow 考慮到為了跟外界環境隔離, 提出了 Connection 概念: 將配置信息放到 Airflow Server 本身配置, 在 DAG 中使用 connection_id 來引用. 例如, 我們配置一個 HiveServer 的 Connection, 使用 liulishuo_hiveserver1 作為 connection_id, 這樣同一個 DAG 文件就可以在測試環境和生成環境調用對應環境中的 HiveServer 服務了. 總結來說, 這就是架構設計模式中的 External Configuration Store Pattern 的標準實現.
那么如果自己開發的 Operator 如何調用這些 Connection 呢? 這里 Airflow 又引入了一個 Hook 的概念:
Hooks are interfaces to external platforms and databases like Hive, S3, MySQL, Postgres, HDFS, and Pig. Hooks implement a common interface when possible, and act as a building block for operators. They also use the airflow.models.Connection model to retrieve hostnames and authentication information. Hooks keep authentication code and information out of pipelines, centralized in the metadata database.
我個人覺得這個概念的引入有點兒臃腫了, 沒有任何意義. 開發者在 Airflow Web 中看到的也是 Connection 的配置, 那么開發 DAG 時第一個想到的就是去找 Connection 相關的 class, 再多一個 Hook的概念有點兒繞.
那么 Operator 中想使用對應的 Connection, 直接根據 connection_id 創建對應的 Hook 就好了(講真, 真繞), 例如, 想使用 HiveServer2 的 Connection, 創建一個 HiveServer2Hook 即可.
# Operator 調用 Connection 示例代碼
class LiulishuoDemoOperator(BaseOperator):
def __init__(self, hive_server2_connection_id, *args, **kwargs):
super(LiulishuoDemoOperator, self).__init__(*args, **kwargs)
self.hive_server2_connection_id = hive_server2_connection_id
def execute(self, context):
hive_server2 = HiveServer2Hook(self.hive_server2_connection_id)
hive_serve2.get_records('SELECT * FROM testdb.table1 LIMIT 20')
# ....
HiveServer2Hook 設計有還有一個貼心之處就是, 在創建 HiveServer2Hook 時根本不涉及真正連接 HiveServer2 的邏輯, 只有真正調用其get_records 等方法時才會真正去連接 HiveServer2, 這樣就給單元測試 mock 帶來很大的方便, 畢竟在 CI 環境中構建一個隔離的專門用于跑自己的 test-case 的 HiveServer2 也不是那么容易的.
def test_operator_with_mock(self):
with mock.patch.object(HiveServer2Hook, 'get_records') as mock_get_records:
# 這里設置 mock 的返回值
mock_get_records.return_value = [['Location: ', 's3://test-bucket/valid_table']]
hive_server_id = 'test-hive-server'
# 這里測試對應的 Operator 代碼
0x03 總結
Airflow 通過精簡的抽象, 將 DAG 開發簡化到了會寫 Python 基本就沒問題的程度, 還是值得點贊的. 自己開發一個 Operator 也是很簡單, 不過自己開發 Operator 也僅僅是技術選型的其中一個方案而已, 復雜邏輯也可以通過暴露一個 Restful API 的形式, 使用 Airflow 提供的 SimpleHttpOperator 調用對應的服務執行任務.
