Activa los DAG de Cloud Composer con Cloud Functions y la API de REST de Airflow

Cloud Composer 1 | Cloud Composer 2

En esta página, se describe cómo usar Cloud Functions para activar los DAG de Cloud Composer en respuesta a eventos.

Nota: También puedes activar los DAG con solo la API de REST de Airflow.

Apache Airflow está diseñado para ejecutar DAG de forma periódica, pero también puedes activar los DAG en respuesta a eventos. Una forma de hacerlo es usar Cloud Functions para activar los DAG de Cloud Composer cuando se produce un evento especificado.

En el ejemplo de esta guía, se ejecuta un DAG cada vez que se produce un cambio en un bucket de Cloud Storage. Los cambios en cualquier objeto de un bucket activan una función. Esta función realiza una solicitud a la API de REST de Airflow del entorno de Cloud Composer. Airflow procesa esta solicitud y ejecuta un DAG. El DAG genera información sobre el cambio.

Nota: En una guía alternativa, se muestra cómo crear una función que active un DAG cuando se envíe un mensaje a un tema de Pub/Sub.

Antes de comenzar

Verifica la configuración de red de tu entorno

Esta solución no funciona en las configuraciones de IP privada ni de los Controles del servicio de VPC, ya que no es posible configurar la conectividad de Cloud Functions al servidor web de Airflow con estos parámetros de configuración.

En Cloud Composer 2, puedes usar otro enfoque: Activar DAG mediante Cloud Functions y mensajes de Pub/Sub

Habilita las API para tu proyecto.

Console

Habilita las API de Cloud Composer and Cloud Functions.

Habilita las API

gcloud

Habilita las APIs de Cloud Composer and Cloud Functions:

gcloud services enable cloudfunctions.googleapis.com composer.googleapis.com

Habilita la API de REST de Airflow

Para Airflow 2, la API de REST estable ya está habilitada de forma predeterminada. Si tu entorno tiene inhabilitada la API estable, habilita la API estable de REST.

Permitir llamadas a la API de REST de Airflow con el control de acceso de Webserver

Cloud Functions puede comunicarse con la API de REST de Airflow con una dirección IPv4 o IPv6.

Si no estás seguro de cuál será el rango de IP de llamada, usa una opción de configuración predeterminada en el Control de acceso del servidor web que es All IP addresses have access (default) para no bloquear accidentalmente tus Cloud Functions.

Crea un bucket de Cloud Storage

Este ejemplo activa un DAG en respuesta a los cambios en un bucket de Cloud Storage. Crea un bucket nuevo para usar en este ejemplo.

Nota: No uses depósitos de entorno de Cloud Composer, ya que esto podría provocar que la función se active varias veces por segundo. Para los fines de este ejemplo, recomendamos crear un bucket distinto.

Obtén la URL del servidor web de Airflow

En este ejemplo, se realizan solicitudes a la API de REST al extremo del servidor web de Airflow. Debes usar la URL del servidor web de Airflow en el código de tu Cloud Function.

Console

1. En la consola de Google Cloud, ve a la página Entornos.

   Ir a Entornos

2. Haz clic en el nombre de tu entorno.

3. En la página Detalles del entorno, ve a la pestaña Configuración del entorno.

4. La URL del servidor web de Airflow aparece en el elemento de la IU web de Airflow.

gcloud

Ejecuta el siguiente comando:

gcloud composer environments describe ENVIRONMENT_NAME \
    --location LOCATION \
    --format='value(config.airflowUri)'

Reemplaza lo siguiente:

• ENVIRONMENT_NAME por el nombre del entorno.
• LOCATION por la región en la que se encuentra el entorno

Sube un DAG a tu entorno

Sube un DAG a tu entorno. En el siguiente ejemplo de DAG, se muestra la configuración de ejecución del DAG recibido. Debes activar este DAG desde una función que crearás más adelante en esta guía.

import datetime

import airflow
from airflow.operators.bash_operator import BashOperator

with airflow.DAG(
    "composer_sample_trigger_response_dag",
    start_date=datetime.datetime(2021, 1, 1),
    # Not scheduled, trigger only
    schedule_interval=None,
) as dag:
    # Print the dag_run's configuration, which includes information about the
    # Cloud Storage object change.
    print_gcs_info = BashOperator(
        task_id="print_gcs_info", bash_command="echo {{ dag_run.conf }}"
    )

Implementa una función de Cloud Functions que active el DAG

Puedes implementar una Cloud Function con tu lenguaje preferido compatible con Cloud Functions o Cloud Run. En este instructivo, se muestra una Cloud Function implementada en Python y Java.

Especifica los parámetros de configuración de Cloud Functions

• Activador Para este ejemplo, selecciona un activador que funcione cuando se cree un objeto nuevo en un bucket o se reemplace un objeto existente.

  • Tipo de activador Cloud Storage

  • Tipo de evento Finalizar/Crear

  • Bucket Selecciona un bucket que debe activar esta función.

  • Volver a intentar en caso de error Te recomendamos inhabilitar esta opción para los fines de este ejemplo. Si usas tu propia función en un entorno de producción, habilita esta opción para controlar errores transitorios.

• Cuenta de servicio del entorno de ejecución, en la sección Configuración del entorno de ejecución, la compilación, las conexiones y la seguridad Usa una de las siguientes opciones, según tus preferencias:

  • Selecciona Cuenta de servicio predeterminada de Compute Engine. Con los permisos predeterminados de IAM, esta cuenta puede ejecutar funciones que acceden a los entornos de Cloud Composer.

  • Crea una cuenta de servicio personalizada que tenga la función de usuario de Composer y especifícala como una cuenta de servicio del entorno de ejecución para esta función. Esta opción sigue el principio de privilegio mínimo.

• Entorno de ejecución y punto de entrada, en el paso Code:

  • (Python) Cuando agregues código para este ejemplo, selecciona el entorno de ejecución Python 3.7 o posterior, y especifica trigger_dag_gcf como punto de entrada.

  • (Java) Cuando agregues código para este ejemplo, selecciona el entorno de ejecución de Java 17 y especifica com.example.Example como punto de entrada.

Agrega requisitos

Python

Especifica las dependencias en el archivo requirements.txt:

google-auth==2.19.1
requests==2.31.0

Java

Agrega las siguientes dependencias a la sección dependencies en la pom.xml generada por la IU de Google Cloud Functions.

    <dependency>
      <groupId>com.google.apis</groupId>
      <artifactId>google-api-services-docs</artifactId>
      <version>v1-rev20210707-1.32.1</version>
    </dependency>
    <dependency>
      <groupId>com.google.api-client</groupId>
      <artifactId>google-api-client</artifactId>
      <version>1.32.1</version>
    </dependency>
    <dependency>
      <groupId>com.google.auth</groupId>
      <artifactId>google-auth-library-credentials</artifactId>
      <version>1.14.0</version>
    </dependency>
    <dependency>
      <groupId>com.google.auth</groupId>
      <artifactId>google-auth-library-oauth2-http</artifactId>
      <version>1.14.0</version>
    </dependency>

Python

Agregar el código para activar los DAG con la API de REST de Airflow Crea un archivo llamado composer2_airflow_rest_api.py y coloca el código para realizar llamadas a la API de REST de Airflow en este archivo.

No modifiques ninguna variable. La Cloud Function importa este archivo desde el archivo main.py.

Nota: A veces, las cuentas de servicio tienen direcciones de correo electrónico de más de 64 caracteres. Para acceder a la API de REST de Airflow con una cuenta de servicio de este tipo, usa la solución alternativa descrita en Accede a la API de REST de Airflow con una cuenta de servicio.

from __future__ import annotations

from typing import Any

import google.auth
from google.auth.transport.requests import AuthorizedSession
import requests

# Following GCP best practices, these credentials should be
# constructed at start-up time and used throughout
# https://cloud.google.com/apis/docs/client-libraries-best-practices
AUTH_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
CREDENTIALS, _ = google.auth.default(scopes=[AUTH_SCOPE])

def make_composer2_web_server_request(
    url: str, method: str = "GET", **kwargs: Any
) -> google.auth.transport.Response:
    """
    Make a request to Cloud Composer 2 environment's web server.
    Args:
      url: The URL to fetch.
      method: The request method to use ('GET', 'OPTIONS', 'HEAD', 'POST', 'PUT',
        'PATCH', 'DELETE')
      **kwargs: Any of the parameters defined for the request function:
                https://github.com/requests/requests/blob/master/requests/api.py
                  If no timeout is provided, it is set to 90 by default.
    """

    authed_session = AuthorizedSession(CREDENTIALS)

    # Set the default timeout, if missing
    if "timeout" not in kwargs:
        kwargs["timeout"] = 90

    return authed_session.request(method, url, **kwargs)

def trigger_dag(web_server_url: str, dag_id: str, data: dict) -> str:
    """
    Make a request to trigger a dag using the stable Airflow 2 REST API.
    https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html

    Args:
      web_server_url: The URL of the Airflow 2 web server.
      dag_id: The DAG ID.
      data: Additional configuration parameters for the DAG run (json).
    """

    endpoint = f"api/v1/dags/{dag_id}/dagRuns"
    request_url = f"{web_server_url}/{endpoint}"
    json_data = {"conf": data}

    response = make_composer2_web_server_request(
        request_url, method="POST", json=json_data
    )

    if response.status_code == 403:
        raise requests.HTTPError(
            "You do not have a permission to perform this operation. "
            "Check Airflow RBAC roles for your account."
            f"{response.headers} / {response.text}"
        )
    elif response.status_code != 200:
        response.raise_for_status()
    else:
        return response.text

Ingresa el siguiente código en el archivo main.py. Reemplaza el valor de la variable web_server_url por la dirección del servidor web de Airflow que obtuviste antes.

# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Trigger a DAG in a Cloud Composer 2 environment in response to an event,
using Cloud Functions.
"""

from typing import Any

import composer2_airflow_rest_api

def trigger_dag_gcf(data, context=None):
    """
    Trigger a DAG and pass event data.

    Args:
      data: A dictionary containing the data for the event. Its format depends
      on the event.
      context: The context object for the event.

    For more information about the arguments, see:
    https://cloud.google.com/functions/docs/writing/background#function_parameters
    """

    # TODO(developer): replace with your values
    # Replace web_server_url with the Airflow web server address. To obtain this
    # URL, run the following command for your environment:
    # gcloud composer environments describe example-environment \
    #  --location=your-composer-region \
    #  --format="value(config.airflowUri)"
    web_server_url = (
        "https://example-airflow-ui-url-dot-us-central1.composer.googleusercontent.com"
    )
    # Replace with the ID of the DAG that you want to run.
    dag_id = 'composer_sample_trigger_response_dag'

    composer2_airflow_rest_api.trigger_dag(web_server_url, dag_id, data)

Java

Ingresa el siguiente código en el archivo Example.java. Reemplaza el valor de la variable webServerUrl por la dirección del servidor web de Airflow que obtuviste antes.

// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.example;

import com.example.Example.GcsEvent;
import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpContent;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpRequestFactory;
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.javanet.NetHttpTransport;
import com.google.api.client.http.json.JsonHttpContent;
import com.google.api.client.json.gson.GsonFactory;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.functions.BackgroundFunction;
import com.google.cloud.functions.Context;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Logger;

/**
 * Cloud Function that triggers an Airflow DAG in response to an event (in
 * this case a Cloud Storage event).
 */
public class Example implements BackgroundFunction<GcsEvent> {
  private static final Logger logger = Logger.getLogger(Example.class.getName());

  // TODO(developer): replace with your values
  // Replace webServerUrl with the Airflow web server address. To obtain this
  // URL, run the following command for your environment:
  // gcloud composer environments describe example-environment \
  //  --location=your-composer-region \
  //  --format="value(config.airflowUri)"
  @Override
  public void accept(GcsEvent event, Context context) throws Exception {
    String webServerUrl = "https://example-airflow-ui-url-dot-us-central1.composer.googleusercontent.com";
    String dagName = "composer_sample_trigger_response_dag";
    String url = String.format("%s/api/v1/dags/%s/dagRuns", webServerUrl, dagName);

    logger.info(String.format("Triggering DAG %s as a result of an event on the object %s.",
      dagName, event.name));
    logger.info(String.format("Triggering DAG via the following URL: %s", url));

    GoogleCredentials googleCredentials = GoogleCredentials.getApplicationDefault()
        .createScoped("https://www.googleapis.com/auth/cloud-platform");
    HttpCredentialsAdapter credentialsAdapter = new HttpCredentialsAdapter(googleCredentials);
    HttpRequestFactory requestFactory =
      new NetHttpTransport().createRequestFactory(credentialsAdapter);

    Map<String, Map<String, String>> json = new HashMap<String, Map<String, String>>();
    Map<String, String> conf = new HashMap<String, String>();
    conf.put("bucket", event.bucket);
    conf.put("name", event.name);
    conf.put("generation", event.generation);
    conf.put("operation", context.eventType());
    json.put("conf", conf);
    HttpContent content = new JsonHttpContent(new GsonFactory(), json);
    HttpRequest request = requestFactory.buildPostRequest(new GenericUrl(url), content);
    request.getHeaders().setContentType("application/json");
    HttpResponse response;
    try {
      response = request.execute();
      int statusCode = response.getStatusCode();
      logger.info("Response code: " + statusCode);
      logger.info(response.parseAsString());
    } catch (HttpResponseException e) {
      // https://cloud.google.com/java/docs/reference/google-http-client/latest/com.google.api.client.http.HttpResponseException
      logger.info("Received HTTP exception");
      logger.info(e.getLocalizedMessage());
      logger.info("- 400 error: wrong arguments passed to Airflow API");
      logger.info("- 401 error: check if service account has Composer User role");
      logger.info("- 403 error: check Airflow RBAC roles assigned to service account");
      logger.info("- 404 error: check Web Server URL");
    } catch (Exception e) {
      logger.info("Received exception");
      logger.info(e.getLocalizedMessage());
    }
  }

  /** Details of the storage event. */
  public static class GcsEvent {
    /** Bucket name. */
    String bucket;
    /** Object name. */
    String name;
    /** Object version. */
    String generation;
  }
}

Prueba la función

Para verificar que tu función y DAG funcionen según lo previsto, haz lo siguiente:

1. Espera hasta que se implemente la función.
2. Sube un archivo a tu bucket de Cloud Storage. Como alternativa, puedes activar la función de forma manual. Para ello, selecciona la acción Probar la función para ella en la consola de Google Cloud.
3. Consulta la página del DAG en la interfaz web de Airflow. El DAG debe tener una DAG activa o ya completada.
4. En la IU de Airflow, verifica los registros de tareas de esta ejecución. Deberías ver que la tarea print_gcs_info genera los datos recibidos de la función a los registros:

Python

[2021-04-04 18:25:44,778] {bash_operator.py:154} INFO - Output:
[2021-04-04 18:25:44,781] {bash_operator.py:158} INFO - Triggered from GCF:
    {bucket: example-storage-for-gcf-triggers, contentType: text/plain,
    crc32c: dldNmg==, etag: COW+26Sb5e8CEAE=, generation: 1617560727904101,
    ... }
[2021-04-04 18:25:44,781] {bash_operator.py:162} INFO - Command exited with
    return code 0h

Java

[2023-02-08, 08:00:09 UTC] {subprocess.py:86} INFO - Output:
[2023-02-08, 08:00:09 UTC] {subprocess.py:93} INFO - {bucket: example-storage-for-gcf-triggers, generation: 1675843189006715, name: file.txt, operation: google.storage.object.create}
[2023-02-08, 08:00:09 UTC] {subprocess.py:97} INFO - Command exited with return code 0

¿Qué sigue?

• Accede a la IU de Airflow
• Accede a la API de REST de Airflow
• Escribir DAG
• Escribe Cloud Functions
• Activadores de Cloud Storage