feat(cli): add upgrade command for CLI to simplify version updates

haystack_enterprise_sdk/cli.py

@@ -6,6 +6,8 @@
 import logging
 import os
 import secrets
+import shutil
+import subprocess
 import sys
 import threading
 import time
@@ -2158,28 +2160,48 @@ def _latest_version() -> Optional[str]:
 def _maybe_print_update_hint() -> None:
     """Print a hint on stderr when a newer release is on PyPI.
 
-    Skipped in CI, off a TTY, and when ``HAYSTACK_ENTERPRISE_NO_UPDATE_CHECK`` is set. Never raises.
+    Skipped in CI, off a TTY, when ``HAYSTACK_ENTERPRISE_NO_UPDATE_CHECK`` is set, and after ``upgrade``
+    (this process still carries the old version). Never raises.
     """
-    if not animations_enabled() or not _stderr_is_tty() or os.environ.get("HAYSTACK_ENTERPRISE_NO_UPDATE_CHECK"):
+    if (
+        not animations_enabled()
+        or not _stderr_is_tty()
+        or os.environ.get("HAYSTACK_ENTERPRISE_NO_UPDATE_CHECK")
+        or "upgrade" in sys.argv[1:]
+    ):
         return
     latest = _latest_version()
     try:
         if latest is None or tuple(map(int, latest.split("."))) <= tuple(map(int, __version__.split("."))):
             return
     except ValueError:  # a pre-release or dev version on either side: don't guess
         return
-    # `uv tool install` (the README's and install.sh's way) puts the environment at .../uv/tools/<name>.
-    prefix = Path(sys.prefix)
-    if prefix.parent.name == "tools" and prefix.parent.parent.name == "uv":
-        upgrade = "uv tool upgrade haystack-enterprise-sdk"
-    else:
-        upgrade = "pip install -U haystack-enterprise-sdk"
     typer.echo(
-        f"A new version of haystack-enterprise-sdk is available: {__version__} -> {latest}. Upgrade with: {upgrade}",
+        f"A new version of haystack-enterprise-sdk is available: {__version__} -> {latest}. "
+        "Run: haystack-enterprise upgrade",
         err=True,
     )
 
 
+@cli_app.command()
+def upgrade() -> None:
+    """Upgrade the CLI to the latest release, the same way it was installed.
+
+    A `uv tool install` (the installer's and README's way) is upgraded in place, keeping its extras.
+    Any other install is left alone and the command to upgrade it is printed instead.
+
+    Example:
+    `haystack-enterprise upgrade`
+    """
+    # `uv tool install` puts the environment at .../uv/tools/<name>.
+    prefix = Path(sys.prefix)
+    if prefix.parent.name == "tools" and prefix.parent.parent.name == "uv" and shutil.which("uv"):
+        # uv reports the result itself: "Updated ... v1.0.2 -> v1.1.0" or "Nothing to upgrade".
+        completed = subprocess.run(["uv", "tool", "upgrade", "haystack-enterprise-sdk"], check=False)  # noqa: S603, S607
+        raise typer.Exit(code=completed.returncode)
+    typer.echo("To upgrade, run: pip install -U haystack-enterprise-sdk")
+
+
 def run_packaged() -> None:
     """Run the packaged CLI.
 

tests/unit/test_cli.py

@@ -839,21 +839,19 @@ def test_newer_version_prints_hint_and_caches(self, capsys: pytest.CaptureFixtur
         assert capsys.readouterr().err.count("1.0.2 -> 1.10.0") == 2
         assert route.call_count == 1
 
-    @pytest.mark.parametrize(
-        "prefix,command",
-        [
-            ("/home/u/.local/share/uv/tools/haystack-enterprise-sdk", "uv tool upgrade haystack-enterprise-sdk"),
-            ("/home/u/project/.venv", "pip install -U haystack-enterprise-sdk"),
-        ],
-    )
     @respx.mock
-    def test_upgrade_command_matches_install(
-        self, prefix: str, command: str, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
-    ) -> None:
-        monkeypatch.setattr("sys.prefix", prefix)
+    def test_hint_points_to_upgrade_command(self, capsys: pytest.CaptureFixture) -> None:
         respx.get(_PYPI_URL).mock(return_value=httpx.Response(200, json={"info": {"version": "1.1.0"}}))
         _maybe_print_update_hint()
-        assert capsys.readouterr().err.strip().endswith(f"Upgrade with: {command}")
+        assert capsys.readouterr().err.strip().endswith("Run: haystack-enterprise upgrade")
+
+    @respx.mock
+    def test_no_hint_after_upgrade(self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture) -> None:
+        monkeypatch.setattr("sys.argv", ["haystack-enterprise", "upgrade"])
+        route = respx.get(_PYPI_URL)
+        _maybe_print_update_hint()
+        assert capsys.readouterr().err == ""
+        assert not route.called
 
     @pytest.mark.parametrize("latest", ["1.0.2", "0.9.9", "2.0.0rc1"])
     @respx.mock
@@ -877,6 +875,41 @@ def test_disabled_skips_network(self, env: str, monkeypatch: pytest.MonkeyPatch)
         assert not route.called
 
 
+class TestUpgradeCommand:
+    UV_PREFIX = "/home/u/.local/share/uv/tools/haystack-enterprise-sdk"
+
+    @pytest.fixture
+    def run(self, monkeypatch: pytest.MonkeyPatch) -> Mock:
+        run = Mock(return_value=Mock(returncode=0))
+        monkeypatch.setattr("haystack_enterprise_sdk.cli.subprocess.run", run)
+        monkeypatch.setattr("haystack_enterprise_sdk.cli.shutil.which", lambda _name: "/usr/bin/uv")
+        return run
+
+    @pytest.mark.parametrize("returncode", [0, 2])
+    def test_uv_tool_install_runs_uv_upgrade(self, run: Mock, monkeypatch: pytest.MonkeyPatch, returncode: int) -> None:
+        monkeypatch.setattr("sys.prefix", self.UV_PREFIX)
+        run.return_value.returncode = returncode
+        result = runner.invoke(cli_app, ["upgrade"])
+        run.assert_called_once_with(["uv", "tool", "upgrade", "haystack-enterprise-sdk"], check=False)
+        assert result.exit_code == returncode
+
+    def test_other_install_prints_pip_command(self, run: Mock, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr("sys.prefix", "/home/u/project/.venv")
+        result = runner.invoke(cli_app, ["upgrade"])
+        assert result.exit_code == 0
+        assert "pip install -U haystack-enterprise-sdk" in result.output
+        run.assert_not_called()
+
+    def test_uv_tool_install_without_uv_on_path_prints_pip_command(
+        self, run: Mock, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setattr("sys.prefix", self.UV_PREFIX)
+        monkeypatch.setattr("haystack_enterprise_sdk.cli.shutil.which", lambda _name: None)
+        result = runner.invoke(cli_app, ["upgrade"])
+        assert "pip install -U haystack-enterprise-sdk" in result.output
+        run.assert_not_called()
+
+
 def test_run_packaged_uses_os_trust_store() -> None:
     with (
         patch("haystack_enterprise_sdk.cli.truststore.inject_into_ssl") as inject,