Upgrade Airflow content to 2.11.0. (#187)

* Upgrade Airflow content to 2.11.0.

This must only be merged into `master` when actual Airflow 2.11 is already running in production,
not before.

Part of https://dfinity.atlassian.net/browse/DRE-496 .

* Fix build pipeline.

Author: DFINITYManu

.github/workflows/test/action.yaml

@@ -8,7 +8,7 @@ runs:
       id: prep
       run: |
         ret=0
-        out=$(make venv/bin venv/bin/mypy venv/bin/ruff venv/bin/pytest venv/lib/*/site-packages/mock 2>&1) || ret=$?
+        out=$(make airflow-setup venv/bin/mypy venv/bin/ruff venv/bin/pytest venv/lib/*/site-packages/mock 2>&1) || ret=$?
         if [ $ret != 0 ]
         then
           echo "$out" >&2

Makefile

@@ -1,37 +1,51 @@
 VENV_DIR = venv
-VENV_BINDIR = $(VENV_DIR)/bin
-.PHONY = test ruff mypy pytest
+VENV_BIN_DIR = $(VENV_DIR)/bin
+AIRFLOW_SETUP = $(VENV_BIN_DIR)/.airflow-2.11.0
+AIRFLOW_CFG = airflow/airflow.cfg
+AIRFLOW_TESTS_HOME = $(PWD)/tests
+AIRFLOW_TESTS_CFG = $(AIRFLOW_TESTS_HOME)/airflow.cfg
+AIRFLOW_TESTS_DB = $(AIRFLOW_TESTS_HOME)/airflow.db
 
-$(VENV_BINDIR):
+.PHONY = airflow-setup test ruff mypy pytest
+
+$(AIRFLOW_CFG):
 	bin/airflow setup
-	touch $(VENV_BINDIR)
 
-$(VENV_DIR)/lib/*/site-packages/mock: $(VENV_BINDIR)
-	$(VENV_BINDIR)/pip3 install mock
+$(AIRFLOW_SETUP): $(AIRFLOW_CFG)
+	touch $(AIRFLOW_SETUP)
+
+airflow-setup: $(AIRFLOW_SETUP)
+
+$(VENV_DIR)/lib/*/site-packages/mock: $(AIRFLOW_SETUP)
+	$(VENV_BIN_DIR)/pip3 install mock
 	touch $(VENV_DIR)/lib/*/site-packages/mock
 
-$(VENV_BINDIR)/pytest: $(VENV_BINDIR)
-	$(VENV_BINDIR)/pip3 install pytest pytest-mock
-	touch $(VENV_BINDIR)/pytest
+$(VENV_BIN_DIR)/pytest: $(AIRFLOW_SETUP)
+	$(VENV_BIN_DIR)/pip3 install pytest pytest-mock
+	touch $(VENV_BIN_DIR)/pytest
+
+$(VENV_BIN_DIR)/mypy: $(AIRFLOW_SETUP)
+	$(VENV_BIN_DIR)/pip3 install mypy types-PyYAML types-requests types-mock
+	touch $(VENV_BIN_DIR)/mypy
 
-$(VENV_BINDIR)/mypy: $(VENV_BINDIR)
-	$(VENV_BINDIR)/pip3 install mypy types-PyYAML types-requests types-mock
-	touch $(VENV_BINDIR)/mypy
+$(VENV_BIN_DIR)/ruff: $(AIRFLOW_SETUP)
+	$(VENV_BIN_DIR)/pip3 install ruff
+	touch $(VENV_BIN_DIR)/ruff
 
-$(VENV_BINDIR)/ruff: $(VENV_BINDIR)
-	$(VENV_BINDIR)/pip3 install ruff
-	touch $(VENV_BINDIR)/ruff
+mypy: $(VENV_BIN_DIR)/mypy
+	PYTHONPATH=$(PWD)/plugins:$(PWD)/shared MYPY_PATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BIN_DIR)/mypy --config=mypy.ini
 
-mypy: $(VENV_BINDIR)/mypy
-	PYTHONPATH=$(PWD)/plugins:$(PWD)/shared MYPY_PATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BINDIR)/mypy --config=mypy.ini
+ruff: $(VENV_BIN_DIR)/ruff
+	PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BIN_DIR)/ruff check shared plugins dags
 
-ruff: $(VENV_BINDIR)/ruff
-	PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BINDIR)/ruff check shared plugins dags
+$(AIRFLOW_TESTS_CFG): $(AIRFLOW_SETUP)
+	mkdir -p $(AIRFLOW_TESTS_HOME)
+	cp -f $(AIRFLOW_CFG) $(AIRFLOW_TESTS_CFG)
 
-tests/airflow.db: $(VENV_BINDIR)
-	AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False AIRFLOW__CORE__LOAD_EXAMPLES=False AIRFLOW__CORE__UNIT_TEST_MODE=True AIRFLOW_HOME=$(PWD)/tests PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BINDIR)/airflow db migrate
+tests/airflow.db: $(AIRFLOW_SETUP) $(AIRFLOW_TESTS_CFG)
+	AIRFLOW__CORE__EXECUTOR=SequentialExecutor  AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False AIRFLOW__CORE__LOAD_EXAMPLES=False AIRFLOW__CORE__UNIT_TEST_MODE=True AIRFLOW_HOME=$(AIRFLOW_TESTS_HOME) PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BIN_DIR)/airflow db migrate
 
-pytest: tests/airflow.db $(VENV_BINDIR)/pytest $(VENV_DIR)/lib/*/site-packages/mock
-	AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False AIRFLOW__CORE__LOAD_EXAMPLES=False AIRFLOW__CORE__UNIT_TEST_MODE=True AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES_REGEXP="(airflow|dfinity)[.].*" AIRFLOW_HOME=$(PWD)/tests PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BINDIR)/pytest -vv tests
+pytest: tests/airflow.db $(VENV_BIN_DIR)/pytest $(VENV_DIR)/lib/*/site-packages/mock
+	AIRFLOW__CORE__EXECUTOR=SequentialExecutor AIRFLOW__DATABASE__LOAD_DEFAULT_CONNECTIONS=False AIRFLOW__CORE__LOAD_EXAMPLES=False AIRFLOW__CORE__UNIT_TEST_MODE=True AIRFLOW__CORE__ALLOWED_DESERIALIZATION_CLASSES_REGEXP="(airflow|dfinity)[.].*" AIRFLOW_HOME=$(PWD)/tests PYTHONPATH=$(PWD)/plugins:$(PWD)/shared $(VENV_BIN_DIR)/pytest -vv tests
 
 test: ruff mypy pytest

bin/airflow

@@ -60,9 +60,10 @@ function setup_airflow_variables_and_connections() {
     "$AIRFLOW_BIN" connections add "github.node_allocation" --conn-type generic --conn-password "$NODE_ALLOC_PAT" --conn-description "Used by the target topology flow. Permits access to the https://github.com/dfinity/node_allocation repository."
 }
 
+AIRFLOW_SETUP="$VENV_BIN_DIR"/.airflow-2.11.0
 if [ "$1" == "setup" ]
 then
-    PYTHON_VERSION=$(python -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
+    PYTHON_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
 
     # Split version
     PYTHON_MAJOR=$(echo "$PYTHON_VERSION" | cut -d. -f1)
@@ -73,12 +74,17 @@ then
         exit 1
     fi
 
-    python -m venv "$VENV_DIR"
-    "$VENV_DIR"/bin/pip3 install "apache-airflow[celery]==2.9.1" \
+    if [ "$PYTHON_MAJOR" -lt 3 ] || { [ "$PYTHON_MAJOR" -eq 3 ] && [ "$PYTHON_MINOR" -lt 11 ]; }; then
+        echo "❌ Python $PYTHON_VERSION is not supported. Please use Python 3.11 or higher."
+        exit 1
+    fi
+
+    python3 -m venv "$VENV_DIR"
+    "$VENV_DIR"/bin/pip3 install "apache-airflow[celery]==2.11.0" \
         apache-airflow-providers-slack[common.sql] \
         apache-airflow-providers-google \
         "apache-airflow[amazon]" \
-        --constraint https://raw.githubusercontent.com/apache/airflow/constraints-2.9.1/constraints-3.11.txt
+        --constraint https://raw.githubusercontent.com/apache/airflow/constraints-2.11.0/constraints-$PYTHON_MAJOR.$PYTHON_MINOR.txt
     mkdir -p "$AIRFLOW_HOME"
     if ! test -f "$AIRFLOW_HOME"/airflow.db
     then
@@ -87,7 +93,9 @@ then
         PATH="$VENV_DIR/bin:$PATH" "$VENV_DIR"/bin/airflow db migrate
     fi
     pushd "$AIRFLOW_HOME"
+    rm -f dags
     ln -sfT ../dags dags
+    rm -f plugins
     ln -sfT ../plugins plugins
     sed -i 's/reload_on_plugin_change.*/reload_on_plugin_change = True/' airflow.cfg
     sed -i 's/load_examples.*/load_examples = False/' airflow.cfg
@@ -102,6 +110,8 @@ then
     else
         echo "Skipping Airflow variables and connections setup."
     fi
+
+    touch "$AIRFLOW_SETUP"
 fi
 
 test -x "$VENV_DIR"/bin/airflow || {
@@ -114,6 +124,11 @@ test -f "$AIRFLOW_HOME"/airflow.cfg || {
     echo "Run this command with 'setup' as its only argument to initialize it." >&2
     exit 32
 }
+if ! test -f "$AIRFLOW_SETUP"
+then
+    echo "Sorry, you have to upgrade your Airflow stack to 2.11.0.  Run '$0 setup' to do so." >&2
+    exit 48
+fi
 
 if [ "$1" == "check-setup" ]
 then
@@ -126,6 +141,7 @@ then
     PATH="$VENV_DIR/bin:$PATH" "$VENV_DIR"/bin/airflow info >&2
     echo "You can now run this command with 'standalone' as its only argument." >&2
     echo "When you do that, note the admin password in the scrollback buffer -- you will need it to log in." >&2
+    echo "Also, please exit any IDEs that have this project open, in order to cause it to refresh its Python configuration." >&2
     exit
 fi
 

dags/test_dag.py

@@ -51,7 +51,7 @@ def finish() -> None:
         >> test_operator.TestTask(
             task_id="task_2",
         )
-        >> TimeDeltaSensorAsync(  # type: ignore
+        >> TimeDeltaSensorAsync(
             task_id="wait_2_minutes",
             delta=datetime.timedelta(minutes=2),
         )

plugins/operators/gsheets_rollout.py

@@ -101,7 +101,6 @@ class GetFeatureRolloutPlan(GoogleSheetsCreateSpreadsheetOperator):
     def execute(self, context: Any) -> RolloutFeatures:  # type: ignore
         hook = GSheetsHook(
             gcp_conn_id=self.gcp_conn_id,
-            delegate_to=self.delegate_to,
             impersonation_chain=self.impersonation_chain,
         )
 

plugins/sensors/ic_os_rollout.py

@@ -5,7 +5,7 @@
 import datetime
 import itertools
 import time
-from typing import Any, Sequence, TypedDict, cast
+from typing import Any, NoReturn, Sequence, TypedDict, cast
 
 import dfinity.dre as dre
 import dfinity.ic_types as ic_types
@@ -25,7 +25,8 @@
 from airflow.sensors.base import BaseSensorOperator
 from airflow.sensors.date_time import DateTimeSensorAsync
 from airflow.serialization.pydantic.dag_run import DagRunPydantic
-from airflow.triggers.temporal import TimeDeltaTrigger
+from airflow.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
+from airflow.utils import timezone
 from airflow.utils.context import Context
 from airflow.utils.state import DagRunState
 
@@ -80,11 +81,18 @@ def __init__(  # type:ignore
         else:
             self.target_time = target_time
 
-    def execute(self, context: Context) -> None:
+    def execute(self, context: Context) -> NoReturn:
         if self.simulate:
             print("Nominally we would sleep, but this is a simulation.  Returning now.")
-            return
-        super().execute(context=context)
+            self.defer(
+                method_name="execute_complete",
+                trigger=DateTimeTrigger(
+                    moment=timezone.parse("1980-01-01"),
+                    end_from_trigger=self.end_from_trigger,
+                ),
+            )
+        else:
+            super().execute(context=context)
 
 
 class WaitForRevisionToBeElected(ICRolloutSensorBaseOperator):