Retry transient database errors in the metastore secrets backend

[1fanwang]

MetastoreBackend reads connections and variables from the metadata database with no retry. Every caller that wraps it (Connection.get_connection_from_secrets, Variable.get_variable_from_secrets, and the Task SDK's server-context lookup) catches Exception, logs it at debug level, and falls through to the next backend. A transient database failure therefore reaches the caller as the secret does not exist.

Why it matters

The Execution API answers 404 for a connection or variable that is present in the database. 404 is permanent, so the Task SDK does not retry and the task fails with The conn_id 'x' isn't defined. The real cause is visible only at debug level.

Two amplifiers. Under [secrets] use_cache = True the resulting None is cached for cache_ttl_seconds (see variable.py:498, which notes "we save None as well"), so one blip poisons the key. The triggerer, Dag processor, and callback supervisor call this backend in-process, so deferred tasks hit the same bogus not-found.

pool_pre_ping does not cover this. It recycles a stale pooled connection, but the failure here lands on the reconnect itself.

The fix

Apply the existing retry_db_transaction under @provide_session, so retries run within one session with a rollback between attempts. That is the stacking already used in renderedtifields.py:241, dagwarning.py:78, and manager.py:701. Both lookups are reads, the budget is the existing [database] max_db_retries, and a missing secret still returns None on the first attempt.

Testing Done

Real Postgres, real psycopg2, real TCP-level failure. Airflow connects through a local forwarder; blip() drops the open connections and refuses exactly the next connection attempt, then serves normally. One failover-shaped blip, so the outcome is deterministic rather than a race against the backoff.

Setup

docker run -d --name af-e2e-pg -e POSTGRES_PASSWORD=airflow \
    -e POSTGRES_USER=airflow -e POSTGRES_DB=airflow -p 55432:5432 postgres:16

export AIRFLOW_HOME=/tmp/af-e2e-metastore
export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://airflow:airflow@127.0.0.1:55433/airflow

python repro.py --setup    # airflow db migrate, then seed e2e_conn and e2e_var
python repro.py

repro.py

from __future__ import annotations

import socket
import subprocess
import sys
import threading
import time

UPSTREAM = ("127.0.0.1", 55432)
LISTEN = ("127.0.0.1", 55433)


class Forwarder:
    """TCP forwarder to the database that can drop and refuse connections on demand."""

    def __init__(self):
        self.refuse_next = 0
        self.live: list[socket.socket] = []
        self._sock = socket.socket()
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._sock.bind(LISTEN)
        self._sock.listen(64)
        threading.Thread(target=self._serve, daemon=True).start()

    def _serve(self):
        while True:
            client, _ = self._sock.accept()
            if self.refuse_next > 0:
                self.refuse_next -= 1
                client.close()
                continue
            upstream = socket.create_connection(UPSTREAM)
            self.live += [client, upstream]
            for a, b in ((client, upstream), (upstream, client)):
                threading.Thread(target=self._pipe, args=(a, b), daemon=True).start()

    @staticmethod
    def _pipe(src, dst):
        try:
            while chunk := src.recv(65536):
                dst.sendall(chunk)
        except OSError:
            pass
        finally:
            src.close()
            dst.close()

    def blip(self, attempts: int = 1) -> None:
        """Drop open connections and refuse the next `attempts` connection attempts."""
        for sock in self.live:
            try:
                sock.close()
            except OSError:
                pass
        self.live = []
        self.refuse_next = attempts


def setup():
    for cmd in (
        ["airflow", "db", "migrate"],
        ["airflow", "connections", "add", "e2e_conn", "--conn-type", "mysql", "--conn-host", "db.example.com"],
        ["airflow", "variables", "set", "e2e_var", "e2e_value"],
    ):
        subprocess.run(cmd, check=True)


def main():
    forwarder = Forwarder()

    from airflow.exceptions import AirflowNotFoundException
    from airflow.models.connection import Connection
    from airflow.models.variable import Variable

    Connection.get_connection_from_secrets("e2e_conn")
    Variable.get_variable_from_secrets("e2e_var")
    print("baseline: connection 'e2e_conn' and variable 'e2e_var' both resolve")

    print("\none refused connection attempt, then look up the connection:")
    forwarder.blip()
    try:
        conn = Connection.get_connection_from_secrets("e2e_conn")
        print(f"  -> OK: conn_id={conn.conn_id} host={conn.host}")
    except AirflowNotFoundException as e:
        print(f"  -> SPURIOUS NOT-FOUND: AirflowNotFoundException: {e}")

    print("\none refused connection attempt, then look up the variable:")
    forwarder.blip()
    value = Variable.get_variable_from_secrets("e2e_var")
    if value is None:
        print("  -> SPURIOUS NOT-FOUND: None (Variable.get raises KeyError -> HTTP 404)")
    else:
        print(f"  -> OK: {value!r}")

    time.sleep(0.2)


if __name__ == "__main__":
    if "--setup" in sys.argv:
        Forwarder()
        time.sleep(0.2)
        setup()
    else:
        main()

Before, on origin/main

$ git checkout origin/main -- airflow-core/src/airflow/secrets/metastore.py
$ python repro.py
baseline: connection 'e2e_conn' and variable 'e2e_var' both resolve

one refused connection attempt, then look up the connection:
  -> SPURIOUS NOT-FOUND: AirflowNotFoundException: The conn_id `e2e_conn` isn't defined

one refused connection attempt, then look up the variable:
  -> SPURIOUS NOT-FOUND: None (Variable.get raises KeyError -> HTTP 404)

The swallowed driver error (debug level)

Unable to retrieve connection from secrets backend (MetastoreBackend). Checking subsequent secrets backend.
Traceback (most recent call last):
  File "sqlalchemy/pool/base.py", line 896, in __connect
    self.dbapi_connection = connection = pool._invoke_creator(self)
  File "sqlalchemy/engine/default.py", line 630, in connect
    return self.loaded_dbapi.connect(*cargs, **cparams)
  File "psycopg2/__init__.py", line 122, in connect
    conn = _connect(dsn, connection_factory=connection_factory, **kwasync)
psycopg2.OperationalError: connection to server at "127.0.0.1", port 55433 failed: server closed the connection unexpectedly
sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) connection to server at "127.0.0.1", port 55433 failed: server closed the connection unexpectedly

After, on this branch

$ git checkout HEAD -- airflow-core/src/airflow/secrets/metastore.py
$ python repro.py
baseline: connection 'e2e_conn' and variable 'e2e_var' both resolve

one refused connection attempt, then look up the connection:
  -> OK: conn_id=e2e_conn host=db.example.com

one refused connection attempt, then look up the variable:
  -> OK: 'e2e_value'

Three consecutive runs each way:

### PRE-FIX (origin/main) — 3 consecutive runs ###
  -> SPURIOUS NOT-FOUND: AirflowNotFoundException: The conn_id `e2e_conn` isn't defined
  -> SPURIOUS NOT-FOUND: None (Variable.get raises KeyError -> HTTP 404)
  --
  -> SPURIOUS NOT-FOUND: AirflowNotFoundException: The conn_id `e2e_conn` isn't defined
  -> SPURIOUS NOT-FOUND: None (Variable.get raises KeyError -> HTTP 404)
  --
  -> SPURIOUS NOT-FOUND: AirflowNotFoundException: The conn_id `e2e_conn` isn't defined
  -> SPURIOUS NOT-FOUND: None (Variable.get raises KeyError -> HTTP 404)
  --

### POST-FIX (this PR) — 3 consecutive runs ###
  -> OK: conn_id=e2e_conn host=db.example.com
  -> OK: 'e2e_value'
  --
  -> OK: conn_id=e2e_conn host=db.example.com
  -> OK: 'e2e_value'
  --
  -> OK: conn_id=e2e_conn host=db.example.com
  -> OK: 'e2e_value'
  --

Regression tests in airflow-core/tests/unit/always/test_secrets_metastore.py drive a real session against a real database and fail on unpatched source.