feat: Add Airflow 3.2.1 (#803)

* test: Add Airflow 3.2.1

Note: The metrics test broke because airflow no longer exports metrics from the webserver role.

Instead, we check completion via the API.

* chore(nix): add grpcurl

Used for interacting with Vector API. Likely not needed, but handy if debugging the vector-aggregator tests for new versions.

* chore(nix): Formatting

* docs(supported-versions): Add 3.2.1, deprecate 3.1.6

* test: Update assertions

Airflow 3.2 changes "migrating" to "migration"

* test: Increase schedule memory limit

Keeps getting OOMKilled in the logging test since Airflow 3.2

* test(common/metrics): Check dag started via the web API.

Later versions of Airflow don't emit metrics from the web server.
These tests also never waited for a success, only checked creation. So that remains for now, and could be improved later.

* docs(getting_started): Bump to Airflow 3.2.1

* docs: Bump postgres requirement to 17

I removed 12, even though it is supported by older versions that we support.

* revert(nix): Restore shell.nix

This file is templated from https://github.com/stackabletech/operator-templating/blob/main/template/shell.nix

Author: NickLarsenNZ

docs/modules/airflow/examples/getting_started/code/airflow.yaml

@@ -5,7 +5,7 @@ metadata:
   name: airflow
 spec:
   image:
-    productVersion: 3.1.6
+    productVersion: 3.2.1
     pullPolicy: IfNotPresent
   clusterConfig:
     loadExamples: true

docs/modules/airflow/pages/required-external-components.adoc

@@ -7,7 +7,7 @@ The {airflow-prerequisites}[Airflow documentation] specifies:
 
 Fully supported for production usage:
 
-* PostgreSQL: 12, 13, 14, 15, 16
+* PostgreSQL: 13, 14, 15, 16, 17
 
 NOTE: The SDP Airflow images do not bundle the MySQL provider for Airflow.
 If you need MySQL or MariaDB as an Airflow back-end, you will need to create a custom image to include this provider.

docs/modules/airflow/partials/supported-versions.adoc

@@ -2,6 +2,7 @@
 // This is a separate file, since it is used by both the direct Airflow-Operator documentation, and the overarching
 // Stackable Platform documentation.
 
-- 3.1.6
+- 3.2.1
+- 3.1.6 (deprecated)
 - 3.0.6 (LTS)
 - 2.9.3 (deprecated)

tests/templates/kuttl/cluster-operation/09-assert.yaml

@@ -5,4 +5,4 @@ kind: TestAssert
 timeout: 30
 commands:
 - script: |
-    kubectl -n $NAMESPACE logs airflow-scheduler-default-0 | grep "Database migrating done!"
+    kubectl -n $NAMESPACE logs airflow-scheduler-default-0 | grep -E "Database migrat(ing|ion) done!"

tests/templates/kuttl/cluster-operation/31-assert.yaml

@@ -5,4 +5,4 @@ kind: TestAssert
 timeout: 30
 commands:
 - script: |
-    kubectl -n $NAMESPACE logs airflow-scheduler-default-0 | grep "Database migrating done!" && exit 1 || exit 0
+    kubectl -n $NAMESPACE logs airflow-scheduler-default-0 | grep -E "Database migrat(ing|ion) done!" && exit 1 || exit 0

tests/templates/kuttl/commons/metrics.py

@@ -25,6 +25,22 @@ def assert_metric(role, role_group, metric):
     return metric in metric_response.text
 
 
+# Check if dag run state is "success", "queued", or "running"
+# TODO: in future, we could wait on it.
+# See: https://airflow.apache.org/docs/apache-airflow/3.1.6/stable-rest-api-ref.html#operation/wait_dag_run_until_finished
+def assert_dag_started(rest_url, headers, dag_id, dag_run_id):
+    dag_run_response = requests.get(
+        f"{rest_url}/dags/{dag_id}/dagRuns/{dag_run_id}", headers=headers
+    )
+    dag_run_state = dag_run_response.json()["state"]
+    print(f"DAG RUN STATE: {dag_run_state}")
+    return (
+        dag_run_state == "success"
+        or dag_run_state == "queued"
+        or dag_run_state == "running"
+    )
+
+
 def metrics_v3(role_group: str) -> None:
     now = datetime.now(timezone.utc)
     ts = now.strftime("%Y-%m-%dT%H:%M:%S.%f") + now.strftime("%z")
@@ -66,10 +82,14 @@ def metrics_v3(role_group: str) -> None:
     response = requests.patch(
         f"{rest_url}/dags/{dag_id}", headers=headers, json={"is_paused": False}
     )
+
     # trigger DAG
     response = requests.post(
         f"{rest_url}/dags/{dag_id}/dagRuns", headers=headers, json=dag_data
     )
+    dag_run_id = response.json()["dag_run_id"]
+
+    print(f"DAG RUN ID: {dag_run_id}")
 
     # Test the DAG in a loop. Each time we call the script a new job will be started: we can avoid
     # or minimize this by looping over the check instead.
@@ -79,24 +99,13 @@ def metrics_v3(role_group: str) -> None:
         assert response.status_code == 200, "DAG run could not be triggered."
         # Wait for the metrics to be consumed by the statsd-exporter
         time.sleep(5)
-        # (disable line-break flake checks)
+        heartbeat_metric = "airflow_scheduler_heartbeat"
+        dag_run_success_count_metric = f"airflow_dagrun_duration_success_{dag_id}_count"
         if (
-            (assert_metric("scheduler", role_group, "airflow_scheduler_heartbeat"))
-            and (
-                assert_metric(
-                    "webserver",
-                    role_group,
-                    "airflow_task_instance_created_BashOperator",
-                )
-            )  # noqa: W503, W504
-            and (
-                assert_metric(
-                    "scheduler",
-                    role_group,
-                    "airflow_dagrun_duration_success_example_trigger_target_dag_count",
-                )
-            )
-        ):  # noqa: W503, W504
+            assert_dag_started(rest_url, headers, dag_id, dag_run_id)
+            and assert_metric("scheduler", role_group, heartbeat_metric)
+            and assert_metric("scheduler", role_group, dag_run_success_count_metric)
+        ):
             break
         time.sleep(10)
         loop += 1

tests/templates/kuttl/logging/41-install-airflow-cluster.yaml.j2

@@ -271,7 +271,7 @@ spec:
           min: 1000m
           max: 2000m
         memory:
-          limit: 1Gi
+          limit: 2Gi
     roleGroups:
       automatic-log-config:
         replicas: 1

tests/templates/kuttl/triggerer/triggerer_metrics.py

@@ -25,6 +25,22 @@ def assert_metric(role, role_group, metric):
     return metric in metric_response.text
 
 
+# Check if dag run state is "success", "queued", or "running"
+# TODO: in future, we could wait on it.
+# See: https://airflow.apache.org/docs/apache-airflow/3.1.6/stable-rest-api-ref.html#operation/wait_dag_run_until_finished
+def assert_dag_started(rest_url, headers, dag_id, dag_run_id):
+    dag_run_response = requests.get(
+        f"{rest_url}/dags/{dag_id}/dagRuns/{dag_run_id}", headers=headers
+    )
+    dag_run_state = dag_run_response.json()["state"]
+    print(f"DAG RUN STATE: {dag_run_state}")
+    return (
+        dag_run_state == "success"
+        or dag_run_state == "queued"
+        or dag_run_state == "running"
+    )
+
+
 def metrics_v3(role_group: str) -> None:
     now = datetime.now(timezone.utc)
     ts = now.strftime("%Y-%m-%dT%H:%M:%S.%f") + now.strftime("%z")
@@ -64,11 +80,16 @@ def metrics_v3(role_group: str) -> None:
     response = requests.patch(
         f"{rest_url}/dags/{dag_id}", headers=headers, json={"is_paused": False}
     )
+
     # trigger DAG
     response = requests.post(
         f"{rest_url}/dags/{dag_id}/dagRuns", headers=headers, json=dag_data
     )
 
+    dag_run_id = response.json()["dag_run_id"]
+
+    print(f"DAG RUN ID: {dag_run_id}")
+
     # Test the DAG in a loop. Each time we call the script a new job will be started: we can avoid
     # or minimize this by looping over the check instead.
     iterations = 4
@@ -77,24 +98,13 @@ def metrics_v3(role_group: str) -> None:
         assert response.status_code == 200, "DAG run could not be triggered."
         # Wait for the metrics to be consumed by the statsd-exporter
         time.sleep(5)
-        # (disable line-break flake checks)
+        heartbeat_metric = "airflow_scheduler_heartbeat"
+        dag_run_success_count_metric = f"airflow_dagrun_duration_success_{dag_id}_count"
         if (
-            (assert_metric("scheduler", role_group, "airflow_scheduler_heartbeat"))
-            and (
-                assert_metric(
-                    "webserver",
-                    role_group,
-                    "airflow_task_instance_created_CoreDeferrableSleepOperator",
-                )
-            )  # noqa: W503, W504
-            and (
-                assert_metric(
-                    "scheduler",
-                    role_group,
-                    "airflow_dagrun_duration_success_core_deferrable_sleep_demo_count",
-                )
-            )
-        ):  # noqa: W503, W504
+            assert_dag_started(rest_url, headers, dag_id, dag_run_id)
+            and assert_metric("scheduler", role_group, heartbeat_metric)
+            and assert_metric("scheduler", role_group, dag_run_success_count_metric)
+        ):
             break
         time.sleep(10)
         loop += 1

tests/test-definition.yaml

@@ -10,11 +10,12 @@ dimensions:
       - 2.9.3
       - 3.0.6
       - 3.1.6
+      - 3.2.1
       # To use a custom image, add a comma and the full name after the product version
       # - x.x.x,oci.stackable.tech/sandbox/airflow:x.x.x-stackable0.0.0-dev
   - name: airflow-latest
     values:
-      - 3.1.6
+      - 3.2.1
       # To use a custom image, add a comma and the full name after the product version
       # - x.x.x,oci.stackable.tech/sandbox/airflow:x.x.x-stackable0.0.0-dev
   - name: opa-latest