This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch changes-3.4.0rc3 in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e07fa18cbb0add546f37fc94046f51560647933c Author: Jason(Zhe-You) Liu <[email protected]> AuthorDate: Thu Nov 5 13:37:40 2026 +0800 [v3-4-test] Fix async engine datetimes being shifted on non-UTC MySQL servers (#76505) (#76517) The async engine skipped the connection setup the sync engine gets, so its MySQL sessions kept the server time zone and stored task heartbeats and audit log rows shifted by the server's UTC offset. (cherry picked from commit 494f65e0f682ab8c09cca83f7776de714d638462) --- airflow-core/src/airflow/settings.py | 1 + airflow-core/tests/unit/core/test_settings.py | 41 ++++++++++++++++++++++++--- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/settings.py b/airflow-core/src/airflow/settings.py index 00a33a83f9d..b374ed85afb 100644 --- a/airflow-core/src/airflow/settings.py +++ b/airflow-core/src/airflow/settings.py @@ -504,6 +504,7 @@ def _configure_async_session() -> None: connect_args=_get_connect_args("async"), engine_args=engine_args, ) + setup_event_handlers(async_engine.sync_engine) AsyncSession = async_sessionmaker( bind=async_engine, class_=SAAsyncSession, diff --git a/airflow-core/tests/unit/core/test_settings.py b/airflow-core/tests/unit/core/test_settings.py index 49a92c24dd7..e6802350e40 100644 --- a/airflow-core/tests/unit/core/test_settings.py +++ b/airflow-core/tests/unit/core/test_settings.py @@ -28,6 +28,7 @@ from unittest import mock from unittest.mock import MagicMock, call, patch import pytest +from sqlalchemy import text from sqlalchemy.engine import Engine, make_url from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.pool import NullPool @@ -295,8 +296,9 @@ class TestMetadataEngineHooks: assert mock_create_engine.call_args[1]["connect_args"] == {"prepare_threshold": None} + @patch("airflow.settings.setup_event_handlers") @patch("airflow.settings.create_async_metadata_engine") - def test_configure_async_session_forwards_the_psycopg3_default(self, mock_create_async_engine): + def test_configure_async_session_forwards_the_psycopg3_default(self, mock_create_async_engine, _): """The docs promise the default on both engines, so assert the async one separately.""" from airflow import settings @@ -307,9 +309,10 @@ class TestMetadataEngineHooks: assert mock_create_async_engine.call_args[1]["connect_args"] == {"prepare_threshold": None} + @patch("airflow.settings.setup_event_handlers") @patch("airflow.settings.create_async_metadata_engine") def test_configure_async_session_delegates_to_create_async_metadata_engine( - self, mock_create_async_engine + self, mock_create_async_engine, mock_setup_event_handlers ): """_configure_async_session() must call create_async_metadata_engine with no pool args for sqlite.""" from airflow import settings @@ -330,9 +333,11 @@ class TestMetadataEngineHooks: assert "connect_args" in call_kwargs[1] # sqlite doesn't support pool size args assert call_kwargs[1]["engine_args"] == {} + mock_setup_event_handlers.assert_called_once_with(mock_create_async_engine.return_value.sync_engine) + @patch("airflow.settings.setup_event_handlers") @patch("airflow.settings.create_async_metadata_engine") - def test_configure_async_session_passes_pool_args_for_non_sqlite(self, mock_create_async_engine): + def test_configure_async_session_passes_pool_args_for_non_sqlite(self, mock_create_async_engine, _): """_configure_async_session() must pass pool configuration for non-sqlite backends.""" from airflow import settings @@ -357,8 +362,9 @@ class TestMetadataEngineHooks: assert engine_args["pool_pre_ping"] is True assert engine_args["max_overflow"] == 5 + @patch("airflow.settings.setup_event_handlers") @patch("airflow.settings.create_async_metadata_engine") - def test_configure_async_session_uses_nullpool_when_pool_disabled(self, mock_create_async_engine): + def test_configure_async_session_uses_nullpool_when_pool_disabled(self, mock_create_async_engine, _): """_configure_async_session() must use NullPool when SQL_ALCHEMY_POOL_ENABLED is False.""" from airflow import settings @@ -748,6 +754,33 @@ def test_async_pool_is_closed_before_process_shutdown(driver): assert diagnostic not in output, output [email protected]_test [email protected]( + ("driver", "query", "expected"), + [ + pytest.param( + "mysql+aiomysql", "SELECT @@session.time_zone", "+00:00", marks=pytest.mark.backend("mysql") + ), + pytest.param("sqlite+aiosqlite", "PRAGMA foreign_keys", 1, marks=pytest.mark.backend("sqlite")), + ], +) +def test_async_engine_connections_get_the_sync_engine_setup(monkeypatch, driver, query, expected): + url = make_url(settings.SQL_ALCHEMY_CONN_ASYNC).set(drivername=driver) + monkeypatch.setattr(settings, "SQL_ALCHEMY_CONN_ASYNC", url.render_as_string(hide_password=False)) + monkeypatch.setattr(settings, "async_engine", settings.async_engine) + monkeypatch.setattr(settings, "AsyncSession", settings.AsyncSession) + settings._configure_async_session() + + async def run(): + try: + async with settings.AsyncSession() as session: + return (await session.execute(text(query))).scalar_one() + finally: + await settings.dispose_async_engine() + + assert asyncio.run(run()) == expected + + @pytest.mark.parametrize( ("sql_alchemy_conn", "expected"), [
