You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "jykae (via GitHub)" <gi...@apache.org> on 2026/06/05 10:48:12 UTC
[PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
jykae opened a new pull request, #68074:
URL: https://github.com/apache/airflow/pull/68074
Extends the existing `migrate-database-job` so a single `helm upgrade` reconciles the Airflow metadata DB in both directions — forward migrate, no-op, or `airflow db downgrade` executed inside the still-running api-server pod (whose source-version image still ships the reverse alembic scripts).
**What the chart does now**
1. The job runs `pre-install,pre-upgrade` (was `post-install,post-upgrade`), so the DB is reconciled *before* the new image rolls out.
2. The default `args` embed a bash entrypoint (`chart/files/db_migrate.sh`) that:
- Reads the alembic revision currently in the DB.
- Reads `AIRFLOW_TARGET_VERSION` (the chart's `.airflowVersion`).
- Decides one of: `noop`, `forward` (`exec airflow db migrate`), or `downgrade` (exec into the running api-server pod and run `airflow db downgrade --to-version X.Y.Z --yes` there).
3. New namespace-scoped RBAC (`Role` + `RoleBinding`) grants the job's existing ServiceAccount `pods get/list` and `pods/exec create/get`, only when both `rbac.create` and `migrateDatabaseJob.enabled` are true.
4. User-supplied `migrateDatabaseJob.command` / `migrateDatabaseJob.args` still win — only the default is changed.
**Why exec into the api-server pod**
Reverse alembic scripts only exist in the source-version image. The new migrations image cannot run a downgrade against the schema of an older Airflow. Exec-ing into the still-running source-version api-server pod sidesteps the need to ship a second image or to coordinate a rolling pod swap.
**Tests**
50 helm tests in `chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py` and `chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py` all pass.
closes: #68072
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Opus 4.7
Generated-by: Claude Opus 4.7 following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "bugraoz93 (via GitHub)" <gi...@apache.org>.
bugraoz93 commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4652210448
> Took me a moment to read and igest the idea. First I was a bit reluctant but then reading more and having some hours of thought I really like the idea.
>
> Some comments besides CI is not getting green - probably you find the bug.
>
> I see this as a very cool improvement the more I think about this and would like to vote for having this. Downgrade with Helm was always a pain. This PR addresses this.
>
> What do other reviewers think about it? @bugraoz93 @jedcunningham @Miretpl @potiuk ?
I very much agree with Jens! This would be a great feature to have. It is a huge LoC, but once you update your environment, you always try to make it work rather than reverting, since sometimes it is harder to return back,. This could help with rolling back updates.
My only concern is the testability of the script. We should test them in one of the integration tests. Unit tests don't reflect what would happen in the environment; even integration tests don't fully, but it would be great comfort if we had it. I haven't taken a deep dive to check if it is easy to achieve. This one `chart/files/db_migrate.py`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Miretpl (via GitHub)" <gi...@apache.org>.
Miretpl commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4698292755
> While I think that it would be "cool" to use public interfaces I think ... the dilemma we are in is that this would get (earliest, we are cutting 3.3.0 in 2 days) into 3.4.0, which means you could only offer to downgrade to 3.3.0 and not earlier. But as we support all down to 3.0.0 we either only have this as a partial feature and leave many (installed base) use case outside.
>
> So for me would be OK like now and aimming to improve to a public interface for next released. Then depending on the version either the "new public" or "old method" is used. Once we drop support for the current versions we could clean-up.
Agreed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387219598
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
Review Comment:
Done. The `kubernetes` import is now lazy: it's wrapped in a module-level `try/except ImportError` guard, and only the downgrade branch calls `_require_kubernetes_client()`, which raises a clear `SystemExit` if the client is missing. The forward/no-op/fresh paths no longer require the `kubernetes` package.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387219415
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -43,7 +43,7 @@ metadata:
{{- end }}
{{- $annotations := dict }}
{{- if .Values.migrateDatabaseJob.useHelmHooks }}
- {{- $_ := set $annotations "helm.sh/hook" "post-install,post-upgrade" }}
+ {{- $_ := set $annotations "helm.sh/hook" "pre-install,pre-upgrade" }}
{{- $_ := set $annotations "helm.sh/hook-weight" "1" }}
{{- $_ := set $annotations "helm.sh/hook-delete-policy" "before-hook-creation,hook-succeeded" }}
{{- end }}
Review Comment:
The ServiceAccount, Role, and RoleBinding are now hook resources (`helm.sh/hook: post-install,pre-upgrade`) with `hook-weight: "-5"`, while the Job runs at `hook-weight: "1"`. Helm applies lower-weighted hook resources first, so the SA and RBAC exist before the Job hook runs on both fresh installs and the first downgrade after adopting this version.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387219780
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,37 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row, or DB unreachable) — ``airflow db migrate``.
Review Comment:
Good catch. Updated the newsfragment to match the implementation: the action is only decided once the DB is reachable, and if it's still unreachable after the connection-wait window (`DB_CONNECT_MAX_WAIT_SECONDS`, default 120) the job exits non-zero and is retried via the Job `backoffLimit` rather than proceeding with a migrate.
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
Review Comment:
Tightened the Role to least privilege:
- `pods/exec` now grants only `get`. The reconciler execs via `connect_get_namespaced_pod_exec`, which is an HTTP GET to the exec subresource, so the API server authorizes it under `get` (not `create`, which a POST-based exec would need). Added a comment documenting this.
- `deployments/scale` and `statefulsets/scale` now grant only `patch` — the reconciler patches replicas to 0 and never reads the scale subresource back. Added a comment as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387477354
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,378 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _int_env(name: str, default: int) -> int:
+ """Return integer env var *name*, or *default* when unset.
+
+ These knobs come from chart config; a typo (e.g. ``"5m"``) should fail with
+ an actionable message rather than an opaque ``ValueError`` traceback deep in
+ the job logs.
+ """
+ raw = os.environ.get(name)
+ if raw is None:
+ return default
+ try:
+ return int(raw)
+ except ValueError:
+ raise SystemExit(f"{name} must be an integer number of seconds, got {raw!r}")
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = _int_env("DB_CONNECT_MAX_WAIT_SECONDS", 120)
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
Review Comment:
Fixed — `discover_api_server_pod()` now takes `release_name` and filters on `release=<release_name>,component=api-server`, so when several Airflow releases share a namespace the downgrade is exec'd strictly in the current release's api-server pod. `RELEASE_NAME` is already required/validated on the downgrade branch, and the tests assert the release-scoped selector.
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py:
##########
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for the RBAC granted to the migrate-database-job.
+
+The migrate-database-job needs ``pods`` + ``pods/exec`` access only on the
+downgrade branch, where it must exec ``airflow db downgrade`` inside the
+still-running api-server pod (whose image still ships the reverse alembic
+scripts). The Role is always rendered so the forward-migrate path remains
+identical for users who only ever upgrade.
+
+Tracked in https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import jmespath
+import pytest
+from chart_utils.helm_template_generator import render_chart
+
+ROLE_TEMPLATE = "templates/rbac/migrate-database-job-role.yaml"
+ROLEBINDING_TEMPLATE = "templates/rbac/migrate-database-job-rolebinding.yaml"
+
+
+class TestMigrateDatabaseJobRBAC:
+ def test_role_and_binding_render_by_default(self):
+ docs = render_chart(show_only=[ROLE_TEMPLATE, ROLEBINDING_TEMPLATE])
+ kinds = sorted(d["kind"] for d in docs)
+ assert kinds == ["Role", "RoleBinding"]
+
+ @pytest.mark.parametrize("rbac_create", [False, True])
+ def test_gated_on_rbac_create(self, rbac_create):
+ docs = render_chart(
+ values={"rbac": {"create": rbac_create}},
+ show_only=[ROLE_TEMPLATE, ROLEBINDING_TEMPLATE],
+ )
+ assert bool(docs) is rbac_create
+
+ def test_role_rules_grant_pods_and_exec(self):
+ docs = render_chart(show_only=[ROLE_TEMPLATE])
+ rules = jmespath.search("rules", docs[0])
+ resources_to_verbs = {tuple(r["resources"]): set(r["verbs"]) for r in rules}
+ assert ("pods",) in resources_to_verbs
+ assert {"get", "list"}.issubset(resources_to_verbs[("pods",)])
+ assert ("pods/exec",) in resources_to_verbs
+ # Exec is authorized under ``create`` or ``get`` depending on the
+ # apiserver/client path, so the role grants both for compatibility.
+ assert resources_to_verbs[("pods/exec",)] == {"create", "get"}
+
+ def test_role_rules_grant_deployments_and_statefulsets_scale(self):
+ # The downgrade branch scales every DB-touching workload to 0 after
+ # running ``airflow db downgrade`` so no OLD code keeps querying the
+ # now-downgraded schema before helm rolls in TARGET pods.
+ docs = render_chart(show_only=[ROLE_TEMPLATE])
+ rules = jmespath.search("rules", docs[0])
+ resources_to_verbs = {tuple(r["resources"]): set(r["verbs"]) for r in rules}
+ assert ("deployments", "statefulsets") in resources_to_verbs
+ assert {"get", "list"}.issubset(resources_to_verbs[("deployments", "statefulsets")])
+ assert ("deployments/scale", "statefulsets/scale") in resources_to_verbs
+ # Only ``patch`` is needed -- the reconciler never reads the scale back.
+ assert resources_to_verbs[("deployments/scale", "statefulsets/scale")] == {"patch"}
+
+ def test_rolebinding_subject_is_migrate_db_job_sa(self):
+ docs = render_chart(show_only=[ROLEBINDING_TEMPLATE])
+ subjects = jmespath.search("subjects", docs[0])
+ assert len(subjects) == 1
+ assert subjects[0]["kind"] == "ServiceAccount"
+ # Matches the name rendered by templates/jobs/migrate-database-job-serviceaccount.yaml
+ # via the "migrateDatabaseJob.serviceAccountName" helper.
+ assert subjects[0]["name"] == "release-name-airflow-migrate-database-job"
+
+ def test_rolebinding_references_role(self):
+ docs = render_chart(show_only=[ROLEBINDING_TEMPLATE])
+ role_ref = jmespath.search("roleRef", docs[0])
+ assert role_ref["kind"] == "Role"
+ assert role_ref["name"] == "release-name-migrate-database-job-role"
+ assert role_ref["apiGroup"] == "rbac.authorization.k8s.io"
Review Comment:
This assertion already matches the rendered output, so no change is needed. Under the chart default (`useStandardNaming=false`) `airflow.fullname` resolves to just the release name, so both the `Role` (`metadata.name`) and the `RoleBinding` `roleRef.name` render as `release-name-migrate-database-job-role` from the same helper — the test passes and the binding stays internally consistent. (`useStandardNaming=true` would shift both names together to `release-name-airflow-...`.)
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -103,8 +103,18 @@ spec:
{{- if .Values.migrateDatabaseJob.command }}
command: {{- tpl (toYaml .Values.migrateDatabaseJob.command) . | nindent 12 }}
{{- end }}
- {{- if .Values.migrateDatabaseJob.args }}
+ {{- if not (kindIs "invalid" .Values.migrateDatabaseJob.args) }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
+ # Bidirectional metadata DB reconciliation: forward migrate, no-op,
+ # fresh install, or exec downgrade in the still-running api-server
+ # pod depending on the relationship between the chart's target
+ # version and the DB's current alembic head.
+ args:
+ - "python"
+ - "-c"
+ - |-
+ {{- .Files.Get "files/db_migrate.py" | nindent 14 }}
Review Comment:
Keeping `{{-` here intentionally — this is the standard chart idiom for embedding a file into a block scalar (`{{- .Files.Get ... | nindent N }}`), used throughout the chart. `nindent` prepends its own newline, so the `{{-` left-trim removes the template's literal indentation/newline and the content lands correctly. Dropping the `-` would instead emit a stray whitespace-only first line inside the `python -c` block. Verified the rendered `|-` block starts directly on the license header with correct 14-space indentation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387335062
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,42 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row) — ``airflow db migrate``.
+* ``target == current`` — no-op.
+* ``target > current`` — ``airflow db migrate`` (existing behaviour).
+* ``target < current`` — ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (so reverse alembic
+ scripts that ship only in the source-version image are available),
+ followed by scaling every DB-touching workload of this release
+ (``api-server``, ``scheduler``, ``triggerer``, ``dag-processor``,
+ ``worker``) to ``0``. Helm then patches those workloads back to their
+ configured ``replicas`` with the TARGET image as the upgrade proceeds,
+ so the cluster comes back up cleanly on the target version without any
+ OLD code running against the now-downgraded schema. A downgrade is
+ therefore a brief outage rather than a zero-downtime rolling update —
+ this is unavoidable when the schema goes backwards.
+
+The action above is only decided once the metadata DB is reachable. If the DB
+is still unreachable after the connection-wait window
+(``DB_CONNECT_MAX_WAIT_SECONDS``, default ``120``), the job exits non-zero and
+is retried via the Job ``backoffLimit`` rather than proceeding with a migrate.
+
+A ``Role`` + ``RoleBinding`` granting ``pods`` (``get``, ``list``),
+``pods/exec`` (``get``), ``deployments`` / ``statefulsets`` (``get``, ``list``)
+and ``deployments/scale`` / ``statefulsets/scale`` (``patch``) to the
+existing ``migrate-database-job`` ``ServiceAccount`` are now rendered when
Review Comment:
The newsfragment says ``pods/exec`` is granted only ``get``, but the new Role template grants both ``get`` and ``create`` for ``pods/exec``. Please update the fragment to match the actual RBAC (or, if intended, change the Role + tests to match the documented verbs).
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,362 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
Review Comment:
Both env var reads use `int(...)` without handling non-integer values. If an operator accidentally sets `DB_CONNECT_MAX_WAIT_SECONDS` or `MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS` to a non-int, the job will crash with `ValueError` before it can emit a helpful error or fall back to defaults. Consider catching `ValueError` and either (a) defaulting to 120/300 with a warning, or (b) raising `SystemExit` with a clear message indicating the env var must be an integer.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,362 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
Review Comment:
Both env var reads use `int(...)` without handling non-integer values. If an operator accidentally sets `DB_CONNECT_MAX_WAIT_SECONDS` or `MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS` to a non-int, the job will crash with `ValueError` before it can emit a helpful error or fall back to defaults. Consider catching `ValueError` and either (a) defaulting to 120/300 with a warning, or (b) raising `SystemExit` with a clear message indicating the env var must be an integer.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,516 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from kubernetes import client as k8s_client
+from sqlalchemy.exc import OperationalError
Review Comment:
This new unit test module imports the Kubernetes python client at import time, but the embedded reconciler is intentionally written to work on images/environments where the Kubernetes client is not installed (until the downgrade branch is selected). Since the tests appear to use `k8s_client` only for `spec=` on `MagicMock`, consider removing this hard dependency by avoiding the `spec` (or guarding the import with a skip) so the test suite remains runnable in environments without `kubernetes` installed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3371312530
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
Review Comment:
Good catch. Added `_resolve_target_rev()` which falls back to the highest mapped version `<= target` (mirroring Airflow's own CLI behaviour), plus a regression test `test_decide_action_resolves_patch_version_via_nearest_lower` that exercises 3.2.2 → rev_320.
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -105,6 +105,16 @@ spec:
{{- end }}
{{- if .Values.migrateDatabaseJob.args }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
Review Comment:
Fixed. Switched to `{{- if not (kindIs "invalid" .Values.migrateDatabaseJob.args) }}` so an explicit `[]` is distinguished from `~` / unset and is no longer overridden by the embedded reconciler.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
Review Comment:
Added `test_decide_action_resolves_patch_version_via_nearest_lower` covering the 3.2.2 → 3.2.0-rev case.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
Review Comment:
Tightened the kubernetes mocks with `spec=k8s_client.CoreV1Api` / `spec=k8s_client.AppsV1Api` where they're used as API stubs (the exec-stream tests keep an unspecced mock because the stubbed `stream` callable accepts the api method as an arg).
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,31 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``pre-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of three actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``target == current`` — no-op.
+* ``target > current`` — ``airflow db migrate`` (existing behaviour).
+* ``target < current`` — ``airflow db downgrade --to-version <target>``
Review Comment:
Updated. The newsfragment now lists four actions including the fresh-install path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3375905676
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
Review Comment:
Please close if you think is resolved 8-D
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387222253
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-6457258157
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387223892
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,85 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ # The reconciler execs via ``connect_get_namespaced_pod_exec`` (an HTTP GET
+ # to the exec subresource), so the API server authorizes it under the
+ # ``get`` verb -- not ``create`` (which is what a POST-based exec needs).
+ verbs:
+ - "get"
Review Comment:
The RBAC verb for `pods/exec` is commonly authorized under `create` in Kubernetes (including for `kubectl exec`-style access), and using only `get` can lead to authorization failures in clusters/apiserver versions that don’t treat exec as `get`. To prevent downgrade failures, grant `create` on `pods/exec` (optionally alongside `get` if you want maximum compatibility across client behaviors) and update the associated tests accordingly.
##########
chart/values.schema.json:
##########
@@ -6731,19 +6731,21 @@
"default": null
},
"args": {
- "description": "Args to use when running migrate database job (templated).",
+ "description": "Args to use when running migrate database job (templated). When unset, the chart embeds a script that reconciles the metadata DB bidirectionally (forward migrate, no-op, or downgrade via the still-running api-server pod).",
Review Comment:
This updated description omits the “fresh install (no alembic row)” path that the reconciler explicitly supports (and that `values.yaml` / the newsfragment describe). To avoid confusing users, update the schema description to include the fresh-install behavior and (optionally) clarify that `[]` is a valid explicit override that disables the embedded reconciler.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,355 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
Review Comment:
`packaging.version.Version(target)` will raise `InvalidVersion` if `AIRFLOW_TARGET_VERSION` is not a PEP 440–parseable version string. Since this value comes from chart config, a misconfigured (or nonstandard) `airflowVersion` would crash the job instead of failing gracefully. Consider catching `InvalidVersion` and treating it as “unknown target” (conservative `forward`), or exiting with a clear error message that points users to fixing `airflowVersion`.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,355 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
Review Comment:
`packaging.version.Version(target)` will raise `InvalidVersion` if `AIRFLOW_TARGET_VERSION` is not a PEP 440–parseable version string. Since this value comes from chart config, a misconfigured (or nonstandard) `airflowVersion` would crash the job instead of failing gracefully. Consider catching `InvalidVersion` and treating it as “unknown target” (conservative `forward`), or exiting with a clear error message that points users to fixing `airflowVersion`.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,511 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from kubernetes import client as k8s_client
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load ``chart/files/db_migrate.py`` as a module.
+
+ The file is normally fed to ``python -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# Shared fixtures
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable ``{version: revision}`` mapping for ``decide_action`` tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+@pytest.fixture
+def patch_engine_returning(db_migrate, monkeypatch):
+ """Factory: patch ``engine.connect()`` so ``get_current_revision()`` returns the given value."""
+
+ def _patch(current_rev):
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+ return _patch
+
+
+@pytest.fixture
+def make_pod():
+ """Factory: build a fake api-server pod with optional ``Ready`` condition."""
+
+ def _build(name, ready=True):
+ return types.SimpleNamespace(
+ metadata=types.SimpleNamespace(name=name),
+ status=types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ ),
+ )
+
+ return _build
+
+
+@pytest.fixture
+def fake_stream():
+ """Factory: build a fake kubernetes exec-stream response with the given exit code."""
+
+ def _build(returncode):
+ resp = mock.MagicMock()
+ # is_open(): True once so the loop body runs, then False to exit.
+ resp.is_open.side_effect = [True, False]
+ resp.peek_stdout.return_value = False
+ resp.peek_stderr.return_value = False
+ resp.returncode = returncode
+ return resp
+
+ return _build
+
+
+@pytest.fixture
+def make_workload():
+ """Factory: build a fake Deployment / StatefulSet object with a ``metadata.name``."""
+
+ def _build(name):
+ return types.SimpleNamespace(metadata=types.SimpleNamespace(name=name))
+
+ return _build
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_propagates_db_error(db_migrate, monkeypatch, patched_revision_map):
+ """``decide_action`` assumes the DB is already reachable.
+
+ The wait is the responsibility of :func:`main` via
+ :func:`_wait_for_db_ready`; ``decide_action`` itself must surface a
+ connectivity failure rather than silently treating it as a fresh
+ install (which would then run ``airflow db migrate`` against a still-
+ unreachable DB and cause ``BackoffLimitExceeded``).
+ """
+
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ with pytest.raises(OperationalError):
+ db_migrate.decide_action("3.1.0")
+
+
+def test_decide_action_fresh_when_no_alembic_row(patched_revision_map, db_migrate, patch_engine_returning):
+ patch_engine_returning(current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ patch_engine_returning(current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_older_than_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ # current=3.0.0 (rev_300), target=3.1.0 (rev_310) -> forward.
+ patch_engine_returning(current_rev="rev_300")
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_is_newer_than_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade the TARGET image cannot resolve a revision newer than
+ its own head; relying on ``ScriptDirectory.walk_revisions`` raised
+ ``RevisionError`` and skipped the downgrade branch entirely. The new
+ implementation uses a version-aware reverse lookup so the comparison
+ works no matter which image is loaded.
+ """
+ # current=3.2.0 (rev_320, newer), target=3.1.0 (older) -> downgrade.
+ patch_engine_returning(current_rev="rev_320")
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+def test_decide_action_forward_when_current_revision_is_unknown(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ """Dev / intermediate alembic rev -> conservative forward, never downgrade."""
+ patch_engine_returning(current_rev="rev_unknown")
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_resolves_patch_version_via_nearest_lower(
+ db_migrate, monkeypatch, patch_engine_returning
+):
+ """Regression test for the Copilot-flagged patch-version mapping bug.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch (e.g. ``3.2.2`` is not
+ a key — its head is the same as ``3.2.0``'s). The resolver must fall back
+ to the highest lower mapped version so a deploy targeting a patch release
+ is correctly classified as a no-op or downgrade against the same rev.
+ """
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {"3.2.0": "rev_320"})
+ patch_engine_returning(current_rev="rev_320")
+ # 3.2.2 must resolve to rev_320 -> noop, not "forward" (the previous bug).
+ assert db_migrate.decide_action("3.2.2") == "noop"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _call_discover_no_retry(db_migrate, namespace):
+ """Invoke the underlying function, bypassing the @retry wrapper backoff."""
+ return db_migrate.discover_api_server_pod.retry_with(stop=lambda _r: True)(namespace)
Review Comment:
`_call_discover_no_retry` is defined but not used anywhere in this new test module. Removing dead helpers keeps the test file easier to maintain and avoids implying there’s a supported “no retry” call pattern when it isn’t exercised.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387222046
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
+ # The downgrade branch scales every DB-touching workload to 0 after running
+ # ``airflow db downgrade`` so no OLD code keeps querying the now-downgraded
+ # schema before helm rolls in TARGET pods. See chart/files/db_migrate.py.
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments"
+ - "statefulsets"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments/scale"
+ - "statefulsets/scale"
+ verbs:
+ - "get"
+ - "patch"
Review Comment:
Tightened the Role to least privilege:
- `pods/exec` now grants only `get`. The reconciler execs via `connect_get_namespaced_pod_exec`, which is an HTTP GET to the exec subresource, so the API server authorizes it under `get` (not `create`, which a POST-based exec would need). Added a comment documenting this.
- `deployments/scale` and `statefulsets/scale` now grant only `patch` — the reconciler patches replicas to 0 and never reads the scale subresource back. Added a comment as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4679401152
> > Took me a moment to read and igest the idea. First I was a bit reluctant but then reading more and having some hours of thought I really like the idea.
> > Some comments besides CI is not getting green - probably you find the bug.
> > I see this as a very cool improvement the more I think about this and would like to vote for having this. Downgrade with Helm was always a pain. This PR addresses this.
> > What do other reviewers think about it? @bugraoz93 @jedcunningham @Miretpl @potiuk ?
>
> I very much agree with Jens! This would be a great feature to have. It is a huge LoC, but once you update your environment, you always try to make it work rather than reverting, since sometimes it is harder to return back,. This could help with rolling back updates. My only concern is the testability of the script. We should test them in one of the integration tests. Unit tests don't reflect what would happen in the environment; even integration tests don't fully, but it would be great comfort if we had it. I haven't taken a deep dive to check if it is easy to achieve. This one `chart/files/db_migrate.py`
As @bugraoz93 pointed out having E2E tests for each path would be really cool to have, currently just existing forward path and fresh install are covered if noticed correct from tests ? can you take on further E2E testing for this as I dont know that much this project test setups
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4694585870
> After a couple of days of thinking about it, I think we should do it differently. At first, I believe that we should have this feature to use the Airflow public interface (if I recall correctly, @jscheffl mentioned something about that here too). If we do not, it can work for 3.2.0 and not work for 3.2.1, as the non-public interface can change, and as a chart supports multiple different versions of Airflow, I think that we should be sure that it will work for all supported versions in the chart. Without that guarantee, this feature can be basically useless when it is needed during the upgrade.
>
> Regarding the `db_migrate.py` script, I think that most of the logic there should be part of the Airflow CLI. In the chart, we could just invoke `airflow db migrate --bidirectional ...` or sth like that, which would do everything, and we would be sure that it will work properly for all Airflow versions.
While I think that it would be "cool" to use public interfaces I think ... the dilemma we are in is that this would get (earliest, we are cutting 3.3.0 in 2 days) into 3.4.0, which means you could only offer to downgrade to 3.3.0 and not earlier. But as we support all down to 3.0.0 we either only have this as a partial feature and leave many (installed base) use case outside.
So for me would be OK like now and aimming to improve to a public interface for next released. Then depending on the version either the "new public" or "old method" is used. Once we drop support for the current versions we could clean-up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4684518276
> > > Took me a moment to read and igest the idea. First I was a bit reluctant but then reading more and having some hours of thought I really like the idea.
> > > Some comments besides CI is not getting green - probably you find the bug.
> > > I see this as a very cool improvement the more I think about this and would like to vote for having this. Downgrade with Helm was always a pain. This PR addresses this.
> > > What do other reviewers think about it? @bugraoz93 @jedcunningham @Miretpl @potiuk ?
> >
> >
> > I very much agree with Jens! This would be a great feature to have. It is a huge LoC, but once you update your environment, you always try to make it work rather than reverting, since sometimes it is harder to return back,. This could help with rolling back updates. My only concern is the testability of the script. We should test them in one of the integration tests. Unit tests don't reflect what would happen in the environment; even integration tests don't fully, but it would be great comfort if we had it. I haven't taken a deep dive to check if it is easy to achieve. This one `chart/files/db_migrate.py`
>
> As @bugraoz93 pointed out having E2E tests for each path would be really cool to have, currently just existing forward path and fresh install are covered if noticed correct from tests ? can you take on further E2E testing for this as I dont know that much this project test setups
@bugraoz93 Would you force/request to have an E2E test in the scope of this PR or also OK to put this to a follow-up PR?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Miretpl (via GitHub)" <gi...@apache.org>.
Miretpl commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3407871141
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -279,6 +279,77 @@ def test_should_set_correct_helm_hooks_weight(self):
annotations = jmespath.search("metadata.annotations", docs[0])
assert annotations["helm.sh/hook-weight"] == "1"
+ def test_default_hooks_run_post_install_pre_upgrade(self):
+ # post-install: the chart's bundled postgres is created as a regular
+ # release resource, so the job needs the DB to exist before it runs.
+ # pre-upgrade: the reconciler must run *before* helm rolls the new
+ # image, so on a downgrade it can exec ``airflow db downgrade`` in the
+ # still-running OLD api-server pod (the OLD image still ships the
+ # reverse alembic scripts; the new one does not).
+ # Tracked in https://github.com/apache/airflow/issues/68072.
Review Comment:
```suggestion
# still-running OLD api-server pod.
```
I think that this should be clear for anyone working with Airflow and Helm and tracking issue is rather not needed as everyone can see blame for the code and what, when and where was introduced.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3371310814
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
Review Comment:
Done. Dropped the issue reference in the docstring.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
Review Comment:
Good idea — I'll open the follow-up PR to expose a public accessor and update this script to prefer it once it's available.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
Review Comment:
Good catch. Switched to a `_REVISION_HEADS_MAP` reverse lookup with a `packaging.version.Version` comparison. If the current revision isn't in the map (dev / intermediate alembic rev) we conservatively forward-migrate rather than risk an incorrect downgrade. Dropped the `ScriptDirectory` dependency entirely.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
Review Comment:
The list comprehension is followed by `(ready or pods)[0]`, and we already raised `RuntimeError("no Running api-server pod...")` above if `pods` was empty, so the slice can't OOB. With the new `@retry` decorator a transient empty result is also retried with backoff.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
Review Comment:
Added `@retry(stop=stop_after_attempt(5), wait=wait_exponential(...), retry=retry_if_exception_type(RuntimeError), reraise=True)`. Added a regression test that exercises the transient-empty case.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(namespace: str, release_name: str, timeout_seconds: int = 300) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+ """
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
Review Comment:
Made it configurable via `migrateDatabaseJob.drainTimeoutSeconds` (default 300) plumbed in as the `MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS` env var. Operators with long worker tasks can raise it. Agree that workers will likely die once the api-server is scaled down, but the value gives a clean knob to tune.
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -105,6 +105,16 @@ spec:
{{- end }}
{{- if .Values.migrateDatabaseJob.args }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
+ # Bidirectional metadata DB reconciliation: forward migrate, no-op, or
+ # exec downgrade in the still-running api-server pod depending on the
+ # relationship between the chart's target version and the DB's current
+ # alembic head. See https://github.com/apache/airflow/issues/68072.
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387221878
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
Review Comment:
Tightened the Role to least privilege:
- `pods/exec` now grants only `get`. The reconciler execs via `connect_get_namespaced_pod_exec`, which is an HTTP GET to the exec subresource, so the API server authorizes it under `get` (not `create`, which a POST-based exec would need). Added a comment documenting this.
- `deployments/scale` and `statefulsets/scale` now grant only `patch` — the reconciler patches replicas to 0 and never reads the scale subresource back. Added a comment as well.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
+ # The downgrade branch scales every DB-touching workload to 0 after running
+ # ``airflow db downgrade`` so no OLD code keeps querying the now-downgraded
+ # schema before helm rolls in TARGET pods. See chart/files/db_migrate.py.
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments"
+ - "statefulsets"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments/scale"
+ - "statefulsets/scale"
+ verbs:
+ - "get"
+ - "patch"
Review Comment:
Tightened the Role to least privilege:
- `pods/exec` now grants only `get`. The reconciler execs via `connect_get_namespaced_pod_exec`, which is an HTTP GET to the exec subresource, so the API server authorizes it under `get` (not `create`, which a POST-based exec would need). Added a comment documenting this.
- `deployments/scale` and `statefulsets/scale` now grant only `patch` — the reconciler patches replicas to 0 and never reads the scale subresource back. Added a comment as well.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
+
+ action = decide_action(target)
+ print(f"[db_migrate] target={target} action={action}", flush=True)
+
+ if action == "noop":
+ print("[db_migrate] DB already at target revision, nothing to do.")
+ return 0
+ if action in {"fresh", "forward"}:
+ return subprocess.run(["airflow", "db", "migrate"], check=False).returncode
+ if action == "downgrade":
+ if not release_name:
+ raise SystemExit("RELEASE_NAME must be set for the downgrade branch")
+ pod = discover_api_server_pod(namespace)
+ print(f"[db_migrate] downgrading via api-server pod {pod}", flush=True)
+ rc = run_downgrade_in_api_server(pod, namespace, target)
+ if rc != 0:
+ return rc
+ # Drain OLD pods before helm rolls in TARGET pods to avoid OLD code
+ # running against the now-downgraded schema (see module docstring).
+ scale_release_workloads_to_zero(namespace, release_name)
+ return 0
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387388192
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,42 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row) — ``airflow db migrate``.
+* ``target == current`` — no-op.
+* ``target > current`` — ``airflow db migrate`` (existing behaviour).
+* ``target < current`` — ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (so reverse alembic
+ scripts that ship only in the source-version image are available),
+ followed by scaling every DB-touching workload of this release
+ (``api-server``, ``scheduler``, ``triggerer``, ``dag-processor``,
+ ``worker``) to ``0``. Helm then patches those workloads back to their
+ configured ``replicas`` with the TARGET image as the upgrade proceeds,
+ so the cluster comes back up cleanly on the target version without any
+ OLD code running against the now-downgraded schema. A downgrade is
+ therefore a brief outage rather than a zero-downtime rolling update —
+ this is unavoidable when the schema goes backwards.
+
+The action above is only decided once the metadata DB is reachable. If the DB
+is still unreachable after the connection-wait window
+(``DB_CONNECT_MAX_WAIT_SECONDS``, default ``120``), the job exits non-zero and
+is retried via the Job ``backoffLimit`` rather than proceeding with a migrate.
+
+A ``Role`` + ``RoleBinding`` granting ``pods`` (``get``, ``list``),
+``pods/exec`` (``get``), ``deployments`` / ``statefulsets`` (``get``, ``list``)
+and ``deployments/scale`` / ``statefulsets/scale`` (``patch``) to the
+existing ``migrate-database-job`` ``ServiceAccount`` are now rendered when
Review Comment:
Good catch — the newsfragment now lists `pods/exec` (`get`, `create`) to match the Role, which grants both verbs so exec works under either the GET-based api-server client path or a classic POST-based exec.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,362 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
Review Comment:
Fixed — both timeout env vars now go through a shared `_int_env` helper that raises `SystemExit` with a clear "must be an integer number of seconds" message instead of letting an opaque `ValueError` bubble up from deep in the job.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,362 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
Review Comment:
Fixed — `MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS` is now parsed via the same `_int_env` helper that exits with an actionable message on a non-integer value.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,516 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from kubernetes import client as k8s_client
+from sqlalchemy.exc import OperationalError
Review Comment:
Fixed — the kubernetes import in the test module is now guarded (falls back to `None`), and the downgrade-branch tests that need `spec=k8s_client...` are skipped via `@pytest.mark.skipif` so the suite stays importable and runnable on images without the kubernetes client, mirroring the reconciler's lazy import.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3371311754
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,78 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+##
+## Tracked in https://github.com/apache/airflow/issues/68072.
Review Comment:
Done.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
Review Comment:
Done — converted to a `patch_engine_returning` fixture factory.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
Review Comment:
Done — converted to a `make_pod` fixture factory.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
Review Comment:
Done — converted to a `fake_stream` fixture factory.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-new"
+ fake_api.list_namespaced_pod.assert_called_once_with(
+ namespace="airflow",
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ )
+
+
+def test_discover_api_server_pod_falls_back_to_non_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [_pod("api-server-old", ready=False)]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-old"
+
+
+def test_discover_api_server_pod_raises_when_none(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = []
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ with pytest.raises(RuntimeError, match="no Running api-server pod"):
+ db_migrate.discover_api_server_pod("airflow")
+
+
+# --------------------------------------------------------------------------
+# run_downgrade_in_api_server
+# --------------------------------------------------------------------------
+
+
+def _fake_stream(returncode):
Review Comment:
Done — converted to a `make_workload` fixture factory.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-new"
+ fake_api.list_namespaced_pod.assert_called_once_with(
+ namespace="airflow",
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ )
+
+
+def test_discover_api_server_pod_falls_back_to_non_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [_pod("api-server-old", ready=False)]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-old"
+
+
+def test_discover_api_server_pod_raises_when_none(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = []
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ with pytest.raises(RuntimeError, match="no Running api-server pod"):
+ db_migrate.discover_api_server_pod("airflow")
+
+
+# --------------------------------------------------------------------------
+# run_downgrade_in_api_server
+# --------------------------------------------------------------------------
+
+
+def _fake_stream(returncode):
+ resp = mock.MagicMock()
+ # is_open(): True once so the loop body runs, then False to exit.
+ resp.is_open.side_effect = [True, False]
+ resp.peek_stdout.return_value = False
+ resp.peek_stderr.return_value = False
+ resp.returncode = returncode
+ return resp
+
+
+def test_run_downgrade_returns_zero_on_success(db_migrate, monkeypatch):
+ resp = _fake_stream(returncode=0)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 0
+
+
+def test_run_downgrade_propagates_nonzero(db_migrate, monkeypatch):
+ resp = _fake_stream(returncode=2)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 2
+
+
+def test_run_downgrade_treats_missing_returncode_as_failure(db_migrate, monkeypatch):
+ # Regression: previous code did ``return returncode or 0`` which silently
+ # reported success when the stream closed without an exit code.
+ resp = _fake_stream(returncode=None)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 1
+
+
+# --------------------------------------------------------------------------
+# scale_release_workloads_to_zero
+# --------------------------------------------------------------------------
+
+
+def _wl(name):
Review Comment:
Done — the workload helper is now the `make_workload` fixture factory used by the scale tests.
##########
chart/values.yaml:
##########
@@ -2121,13 +2121,11 @@ migrateDatabaseJob:
command: ~
# Args to use when running the migrate database job (templated).
- args:
- - "bash"
- - "-c"
- - >-
- exec \
-
- airflow db migrate
+ # When unset, the chart runs a bidirectional reconciliation script
+ # (forward migrate, no-op, or exec downgrade against the running api-server)
+ # depending on whether the chart is being upgraded or downgraded.
+ # See https://github.com/apache/airflow/issues/68072.
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4646180992
Thanks for the careful review @jscheffl (and Copilot)! Just pushed b9824f9 addressing all the comments. Highlights:
- **`decide_action` rewritten** to use a `_REVISION_HEADS_MAP` reverse lookup + `packaging.version.Version` comparison, with a nearest-lower fallback so patch versions like `3.2.2` resolve to `3.2.0`s head (matches Airflow CLIs own behaviour). Dropped the `ScriptDirectory.walk_revisions` plumbing entirely — much less coupling.
- **`@retry`** on `discover_api_server_pod` (5 attempts, exponential backoff) so an api-server rollover doesnt fail the job.
- **`drainTimeoutSeconds`** chart value (default 300) plumbed in as `MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS` for operators with long-running worker tasks.
- **`args: []`** is now distinguished from `args: ~` (`kindIs "invalid"` check) so the empty-list override is respected.
- **Tests** restructured around `patch_engine_returning` / `make_pod` / `fake_stream` / `make_workload` fixture factories, with a new `test_decide_action_resolves_patch_version_via_nearest_lower` regression test for the patch-version mapping bug.
- **Newsfragment** updated to cover all four actions (added the fresh-install path), ticket references dropped throughout.
- **`values.schema.json`** updated for the new `drainTimeoutSeconds` key.
Follow-up: Ill open a separate PR to expose a public accessor in `airflow.utils.db` so this script can drop the `_REVISION_HEADS_MAP` underscore once 3.3 lands.
Full chart test suite (115 tests across `test_basic_helm_chart` + `test_migrate_database_job*` + `test_db_migrate_script`) passes locally.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387411342
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -103,8 +103,18 @@ spec:
{{- if .Values.migrateDatabaseJob.command }}
command: {{- tpl (toYaml .Values.migrateDatabaseJob.command) . | nindent 12 }}
{{- end }}
- {{- if .Values.migrateDatabaseJob.args }}
+ {{- if not (kindIs "invalid" .Values.migrateDatabaseJob.args) }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
+ # Bidirectional metadata DB reconciliation: forward migrate, no-op,
+ # fresh install, or exec downgrade in the still-running api-server
+ # pod depending on the relationship between the chart's target
+ # version and the DB's current alembic head.
+ args:
+ - "python"
+ - "-c"
+ - |-
+ {{- .Files.Get "files/db_migrate.py" | nindent 14 }}
Review Comment:
The `{{- ... }}` whitespace-trimming form inside a YAML block scalar is fragile and can inadvertently alter required newlines/indentation depending on surrounding whitespace. Prefer `{{ ... }}` here (no left-trim) so the rendered `|-` block consistently starts on the next line with the correct indentation.
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py:
##########
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for the RBAC granted to the migrate-database-job.
+
+The migrate-database-job needs ``pods`` + ``pods/exec`` access only on the
+downgrade branch, where it must exec ``airflow db downgrade`` inside the
+still-running api-server pod (whose image still ships the reverse alembic
+scripts). The Role is always rendered so the forward-migrate path remains
+identical for users who only ever upgrade.
+
+Tracked in https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import jmespath
+import pytest
+from chart_utils.helm_template_generator import render_chart
+
+ROLE_TEMPLATE = "templates/rbac/migrate-database-job-role.yaml"
+ROLEBINDING_TEMPLATE = "templates/rbac/migrate-database-job-rolebinding.yaml"
+
+
+class TestMigrateDatabaseJobRBAC:
+ def test_role_and_binding_render_by_default(self):
+ docs = render_chart(show_only=[ROLE_TEMPLATE, ROLEBINDING_TEMPLATE])
+ kinds = sorted(d["kind"] for d in docs)
+ assert kinds == ["Role", "RoleBinding"]
+
+ @pytest.mark.parametrize("rbac_create", [False, True])
+ def test_gated_on_rbac_create(self, rbac_create):
+ docs = render_chart(
+ values={"rbac": {"create": rbac_create}},
+ show_only=[ROLE_TEMPLATE, ROLEBINDING_TEMPLATE],
+ )
+ assert bool(docs) is rbac_create
+
+ def test_role_rules_grant_pods_and_exec(self):
+ docs = render_chart(show_only=[ROLE_TEMPLATE])
+ rules = jmespath.search("rules", docs[0])
+ resources_to_verbs = {tuple(r["resources"]): set(r["verbs"]) for r in rules}
+ assert ("pods",) in resources_to_verbs
+ assert {"get", "list"}.issubset(resources_to_verbs[("pods",)])
+ assert ("pods/exec",) in resources_to_verbs
+ # Exec is authorized under ``create`` or ``get`` depending on the
+ # apiserver/client path, so the role grants both for compatibility.
+ assert resources_to_verbs[("pods/exec",)] == {"create", "get"}
+
+ def test_role_rules_grant_deployments_and_statefulsets_scale(self):
+ # The downgrade branch scales every DB-touching workload to 0 after
+ # running ``airflow db downgrade`` so no OLD code keeps querying the
+ # now-downgraded schema before helm rolls in TARGET pods.
+ docs = render_chart(show_only=[ROLE_TEMPLATE])
+ rules = jmespath.search("rules", docs[0])
+ resources_to_verbs = {tuple(r["resources"]): set(r["verbs"]) for r in rules}
+ assert ("deployments", "statefulsets") in resources_to_verbs
+ assert {"get", "list"}.issubset(resources_to_verbs[("deployments", "statefulsets")])
+ assert ("deployments/scale", "statefulsets/scale") in resources_to_verbs
+ # Only ``patch`` is needed -- the reconciler never reads the scale back.
+ assert resources_to_verbs[("deployments/scale", "statefulsets/scale")] == {"patch"}
+
+ def test_rolebinding_subject_is_migrate_db_job_sa(self):
+ docs = render_chart(show_only=[ROLEBINDING_TEMPLATE])
+ subjects = jmespath.search("subjects", docs[0])
+ assert len(subjects) == 1
+ assert subjects[0]["kind"] == "ServiceAccount"
+ # Matches the name rendered by templates/jobs/migrate-database-job-serviceaccount.yaml
+ # via the "migrateDatabaseJob.serviceAccountName" helper.
+ assert subjects[0]["name"] == "release-name-airflow-migrate-database-job"
+
+ def test_rolebinding_references_role(self):
+ docs = render_chart(show_only=[ROLEBINDING_TEMPLATE])
+ role_ref = jmespath.search("roleRef", docs[0])
+ assert role_ref["kind"] == "Role"
+ assert role_ref["name"] == "release-name-migrate-database-job-role"
+ assert role_ref["apiGroup"] == "rbac.authorization.k8s.io"
Review Comment:
The expected Role name in this test does not match the template: `migrate-database-job-rolebinding.yaml` sets `roleRef.name` to `{{ include "airflow.fullname" . }}-migrate-database-job-role`, which (with default naming) typically expands to `release-name-airflow-migrate-database-job-role`. Update the assertion to match the rendered name (or compute it from the same naming helper if the test harness supports it).
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,378 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _int_env(name: str, default: int) -> int:
+ """Return integer env var *name*, or *default* when unset.
+
+ These knobs come from chart config; a typo (e.g. ``"5m"``) should fail with
+ an actionable message rather than an opaque ``ValueError`` traceback deep in
+ the job logs.
+ """
+ raw = os.environ.get(name)
+ if raw is None:
+ return default
+ try:
+ return int(raw)
+ except ValueError:
+ raise SystemExit(f"{name} must be an integer number of seconds, got {raw!r}")
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = _int_env("DB_CONNECT_MAX_WAIT_SECONDS", 120)
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
Review Comment:
`discover_api_server_pod()` selects pods only by `component=api-server` within the namespace. If multiple Airflow releases exist in the same namespace, this can pick the wrong release’s api-server pod and exec the downgrade in the wrong place. Pass `release_name` into `discover_api_server_pod()` and include it in the selector (e.g., `release=<release_name>,component=api-server`) to scope the exec strictly to the current Helm release.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4679393978
>
as @bugraoz93 pointed out having E2E tests for each path would be really cool to have, currently just existing forward path and fresh install are covered if noticed correct from tests ? can you take on further E2E testing for this as I dont know that much this project test setups
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "github-actions[bot] (via GitHub) via commits" <co...@airflow.apache.org>.
github-actions[bot] closed pull request #68074: Chart: Reconcile metadata DB bidirectionally on helm upgrade
URL: https://github.com/apache/airflow/pull/68074
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3365404385
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -105,6 +105,16 @@ spec:
{{- end }}
{{- if .Values.migrateDatabaseJob.args }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
Review Comment:
The args override check uses Helm truthiness. If a user sets migrateDatabaseJob.args to an empty list (valid per schema) to intentionally pass no args (e.g. rely on image defaults, or pair with a custom command), Helm treats [] as false and the chart will unexpectedly inject the embedded db_migrate.py args.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
Review Comment:
decide_action() treats any target version not present as an exact key in _REVISION_HEADS_MAP as "unknown" and forces the forward-migrate branch. _REVISION_HEADS_MAP intentionally does not list every patch release (it falls back to the nearest lower mapped version in Airflow's own CLI), so this will break common targets like 3.2.2 and prevent downgrades/no-ops from being detected correctly.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
Review Comment:
The new script relies on a "nearest lower" mapping behavior for patch versions that aren't explicit keys in _REVISION_HEADS_MAP (e.g. 3.2.2 -> 3.2.0 head). There isn't a regression test covering this, so a future refactor could reintroduce the exact-key-only bug and silently break default chart versions.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
Review Comment:
These tests create several unspecced MagicMocks. Using autospecced mocks helps catch typos and API drift (and matches Airflow's test guidelines about avoiding free-form mocks).
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,31 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``pre-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of three actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``target == current`` — no-op.
+* ``target > current`` — ``airflow db migrate`` (existing behaviour).
+* ``target < current`` — ``airflow db downgrade --to-version <target>``
Review Comment:
This newsfragment says the reconciler picks one of three actions, but the shipped script also has a distinct "fresh install" path when there is no current alembic revision. The release note should include that case (and update the count) so operators understand the full behavior.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3365334043
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -105,6 +105,16 @@ spec:
{{- end }}
{{- if .Values.migrateDatabaseJob.args }}
args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }}
+ {{- else }}
+ # Bidirectional metadata DB reconciliation: forward migrate, no-op, or
+ # exec downgrade in the still-running api-server pod depending on the
+ # relationship between the chart's target version and the DB's current
+ # alembic head. See https://github.com/apache/airflow/issues/68072.
Review Comment:
Same here: Would not point to the issue ticket.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
Review Comment:
Would be great to open a PR that exposes the needed information on a public interface in parallel such that this can land in Airflow 3.3. Then this script could attempt first to take the public method and fall-back on private member until (somewhen in future) support for Airflow 3.3 is cleaned.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
Review Comment:
Handling of case that no pod is ready is missing, gets an array-out-of-bounds.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
Review Comment:
Why are you not looking into `_REVISION_HEADS_MAP` like above? Assuming that only "real" versions are listed I _think_ this is sufficient, if somebody has a version in between or build manually from main, I think would be OK not to support all side-cases. Would reduce dependability in internal Airflow details.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
Review Comment:
I think reference to a past issue is not really needed. Context from above is sufficient.
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,78 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+##
+## Tracked in https://github.com/apache/airflow/issues/68072.
Review Comment:
Also here propose to remove ticket number
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-new"
+ fake_api.list_namespaced_pod.assert_called_once_with(
+ namespace="airflow",
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ )
+
+
+def test_discover_api_server_pod_falls_back_to_non_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [_pod("api-server-old", ready=False)]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-old"
+
+
+def test_discover_api_server_pod_raises_when_none(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = []
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ with pytest.raises(RuntimeError, match="no Running api-server pod"):
+ db_migrate.discover_api_server_pod("airflow")
+
+
+# --------------------------------------------------------------------------
+# run_downgrade_in_api_server
+# --------------------------------------------------------------------------
+
+
+def _fake_stream(returncode):
Review Comment:
Why not making this a fixture?
##########
chart/values.yaml:
##########
@@ -2121,13 +2121,11 @@ migrateDatabaseJob:
command: ~
# Args to use when running the migrate database job (templated).
- args:
- - "bash"
- - "-c"
- - >-
- exec \
-
- airflow db migrate
+ # When unset, the chart runs a bidirectional reconciliation script
+ # (forward migrate, no-op, or exec downgrade against the running api-server)
+ # depending on whether the chart is being upgraded or downgraded.
+ # See https://github.com/apache/airflow/issues/68072.
Review Comment:
Also here... ticket ref not needed.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
+ pod = types.SimpleNamespace()
+ pod.metadata = types.SimpleNamespace(name=name)
+ pod.status = types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ )
+ return pod
+
+
+def test_discover_api_server_pod_prefers_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [
+ _pod("api-server-old", ready=False),
+ _pod("api-server-new", ready=True),
+ ]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-new"
+ fake_api.list_namespaced_pod.assert_called_once_with(
+ namespace="airflow",
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ )
+
+
+def test_discover_api_server_pod_falls_back_to_non_ready(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = [_pod("api-server-old", ready=False)]
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ assert db_migrate.discover_api_server_pod("airflow") == "api-server-old"
+
+
+def test_discover_api_server_pod_raises_when_none(db_migrate, monkeypatch):
+ fake_api = mock.MagicMock()
+ fake_api.list_namespaced_pod.return_value.items = []
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: fake_api)
+ with pytest.raises(RuntimeError, match="no Running api-server pod"):
+ db_migrate.discover_api_server_pod("airflow")
+
+
+# --------------------------------------------------------------------------
+# run_downgrade_in_api_server
+# --------------------------------------------------------------------------
+
+
+def _fake_stream(returncode):
+ resp = mock.MagicMock()
+ # is_open(): True once so the loop body runs, then False to exit.
+ resp.is_open.side_effect = [True, False]
+ resp.peek_stdout.return_value = False
+ resp.peek_stderr.return_value = False
+ resp.returncode = returncode
+ return resp
+
+
+def test_run_downgrade_returns_zero_on_success(db_migrate, monkeypatch):
+ resp = _fake_stream(returncode=0)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 0
+
+
+def test_run_downgrade_propagates_nonzero(db_migrate, monkeypatch):
+ resp = _fake_stream(returncode=2)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 2
+
+
+def test_run_downgrade_treats_missing_returncode_as_failure(db_migrate, monkeypatch):
+ # Regression: previous code did ``return returncode or 0`` which silently
+ # reported success when the stream closed without an exit code.
+ resp = _fake_stream(returncode=None)
+ monkeypatch.setattr(db_migrate.k8s_config, "load_incluster_config", lambda: None)
+ monkeypatch.setattr(db_migrate.client, "CoreV1Api", lambda: mock.MagicMock())
+ monkeypatch.setattr(db_migrate, "stream", lambda *a, **kw: resp)
+ assert db_migrate.run_downgrade_in_api_server("p", "ns", "3.0.0") == 1
+
+
+# --------------------------------------------------------------------------
+# scale_release_workloads_to_zero
+# --------------------------------------------------------------------------
+
+
+def _wl(name):
Review Comment:
Why not making this a fixture?
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
Review Comment:
Why not making this a fixture?
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
+ """Patch ``engine.connect()`` so MigrationContext.get_current_revision() returns *current_rev*."""
+
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+
+def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head):
+ """Patch ScriptDirectory.from_config so walk_revisions("base", X) yields ancestors_by_head[X]."""
+
+ class _Rev:
+ def __init__(self, revision):
+ self.revision = revision
+
+ class _ScriptDir:
+ def walk_revisions(self, base, head):
+ assert base == "base"
+ return [_Rev(r) for r in ancestors_by_head[head]]
+
+ monkeypatch.setattr(db_migrate.ScriptDirectory, "from_config", staticmethod(lambda _cfg: _ScriptDir()))
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_fresh_when_db_unreachable(db_migrate, monkeypatch, patched_revision_map):
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_fresh_when_no_alembic_row(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(db_migrate, monkeypatch, patched_revision_map):
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_ancestor_of_target(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ # current=3.0.0, target=3.1.0 -> rev_300 is in target's ancestor set.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_300")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_not_in_target_ancestors(
+ db_migrate, monkeypatch, patched_revision_map
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade, the TARGET image's ScriptDirectory does NOT contain
+ ``current_rev`` (current is newer than target). The previous implementation
+ called ``walk_revisions("base", current_rev)`` which raised
+ ``RevisionError`` and never reached the downgrade branch. The fix walks to
+ ``target_rev`` (always present in target's scripts) and checks membership.
+ """
+ # current=3.2.0 (newer), target=3.1.0 (older). rev_320 is unknown to target.
+ _patch_engine_returning(db_migrate, monkeypatch, current_rev="rev_320")
+ _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head={"rev_310": ["rev_300", "rev_310"]})
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _pod(name, ready=True):
Review Comment:
Why not making this a fixture?
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,305 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+The decide_action regression test exists because the first iteration walked
+the target image's ScriptDirectory from base to ``current_rev``, which raises
+``RevisionError`` in the actual downgrade case where ``current_rev`` is newer
+than anything the target image knows about. See
+https://github.com/apache/airflow/issues/68072.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load chart/files/db_migrate.py as a module.
+
+ The file is normally fed to ``python3 -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable {version: revision} mapping for decide_action tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+def _patch_engine_returning(db_migrate, monkeypatch, current_rev):
Review Comment:
Why not making this a fixture?
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(namespace: str, release_name: str, timeout_seconds: int = 300) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+ """
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
Review Comment:
I think it is fair and usual that worker Pods might have a longer job running and might need to scale down for an hour. Can you make it possible to set a higher timeout or ignore longer worker runtime and leave them running? Whereas ... probably the workload "dies" anyway is api server is scaled-down, then no heartbeats can be sent...
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``."""
+ target_rev = _REVISION_HEADS_MAP.get(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ try:
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+ except OperationalError:
+ # DB unreachable -> treat as fresh install; ``airflow db migrate`` will
+ # re-surface a real connectivity error with a clearer message.
+ return "fresh"
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Walk the TARGET image's revision graph from base to target_rev. Target's
+ # ScriptDirectory only contains revisions up to target's heads, so we
+ # cannot walk from an unknown current_rev (the downgrade case): the call
+ # would raise ``RevisionError``. Instead, build the ancestor set of target
+ # and check membership.
+ cfg = Config()
+ cfg.set_main_option(
+ "script_location",
+ os.path.join(os.path.dirname(airflow.__file__), "migrations"),
+ )
+ script = ScriptDirectory.from_config(cfg)
+ ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)}
+ if current_rev in ancestors_of_target:
+ # current is an ancestor of target -> moving forward.
+ return "forward"
+ # current is not in target's history -> must be newer than target.
+ return "downgrade"
+
+
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
Review Comment:
As API Server might be "rolled-over" can you add a `@retry()` decorator to this for resiliency?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3371310966
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,247 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Reference: https://github.com/apache/airflow/issues/68072
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.config import Config
+from alembic.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from sqlalchemy.exc import OperationalError
+
+import airflow
+from airflow.settings import engine
+
+# NOTE: _REVISION_HEADS_MAP is a private symbol in airflow.utils.db. Tracked in
+# #68072 to expose a public accessor; using the private name is the only way
+# today to map a target version string to an alembic revision.
Review Comment:
Good idea — I'll open the follow-up PR to expose a public accessor and update this script to prefer it once it's available.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4674833233
> > What do other reviewers think about it? @bugraoz93 @jedcunningham @Miretpl @potiuk ?
>
> Very much needed IMO. If we have it, it will make manual chart adjustments for downgrade procedures or doing different types of stuff around it not needed. It would be nice to have it also for the 1.2x line.
Let's attempt backporting.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387293319
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,85 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ # The reconciler execs via ``connect_get_namespaced_pod_exec`` (an HTTP GET
+ # to the exec subresource), so the API server authorizes it under the
+ # ``get`` verb -- not ``create`` (which is what a POST-based exec needs).
+ verbs:
+ - "get"
Review Comment:
Good point on compatibility. Reverted to granting both `create` and `get` on `pods/exec` so the downgrade exec works regardless of whether the apiserver/client path authorizes it as a POST (`create`) or a GET-based upgrade (`get`). Updated the role comment and the RBAC test assertion accordingly.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,355 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
Review Comment:
Fixed. `_resolve_target_rev` now wraps `Version(target)` in `try/except InvalidVersion` and returns `None` for an unparseable `AIRFLOW_TARGET_VERSION`, so `decide_action` treats it as an unknown target and falls back to a conservative forward migrate instead of crashing the job. Added `test_decide_action_unparseable_target_falls_back_to_forward` as a regression test.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,355 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
Review Comment:
Same fix as the sibling comment: `_resolve_target_rev` now catches `InvalidVersion` and returns `None` (→ conservative forward) so a misconfigured `airflowVersion` no longer crashes the job. Covered by a new regression test.
##########
chart/values.schema.json:
##########
@@ -6731,19 +6731,21 @@
"default": null
},
"args": {
- "description": "Args to use when running migrate database job (templated).",
+ "description": "Args to use when running migrate database job (templated). When unset, the chart embeds a script that reconciles the metadata DB bidirectionally (forward migrate, no-op, or downgrade via the still-running api-server pod).",
Review Comment:
Updated the `migrateDatabaseJob.args` schema description to include the fresh-install path (no alembic row) and to clarify that an explicit list — including `[]` — overrides and disables the embedded reconciler, matching `values.yaml` and the newsfragment.
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,511 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Unit tests for the embedded bidirectional reconciler ``chart/files/db_migrate.py``.
+
+These exercise the runtime behaviour of the script itself (the helm template
+tests only check that the script is embedded with the right env/args/RBAC).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import types
+from unittest import mock
+
+import pytest
+from kubernetes import client as k8s_client
+from sqlalchemy.exc import OperationalError
+
+DB_MIGRATE_PATH = pathlib.Path(__file__).resolve().parents[3] / "files" / "db_migrate.py"
+
+
+@pytest.fixture(scope="module")
+def db_migrate():
+ """Load ``chart/files/db_migrate.py`` as a module.
+
+ The file is normally fed to ``python -c`` by the helm job rather than
+ imported, so it lives outside the chart's Python package tree.
+ """
+ spec = importlib.util.spec_from_file_location("chart_db_migrate", DB_MIGRATE_PATH)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+# --------------------------------------------------------------------------
+# Shared fixtures
+# --------------------------------------------------------------------------
+
+
+@pytest.fixture
+def patched_revision_map(db_migrate, monkeypatch):
+ """Install a stable ``{version: revision}`` mapping for ``decide_action`` tests."""
+ fake_map = {
+ "3.0.0": "rev_300",
+ "3.1.0": "rev_310",
+ "3.2.0": "rev_320",
+ }
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", fake_map)
+ return fake_map
+
+
+@pytest.fixture
+def patch_engine_returning(db_migrate, monkeypatch):
+ """Factory: patch ``engine.connect()`` so ``get_current_revision()`` returns the given value."""
+
+ def _patch(current_rev):
+ class _Ctx:
+ def get_current_revision(self):
+ return current_rev
+
+ class _Conn:
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *_):
+ return False
+
+ monkeypatch.setattr(db_migrate.engine, "connect", lambda: _Conn())
+ monkeypatch.setattr(
+ db_migrate.MigrationContext,
+ "configure",
+ staticmethod(lambda _conn: _Ctx()),
+ )
+
+ return _patch
+
+
+@pytest.fixture
+def make_pod():
+ """Factory: build a fake api-server pod with optional ``Ready`` condition."""
+
+ def _build(name, ready=True):
+ return types.SimpleNamespace(
+ metadata=types.SimpleNamespace(name=name),
+ status=types.SimpleNamespace(
+ conditions=[types.SimpleNamespace(type="Ready", status="True" if ready else "False")]
+ ),
+ )
+
+ return _build
+
+
+@pytest.fixture
+def fake_stream():
+ """Factory: build a fake kubernetes exec-stream response with the given exit code."""
+
+ def _build(returncode):
+ resp = mock.MagicMock()
+ # is_open(): True once so the loop body runs, then False to exit.
+ resp.is_open.side_effect = [True, False]
+ resp.peek_stdout.return_value = False
+ resp.peek_stderr.return_value = False
+ resp.returncode = returncode
+ return resp
+
+ return _build
+
+
+@pytest.fixture
+def make_workload():
+ """Factory: build a fake Deployment / StatefulSet object with a ``metadata.name``."""
+
+ def _build(name):
+ return types.SimpleNamespace(metadata=types.SimpleNamespace(name=name))
+
+ return _build
+
+
+# --------------------------------------------------------------------------
+# decide_action
+# --------------------------------------------------------------------------
+
+
+def test_decide_action_unknown_target_falls_back_to_forward(db_migrate, monkeypatch):
+ # Empty map: any target is unknown -> conservative forward migrate.
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {})
+ assert db_migrate.decide_action("9.9.9") == "forward"
+
+
+def test_decide_action_propagates_db_error(db_migrate, monkeypatch, patched_revision_map):
+ """``decide_action`` assumes the DB is already reachable.
+
+ The wait is the responsibility of :func:`main` via
+ :func:`_wait_for_db_ready`; ``decide_action`` itself must surface a
+ connectivity failure rather than silently treating it as a fresh
+ install (which would then run ``airflow db migrate`` against a still-
+ unreachable DB and cause ``BackoffLimitExceeded``).
+ """
+
+ def _raise(*_a, **_kw):
+ raise OperationalError("SELECT 1", {}, Exception("unreachable"))
+
+ monkeypatch.setattr(db_migrate.engine, "connect", _raise)
+ with pytest.raises(OperationalError):
+ db_migrate.decide_action("3.1.0")
+
+
+def test_decide_action_fresh_when_no_alembic_row(patched_revision_map, db_migrate, patch_engine_returning):
+ patch_engine_returning(current_rev=None)
+ assert db_migrate.decide_action("3.1.0") == "fresh"
+
+
+def test_decide_action_noop_when_current_equals_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ patch_engine_returning(current_rev="rev_310")
+ assert db_migrate.decide_action("3.1.0") == "noop"
+
+
+def test_decide_action_forward_when_current_is_older_than_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ # current=3.0.0 (rev_300), target=3.1.0 (rev_310) -> forward.
+ patch_engine_returning(current_rev="rev_300")
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_downgrade_when_current_is_newer_than_target(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ """Regression test for the original blocker.
+
+ On a real downgrade the TARGET image cannot resolve a revision newer than
+ its own head; relying on ``ScriptDirectory.walk_revisions`` raised
+ ``RevisionError`` and skipped the downgrade branch entirely. The new
+ implementation uses a version-aware reverse lookup so the comparison
+ works no matter which image is loaded.
+ """
+ # current=3.2.0 (rev_320, newer), target=3.1.0 (older) -> downgrade.
+ patch_engine_returning(current_rev="rev_320")
+ assert db_migrate.decide_action("3.1.0") == "downgrade"
+
+
+def test_decide_action_forward_when_current_revision_is_unknown(
+ patched_revision_map, db_migrate, patch_engine_returning
+):
+ """Dev / intermediate alembic rev -> conservative forward, never downgrade."""
+ patch_engine_returning(current_rev="rev_unknown")
+ assert db_migrate.decide_action("3.1.0") == "forward"
+
+
+def test_decide_action_resolves_patch_version_via_nearest_lower(
+ db_migrate, monkeypatch, patch_engine_returning
+):
+ """Regression test for the Copilot-flagged patch-version mapping bug.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch (e.g. ``3.2.2`` is not
+ a key — its head is the same as ``3.2.0``'s). The resolver must fall back
+ to the highest lower mapped version so a deploy targeting a patch release
+ is correctly classified as a no-op or downgrade against the same rev.
+ """
+ monkeypatch.setattr(db_migrate, "_REVISION_HEADS_MAP", {"3.2.0": "rev_320"})
+ patch_engine_returning(current_rev="rev_320")
+ # 3.2.2 must resolve to rev_320 -> noop, not "forward" (the previous bug).
+ assert db_migrate.decide_action("3.2.2") == "noop"
+
+
+# --------------------------------------------------------------------------
+# discover_api_server_pod
+# --------------------------------------------------------------------------
+
+
+def _call_discover_no_retry(db_migrate, namespace):
+ """Invoke the underlying function, bypassing the @retry wrapper backoff."""
+ return db_migrate.discover_api_server_pod.retry_with(stop=lambda _r: True)(namespace)
Review Comment:
Removed the unused `_call_discover_no_retry` helper.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Miretpl (via GitHub)" <gi...@apache.org>.
Miretpl commented on PR #68074:
URL: https://github.com/apache/airflow/pull/68074#issuecomment-4650776675
> What do other reviewers think about it? @bugraoz93 @jedcunningham @Miretpl @potiuk ?
Very much needed IMO. If we have it, it will make manual chart adjustments for downgrade procedures or doing different types of stuff around it not needed. It would be nice to have it also for the 1.2x line.
I will check the PR in a short period of time, probably.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3376369384
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -43,7 +43,7 @@ metadata:
{{- end }}
{{- $annotations := dict }}
{{- if .Values.migrateDatabaseJob.useHelmHooks }}
- {{- $_ := set $annotations "helm.sh/hook" "post-install,post-upgrade" }}
+ {{- $_ := set $annotations "helm.sh/hook" "pre-install,pre-upgrade" }}
{{- $_ := set $annotations "helm.sh/hook-weight" "1" }}
{{- $_ := set $annotations "helm.sh/hook-delete-policy" "before-hook-creation,hook-succeeded" }}
{{- end }}
Review Comment:
The migrate-database Job is now a `pre-install,pre-upgrade` hook, but it still runs under `serviceAccountName: {{ include "migrateDatabaseJob.serviceAccountName" . }}`. Helm does not create non-hook resources (including the migrate-database-job ServiceAccount and these new Role/RoleBindings) before running pre-install/pre-upgrade hooks, so on a fresh install the Job will start before its ServiceAccount exists, and on the first downgrade after adopting this chart version the Job may run before the new RBAC is created. To make the hook reliable, the ServiceAccount (and any required Role/RoleBinding) need to be created as hook resources with a lower hook weight than the Job, or the Job hook timing needs to be adjusted so required RBAC exists before it runs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Miretpl (via GitHub)" <gi...@apache.org>.
Miretpl commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3407864073
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,88 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
Review Comment:
This will be only needed when the proper args value will be set. Could you add the condition for that too? We should not generate things which are not needed.
##########
chart/values.yaml:
##########
@@ -2128,13 +2128,18 @@ migrateDatabaseJob:
command: ~
# Args to use when running the migrate database job (templated).
- args:
- - "bash"
- - "-c"
- - >-
- exec \
+ # When unset (the default), the chart runs a bidirectional reconciliation
+ # script that picks one of fresh install, forward migrate, no-op, or exec
+ # downgrade against the running api-server, depending on whether the chart
+ # is being installed, upgraded, or downgraded. Set to an explicit list
+ # (including ``[]``) to override.
+ args: ~
Review Comment:
Due to the discussion about potential instability of the feature, I would leave the default for being forward only, making sure that users are aware of the downsides of the current solution (we would need also doc regarding it).
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -372,7 +444,13 @@ def test_command_and_args_overrides(self, command, args):
)
assert command == jmespath.search("spec.template.spec.containers[0].command", docs[0])
- assert args == jmespath.search("spec.template.spec.containers[0].args", docs[0])
+ if args is None:
+ # When args is unset, the chart's bidirectional default kicks in
+ # (covered in detail by other tests in this class).
+ rendered_args = jmespath.search("spec.template.spec.containers[0].args", docs[0])
+ assert rendered_args[:2] == ["python", "-c"]
+ else:
+ assert args == jmespath.search("spec.template.spec.containers[0].args", docs[0])
Review Comment:
I think that these should be two separate tests to make debugging easier and provide more clarity too.
##########
chart/templates/jobs/migrate-database-job-serviceaccount.yaml:
##########
@@ -35,8 +35,17 @@ metadata:
{{- if or .Values.labels .Values.migrateDatabaseJob.labels }}
{{- mustMerge .Values.migrateDatabaseJob.labels .Values.labels | toYaml | nindent 4 }}
{{- end }}
+ {{- $annotations := dict }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ {{- $_ := set $annotations "helm.sh/hook" "post-install,pre-upgrade" }}
+ {{- $_ := set $annotations "helm.sh/hook-weight" "-5" }}
Review Comment:
```suggestion
{{- $_ := set $annotations "helm.sh/hook-weight" "0" }}
```
Should be enough
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -279,6 +279,77 @@ def test_should_set_correct_helm_hooks_weight(self):
annotations = jmespath.search("metadata.annotations", docs[0])
assert annotations["helm.sh/hook-weight"] == "1"
+ def test_default_hooks_run_post_install_pre_upgrade(self):
+ # post-install: the chart's bundled postgres is created as a regular
+ # release resource, so the job needs the DB to exist before it runs.
+ # pre-upgrade: the reconciler must run *before* helm rolls the new
+ # image, so on a downgrade it can exec ``airflow db downgrade`` in the
+ # still-running OLD api-server pod (the OLD image still ships the
+ # reverse alembic scripts; the new one does not).
+ # Tracked in https://github.com/apache/airflow/issues/68072.
+ docs = render_chart(
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ annotations = jmespath.search("metadata.annotations", docs[0])
+ assert annotations["helm.sh/hook"] == "post-install,pre-upgrade"
+
+ def test_default_args_contain_bidirectional_migrate_script(self):
+ docs = render_chart(
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ args = jmespath.search("spec.template.spec.containers[0].args", docs[0])
+ assert args[:2] == ["python", "-c"]
+ # Confirm the embedded script is the bidirectional reconciler.
+ assert "AIRFLOW_TARGET_VERSION" in args[2]
Review Comment:
I understand it in the envs, why it is in args?
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,384 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from packaging.version import InvalidVersion, Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# The kubernetes client is only used by the downgrade branch (exec ``airflow db
+# downgrade`` in the api-server pod, then scale DB-touching workloads to zero).
+# The base Airflow image does not ship the kubernetes client unless the
+# cncf.kubernetes provider is installed, so importing it unconditionally would
+# break the far more common forward / no-op / fresh paths. Import it lazily and
+# only fail when a downgrade is actually selected.
+try:
+ from kubernetes import client, config as k8s_config
+ from kubernetes.stream import stream
+
+ KUBERNETES_IMPORT_ERROR: ImportError | None = None
+except ImportError as exc: # pragma: no cover - only on images without the k8s client
+ client = k8s_config = stream = None # type: ignore[assignment]
+ KUBERNETES_IMPORT_ERROR = exc
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _int_env(name: str, default: int) -> int:
+ """Return integer env var *name*, or *default* when unset.
+
+ These knobs come from chart config; a typo (e.g. ``"5m"``) should fail with
+ an actionable message rather than an opaque ``ValueError`` traceback deep in
+ the job logs.
+ """
+ raw = os.environ.get(name)
+ if raw is None:
+ return default
+ try:
+ return int(raw)
+ except ValueError:
+ raise SystemExit(f"{name} must be an integer number of seconds, got {raw!r}")
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ try:
+ target_v = Version(target)
+ except InvalidVersion:
+ # ``AIRFLOW_TARGET_VERSION`` comes from chart config; a nonstandard or
+ # misconfigured ``airflowVersion`` (e.g. a dev build) is not PEP 440
+ # parseable. Treat it as an unknown target so :func:`decide_action`
+ # falls back to a conservative forward migrate instead of crashing.
+ return None
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = _int_env("DB_CONNECT_MAX_WAIT_SECONDS", 120)
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str, release_name: str) -> str:
+ """Return the name of a Running api-server pod for *release_name* in *namespace*.
+
+ Scoped to the release via the ``release`` label so that, when several
+ Airflow releases share a namespace, the downgrade is exec'd strictly in the
+ current release's api-server pod and never another release's.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector=f"release={release_name},component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(
+ f"no Running api-server pod found for release {release_name} in namespace {namespace}"
+ )
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = _int_env("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", 300)
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def _require_kubernetes_client() -> None:
+ """Fail clearly if the downgrade branch runs on an image without the k8s client.
+
+ The forward / no-op / fresh paths never import kubernetes, so only the
+ downgrade branch needs to assert it is available (see the guarded import
+ at the top of this module).
+ """
+ if KUBERNETES_IMPORT_ERROR is not None:
+ raise SystemExit(
+ "the chart downgrade path needs the 'kubernetes' python client, which is not "
+ "installed in this image. Install the cncf.kubernetes provider, or set "
+ "migrateDatabaseJob.args explicitly to supply your own migration command."
+ )
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
Review Comment:
Code should be self-explanatory, and the comments themselves should rather not tell what the code is doing but rather why it was written that way, not the other (e.g. some history context, edge case, etc.). This case is, at least for me, self-explanatory as to why we need to wait for the DB to be ready. I think that most of the comments in this file are just not needed and are the noise generated by AI, which makes, from my perspective, reviewing harder, cause I actually read the same thing over and over again. Due to that, I didn't finish reviewing this file - I will do it in the next iteration.
Could you please review all of the comments in this PR in terms of whether they are actually bringing any value which cannot be taken from the actual code?
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -279,6 +279,77 @@ def test_should_set_correct_helm_hooks_weight(self):
annotations = jmespath.search("metadata.annotations", docs[0])
assert annotations["helm.sh/hook-weight"] == "1"
+ def test_default_hooks_run_post_install_pre_upgrade(self):
+ # post-install: the chart's bundled postgres is created as a regular
+ # release resource, so the job needs the DB to exist before it runs.
+ # pre-upgrade: the reconciler must run *before* helm rolls the new
+ # image, so on a downgrade it can exec ``airflow db downgrade`` in the
+ # still-running OLD api-server pod (the OLD image still ships the
+ # reverse alembic scripts; the new one does not).
+ # Tracked in https://github.com/apache/airflow/issues/68072.
+ docs = render_chart(
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ annotations = jmespath.search("metadata.annotations", docs[0])
+ assert annotations["helm.sh/hook"] == "post-install,pre-upgrade"
+
+ def test_default_args_contain_bidirectional_migrate_script(self):
+ docs = render_chart(
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ args = jmespath.search("spec.template.spec.containers[0].args", docs[0])
+ assert args[:2] == ["python", "-c"]
+ # Confirm the embedded script is the bidirectional reconciler.
+ assert "AIRFLOW_TARGET_VERSION" in args[2]
+ assert "airflow db downgrade" in args[2]
+ assert "airflow db migrate" in args[2]
+
+ def test_user_supplied_args_still_win(self):
+ custom_args = ["bash", "-c", "exec airflow db migrate"]
+ docs = render_chart(
+ values={"migrateDatabaseJob": {"args": custom_args}},
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ assert jmespath.search("spec.template.spec.containers[0].args", docs[0]) == custom_args
+
+ def test_user_supplied_command_still_wins(self):
+ custom_command = ["/bin/my-entrypoint"]
+ docs = render_chart(
+ values={"migrateDatabaseJob": {"command": custom_command}},
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) == custom_command
+
+ def test_airflow_target_version_env_uses_airflow_version(self):
+ docs = render_chart(
+ values={"airflowVersion": "3.1.8"},
+ show_only=["templates/jobs/migrate-database-job.yaml"],
+ )
+ env = jmespath.search("spec.template.spec.containers[0].env", docs[0])
+ target_var = next((e for e in env if e["name"] == "AIRFLOW_TARGET_VERSION"), None)
+ assert target_var is not None
+ assert target_var["value"] == "3.1.8"
Review Comment:
```suggestion
assert jmespath.search("spec.template.spec.containers[0].env[?name=='AIRFLOW_TARGET_VERSION'] | [0].value", docs[0]) == "3.1.8"
```
If I remember notation correctly.
*the same applies to the other tests*
##########
chart/templates/rbac/migrate-database-job-rolebinding.yaml:
##########
@@ -0,0 +1,52 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+##############################################
+## Airflow Migrate Database Job Role Binding
+##############################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
Review Comment:
Same comment as for the roles.
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -352,16 +423,17 @@ def test_job_ttl_after_finished_nil(self):
assert "ttlSecondsAfterFinished" not in spec
def test_default_command_and_args_airflow_version(self):
+ # Default args now embed the bidirectional reconciliation script.
+ # The "args still win" / "command still wins" / "script body present"
Review Comment:
What does it mean?
##########
chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py:
##########
@@ -0,0 +1,554 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Could you please review this file in terms of consistency with the whole Helm chart test suite?
##########
chart/templates/rbac/migrate-database-job-rolebinding.yaml:
##########
@@ -0,0 +1,52 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+##############################################
+## Airflow Migrate Database Job Role Binding
+##############################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-rolebinding
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
Review Comment:
```suggestion
helm.sh/hook-weight: "0"
```
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -352,16 +423,17 @@ def test_job_ttl_after_finished_nil(self):
assert "ttlSecondsAfterFinished" not in spec
def test_default_command_and_args_airflow_version(self):
+ # Default args now embed the bidirectional reconciliation script.
+ # The "args still win" / "command still wins" / "script body present"
+ # cases above cover the override semantics in detail.
docs = render_chart(
show_only=["templates/jobs/migrate-database-job.yaml"],
)
assert jmespath.search("spec.template.spec.containers[0].command", docs[0]) is None
- assert jmespath.search("spec.template.spec.containers[0].args", docs[0]) == [
- "bash",
- "-c",
- "exec \\\nairflow db migrate",
- ]
+ args = jmespath.search("spec.template.spec.containers[0].args", docs[0])
+ assert args[:2] == ["python", "-c"]
+ assert "airflow db migrate" in args[2]
Review Comment:
I think it is better to test the whole content of the command. Without that, we can miss something in the future when there will be some change.
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py:
##########
@@ -279,6 +279,77 @@ def test_should_set_correct_helm_hooks_weight(self):
annotations = jmespath.search("metadata.annotations", docs[0])
assert annotations["helm.sh/hook-weight"] == "1"
+ def test_default_hooks_run_post_install_pre_upgrade(self):
+ # post-install: the chart's bundled postgres is created as a regular
+ # release resource, so the job needs the DB to exist before it runs.
+ # pre-upgrade: the reconciler must run *before* helm rolls the new
+ # image, so on a downgrade it can exec ``airflow db downgrade`` in the
+ # still-running OLD api-server pod (the OLD image still ships the
+ # reverse alembic scripts; the new one does not).
+ # Tracked in https://github.com/apache/airflow/issues/68072.
Review Comment:
```suggestion
# still-running OLD api-server pod.
```
I think that this should be clear for anyone working with Airflow and Helm
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py:
##########
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Should be under the `security` test set, not `airflow_aux`, to keep consistency within the tests.
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,42 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row) — ``airflow db migrate``.
+* ``target == current`` — no-op.
+* ``target > current`` — ``airflow db migrate`` (existing behaviour).
+* ``target < current`` — ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (so reverse alembic
+ scripts that ship only in the source-version image are available),
+ followed by scaling every DB-touching workload of this release
+ (``api-server``, ``scheduler``, ``triggerer``, ``dag-processor``,
+ ``worker``) to ``0``. Helm then patches those workloads back to their
+ configured ``replicas`` with the TARGET image as the upgrade proceeds,
+ so the cluster comes back up cleanly on the target version without any
+ OLD code running against the now-downgraded schema. A downgrade is
+ therefore a brief outage rather than a zero-downtime rolling update —
+ this is unavoidable when the schema goes backwards.
+
+The action above is only decided once the metadata DB is reachable. If the DB
+is still unreachable after the connection-wait window
+(``DB_CONNECT_MAX_WAIT_SECONDS``, default ``120``), the job exits non-zero and
+is retried via the Job ``backoffLimit`` rather than proceeding with a migrate.
+
+A ``Role`` + ``RoleBinding`` granting ``pods`` (``get``, ``list``),
+``pods/exec`` (``get``, ``create``), ``deployments`` / ``statefulsets`` (``get``, ``list``)
+and ``deployments/scale`` / ``statefulsets/scale`` (``patch``) to the
+existing ``migrate-database-job`` ``ServiceAccount`` are now rendered when
Review Comment:
Implementation details which could be rather documentation than part of the release notes.
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,88 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
Review Comment:
```suggestion
helm.sh/hook-weight: "0"
```
##########
chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py:
##########
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Could you please review this file in terms of consistency and test duplication within the whole Helm chart test suite?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3386801280
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,37 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row, or DB unreachable) — ``airflow db migrate``.
Review Comment:
This release note says “DB unreachable” is treated like a fresh install and runs `airflow db migrate`, but `chart/files/db_migrate.py` currently returns exit code 1 after the DB wait window when the DB is still unreachable. Please update the newsfragment wording to match the implemented behavior (or adjust the script if the intended behavior is actually to proceed with migrate when DB is unreachable, though that would likely just fail and backoff).
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
Review Comment:
The module imports the Kubernetes Python client at import-time, which will cause the migrate job to fail immediately (even for the forward/no-op paths) on images that don’t include `kubernetes` (common unless the Kubernetes provider/deps are installed). To avoid turning an optional downgrade-only dependency into a hard runtime requirement, move the `kubernetes` imports (and `stream`) inside the downgrade-only functions (`discover_api_server_pod`, `run_downgrade_in_api_server`, `scale_release_workloads_to_zero`) or guard them with a `try/except ImportError` that only raises if the selected action is `downgrade`.
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
+ # The downgrade branch scales every DB-touching workload to 0 after running
+ # ``airflow db downgrade`` so no OLD code keeps querying the now-downgraded
+ # schema before helm rolls in TARGET pods. See chart/files/db_migrate.py.
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments"
+ - "statefulsets"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - "apps"
+ resources:
+ - "deployments/scale"
+ - "statefulsets/scale"
+ verbs:
+ - "get"
+ - "patch"
Review Comment:
The Role grants extra verbs on subresources that don’t appear to be used by the script: `pods/exec` typically only needs `create` for exec, and the code patches `/scale` but doesn’t call `get` on the scale subresource. Dropping unnecessary verbs reduces the blast radius of this ServiceAccount if compromised. If `get` is required for your specific client/server behavior, consider adding a short comment explaining why it’s needed.
##########
chart/templates/rbac/migrate-database-job-role.yaml:
##########
@@ -0,0 +1,82 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+#######################################
+## Airflow Migrate Database Job Role
+##
+## Grants pods + pods/exec in the release namespace so the migrate-database-job
+## can exec ``airflow db downgrade`` inside the still-running api-server pod on
+## a chart downgrade. The forward-migrate branch never uses these permissions.
+#######################################
+{{- if and .Values.rbac.create .Values.migrateDatabaseJob.enabled }}
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+ name: {{ include "airflow.fullname" . }}-migrate-database-job-role
+ namespace: "{{ .Release.Namespace }}"
+ labels:
+ tier: airflow
+ component: run-airflow-migrations
+ release: {{ .Release.Name }}
+ chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+ heritage: {{ .Release.Service }}
+ {{- with .Values.labels }}
+ {{- toYaml . | nindent 4 }}
+ {{- end }}
+ {{- if .Values.migrateDatabaseJob.useHelmHooks }}
+ annotations:
+ helm.sh/hook: post-install,pre-upgrade
+ helm.sh/hook-weight: "-5"
+ helm.sh/hook-delete-policy: before-hook-creation,hook-succeeded
+ {{- end }}
+rules:
+ - apiGroups:
+ - ""
+ resources:
+ - "pods"
+ verbs:
+ - "get"
+ - "list"
+ - apiGroups:
+ - ""
+ resources:
+ - "pods/exec"
+ verbs:
+ - "create"
+ - "get"
Review Comment:
The Role grants extra verbs on subresources that don’t appear to be used by the script: `pods/exec` typically only needs `create` for exec, and the code patches `/scale` but doesn’t call `get` on the scale subresource. Dropping unnecessary verbs reduces the blast radius of this ServiceAccount if compromised. If `get` is required for your specific client/server behavior, consider adding a short comment explaining why it’s needed.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
Review Comment:
The new unit tests cover helper functions well, but `main()` (env validation, DB-unreachable exit behavior, and the action-based orchestration including short-circuiting on nonzero downgrade rc) is not directly tested. Adding focused tests with monkeypatching/mocking for `_wait_for_db_ready`, `decide_action`, `subprocess.run`, and the downgrade helpers would make regressions in the overall control flow much easier to catch.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
Review Comment:
The new unit tests cover helper functions well, but `main()` (env validation, DB-unreachable exit behavior, and the action-based orchestration including short-circuiting on nonzero downgrade rc) is not directly tested. Adding focused tests with monkeypatching/mocking for `_wait_for_db_ready`, `decide_action`, `subprocess.run`, and the downgrade helpers would make regressions in the overall control flow much easier to catch.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
+
+ action = decide_action(target)
+ print(f"[db_migrate] target={target} action={action}", flush=True)
+
+ if action == "noop":
+ print("[db_migrate] DB already at target revision, nothing to do.")
+ return 0
+ if action in {"fresh", "forward"}:
+ return subprocess.run(["airflow", "db", "migrate"], check=False).returncode
+ if action == "downgrade":
+ if not release_name:
+ raise SystemExit("RELEASE_NAME must be set for the downgrade branch")
+ pod = discover_api_server_pod(namespace)
+ print(f"[db_migrate] downgrading via api-server pod {pod}", flush=True)
+ rc = run_downgrade_in_api_server(pod, namespace, target)
+ if rc != 0:
+ return rc
+ # Drain OLD pods before helm rolls in TARGET pods to avoid OLD code
+ # running against the now-downgraded schema (see module docstring).
+ scale_release_workloads_to_zero(namespace, release_name)
+ return 0
Review Comment:
The new unit tests cover helper functions well, but `main()` (env validation, DB-unreachable exit behavior, and the action-based orchestration including short-circuiting on nonzero downgrade rc) is not directly tested. Adding focused tests with monkeypatching/mocking for `_wait_for_db_ready`, `decide_action`, `subprocess.run`, and the downgrade helpers would make regressions in the overall control flow much easier to catch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387222498
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+from tenacity import (
+ retry,
+ retry_if_exception_type,
+ stop_after_attempt,
+ stop_after_delay,
+ wait_exponential,
+ wait_fixed,
+)
+
+from airflow.settings import engine
+
+# ``_REVISION_HEADS_MAP`` is private today; a public accessor is being tracked
+# upstream so the chart can drop the leading underscore in a future release.
+from airflow.utils.db import _REVISION_HEADS_MAP
+
+
+def _resolve_target_rev(target: str) -> str | None:
+ """Return the alembic head for *target*, falling back to the nearest lower mapped version.
+
+ ``_REVISION_HEADS_MAP`` does not list every patch release — patches often
+ share the head of their minor's first release. Mirror Airflow's own CLI
+ behaviour by picking the highest mapped version that is ``<= target``.
+ """
+ if target in _REVISION_HEADS_MAP:
+ return _REVISION_HEADS_MAP[target]
+ target_v = Version(target)
+ candidates = [(Version(v), rev) for v, rev in _REVISION_HEADS_MAP.items() if Version(v) <= target_v]
+ if not candidates:
+ return None
+ return max(candidates, key=lambda pair: pair[0])[1]
+
+
+def _db_connect_stop(retry_state):
+ # Evaluate ``DB_CONNECT_MAX_WAIT_SECONDS`` at retry time (not import time)
+ # so that operators -- and unit tests -- can tune the wait window without
+ # reloading the module. ``/entrypoint`` skips its DB-wait for non-airflow
+ # commands (we run as ``python -c ...``), so on a fresh install with a
+ # bundled postgres still starting, the first connect attempt races the DB.
+ # Default 120s matches the entrypoint's
+ # ``CONNECTION_CHECK_MAX_COUNT`` * ``CONNECTION_CHECK_SLEEP_TIME``.
+ delay = int(os.environ.get("DB_CONNECT_MAX_WAIT_SECONDS", "120"))
+ return stop_after_delay(delay)(retry_state)
+
+
+@retry(
+ stop=_db_connect_stop,
+ wait=wait_fixed(3),
+ retry=retry_if_exception_type(OperationalError),
+ reraise=True,
+)
+def _wait_for_db_ready() -> None:
+ """Block until the metadata DB accepts a connection.
+
+ Called once at the top of :func:`main` so that *every* downstream step
+ (``decide_action``, ``airflow db migrate`` subprocess, ``kubernetes``
+ api-server pod exec) can assume the DB is reachable. Without this, the
+ subprocess branch had no DB-wait of its own and would fail immediately
+ on a fresh install where the bundled postgres was still initialising,
+ causing ``BackoffLimitExceeded`` on the Job.
+ """
+ with engine.connect() as conn:
+ conn.execute(text("SELECT 1"))
+
+
+def decide_action(target: str) -> str:
+ """Return one of ``noop``, ``forward``, ``downgrade``, ``fresh``.
+
+ Assumes the DB is already reachable -- :func:`_wait_for_db_ready` must
+ have been called first.
+ """
+ target_rev = _resolve_target_rev(target)
+ if target_rev is None:
+ # Unknown target version (e.g. dev build). Be conservative: forward only.
+ return "forward"
+
+ with engine.connect() as conn:
+ current_rev = MigrationContext.configure(conn).get_current_revision()
+
+ if current_rev is None:
+ return "fresh"
+ if current_rev == target_rev:
+ return "noop"
+
+ # Reverse-lookup current_rev to determine which version it belongs to,
+ # then compare versions. If current_rev isn't mapped (dev/intermediate
+ # alembic revision) be conservative and forward-migrate rather than risk
+ # an incorrect downgrade.
+ rev_to_version = {rev: ver for ver, rev in _REVISION_HEADS_MAP.items()}
+ current_version = rev_to_version.get(current_rev)
+ if current_version is None:
+ return "forward"
+ if Version(current_version) > Version(target):
+ return "downgrade"
+ return "forward"
+
+
+@retry(
+ stop=stop_after_attempt(5),
+ wait=wait_exponential(multiplier=2, min=2, max=30),
+ retry=retry_if_exception_type(RuntimeError),
+ reraise=True,
+)
+def discover_api_server_pod(namespace: str) -> str:
+ """Return the name of a Running api-server pod in *namespace*.
+
+ Retries on ``RuntimeError`` so a rolling restart of the api-server
+ deployment (no Running pod for a few seconds) does not fail the job.
+ """
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ pods = api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector="component=api-server",
+ field_selector="status.phase=Running",
+ ).items
+ if not pods:
+ raise RuntimeError(f"no Running api-server pod found in namespace {namespace}")
+
+ # Prefer Ready pods so we don't pick one mid-rollout.
+ ready = [
+ p for p in pods if any(c.type == "Ready" and c.status == "True" for c in (p.status.conditions or []))
+ ]
+ return (ready or pods)[0].metadata.name
+
+
+def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int:
+ """Exec ``airflow db downgrade`` in the api-server pod via the Kubernetes API."""
+ k8s_config.load_incluster_config()
+ api = client.CoreV1Api()
+ command = ["airflow", "db", "downgrade", "--to-version", target_version, "--yes"]
+
+ resp = stream(
+ api.connect_get_namespaced_pod_exec,
+ pod_name,
+ namespace,
+ container="api-server",
+ command=command,
+ stderr=True,
+ stdin=False,
+ stdout=True,
+ tty=False,
+ _preload_content=False,
+ )
+ while resp.is_open():
+ resp.update(timeout=1)
+ if resp.peek_stdout():
+ sys.stdout.write(resp.read_stdout())
+ sys.stdout.flush()
+ if resp.peek_stderr():
+ sys.stderr.write(resp.read_stderr())
+ sys.stderr.flush()
+ returncode = resp.returncode
+ resp.close()
+ # Treat an unknown exit code as failure rather than success: if the stream
+ # closes without populating returncode we cannot prove the downgrade ran.
+ if returncode is None:
+ sys.stderr.write("[db_migrate] downgrade exec stream closed without an exit code\n")
+ return 1
+ return returncode
+
+
+# Components that talk to the metadata DB and are managed by this helm release.
+# A downgrade-then-scale-back sequence must drain all of them so no OLD code
+# keeps talking to the now-downgraded schema before helm rolls in TARGET pods.
+_DB_TOUCHING_COMPONENTS = (
+ "api-server",
+ "scheduler",
+ "triggerer",
+ "dag-processor",
+ "worker",
+)
+
+
+def scale_release_workloads_to_zero(
+ namespace: str, release_name: str, timeout_seconds: int | None = None
+) -> None:
+ """Scale all DB-touching workloads of this release to 0 and wait for drain.
+
+ Helm will patch the same Deployments/StatefulSets back to ``replicas: N``
+ when it applies the post-hook manifests, so we deliberately do NOT scale
+ them up again here.
+
+ The drain deadline defaults to ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` (or
+ 300s if unset). Operators with long-running worker tasks can raise this
+ via the ``migrateDatabaseJob.drainTimeoutSeconds`` chart value.
+ """
+ if timeout_seconds is None:
+ timeout_seconds = int(os.environ.get("MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS", "300"))
+ k8s_config.load_incluster_config()
+ apps = client.AppsV1Api()
+ core = client.CoreV1Api()
+
+ selector = f"release={release_name},component in ({','.join(_DB_TOUCHING_COMPONENTS)})"
+ scale_body = {"spec": {"replicas": 0}}
+
+ deployments = apps.list_namespaced_deployment(namespace, label_selector=selector).items
+ statefulsets = apps.list_namespaced_stateful_set(namespace, label_selector=selector).items
+
+ for d in deployments:
+ print(f"[db_migrate] scaling Deployment/{d.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_deployment_scale(d.metadata.name, namespace, scale_body)
+ for s in statefulsets:
+ print(f"[db_migrate] scaling StatefulSet/{s.metadata.name} to 0", flush=True)
+ apps.patch_namespaced_stateful_set_scale(s.metadata.name, namespace, scale_body)
+
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ remaining = core.list_namespaced_pod(namespace, label_selector=selector).items
+ if not remaining:
+ print("[db_migrate] all DB-touching pods drained.", flush=True)
+ return
+ print(
+ f"[db_migrate] waiting for {len(remaining)} pod(s) to terminate...",
+ flush=True,
+ )
+ time.sleep(2)
+ raise TimeoutError(f"DB-touching pods did not drain within {timeout_seconds}s after scale-to-zero")
+
+
+def main() -> int:
+ target = os.environ.get("AIRFLOW_TARGET_VERSION")
+ namespace = os.environ.get("POD_NAMESPACE")
+ release_name = os.environ.get("RELEASE_NAME")
+ if not target:
+ raise SystemExit("AIRFLOW_TARGET_VERSION must be set")
+ if not namespace:
+ raise SystemExit("POD_NAMESPACE must be set")
+
+ # Block until the DB is reachable so that BOTH ``decide_action`` and the
+ # subsequent ``airflow db migrate`` subprocess (which does not go through
+ # ``/entrypoint`` and therefore has no DB-wait of its own) succeed on a
+ # fresh install where the bundled postgres may still be initialising.
+ try:
+ _wait_for_db_ready()
+ except OperationalError as exc:
+ print(f"[db_migrate] DB unreachable after wait window: {exc}", flush=True)
+ return 1
+
+ action = decide_action(target)
+ print(f"[db_migrate] target={target} action={action}", flush=True)
+
+ if action == "noop":
+ print("[db_migrate] DB already at target revision, nothing to do.")
+ return 0
+ if action in {"fresh", "forward"}:
+ return subprocess.run(["airflow", "db", "migrate"], check=False).returncode
+ if action == "downgrade":
+ if not release_name:
+ raise SystemExit("RELEASE_NAME must be set for the downgrade branch")
+ pod = discover_api_server_pod(namespace)
+ print(f"[db_migrate] downgrading via api-server pod {pod}", flush=True)
+ rc = run_downgrade_in_api_server(pod, namespace, target)
+ if rc != 0:
+ return rc
+ # Drain OLD pods before helm rolls in TARGET pods to avoid OLD code
+ # running against the now-downgraded schema (see module docstring).
+ scale_release_workloads_to_zero(namespace, release_name)
+ return 0
Review Comment:
Done. Added focused `main()` tests in `test_db_migrate_script.py` that monkeypatch `_wait_for_db_ready`, `decide_action`, `subprocess.run` and the downgrade helpers, covering: missing `AIRFLOW_TARGET_VERSION`/`POD_NAMESPACE` validation, the DB-unreachable exit (returns 1), noop/forward/fresh dispatch, downgrade orchestration (exec then scale), short-circuit on a nonzero downgrade rc, and the missing-`kubernetes`-client path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Chart: Reconcile metadata DB bidirectionally on helm upgrade [airflow]
Posted by "jykae (via GitHub)" <gi...@apache.org>.
jykae commented on code in PR #68074:
URL: https://github.com/apache/airflow/pull/68074#discussion_r3387219415
##########
chart/templates/jobs/migrate-database-job.yaml:
##########
@@ -43,7 +43,7 @@ metadata:
{{- end }}
{{- $annotations := dict }}
{{- if .Values.migrateDatabaseJob.useHelmHooks }}
- {{- $_ := set $annotations "helm.sh/hook" "post-install,post-upgrade" }}
+ {{- $_ := set $annotations "helm.sh/hook" "pre-install,pre-upgrade" }}
{{- $_ := set $annotations "helm.sh/hook-weight" "1" }}
{{- $_ := set $annotations "helm.sh/hook-delete-policy" "before-hook-creation,hook-succeeded" }}
{{- end }}
Review Comment:
The ServiceAccount, Role, and RoleBinding are now hook resources (`helm.sh/hook: post-install,pre-upgrade`) with `hook-weight: "-5"`, while the Job runs at `hook-weight: "1"`. Helm applies lower-weighted hook resources first, so the SA and RBAC exist before the Job hook runs on both fresh installs and the first downgrade after adopting this version.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/files/db_migrate.py:
##########
@@ -0,0 +1,326 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bidirectional Airflow metadata DB reconciliation for the helm chart.
+
+Decides at runtime whether the helm release wants a forward migrate, a
+downgrade, or a no-op, and runs the right command:
+
+* target == current -> no-op (idempotent check)
+* target > current -> ``airflow db migrate`` inside this job's container
+ (uses the TARGET image, which ships forward scripts).
+* target < current -> ``airflow db downgrade --to-version <target>``
+ executed inside the still-running api-server pod (the OLD image still
+ ships the reverse scripts), followed by scaling every DB-touching
+ workload (api-server, scheduler, triggerer, dag-processor, worker) to
+ zero so that no OLD pod keeps talking to the now-downgraded schema. Helm
+ then patches those workloads back to ``replicas: N`` with the TARGET
+ image as the upgrade proceeds, so the cluster comes back up cleanly on
+ the target version. This means a downgrade trades the otherwise-broken
+ rolling-update window for a brief outage (which is unavoidable when the
+ schema goes backwards).
+
+Required env:
+
+* ``AIRFLOW_TARGET_VERSION`` - the version the chart is being upgraded/installed to.
+* ``POD_NAMESPACE`` - release namespace, injected via downward API.
+* ``RELEASE_NAME`` - the helm release name, used to scope the scale-down to
+ only the workloads owned by this release.
+
+Optional env:
+
+* ``MIGRATE_JOB_DRAIN_TIMEOUT_SECONDS`` - how long to wait for DB-touching
+ pods to terminate after scale-to-zero on a downgrade. Defaults to 300.
+ Increase when long-running worker tasks need more time to wind down.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+
+from alembic.migration import MigrationContext
+from kubernetes import client, config as k8s_config
+from kubernetes.stream import stream
+from packaging.version import Version
Review Comment:
Done. The `kubernetes` import is now lazy: it's wrapped in a module-level `try/except ImportError` guard, and only the downgrade branch calls `_require_kubernetes_client()`, which raises a clear `SystemExit` if the client is missing. The forward/no-op/fresh paths no longer require the `kubernetes` package.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
##########
chart/newsfragments/68074.significant.rst:
##########
@@ -0,0 +1,37 @@
+Helm chart now reconciles the Airflow metadata DB bidirectionally on ``helm upgrade``
+
+The ``migrate-database-job`` is now run as a ``post-install,pre-upgrade`` hook
+(was ``post-install,post-upgrade``) and its default args run a reconciliation
+script that picks one of four actions depending on the relationship between
+the chart's ``airflowVersion`` and the alembic head currently in the database:
+
+* ``fresh install`` (no alembic row, or DB unreachable) — ``airflow db migrate``.
Review Comment:
Good catch. Updated the newsfragment to match the implementation: the action is only decided once the DB is reachable, and if it's still unreachable after the connection-wait window (`DB_CONNECT_MAX_WAIT_SECONDS`, default 120) the job exits non-zero and is retried via the Job `backoffLimit` rather than proceeding with a migrate.
---
Drafted-by: GitHub Copilot (Claude Opus 4.8) (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org