Question

我正在尝试在许多DAG中导入脚本，以多次调用同一操作。应用这种解决方案的最佳方法是什么？

现在我的文件夹结构为：

dags/
|-- some_dags_folder/
|---- some_dag.py
|-- other_dags_folder/
|---- another_dag.py
|-- utils/
|---- util_slack.py

当我尝试导入util_slack文件时，将以下内容放入DAG代码中，例如，假设该代码来自some_dag.py：

from ..utils.util_slack import some_function

将所有内容放入Airflow之后，出现以下错误：

Broken DAG: [/usr/local/airflow/dags/some_dags_folder/some_dag.py] attempted relative import with no known parent package

util_slack脚本是发送成功消息或失败消息的文件，看起来像这样

from airflow.contrib.operators.slack_webhook_operator import SlackWebhookOperator
from airflow.hooks.base_hook import BaseHook

CHANNEL = BaseHook.get_connection('Slack').login
TOKEN = BaseHook.get_connection('Slack').password

def slack_success(context):
    ...
    alterHook = SlackWebhookOperator(...)
    return alterHook.execut(context=context)

def slack_fail(context):
    ...
    alterHook = SlackWebhookOperator(...)
    return alterHook.execut(context=context)

这个想法是，我可以将util_slack模块或任何其他自制模块导入多个DAG，并以

的方式调用我需要的功能。

...
from ..utils.util_slack import slack_success

...

def task_success(context):
    return slack_success(context)
...
some_task_in_dag = SSHOperator(
    ...
    on_success_callback=task_success
    ...)

这是最好的方法还是创建像https://airflow.apache.org/plugins.html所示的自定义插件会更好？

Answer 1

不确定plugins是否适合您的情况。 Plugins将外部功能集成到Airflow核心（例如自定义端点，自定义登录/身份验证等）。

下面是我的方法。目前，我有许多与ClickHouse一起使用的任务。因此，我需要在不同的DAG's中获得连接/截断/插入/复制/等。结构示例：

 dags
    ├── lib  # you can choose any your favorite name(utils, tools etc)
    │   ├── ... just another common package / module
    │   ├── default.py
    │   ├── configurator.py
    │   └── telegram.py
    └── # dag1, dag2...dag_n

default.py -只是默认的DAG参数

from lib.telegram import send_message

def on_success_callback(context):
    pass


def on_failure_callback(context):
    config = get_main_config()
    if not config.get('NOTIFY_ON_FAILURE'):
        return
    send_message('failed blabla')


def get_main_config():
    # I use variable with key 'MAIN_CONFIG' to store some common settings for all dags
    return Variable.get('MAIN_CONFIG', deserialize_json=True)


def get_default_args():
    return {
        'email_on_failure': False,
        'email_on_retry': False,
        'on_failure_callback': on_failure_callback,
        'on_success_callback': on_success_callback,
        # etc...
    }

configurator.py -所有必需的初始化都集中在一个位置。我使用inject，但是您可以使用任何工具/方法，这只是一个例子。

from lib.default import get_main_config
from airflow.hooks.base_hook import BaseHook


class InstancesPool:
    def __init__(self, slack_connection, db_connection):
        self._db_connection = db_connection
        self._slack_connection = slack_connection

    def get_slack_connection(self):
        return self._slack_connection

    def get_db_connection():
        return self._db_connection


class DbConnection:
    # just an example
    def __init__(self, user, password):
        pass


def configure():
    config = get_main_config()

    return InstancesPool(
        BaseHook.get_connection('Slack'),
        DbConnection(config['DB_USER'], config['DB_PASSWORD'])
    )

这样，您将不会在导入或初始化方面遇到问题。您只需致电：

from lib.configurator import configure


def my_task(ds, **kwargs):
    pool = configure()
    # pool.get_slack_connection() etc...

希望这会有所帮助。

如何在不同的DAG中使用通用脚本？

1 个答案: