Airflow自定义插件
Airflow之所以受欢迎的一个重要因素就是它的插件机制。Python成熟类库可以很方便的引入各种插件。在我们实际工作中，必然会遇到官方的一些插件不足够满足需求的时候。这时候，我们可以编写自己的插件。不需要你了解内部原理，甚至不需要很熟悉Python， 反正我连蒙带猜写的。
插件分类
Airflow的插件分为Operator和Sensor两种。Operator是具体要执行的任务插件， Sensor则是条件传感器，当我需要设定某些依赖的时候可以通过不同的sensor来感知条件是否满足。
Airflow对插件提供的支持
插件肯定是Python文件了，系统必然需要加载才能执行。Airflow提供了一个简单插件管理器，会扫描$AIRFLOW_HOME/plugins加载我们的插件。
所以，我们只需要将写好的插件放入这个目录下就可以了。
插件语法
Operator和Sensor都声明了需要的参数，Operator通过调用execute来执行， sensor通过poke来确认。以Operator为例子。
插件的使用过程为：
dag -> operator -> hook
Hook就是任务执行的具体操作了。
Operator通过继承BaseOperator实现对dag相关属性的绑定， Hook通过继承BaseHook实现对系统配置和资源获取的一些封装。
自定义一个通知插件NotifyOperator
前文https://www.cnblogs.com/woshimrf/p/airflow-dag.html 提到我们通过自定义通知实现多功能任务告警，以下就是一个demo。
文件结构如下：
plugins
│ ├── hooks
│ └── operators
NotifyOperator
首先，在operators目录下创建一个Operator.
# -*- coding: utf-8 -*-
#
from hooks.notify_hook import NotifyHook
from airflow.operators.bash_operator import BaseOperator
class NotifyOperator(BaseOperator):
"""
使用通知服务发送通知
:param message: 内容
:type message: str or dict
:param receivers: 英文逗号分割的罗盘账号
:type receivers: str
:param subject: 邮件主题
:type subject: str
"""
template_fields = ('message', 'subject')
@apply_defaults
def __init__(self,
subject=None,
message=None,
receivers=None,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.message = message
self.receivers = receivers
self.subject = subject
def execute(self, context):
self.log.info('Sending notify message. receivers:{} message:{}'.format(self.receivers, self.message))
hook = NotifyHook(
subject=self.subject,
message=self.message,
receivers=self.receivers
)
hook.send()
- 继承BaseOperator
- 引入NotifyHook， 这个还没创建，等下创建
- template_fields， 想要使用模板变量替换，比如{{ds}}, 字段必须声明到template_fields
- Operator执行的时候会调用execute方法， 这个就是执行的内容
上面可以看出，operator就是接口声明。
NotifyHook
在hooks目录下创建NotifyHook
# -*- coding: utf-8 -*-
#
import json
import requests
from airflow import AirflowException
from airflow.hooks.http_hook import HttpHook
class NotifyHook(HttpHook):
"""
使用通知服务发送通知
:param send_type: 通知类型选填 MAIL,DINGDING,SMS，选填多个时中间用英文逗号隔开
:type send_type: str
:param message: 内容
:type message: str or dict
:param receivers: 英文逗号分割的账号
:type receivers: str
:param subject: 邮件主题
:type subject: str
"""
def __init__(self,
notify_conn_id='notify_default',
send_type='MAIL',
subject=None,
message=None,
receivers=None,
*args,
**kwargs
):
super().__init__(http_conn_id=notify_conn_id, *args, **kwargs)
self.send_type = send_type
self.message = message
self.subject = subject
self.receivers = receivers
def _build_message(self):
"""
构建data
"""
data = {
"content": self.message,
"contentType": "HTML",
"receivers": self.receivers,
"sendType": self.send_type,
"sender": '【Airflow】',
"subject": '【Airflow】' + self.subject
}
return json.dumps(data)
def get_conn(self, headers=None):
"""
Overwrite HttpHook get_conn because just need base_url and headers and
not don't need generic params
:param headers: additional headers to be passed through as a dictionary
:type headers: dict
"""
self.base_url = 'http://notify.ryan-miao.com'
session = requests.Session()
if headers:
session.headers.update(headers)
return session
def send(self):
"""
Send Notify message
"""
data = self._build_message()
self.log.info('Sending message: %s', data)
resp = self.run(endpoint='/api/v2/notify/send',
data=data,
headers={'Content-Type': 'application/json',
'app-id': 'ryan',
'app-key': '123456'})
if int(resp.json().get('retCode')) != 0:
raise AirflowException('Send notify message failed, receive error '
'message %s', resp.text)
self.log.info('Success Send notify message')
- 这里使用的我自己的通知服务api调用。因为是http请求，所以直接继承HttpHook来发送请求就可以了。
- http_conn_id是用来读取数据库中connection里配置的host的，这里直接覆盖，固定我们通知服务的地址。
- 通过抛出异常的方式来终止服务
如何使用
将上面两个文件放到airflow对应的plugins目录下， airflow就自动加载了。然后，当做任务类型使用
from operators.notify_operator import NotifyOperator
notification = NotifyOperator(
task_id="we_are_done",
subject='发送邮件',
message='content',
receivers='ryanmiao'
)
也可以直接执行。比如，我们前面提到任务失败告警可以自定义通知。
from operators.notify_operator import NotifyOperator
def mail_failure_callback(receivers):
"""
失败后邮件通知
:receivers 接收人，多个接收人用英文逗号分开
"""
def mail_back(context):
subject="【执行失败】DAG {} TASK {} ds {}".format(
context['task_instance'].dag_id,
context['task_instance'].task_id,
context['ds'])
message="【执行失败】DAG: {};<br> TASK: {} <br>; ds {} <br>； 原因： {} .<br>"
"查看地址： http://airflow.ryan-miao.com/admin/airflow/tree?dag_id={}"
.format(
context['task_instance'].dag_id,
context['task_instance'].task_id,
context['ds'],
context['exception'],
context['task_instance'].dag_id)
return NotifyOperator(
task_id="mail_failed_notify_callback",
subject=subject,
message=message,
receivers=receivers
).execute(context)
return mail_back
default_args = {
'owner': 'ryanmiao',
'depends_on_past': False,
'start_date': datetime(2019, 5, 1, 9),
'on_failure_callback': mail_failure_callback(receivers='ryanmiao'),
'retries': 0
}
dag = DAG(
'example', default_args=default_args, schedule_interval=None)
自定义一个RDBMS2Hive插件
我们任务调度有个常见的服务是数据抽取到Hive，现在来制作这个插件，可以从关系数据库中读取数据，然后存储到hive。这样，用户只要在airflow配置一下要抽数的database, table和目标hive table就可以实现每天数据入库了。
异构数据传输转换工具很多， 最简单的就是使用原生的dump工具，将数据dump下来，然后import到另一个数据库里。
比如postgres dump
将${sql}查询的列导出到文件${export_data_file}
psql -h$SRC_HOST_IP -U$SRC_USER_NAME -d$SRC_DB -p$SRC_HOST_PORT -c "copy (${sql}) to '${export_data_file}' WITH NULL AS ''"
然后导入hive
LOAD DATA LOCAL INPATH '${export_data_file}' INTO TABLE $TAR_TABLENAME PARTITION (BIZDATE='$BIZ_DATE')
对postgres来说，copy是最快的方案了， 但可能会遇到 ,
等各种转义符号，导出的txt文件或者cvs文件格式就会混乱，需要做对应符号转义处理。
同样， mysql 可以直接把数据查询出来
cat search.sql | mysql -h"$SRC_HOST_IP" -u"$SRC_USER_NAME" -p"$SRC_USER_PWD" -P"$SRC_HOST_PORT" -D"$SRC_DB" --default-character-set=${mysql_charset} -N -s | sed "s/NULL/\\N/ig;s/\\\\n//ig" > result.txt
上述这些命令行的好处就是快，不好的地方在于shell命令的脆弱性和错误处理。最终，选择了集成化的数据转换工具datax. datax是阿里巴巴开源的一款异构数据源同步工具， 虽然看起来不怎么更新了，但简单使用还是可以的。https://github.com/alibaba/DataX
datax的用法相对简单，按照文档配置一下读取数据源和目标数据源，然后执行调用就可以了。可以当做命令行工具来使用。
结合airflow，可以自己实现datax插件。通过读取connections拿到数据源链接配置，然后生成datax的配置文件json，最后调用datax执行。下面是一个从pg或者mysql读取数据，导入hive的插件实现。
主要思路是：
- hdfs创建一个目录
- 生成datax配置文件
- datax执行配置文件，将数据抽取到hdfs
- hive命令行load hdfs
RDBMS2HiveOperator
# -*- coding: utf-8 -*-
#
#
"""
postgres或者mysql 入库到hdfs
"""
import os
import signal
from hooks.rdbms_to_hive_hook import RDBMS2HiveHook
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
class RDBMS2HiveOperator(BaseOperator):
"""
传输pg到hive
https://github.com/alibaba/DataX
:param conn_id: pg连接id
:param query_sql : pg查询语句
:param split_pk : pg分割主键， NONE表示不分割，指定后可以多线程分割，加快传输
:param hive_db : hive的db
:param hive_table: hive的table
:param hive_table_column column数组， column={name:a, type: int} 或者逗号分割的字符串， column=a,b,c
:param hive_table_partition 分区bizdate值
"""
template_fields = ('query_sql', 'hive_db', 'hive_table','hive_table_partition')
ui_color = '#edd5f1'
@apply_defaults
def __init__(self,
conn_id,
query_sql,
hive_db,
hive_table,
hive_table_column,
hive_table_partition,
split_pk=None,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.conn_id = conn_id
self.query_sql = query_sql
self.split_pk = split_pk
self.hive_db = hive_db
self.hive_table = hive_table
self.hive_table_column = hive_table_column
self.hive_table_partition = hive_table_partition
def execute(self, context):
"""
Execute
"""
task_id = context['task_instance'].dag_id + "#" + context['task_instance'].task_id
self.hook = RDBMS2HiveHook(
task_id = task_id,
conn_id = self.conn_id,
query_sql = self.query_sql,
split_pk=self.split_pk,
hive_db=self.hive_db,
hive_table=self.hive_table,
hive_table_column=self.hive_table_column,
hive_table_partition=self.hive_table_partition
)
self.hook.execute(context=context)
def on_kill(self):
self.log.info('Sending SIGTERM signal to bash process group')
os.killpg(os.getpgid(self.hook.sp.pid), signal.SIGTERM)
RDBMS2HiveHook
# -*- coding: utf-8 -*-
#
"""
datax入库hive
"""
import subprocess
import uuid
import json
import os
from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook
class RDBMS2HiveHook(BaseHook):
"""
Datax执行器
"""
def __init__(self,
task_id,
conn_id,
query_sql,
hive_db,
hive_table,
hive_table_column,
hive_table_partition,
split_pk=None):
self.task_id = task_id
self.conn = self.get_connection(conn_id)
self.query_sql = query_sql
self.split_pk = split_pk
self.hive_db = hive_db
self.hive_table = hive_table
self.hive_table_partition = hive_table_partition
self.log.info("Using connection to: {}:{}/{}".format(self.conn.host, self.conn.port, self.conn.schema))
self.hive_table_column = hive_table_column
if isinstance(hive_table_column, str):
self.hive_table_column = []
cl = hive_table_column.split(',')
for item in cl:
hive_table_column_item = {
"name": item,
"type": "string"
}
self.hive_table_column.append(hive_table_column_item)
def Popen(self, cmd, **kwargs):
"""
Remote Popen
:param cmd: command to remotely execute
:param kwargs: extra arguments to Popen (see subprocess.Popen)
:return: handle to subprocess
"""
self.sp = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
**kwargs)
for line in iter(self.sp.stdout):
self.log.info(line.strip().decode('utf-8'))
self.sp.wait()
self.log.info("Command exited with return code %s", self.sp.returncode)
if self.sp.returncode:
raise AirflowException("Execute command failed")
def generate_setting(self):
"""
datax速度等设置
"""
self.setting= {
"speed": {
"byte": 104857600
},
"errorLimit": {
"record": 0,
"percentage": 0.02
}
}
return self.setting
def generate_reader(self):
"""
datax reader
"""
conn_type = 'mysql'
reader_name = 'mysqlreader'
if(self.conn.conn_type == 'postgres'):
conn_type = 'postgresql'
reader_name = 'postgresqlreader'
self.jdbcUrl = "jdbc:"+conn_type+"://"+self.conn.host.strip()+":"+str(self.conn.port)+"/"+ self.conn.schema.strip()
self.reader = {
"name": reader_name,
"parameter": {
"username": self.conn.login.strip(),
"password": self.conn.password.strip(),
"connection": [
{
"querySql": [
self.query_sql
],
"jdbcUrl": [
self.jdbcUrl
]
}
]
}
}
return self.reader
def generate_writer(self):
"""
datax hdafs writer
"""
self.file_type = "text"
self.hdfs_path = "/datax/"+self.hive_db+"/"+self.hive_table+"/"+self.hive_table_partition
self.log.info("临时存储目录：{}".format(self.hdfs_path))
self.writer = {
"name": "hdfswriter",
"parameter": {
"defaultFS": "hdfs://nameservice1",
"hadoopConfig": {
"dfs.nameservices": "nameservice1",
"dfs.ha.automatic-failover.enabled.nameservice1": True,
"ha.zookeeper.quorum": "bigdata2-prod-nn01.ryan-miao.com:2181,bigdata2-prod-nn02.ryan-miao.com:2181,bigdata2-prod-nn03.ryan-miao.com:2181",
"dfs.ha.namenodes.nameservice1": "namenode117,namenode124",
"dfs.namenode.rpc-address.nameservice1.namenode117": "bigdata2-prod-nn01.ryan-miao.com:8020",
"dfs.namenode.servicerpc-address.nameservice1.namenode117": "bigdata2-prod-nn01.ryan-miao.com:8022",
"dfs.namenode.http-address.nameservice1.namenode117": "bigdata2-prod-nn01.ryan-miao.com:50070",
"dfs.namenode.https-address.nameservice1.namenode117": "bigdata2-prod-nn01.ryan-miao.com:50470",
"dfs.namenode.rpc-address.nameservice1.namenode124": "bigdata2-prod-nn02.ryan-miao.com:8020",
"dfs.namenode.servicerpc-address.nameservice1.namenode124": "bigdata2-prod-nn02.ryan-miao.com:8022",
"dfs.namenode.http-address.nameservice1.namenode124": "bigdata2-prod-nn02.ryan-miao.com:50070",
"dfs.namenode.https-address.nameservice1.namenode124": "bigdata2-prod-nn02.ryan-miao.com:50470",
"dfs.replication": 3,
"dfs.client.failover.proxy.provider.nameservice1": "org.apache.hadoop.hdfs.server.namenode.ha.ConfiguredFailoverProxyProvider"
},
"fileType": self.file_type,
"path": self.hdfs_path,
"fileName": self.task_id,
"column": self.hive_table_column,
"writeMode": "nonConflict",
"fieldDelimiter": " "
}
}
return self.writer
def generate_config(self):
content = [{
"reader": self.generate_reader(),
"writer": self.generate_writer()
}]
job = {
"setting": self.generate_setting(),
"content": content
}
config = {
"job": job
}
self.target_json = json.dumps(config)
# write json to file
self.json_file= '/tmp/datax_json_'+self.task_id+ uuid.uuid1().hex
# 打开一个文件
fo = open(self.json_file, "w")
fo.write(self.target_json)
fo.close()
self.log.info("write config json {}".format(self.json_file))
return self.json_file
def execute(self, context):
self.generate_config()
# check hdfs_path
hdfs_path = self.hdfs_path
if(not hdfs_path.startswith('/datax/')):
raise AirflowException("hdfs路径填写错误，不在/datax目录下")
# 创建目录
cmd = ['hadoop', 'fs', '-mkdir', '-p', hdfs_path]
self.Popen(cmd)
# 删除文件
if(not hdfs_path.startswith('/datax/')):
raise AirflowException("hdfs路径填写错误，不在/datax目录下")
files_path = hdfs_path+"/*";
try:
cmd = ['hadoop', 'fs', '-rm', files_path]
self.Popen(cmd)
except Exception:
self.log.info('ignore err, just make sure the dir is clean')
pass
# 上传文件
datax_home = '/data/opt/datax/bin'
cmd = [ 'python', datax_home + '/datax.py', self.json_file]
self.Popen(cmd)
# 删除配置文件
os.remove(self.json_file)
# hive加载
#hive load data from hdfs
hql = "LOAD DATA INPATH '"+ hdfs_path + "' OVERWRITE INTO TABLE "
+ self.hive_db+"."+self.hive_table + " PARTITION (bizdate="+ self.hive_table_partition +")"
cmd = ['hive', '-e', """ + hql + """]
self.Popen(cmd)
如何使用
- admin登录airflow
- 配置connection, 配置pg或者mysql的数据库
- 修改hdfs集群配置信息
- 创建一个DAG
from airflow import DAG
from operators.rdbms_to_hive_operator import RDBMS2HiveOperator
from datetime import datetime, timedelta
from dag_utils import compass_utils
default_args = {
'owner': 'ryanmiao',
'depends_on_past': False,
'start_date': datetime(2019, 5, 1, 9),
'on_failure_callback': compass_utils.failure_callback(dingding_conn_id='dingding_bigdata', receivers='ryanmiao'),
# 'on_success_callback': compass_utils.success_callback(dingding_conn_id='dingding_bigdata', receivers='ryanmiao'),
'retries': 0
}
dag = DAG(
'example_pg2hive', default_args=default_args, schedule_interval=None)
# CREATE TABLE test.pg2hive_test(
# ftime int,
# raw_cp_count int,
# raw_to_delete_cp_count bigint,
# create_time timestamp
# )
# COMMENT '这个是测试datax表'
# PARTITIONED BY (bizdate int)
# ROW FORMAT DELIMITED
# FIELDS TERMINATED BY ' '
# LINES TERMINATED BY '
'
# STORED AS TEXTFILE;
hive_table_column = "ftime,raw_cp_count,raw_to_delete_cp_count,create_time"
t1 = RDBMS2HiveOperator(
task_id='pg2hive',
conn_id='pg_rdb_poi',
query_sql='select ftime, raw_cp_count, raw_to_delete_cp_count, create_time from tbl_poi_report limit 1000',
hive_db='test',
hive_table='pg2hive_test',
hive_table_column=hive_table_column,
hive_table_partition="{{ ds_nodash }}",
dag=dag
)