Chart: Reconcile metadata DB bidirectionally on helm upgrade#68074
Open
jykae wants to merge 12 commits into
Open
Chart: Reconcile metadata DB bidirectionally on helm upgrade#68074jykae wants to merge 12 commits into
jykae wants to merge 12 commits into
Conversation
…e-job Embed a bash entrypoint that picks the right action at runtime — noop, forward migrate, or downgrade into the still-running api-server pod — based on the alembic revision in the metadata DB versus the target airflow_version declared in the chart. This is the building block consumed in the next commit, which switches the migrate-database-job to use it by default. Related: apache#68072
…ally A helm upgrade across an Airflow version now does the right thing in both directions automatically: * On upgrade, the job runs pre-install/pre-upgrade and forward-migrates the metadata DB before the new images start rolling out. * On rollback (helm upgrade to an older airflowVersion), the job detects the version regression and executes 'airflow db downgrade' inside the still-running api-server pod, which still ships the reverse alembic scripts for its own version. The new behaviour is wired up by switching migrateDatabaseJob.args from a baked-in 'airflow db migrate' command to the embedded reconciler script. Users who already supply their own command/args are unaffected: their override still wins. Two new envs feed the script: AIRFLOW_TARGET_VERSION (from .airflowVersion) and POD_NAMESPACE (downward API). Related: apache#68072
The bidirectional reconciler shells into the running api-server pod via the Kubernetes API to run 'airflow db downgrade' there (where the reverse alembic scripts still exist). Without these RBAC rules the downgrade branch would fail with a 403 on the in-cluster ServiceAccount the job already uses. Gated on .Values.rbac.create and .Values.migrateDatabaseJob.enabled so no extra resources are rendered for clusters that opt out of either. Related: apache#68072
* New tests verify pre-install/pre-upgrade hook annotations, embedded default script presence, AIRFLOW_TARGET_VERSION + POD_NAMESPACE envs, and that user-supplied command/args still win over the script default. * New RBAC test file covers the Role + RoleBinding rendering, rule contents (pods get/list, pods/exec create/get), namespace scoping, rbac.create / migrateDatabaseJob.enabled gating, and the ServiceAccount subject wired in by the existing helper. * Pre-existing tests that asserted the previous fixed default args were relaxed to acknowledge the new bidirectional default; explicit-override semantics are still covered. Related: apache#68072
2 tasks
Replaces files/db_migrate.sh (bash wrapper around three python3 heredocs) with files/db_migrate.py - a single Python entrypoint invoked as 'python3 -c ...'. Same decision logic and same external behaviour: noop / forward (airflow db migrate) / downgrade (kubectl-style exec into the still-running api-server pod). Tracked in apache#68072.
A rolling helm upgrade in the downgrade direction would leave the OLD
api-server / scheduler / triggerer / dag-processor / worker pods talking
to a now-downgraded schema (the reverse alembic scripts can drop columns
the OLD code still queries). To avoid that broken window, the downgrade
branch of the migrate-database-job now:
1. execs 'airflow db downgrade' inside the still-running api-server pod
(unchanged), then
2. scales every DB-touching Deployment and StatefulSet of this release
to 0 and waits for the pods to terminate.
Helm then patches the same 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. This trades the otherwise-broken
rolling update for a brief outage, which is unavoidable when the schema
goes backwards.
Required additions:
- env: RELEASE_NAME to scope the scale-to-zero to this release.
- RBAC: deployments + statefulsets (get, list) and
deployments/scale + statefulsets/scale (get, patch) on the existing
migrate-database-job Role.
Tracked in apache#68072.
* Walk target image's revision graph from base to target_rev instead of from base to current_rev. On a downgrade the TARGET image's ScriptDirectory does not contain current_rev, so the previous code raised alembic RevisionError before ever deciding 'downgrade'. Check membership of current_rev in target's ancestors instead. * Narrow the broad 'except Exception' around engine.connect() to sqlalchemy.exc.OperationalError so config / programming errors are not silently reported as a fresh install. * Move airflow / alembic / kubernetes / sqlalchemy imports to module top per Airflow code-style rules. * Treat None exit code from the kubernetes exec stream as failure instead of silently reporting success. * Use subprocess.run(..., check=False).returncode in place of subprocess.call for cleaner signal propagation.
The existing helm tests only verify that the script is embedded with the right env vars / args / RBAC. They never exercise decide_action, discover_api_server_pod, run_downgrade_in_api_server, or scale_release_workloads_to_zero. Adds: * decide_action: unknown target -> forward, OperationalError -> fresh, no alembic row -> fresh, current==target -> noop, current is ancestor of target -> forward, current is NOT in target's ancestors -> downgrade (regression test for the original blocker where the script walked from base to an unknown current_rev and raised alembic RevisionError). * discover_api_server_pod: prefers Ready, falls back to non-Ready, raises when no api-server pod is Running. * run_downgrade_in_api_server: zero / non-zero pass through, None returncode is treated as failure (regression for the silent-success bug from `return returncode or 0`). * scale_release_workloads_to_zero: patches every matched Deployment and StatefulSet to replicas=0, waits for drain, raises TimeoutError when pods never terminate.
jykae
requested review from
bugraoz93,
hussein-awala,
jedcunningham and
jscheffl
as code owners
…ed set
The new Role and RoleBinding rendered by templates/rbac/migrate-database-job-role.yaml
and templates/rbac/migrate-database-job-rolebinding.yaml weren't listed in the
expected set used by test_deployments_with_rbac_{no_sa,with_sa}, so the assertion
'rendered manifests == expected list' fired in the K8s 1.30 helm-tests job.
…chart expected sets
jscheffl
reviewed
Jun 5, 2026
Contributor
jscheffl
left a comment
jscheffl
left a comment
There was a problem hiding this comment.
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 ?
| * ``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 |
Contributor
There was a problem hiding this comment.
I think reference to a past issue is not really needed. Context from above is sufficient.
Comment on lines
+64
to
+66
| # 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. |
Contributor
There was a problem hiding this 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.
Comment on lines
+101
to
+102
| ancestors_of_target = {rev.revision for rev in script.walk_revisions("base", target_rev)} | ||
| if current_rev in ancestors_of_target: |
Contributor
There was a problem hiding this 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.
Comment on lines
+122
to
+125
| 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 |
Contributor
There was a problem hiding this comment.
Handling of case that no pod is ready is missing, gets an array-out-of-bounds.
| return (ready or pods)[0].metadata.name | ||
|
|
||
|
|
||
| def run_downgrade_in_api_server(pod_name: str, namespace: str, target_version: str) -> int: |
Contributor
There was a problem hiding this comment.
As API Server might be "rolled-over" can you add a @retry() decorator to this for resiliency?
| ) | ||
|
|
||
|
|
||
| def _patch_script_dir(db_migrate, monkeypatch, ancestors_by_head): |
Contributor
There was a problem hiding this comment.
Why not making this a fixture?
| # -------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _pod(name, ready=True): |
Contributor
There was a problem hiding this comment.
Why not making this a fixture?
| # -------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _fake_stream(returncode): |
Contributor
There was a problem hiding this comment.
Why not making this a fixture?
| # -------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _wl(name): |
Contributor
There was a problem hiding this comment.
Why not making this a fixture?
| # 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. |
Contributor
There was a problem hiding this comment.
Also here... ticket ref not needed.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Airflow Helm chart’s existing migrate-database-job so a single helm upgrade can reconcile the metadata DB schema in either direction (forward migrate, no-op, or downgrade) based on the chart’s target airflowVersion, running as a pre-install,pre-upgrade hook and adding the RBAC needed to exec into the running api-server and scale DB-touching workloads to zero on downgrades.
Changes:
- Switch the migrations Job hook to
pre-install,pre-upgradeand change the default container args to run an embedded Python reconciler (chart/files/db_migrate.py) viapython3 -c. - Add a namespace-scoped
Role/RoleBindinggrantingpods,pods/exec, anddeployments|statefulsets/*scale permissions to support downgrade+drain. - Add Helm template tests and unit tests for the embedded reconciler; update RBAC/security expectations accordingly.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| chart/values.yaml | Removes the hard-coded default migration args and documents the new “unset => embedded reconciler” behavior. |
| chart/values.schema.json | Aligns schema docs/defaults with args: null so the chart can inject its default script. |
| chart/templates/jobs/migrate-database-job.yaml | Moves hook earlier and embeds db_migrate.py as default args; injects env vars used by the script. |
| chart/files/db_migrate.py | New runtime reconciler that decides forward/no-op/downgrade and executes downgrade via api-server pod + drains DB-touching workloads. |
| chart/templates/rbac/migrate-database-job-role.yaml | New Role granting pod exec + workload scale permissions needed for downgrade/drain. |
| chart/templates/rbac/migrate-database-job-rolebinding.yaml | New RoleBinding attaching the Role to the existing migrations ServiceAccount. |
| chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py | Updates/extends template assertions for new hooks, default args, and env vars. |
| chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py | New Helm tests validating the new Role/RoleBinding and their gating. |
| chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py | New unit tests for the embedded reconciler logic and Kubernetes/Alembic interactions. |
| chart/tests/helm_tests/security/test_rbac.py | Updates the security RBAC inventory to include the new migrations Role/RoleBinding. |
| chart/tests/helm_tests/airflow_aux/test_basic_helm_chart.py | Updates expected rendered resource list to include the new RBAC objects. |
| chart/newsfragments/68074.significant.rst | Adds release notes describing hook change, new reconciler behavior, and RBAC additions. |
Comment on lines
+70
to
+75
| 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" |
Comment on lines
106
to
+108
| {{- if .Values.migrateDatabaseJob.args }} | ||
| args: {{- tpl (toYaml .Values.migrateDatabaseJob.args) . | nindent 12 }} | ||
| {{- else }} |
Comment on lines
+108
to
+112
| 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" | ||
|
|
Comment on lines
+172
to
+179
| 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) |
Comment on lines
+5
to
+10
| 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>`` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Extends the existing
migrate-database-jobso a singlehelm upgradereconciles the Airflow metadata DB in both directions — forward migrate, no-op, or downgrade — based on the chart'sairflowVersion.closes: #68072
What the chart does now
The job runs as a
pre-install,pre-upgradehook (waspost-install,post-upgrade) so the schema is reconciled before helm rolls in the target image.The default
argsinvoke a pure-Python reconciler (chart/files/db_migrate.py) embedded viapython3 -c "<file contents>". It picks one of four actions:alembic_versionrowairflow db migrate(fresh install)airflow db migratein the job container (TARGET image)airflow db downgrade --to-version <target> --yesinside the still-running api-server pod, then (b) scale every DB-touching workload of this release to 0 and wait for drainNew namespace-scoped
Role+RoleBindingfor the existingmigrate-database-jobServiceAccount, gated onrbac.create && migrateDatabaseJob.enabled:podsget, listpods/execcreate, getdeployments,statefulsetsget, listdeployments/scale,statefulsets/scaleget, patchUser-supplied
migrateDatabaseJob.command/migrateDatabaseJob.argsstill take precedence — only the default changed.Why the downgrade branch looks the way it does
Why exec into the api-server pod? Reverse alembic scripts only exist in the image of the version that introduced them. The target image cannot run a downgrade because it doesn't carry the scripts that need to be reversed. The still-running api-server pod does carry them.
Why scale every DB-touching workload to 0 right after? A rolling helm upgrade in the downgrade direction would leave OLD api-server / scheduler / triggerer / dag-processor / worker pods talking to the now-downgraded schema (reverse alembic scripts can drop columns the OLD code still queries). Draining those workloads before the hook returns means helm comes back and patches the same Deployments/StatefulSets to TARGET image + their configured
replicas, so the cluster comes up cleanly on the target version with no OLD pod ever running against the downgraded schema.This makes a downgrade a brief outage rather than a zero-downtime rolling update — which is unavoidable when the schema goes backwards.
Forward migrate and no-op paths are unchanged for users.
Deviations from the design sketch in #68072
The issue described the high-level shape; a few details changed during implementation. None of them alter the user-visible contract (same job, same value keys, same ServiceAccount).
caseovercurrent/targetchart/files/db_migrate.py), embedded viapython3 -c "<file>"airflow.utils.db._REVISION_HEADS_MAP,MigrationContext.configure(conn).get_current_revision(), andScriptDirectory.walk_revisions(...)directly — no string-parsing ofairflow db check-migrationsoutput, and the direction decision walks the actual revision graph instead of comparing version stringskubectl exec ... -- airflow db downgrade ...CoreV1Api.connect_get_namespaced_pod_execvia the kubernetes Python clientkubectlbinary or pin its versionreplicas, so the cluster comes up cleanlypods, pods/execonlypodsget/list,pods/execcreate/get,deployments/statefulsetsget/list,deployments/scale/statefulsets/scaleget/patchappVersionbumps"airflow.utils.db._REVISION_HEADS_MAPfrom the target image_prefix is acknowledged as private — exposing a public accessor is a follow-upAlternatives considered
Discovery of
currentQueries
alembic_versionand maps to aairflow.utils.db._REVISION_HEADS_MAPrevision._REVISION_HEADS_MAPis currently private — exposing a public accessor is tracked in #68072 and would let this script drop the underscore import.Tests
66 helm + runtime tests in
chart/tests/helm_tests/airflow_aux/:Helm template (52 tests) —
test_migrate_database_job.py,test_migrate_database_job_rbac.py:pre-install,pre-upgradepython3 -c "..."withAIRFLOW_TARGET_VERSION,airflow db migrate,airflow db downgradeall present)migrateDatabaseJob.command/argsoverrides still winAIRFLOW_TARGET_VERSION,POD_NAMESPACE(downward API),RELEASE_NAMErbac.createpods+pods/execanddeployments/statefulsets+*/scaleRuntime unit tests (14 tests) —
test_db_migrate_script.py, loadschart/files/db_migrate.pyas a module and mocks alembic + kubernetes:decide_action: unknown target → forward;OperationalError→ fresh; no alembic row → fresh; current==target → noop; current is ancestor of target → forward; current NOT in target's ancestors → downgradediscover_api_server_pod: prefers Ready, falls back to non-Ready, raises when nonerun_downgrade_in_api_server: passes through returncode 0 / 2; treatsNonereturncode as failurescale_release_workloads_to_zero: patches Deployments + StatefulSets toreplicas=0with the right selector; raisesTimeoutErroron stuck drainAll pass:
uv run --project chart pytest chart/tests/helm_tests/airflow_aux/test_migrate_database_job.py chart/tests/helm_tests/airflow_aux/test_migrate_database_job_rbac.py chart/tests/helm_tests/airflow_aux/test_db_migrate_script.py -q→66 passed.A kind-based end-to-end scenario (install 3.0.x → upgrade 3.1.x → downgrade 3.0.x) listed in #68072 is still a follow-up.
Backward compatibility
post-install,post-upgradetopre-install,pre-upgrade. Functionally equivalent for forward migrations; removes the race with thewaitForMigrationsinitContainer on new pods.helm upgradehalf-succeeds and leaves the cluster broken; after this change it completes cleanly (with a brief outage).Was generative AI tooling used to co-author this PR?
Generated-by: Claude Opus 4.7 following the guidelines