feat: Add Airflow 3.2.1

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)

shell.nix

@@ -1,11 +1,21 @@
 let
-  self = import ./. {};
+  self = import ./. { };
   inherit (self) sources pkgs meta;
 
-  beku = pkgs.callPackage (sources."beku.py" + "/beku.nix") {};
-  cargoDependencySetOfCrate = crate: [ crate ] ++ pkgs.lib.concatMap cargoDependencySetOfCrate (crate.dependencies ++ crate.buildDependencies);
-  cargoDependencySet = pkgs.lib.unique (pkgs.lib.flatten (pkgs.lib.mapAttrsToList (crateName: crate: cargoDependencySetOfCrate crate.build) self.cargo.workspaceMembers));
-in pkgs.mkShell rec {
+  beku = pkgs.callPackage (sources."beku.py" + "/beku.nix") { };
+  cargoDependencySetOfCrate =
+    crate:
+    [ crate ]
+    ++ pkgs.lib.concatMap cargoDependencySetOfCrate (crate.dependencies ++ crate.buildDependencies);
+  cargoDependencySet = pkgs.lib.unique (
+    pkgs.lib.flatten (
+      pkgs.lib.mapAttrsToList (
+        crateName: crate: cargoDependencySetOfCrate crate.build
+      ) self.cargo.workspaceMembers
+    )
+  );
+in
+pkgs.mkShell rec {
   name = meta.operator.name;
 
   packages = with pkgs; [
@@ -24,21 +34,25 @@ in pkgs.mkShell rec {
   buildInputs = pkgs.lib.unique (pkgs.lib.concatMap (crate: crate.buildInputs) cargoDependencySet);
 
   # build time dependencies
-  nativeBuildInputs = pkgs.lib.unique (pkgs.lib.concatMap (crate: crate.nativeBuildInputs) cargoDependencySet ++ (with pkgs; [
-    beku
-    docker
-    gettext # for the proper envsubst
-    git
-    jq
-    kind
-    kubectl
-    kubernetes-helm
-    kuttl
-    nix # this is implied, but needed in the pure env
-    # tilt already defined in default.nix
-    which
-    yq-go
-  ]));
+  nativeBuildInputs = pkgs.lib.unique (
+    pkgs.lib.concatMap (crate: crate.nativeBuildInputs) cargoDependencySet
+    ++ (with pkgs; [
+      beku
+      docker
+      gettext # for the proper envsubst
+      git
+      jq
+      kind
+      kubectl
+      kubernetes-helm
+      kuttl
+      nix # this is implied, but needed in the pure env
+      # tilt already defined in default.nix
+      which
+      yq-go
+      grpcurl # for interacting with the Vector API
+    ])
+  );
 
   LIBCLANG_PATH = "${pkgs.libclang.lib}/lib";
   BINDGEN_EXTRA_CLANG_ARGS = "-I${pkgs.glibc.dev}/include -I${pkgs.clang}/resource-root/include";

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