Airflow:如何从 PostgreOperator 推送 xcom 值?
Airflow: How to push xcom value from PostgreOperator?
我正在使用 Airflow 1.8.1,我想推送来自 PostgreOperator 的 sql 请求的结果。
这是我的任务:
check_task = PostgresOperator(
task_id='check_task',
postgres_conn_id='conx',
sql="check_task.sql",
xcom_push=True,
dag=dag)
def py_is_first_execution(**kwargs):
value = kwargs['ti'].xcom_pull(task_ids='check_task')
print 'count ----> ', value
if value == 0:
return 'next_task'
else:
return 'end-flow'
check_branch = BranchPythonOperator(
task_id='is-first-execution',
python_callable=py_is_first_execution,
provide_context=True,
dag=dag)
这是我的 sql 脚本:
select count(1) from table
当我从 check_task
检查 xcom 值时,它检索 none
值。
如果我是正确的,airflow 会在查询 return 值时自动推送到 xcom。但是,当您查看 postgresoperator 的代码时,您会发现它有一个调用 PostgresHook(dbapi_hook 的扩展名)的 运行 方法的执行方法。这两种方法都不 return 任何东西,因此它不会向 xcom 推送任何东西。
我们为解决此问题所做的工作是创建一个 CustomPostgresSelectOperator,它是 PostgresOperator 的副本,但不是 'hook.run(..)' 而是 'return hook.get_records(..)'.
希望对你有所帮助。
最后,我在 $AIRFLOW_HOME/plugins
下的插件管理器中创建了一个新的传感器 ExecuteSqlOperator
。
我以CheckOperator
为例,修改了返回值:这个运算符的基本运行和我需要的完全相反
这是默认值 ExecuteSqlOperator
:
CheckOperator
这是我定制的 SqlSensor
:ReverseSqlSensor
class SqlExecuteOperator(BaseOperator):
"""
Performs checks against a db. The ``CheckOperator`` expects
a sql query that will return a single row.
Note that this is an abstract class and get_db_hook
needs to be defined. Whereas a get_db_hook is hook that gets a
single record from an external source.
:param sql: the sql to be executed
:type sql: string
"""
template_fields = ('sql',)
template_ext = ('.hql', '.sql',)
ui_color = '#fff7e6'
@apply_defaults
def __init__(
self, sql,
conn_id=None,
*args, **kwargs):
super(SqlExecuteOperator, self).__init__(*args, **kwargs)
self.conn_id = conn_id
self.sql = sql
def execute(self, context=None):
logging.info('Executing SQL statement: ' + self.sql)
records = self.get_db_hook().get_first(self.sql)
logging.info("Record: " + str(records))
records_int = int(records[0])
print (records_int)
return records_int
def get_db_hook(self):
return BaseHook.get_hook(conn_id=self.conn_id)
我正在使用 Airflow 1.8.1,我想推送来自 PostgreOperator 的 sql 请求的结果。
这是我的任务:
check_task = PostgresOperator(
task_id='check_task',
postgres_conn_id='conx',
sql="check_task.sql",
xcom_push=True,
dag=dag)
def py_is_first_execution(**kwargs):
value = kwargs['ti'].xcom_pull(task_ids='check_task')
print 'count ----> ', value
if value == 0:
return 'next_task'
else:
return 'end-flow'
check_branch = BranchPythonOperator(
task_id='is-first-execution',
python_callable=py_is_first_execution,
provide_context=True,
dag=dag)
这是我的 sql 脚本:
select count(1) from table
当我从 check_task
检查 xcom 值时,它检索 none
值。
如果我是正确的,airflow 会在查询 return 值时自动推送到 xcom。但是,当您查看 postgresoperator 的代码时,您会发现它有一个调用 PostgresHook(dbapi_hook 的扩展名)的 运行 方法的执行方法。这两种方法都不 return 任何东西,因此它不会向 xcom 推送任何东西。 我们为解决此问题所做的工作是创建一个 CustomPostgresSelectOperator,它是 PostgresOperator 的副本,但不是 'hook.run(..)' 而是 'return hook.get_records(..)'.
希望对你有所帮助。
最后,我在 $AIRFLOW_HOME/plugins
下的插件管理器中创建了一个新的传感器 ExecuteSqlOperator
。
我以CheckOperator
为例,修改了返回值:这个运算符的基本运行和我需要的完全相反
这是默认值 ExecuteSqlOperator
:
CheckOperator
这是我定制的 SqlSensor
:ReverseSqlSensor
class SqlExecuteOperator(BaseOperator):
"""
Performs checks against a db. The ``CheckOperator`` expects
a sql query that will return a single row.
Note that this is an abstract class and get_db_hook
needs to be defined. Whereas a get_db_hook is hook that gets a
single record from an external source.
:param sql: the sql to be executed
:type sql: string
"""
template_fields = ('sql',)
template_ext = ('.hql', '.sql',)
ui_color = '#fff7e6'
@apply_defaults
def __init__(
self, sql,
conn_id=None,
*args, **kwargs):
super(SqlExecuteOperator, self).__init__(*args, **kwargs)
self.conn_id = conn_id
self.sql = sql
def execute(self, context=None):
logging.info('Executing SQL statement: ' + self.sql)
records = self.get_db_hook().get_first(self.sql)
logging.info("Record: " + str(records))
records_int = int(records[0])
print (records_int)
return records_int
def get_db_hook(self):
return BaseHook.get_hook(conn_id=self.conn_id)