Fix Breeze OpenLineage integration gate stuck on stale Postgres versions

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)

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()