Activa DAG de Cloud Composer con funciones de Cloud Run y la API de REST de Airflow

Cloud Composer 1 | Cloud Composer 2 | Cloud Composer 3

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

Apache Airflow está diseñado para ejecutar DAG de forma periódica, pero también puedes para activar DAG en respuesta a eventos. Una forma de hacer esto es con Cloud Run Functions para activar Los DAG de Cloud Composer cuando ocurre 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.

Antes de comenzar

Verifica la configuración de red de tu entorno

Esta solución no funciona en las configuraciones de IP privada y Controles del servicio de VPC porque no es posible configurar la conectividad de las funciones de Cloud Run al servidor web de Airflow en estas configuraciones.

En Cloud Composer 2, puedes usar otro enfoque: Activa DAG con funciones de Cloud Run y mensajes de Pub/Sub

Habilita las API para tu proyecto.

Console

Enable the Cloud Composer and Cloud Run functions APIs.

Enable the APIs

gcloud

Enable the Cloud Composer and Cloud Run functions APIs: 

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 las llamadas a la API de REST de Airflow con el control de acceso al servidor web

Las funciones de Cloud Run pueden 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 funciones de Cloud Run.

Cree 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.

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. Usas la URL del servidor web de Airflow en el código de 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 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 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 Run o Cloud Run. En este instructivo, se muestra una función de Cloud Functions 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 Sección Configuración del entorno de ejecución, la compilación, las conexiones y la seguridad. Usa uno de las siguientes opciones, según tus preferencias:

    • Selecciona Cuenta de servicio predeterminada de Compute Engine. Con los permisos de IAM predeterminadas, esta cuenta puede ejecutar funciones que acceden a 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 Código:

    • (Python) Cuando agregues código para esto Por ejemplo, selecciona el entorno de ejecución de 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.32.2

Java

Agrega las siguientes dependencias a la sección dependencies en el pom.xml que genera 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

Agrega el código para activar los DAG con la API de REST de Airflow. Crea un archivo llamado composer2_airflow_rest_api.py y agrega el código para crear la API de REST de Airflow. llamadas a este archivo.

No cambies ninguna variable. La función de Cloud Functions importa este archivo desde el archivo main.py.

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

Coloca el siguiente código en el archivo Example.java. Reemplaza el valor del La variable webServerUrl con la dirección del servidor web de Airflow que que obtuviste anteriormente.


// 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 la función se implemente.
  2. Sube un archivo a tu bucket de Cloud Storage. Como alternativa, puedes activar la función de forma manual si seleccionas la acción Probar la función 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?