airflow_prefect.py

from airflow.hooks.base import BaseHook
import pendulum
import requests
import uuid


def get_prefect_event(context):
    prefect_event = [
        {
            "id": str(uuid.uuid4()),
            "occurred": pendulum.now("UTC").to_iso8601_string(),
            "event": "airflow.dag",
            "resource": {
                "prefect.resource.id": f"airflow.dag.{context.get('task_instance').dag_id}",
                "prefect.resource.role": "airflow",
            },
            "related": [
                {
                    "prefect.resource.id": "airflow.dag.*",
                    "prefect.resource.role": "airflow",
                }
            ],
            "payload": {
                "logical_date": context["ds"],
                "log_url": context.get("task_instance").log_url,
                "mark_success_url": context.get("task_instance").mark_success_url,
                "dag": context.get("task_instance").dag_id,
                "dag_run_id": context["run_id"],
                "task": context["task_instance_key_str"],
                "final_state": context.get("task_instance").state,
            },
        }
    ]
    return prefect_event


def send_prefect_event_callback(context):
    prefect_event = get_prefect_event(context)
    conn = BaseHook.get_connection("prefect")
    api_key = conn.password
    api_url = conn.login
    url = f"{api_url}/events"
    requests.post(
        url=url,
        json=prefect_event,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        },
    )