airflow系列之任务远程触发调用

1. 目标实现

我们希望通过消息队列或者API方式触发调用Airflow任务

2. 鉴权方式更改

通过如下命令查看

airflow config get-value api auth_backends

如果是session，在values.yaml中改为Basic Auth

airflowVersion: 2.5.3
defaultAirflowTag: 2.5.3.2
config:
  core:
    dags_folder: /opt/airflow/dags/repo/dags
    hostname_callable: airflow.utils.net.get_host_ip_address
  api:
    # 新增该配置
    auth_backends: airflow.api.auth.backend.basic_auth

3. 实现代码

/**
 * Airflow 工具类（使用 Basic Auth）
 */
public class AirflowUtils {
    private static final String AIRFLOW_URL = "https://airflow.sfxs.com";
    private static final String USERNAME = "admin";
    private static final String PASSWORD = "公众号算法小生";

    /**
     * 触发 Airflow DAG 运行（使用 Basic Auth）
     * @param dagId      DAG 的 ID
     * @param conf       DAG 参数，JSON 格式
     * @param dagRunId   DAG 运行的唯一标识（可选，建议包含时间戳）
     * @return 触发结果（JSON 字符串）
     * @throws RuntimeException 如果请求失败
     */
    public static String runDag(String dagId, Map<String, Object> conf, String dagRunId) {
        String url = AIRFLOW_URL + "/api/v1/dags/" + dagId + "/dagRuns";

        // 设置 Basic Auth
        HttpHeaders headers = new HttpHeaders();
        String credentials = USERNAME + ":" + PASSWORD;
        String encodedCredentials = Base64.getEncoder().encodeToString(credentials.getBytes());
        headers.set("Authorization", "Basic " + encodedCredentials);
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 构造请求体
        Map<String, Object> body = new JSONObject();
        body.put("conf", conf);
        if (StringUtils.isNotBlank(dagRunId)) {
            body.put("dag_run_id", dagRunId);
        }

        try {
            return RequestUtils.postForEntity(url, body, headers);
        } catch (Exception e) {
            throw new RuntimeException("Failed to trigger DAG: " + dagId + ", Error: " + e.getMessage(), e);
        }
    }

    /**
     * 查询指定 DAG（使用 Basic Auth）
     *
     * @param dagId DAG 的 ID（精确匹配）
     * @return DAG 详情（JSON 字符串）
     * @throws RuntimeException 如果请求失败
     */
    public static String getDag(String dagId) {
        String url = AIRFLOW_URL + "/api/v1/dags";

        // 设置 Basic Auth
        HttpHeaders headers = new HttpHeaders();
        String credentials = USERNAME + ":" + PASSWORD;
        String encodedCredentials = Base64.getEncoder().encodeToString(credentials.getBytes());
        headers.set("Authorization", "Basic " + encodedCredentials);
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 设置查询参数 dag_id_pattern
        Map<String, String> params = new HashMap<>();
        params.put("dag_id_pattern", dagId);

        try {
            return RequestUtils.getContent(url, params, headers);
        } catch (Exception e) {
            throw new RuntimeException("Failed to get DAG: " + dagId + ", Error: " + e.getMessage(), e);
        }
    }
}

4. 测试验证

@Test
public void testAirflow() {
    //AirflowUtils.getDag("test_failure_notification");
    //AirflowUtils.runDag("test_failure_notification", null, null);
    String message = "{\"dag_id\":\"community_summary_dag\",\"conf\":{\"city\":\"临沂市\"}}";
    JSONObject jsonObject = JSONObject.parseObject(message);
    String dagId = jsonObject.getString("dag_id");
    JSONObject conf = jsonObject.getJSONObject("conf");

    if (StringUtils.isNotBlank(dagId)) {
        AirflowUtils.runDag(dagId, conf, null);
    }
}

OK, 运行后可以airflow中看到任务正在运行了~