From d5ae7dd8655645607bb43f9e25b0b6ad8c1e5c9f Mon Sep 17 00:00:00 2001 From: Aaryan Mahajan Date: Thu, 2 Jul 2026 21:48:10 +0530 Subject: [PATCH 1/2] Fix Breeze OpenLineage integration gate stuck on stale Postgres versions The gate in enter_shell() hardcoded PostgreSQL 12/13/14 as the only versions allowed with --integration openlineage, but Breeze's actual matrix has moved to 14-18 (12 isn't even a valid choice anymore). Derive the allowed set from CURRENT_POSTGRES_VERSIONS in global_constants.py instead, so the gate can't drift out of sync with the matrix again. closes: #69233 --- .../utils/docker_command_utils.py | 8 +- dev/breeze/tests/test_docker_command_utils.py | 75 +++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py index 76bc3f0e7a0f6..186b1b0956ee4 100644 --- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py +++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py @@ -52,6 +52,7 @@ from airflow_breeze.global_constants import ( ALLOWED_CELERY_BROKERS, ALLOWED_DEBIAN_VERSIONS, + CURRENT_POSTGRES_VERSIONS, DEFAULT_PYTHON_MAJOR_MINOR_VERSION, DOCKER_DEFAULT_PLATFORM, KNOWN_DOCKER_COMPOSE_PROJECT_NAMES, @@ -1030,9 +1031,12 @@ def enter_shell( console_print("\n[warn]MySQL use MariaDB client binaries on ARM architecture.[/]\n") if "openlineage" in shell_params.integration or "all" in shell_params.integration: - if shell_params.backend != "postgres" or shell_params.postgres_version not in ["12", "13", "14"]: + if ( + shell_params.backend != "postgres" + or shell_params.postgres_version not in CURRENT_POSTGRES_VERSIONS + ): console_print( - "\n[error]Only PostgreSQL 12, 13, and 14 are supported " + f"\n[error]Only PostgreSQL {', '.join(CURRENT_POSTGRES_VERSIONS)} are supported " "as a backend with OpenLineage integration via Breeze[/]\n" ) sys.exit(1) diff --git a/dev/breeze/tests/test_docker_command_utils.py b/dev/breeze/tests/test_docker_command_utils.py index 4adaab3bcca06..370efce835ad8 100644 --- a/dev/breeze/tests/test_docker_command_utils.py +++ b/dev/breeze/tests/test_docker_command_utils.py @@ -22,12 +22,14 @@ import pytest +from airflow_breeze.global_constants import ALLOWED_POSTGRES_VERSIONS, CURRENT_POSTGRES_VERSIONS from airflow_breeze.utils.docker_command_utils import ( autodetect_docker_context, bring_all_compose_projects_down, check_docker_compose_version, check_docker_version, discover_running_compose_projects, + enter_shell, is_known_breeze_compose_project, ) @@ -389,3 +391,76 @@ def test_bring_all_compose_projects_down_preserve_volumes(mock_discover, mock_ru down_call = next(c for c in mock_run_command.call_args_list if c.args[0][:2] == ["docker", "compose"]) assert "--volumes" not in down_call.args[0] assert "--remove-orphans" in down_call.args[0] + + +def _shell_params_for_openlineage(backend: str, postgres_version: str) -> mock.MagicMock: + shell_params = mock.MagicMock() + shell_params.use_airflow_version = None + shell_params.restart = False + shell_params.include_mypy_volume = False + shell_params.quiet = True + shell_params.project_name = None + shell_params.tty = "disabled" + shell_params.command_passed = None + shell_params.integration = ("openlineage",) + shell_params.backend = backend + shell_params.postgres_version = postgres_version + return shell_params + + +@mock.patch("airflow_breeze.utils.docker_command_utils.fix_ownership_using_docker") +@mock.patch("airflow_breeze.utils.docker_command_utils.cleanup_python_generated_files") +@mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") +@mock.patch("airflow_breeze.utils.docker_command_utils.console_print") +@mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +@pytest.mark.parametrize("postgres_version", CURRENT_POSTGRES_VERSIONS) +def test_enter_shell_openlineage_allows_current_postgres_versions( + mock_run_command, + mock_console_print, + _mock_read_cache, + _mock_cleanup, + _mock_fix_ownership, + postgres_version, +): + mock_run_command.return_value.returncode = 0 + shell_params = _shell_params_for_openlineage("postgres", postgres_version) + enter_shell(shell_params) + mock_run_command.assert_called_once() + + +@mock.patch("airflow_breeze.utils.docker_command_utils.fix_ownership_using_docker") +@mock.patch("airflow_breeze.utils.docker_command_utils.cleanup_python_generated_files") +@mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") +@mock.patch("airflow_breeze.utils.docker_command_utils.console_print") +@mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +@pytest.mark.parametrize("postgres_version", set(ALLOWED_POSTGRES_VERSIONS) - set(CURRENT_POSTGRES_VERSIONS)) +def test_enter_shell_openlineage_rejects_stale_postgres_versions( + mock_run_command, + mock_console_print, + _mock_read_cache, + _mock_cleanup, + _mock_fix_ownership, + postgres_version, +): + shell_params = _shell_params_for_openlineage("postgres", postgres_version) + with pytest.raises(SystemExit) as exc_info: + enter_shell(shell_params) + assert exc_info.value.code == 1 + error_message = mock_console_print.call_args[0][0] + assert all(version in error_message for version in CURRENT_POSTGRES_VERSIONS) + mock_run_command.assert_not_called() + + +@mock.patch("airflow_breeze.utils.docker_command_utils.fix_ownership_using_docker") +@mock.patch("airflow_breeze.utils.docker_command_utils.cleanup_python_generated_files") +@mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") +@mock.patch("airflow_breeze.utils.docker_command_utils.console_print") +@mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +def test_enter_shell_openlineage_rejects_non_postgres_backend( + mock_run_command, mock_console_print, _mock_read_cache, _mock_cleanup, _mock_fix_ownership +): + shell_params = _shell_params_for_openlineage("mysql", CURRENT_POSTGRES_VERSIONS[0]) + with pytest.raises(SystemExit) as exc_info: + enter_shell(shell_params) + assert exc_info.value.code == 1 + mock_run_command.assert_not_called() From bb38659afd60f9eac343fa2ba83dedcbb3682869 Mon Sep 17 00:00:00 2001 From: Aaryan Mahajan Date: Sun, 19 Jul 2026 17:24:37 +0400 Subject: [PATCH 2/2] Cover --integration all in OpenLineage Postgres gate tests Reviewer feedback on apache/airflow#69264: the gate also triggers on "all" (not just "openlineage"), so the tests should exercise that path too. --- dev/breeze/tests/test_docker_command_utils.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/dev/breeze/tests/test_docker_command_utils.py b/dev/breeze/tests/test_docker_command_utils.py index 370efce835ad8..76293c350b0d6 100644 --- a/dev/breeze/tests/test_docker_command_utils.py +++ b/dev/breeze/tests/test_docker_command_utils.py @@ -393,7 +393,9 @@ def test_bring_all_compose_projects_down_preserve_volumes(mock_discover, mock_ru assert "--remove-orphans" in down_call.args[0] -def _shell_params_for_openlineage(backend: str, postgres_version: str) -> mock.MagicMock: +def _shell_params_for_openlineage( + backend: str, postgres_version: str, integration: tuple[str, ...] = ("openlineage",) +) -> mock.MagicMock: shell_params = mock.MagicMock() shell_params.use_airflow_version = None shell_params.restart = False @@ -402,7 +404,7 @@ def _shell_params_for_openlineage(backend: str, postgres_version: str) -> mock.M shell_params.project_name = None shell_params.tty = "disabled" shell_params.command_passed = None - shell_params.integration = ("openlineage",) + shell_params.integration = integration shell_params.backend = backend shell_params.postgres_version = postgres_version return shell_params @@ -413,6 +415,7 @@ def _shell_params_for_openlineage(backend: str, postgres_version: str) -> mock.M @mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") @mock.patch("airflow_breeze.utils.docker_command_utils.console_print") @mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +@pytest.mark.parametrize("integration", [("openlineage",), ("all",)]) @pytest.mark.parametrize("postgres_version", CURRENT_POSTGRES_VERSIONS) def test_enter_shell_openlineage_allows_current_postgres_versions( mock_run_command, @@ -421,9 +424,10 @@ def test_enter_shell_openlineage_allows_current_postgres_versions( _mock_cleanup, _mock_fix_ownership, postgres_version, + integration, ): mock_run_command.return_value.returncode = 0 - shell_params = _shell_params_for_openlineage("postgres", postgres_version) + shell_params = _shell_params_for_openlineage("postgres", postgres_version, integration) enter_shell(shell_params) mock_run_command.assert_called_once() @@ -433,6 +437,7 @@ def test_enter_shell_openlineage_allows_current_postgres_versions( @mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") @mock.patch("airflow_breeze.utils.docker_command_utils.console_print") @mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +@pytest.mark.parametrize("integration", [("openlineage",), ("all",)]) @pytest.mark.parametrize("postgres_version", set(ALLOWED_POSTGRES_VERSIONS) - set(CURRENT_POSTGRES_VERSIONS)) def test_enter_shell_openlineage_rejects_stale_postgres_versions( mock_run_command, @@ -441,8 +446,9 @@ def test_enter_shell_openlineage_rejects_stale_postgres_versions( _mock_cleanup, _mock_fix_ownership, postgres_version, + integration, ): - shell_params = _shell_params_for_openlineage("postgres", postgres_version) + shell_params = _shell_params_for_openlineage("postgres", postgres_version, integration) with pytest.raises(SystemExit) as exc_info: enter_shell(shell_params) assert exc_info.value.code == 1 @@ -456,10 +462,11 @@ def test_enter_shell_openlineage_rejects_stale_postgres_versions( @mock.patch("airflow_breeze.utils.docker_command_utils.read_from_cache_file", return_value="1") @mock.patch("airflow_breeze.utils.docker_command_utils.console_print") @mock.patch("airflow_breeze.utils.docker_command_utils.run_command") +@pytest.mark.parametrize("integration", [("openlineage",), ("all",)]) def test_enter_shell_openlineage_rejects_non_postgres_backend( - mock_run_command, mock_console_print, _mock_read_cache, _mock_cleanup, _mock_fix_ownership + mock_run_command, mock_console_print, _mock_read_cache, _mock_cleanup, _mock_fix_ownership, integration ): - shell_params = _shell_params_for_openlineage("mysql", CURRENT_POSTGRES_VERSIONS[0]) + shell_params = _shell_params_for_openlineage("mysql", CURRENT_POSTGRES_VERSIONS[0], integration) with pytest.raises(SystemExit) as exc_info: enter_shell(shell_params) assert exc_info.value.code == 1