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..76293c350b0d6 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,83 @@ 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, integration: tuple[str, ...] = ("openlineage",) +) -> 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 = integration + 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("integration", [("openlineage",), ("all",)]) +@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, + integration, +): + mock_run_command.return_value.returncode = 0 + shell_params = _shell_params_for_openlineage("postgres", postgres_version, integration) + 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("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, + mock_console_print, + _mock_read_cache, + _mock_cleanup, + _mock_fix_ownership, + postgres_version, + integration, +): + 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 + 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") +@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, integration +): + 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 + mock_run_command.assert_not_called()