Merge pull request #79 from godatadriven/kg-deferrable-operators

Add deferrable operator examples

Author: krisgeus

.whirl.env

@@ -1,3 +1,4 @@
 AIRFLOW_VERSION=2.3.2
 PYTHON_VERSION=3.9
 AIRFLOW__API__AUTH_BACKEND=airflow.api.auth.backend.basic_auth
+MINIMAL_AIRFLOW_VERSION=2.2.5

docker/airflow-python/entrypoint.sh

@@ -17,9 +17,9 @@ if [[ ${AIRFLOW_COMMAND} == "scheduler" || ${AIRFLOW_COMMAND} == "singlemachine"
   echo "y" | airflow db reset
   airflow users create --username admin --password admin --firstname Anonymous --lastname Admin --role Admin --email admin@example.org
 else
-  if [[ ${AIRFLOW_COMMAND} == "webserver" ]]; then
+  if [[ ${AIRFLOW_COMMAND} == "webserver" || ${AIRFLOW_COMMAND} == "triggerer" ]]; then
     echo "wait a bit more to let the scheduler do the database reset."
-    sleep 15
+    sleep 30
   fi
 fi
 
@@ -41,8 +41,7 @@ done
 
 if [[ ${AIRFLOW_COMMAND} == "singlemachine" ]]; then
   nohup /entrypoint scheduler -D &
-  # echo  "wait a while for the scheduler to be started"
-  # sleep 15
+  nohup /entrypoint triggerer -D &
   /entrypoint webserver -p 5000
 else
   /entrypoint "${@}"

envs/airflow-with-mockserver/.whirl.env

@@ -0,0 +1,14 @@
+# postgres env vars
+POSTGRES_HOST=postgresdb
+POSTGRES_PORT=5432
+POSTGRES_PASSWORD=pAssw0rd
+POSTGRES_USER=airflow
+POSTGRES_DB=airflow
+
+# Airflow variables
+AIRFLOW__CORE__EXPOSE_CONFIG=True
+AIRFLOW__WEBSERVER__EXPOSE_CONFIG=True
+AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
+AIRFLOW__CORE__LOAD_EXAMPLES=False
+AIRFLOW__CORE__FERNET_KEY=W5gmA+dp84hkZEzpxPw4LTmhbXA1uVxKZsgIfay8wno=
+AIRFLOW__CORE__SQL_ALCHEMY_CONN=postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${POSTGRES_HOST}:${POSTGRES_PORT}/${POSTGRES_DB}

envs/airflow-with-mockserver/docker-compose.yml

@@ -0,0 +1,75 @@
+version: '3'
+
+services:
+  webserver:
+    image: docker-whirl-airflow:py-${PYTHON_VERSION}-local
+    command: ["webserver", "-p", "5000"]
+    ports:
+      - '5000:5000'  # HTTP (Airflow Web UI)
+    env_file:
+      - .whirl.env
+    environment:
+      - WHIRL_SETUP_FOLDER
+      - AIRFLOW__API__AUTH_BACKEND
+    volumes:
+      - ${DAG_FOLDER}:/opt/airflow/dags
+      - ${ENVIRONMENT_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/env.d/
+      - ${DAG_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/dag.d/
+    depends_on:
+      - mockserver
+      - postgresdb
+
+  scheduler:
+    image: docker-whirl-airflow:py-${PYTHON_VERSION}-local
+    command: ["scheduler"]
+    env_file:
+      - .whirl.env
+    environment:
+      - WHIRL_SETUP_FOLDER
+      - AIRFLOW__API__AUTH_BACKEND
+    volumes:
+      - ${DAG_FOLDER}:/opt/airflow/dags
+      - ${ENVIRONMENT_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/env.d/
+      - ${DAG_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/dag.d/
+    depends_on:
+      - mockserver
+      - postgresdb
+
+  triggerer:
+    image: docker-whirl-airflow:py-${PYTHON_VERSION}-local
+    command: ["triggerer"]
+    env_file:
+      - .whirl.env
+    environment:
+      - WHIRL_SETUP_FOLDER
+      - AIRFLOW__API__AUTH_BACKEND
+    volumes:
+      - ${DAG_FOLDER}:/opt/airflow/dags
+      - ${ENVIRONMENT_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/env.d/
+      - ${DAG_FOLDER}/whirl.setup.d:${WHIRL_SETUP_FOLDER}/dag.d/
+    depends_on:
+      - mockserver
+      - postgresdb
+
+
+  mockserver:
+    image: jamesdbloom/mockserver:mockserver-5.13.2
+    ports:
+      - 1080:1080
+      - 1081:1081
+    environment:
+      - LOG_LEVEL=ERROR
+      - SERVER_PORT=1080,1081
+
+  postgresdb:
+    image: postgres:13
+    ports:
+      - 5432:5432
+    environment:
+      - POSTGRES_HOST=postgresdb
+      - POSTGRES_PORT
+      - POSTGRES_PASSWORD
+      - POSTGRES_USER
+      - POSTGRES_DB
+
+

envs/airflow-with-mockserver/whirl.setup.d/.gitkeep

@@ -0,0 +1 @@
+.gitkeep

examples/airflow-deferrable-operator-custom/.airflowignore

@@ -0,0 +1,2 @@
+whirl.setup.d
+.whirl.env

examples/airflow-deferrable-operator-custom/.whirl.env

@@ -0,0 +1,3 @@
+WHIRL_ENVIRONMENT=airflow-with-mockserver
+AIRFLOW__LOGGING__LOGGING_LEVEL=DEBUG
+MINIMAL_AIRFLOW_VERSION=2.3.0

examples/airflow-deferrable-operator-custom/custom/__init__.py

@@ -0,0 +1,15 @@
+from datetime import timedelta
+
+from airflow.sensors.base import BaseSensorOperator
+
+from custom.triggers.api_check_trigger import ApiCheckTrigger
+
+
+class WaitForStartedStatusSensor(BaseSensorOperator):
+    def execute(self, context):
+        self.defer(trigger=ApiCheckTrigger(url="http://mockserver:1080/testapi", expected_status="Started"), method_name="execute_complete", timeout=timedelta(minutes=2))
+
+    def execute_complete(self, context, event=None):
+        # We have no more work to do here. Mark as complete.
+        print(event)
+        return

examples/airflow-deferrable-operator-custom/custom/operators/__init__.py

@@ -0,0 +1,33 @@
+import aiohttp
+import asyncio
+from typing import Any, Dict, Tuple
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+class ApiCheckTrigger(BaseTrigger):
+    """
+    A trigger that calls a Rest API until a certain status is returned.
+    """
+
+    def __init__(self, url:str, expected_status: int):
+        super().__init__()
+        self.url = url
+        self.expected_status = expected_status
+
+    def serialize(self) -> Tuple[str, Dict[str, Any]]:
+        return ("custom.triggers.api_check_trigger.ApiCheckTrigger", {"url": self.url, "expected_status": self.expected_status})
+
+    async def run(self):
+        """
+        Simple time delay loop until the relevant status is found.
+        """
+        current_status = "Init"
+        while self.expected_status != current_status:
+            await asyncio.sleep(1)
+            async with aiohttp.ClientSession() as session:
+                async with session.get(self.url) as response:
+                    api_call_result = await response.json()
+                    current_status = api_call_result["status"]
+        
+        # Send our single event and then we're done
+        yield TriggerEvent(api_call_result)

examples/airflow-deferrable-operator-custom/custom/operators/api_check_operator.py

@@ -0,0 +1,24 @@
+"""
+Example DAG demonstrating ``TimeDeltaSensorAsync``, a drop in replacement for ``TimeDeltaSensor`` that
+defers and doesn't occupy a worker slot while it waits
+"""
+
+from datetime import timedelta, datetime
+
+from airflow import DAG
+from airflow.operators.empty import EmptyOperator
+from airflow.sensors.time_delta import TimeDeltaSensorAsync
+
+from custom.operators.api_check_operator import WaitForStartedStatusSensor
+
+with DAG(
+    dag_id="example_custom_sensor_async",
+    schedule_interval=None,
+    start_date=datetime.now() - timedelta(minutes=20),
+    catchup=False
+) as dag:
+
+    wait = WaitForStartedStatusSensor(task_id="wait")
+
+    finish = EmptyOperator(task_id="finish")
+    wait >> finish

examples/airflow-deferrable-operator-custom/custom/triggers/__init__.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+
+echo "============================="
+echo "== Configure API Responses =="
+echo "============================="
+
+# Creating a expectation for our mockserver to respond to a specific api rest call with a pendig status for the first 30 invocations
+# For docs on creating expectations see: http://www.mock-server.com/mock_server/creating_expectations.html
+curl -v -X PUT "http://mockserver:1080/mockserver/expectation" -d '{
+  "id": "pending_expectation",
+  "httpRequest": {
+    "path": "/testapi"
+  },
+  "httpResponse": {
+    "statusCode": 200,
+    "headers": {
+      "content-type": [
+        "application/json"
+      ]
+    },
+    "body": {
+      "type": "JSON",
+      "json": "{\"status\": \"Pending\"}"
+    }
+  },
+  "times": {
+    "remainingTimes": 30,
+    "unlimited": false
+  },
+  "timeToLive": {
+    "unlimited": true
+  },
+  "priority": 10
+}'
+
+curl -v -X PUT "http://mockserver:1080/mockserver/expectation" -d '{
+  "id": "started_expectation",
+  "httpRequest": {
+    "path": "/testapi"
+  },
+  "httpResponse": {
+    "statusCode": 200,
+    "headers": {
+      "content-type": [
+        "application/json"
+      ]
+    },
+    "body": {
+      "type": "JSON",
+      "json": "{\"status\": \"Started\"}"
+    }
+  },
+  "priority": 0
+}'
+
+pip install aiohttp

examples/airflow-deferrable-operator-custom/custom/triggers/api_check_trigger.py

@@ -0,0 +1,2 @@
+whirl.setup.d
+.whirl.env

examples/airflow-deferrable-operator-custom/dag.py

@@ -0,0 +1,3 @@
+WHIRL_ENVIRONMENT=just-airflow
+AIRFLOW__LOGGING__LOGGING_LEVEL=DEBUG
+MINIMAL_AIRFLOW_VERSION=2.3.0

examples/airflow-deferrable-operator-custom/whirl.setup.d/.gitkeep

@@ -0,0 +1,22 @@
+"""
+Example DAG demonstrating ``TimeDeltaSensorAsync``, a drop in replacement for ``TimeDeltaSensor`` that
+defers and doesn't occupy a worker slot while it waits
+"""
+
+from datetime import timedelta, datetime
+
+from airflow import DAG
+from airflow.operators.empty import EmptyOperator
+from airflow.sensors.time_delta import TimeDeltaSensorAsync
+
+with DAG(
+    dag_id="example_time_delta_sensor_async",
+    schedule_interval=None,
+    start_date=datetime.now() - timedelta(minutes=20),
+    catchup=False
+) as dag:
+
+    wait = TimeDeltaSensorAsync(task_id="wait", delta=timedelta(seconds=120))
+
+    finish = EmptyOperator(task_id="finish")
+    wait >> finish

examples/airflow-deferrable-operator-custom/whirl.setup.d/01_add_connection_api_and_mockdata.sh

@@ -227,9 +227,23 @@ check_dagrun_result() {
   fi
 }
 
+function version_lt() {
+  local VER1="${1}"
+  local VER2="${2}"
+  test "$(echo "${VER1} ${VER2}" | tr " " "\n" | sort -rV | head -n 1)" != "${VER1}";
+}
+
 start() {
     echo "Starting airflow local run for environment ${WHIRL_ENVIRONMENT}"
 
+    if [ "${CI_MODE}" == true ]; then
+      if version_lt "${AIRFLOW_VERSION}" "${MINIMAL_AIRFLOW_VERSION}"; then
+        echo "${AIRFLOW_VERSION} less then minimal version ${MINIMAL_AIRFLOW_VERSION}";
+        echo "Skipping running in CI MODE....";
+        exit 0;
+      fi
+    fi
+
     # Possible cleanup before starting
     run_compose_setup_scripts