cylc play (restart): fix bug affecting run host reinvocation after interactive upgrade

cylc/flow/scheduler_cli.py

@@ -390,11 +390,7 @@ async def scheduler_cli(
 
     # check the workflow can be safely restarted with this version of Cylc
     db_file = Path(get_workflow_srv_dir(workflow_id), 'db')
-    if not _version_check(
-        db_file,
-        options.upgrade,
-        options.downgrade,
-    ):
+    if not _version_check(db_file, options):
         sys.exit(1)
 
     # upgrade the workflow DB (after user has confirmed upgrade)
@@ -404,7 +400,7 @@ async def scheduler_cli(
     _print_startup_message(options)
 
     # re-execute on another host if required
-    _distribute(options.host, workflow_id_raw, workflow_id, options.color)
+    _distribute(workflow_id_raw, workflow_id, options)
 
     # setup the scheduler
     # NOTE: asyncio.run opens an event loop, runs your coro,
@@ -474,8 +470,7 @@ async def _resume(workflow_id, options):
 
 def _version_check(
     db_file: Path,
-    can_upgrade: bool,
-    can_downgrade: bool
+    options: 'Values',
 ) -> bool:
     """Check the workflow can be safely restarted with this version of Cylc."""
     if not db_file.is_file():
@@ -491,7 +486,7 @@ def _version_check(
     )):
         if this < that:
             # restart would REDUCE the Cylc version
-            if can_downgrade:
+            if options.downgrade:
                 # permission to downgrade given in CLI flags
                 LOG.warning(
                     'Restarting with an older version of Cylc'
@@ -517,7 +512,7 @@ def _version_check(
             return False
         elif itt < 2 and this > that:
             # restart would INCREASE the Cylc version in a big way
-            if can_upgrade:
+            if options.upgrade:
                 # permission to upgrade given in CLI flags
                 LOG.warning(
                     'Restarting with a newer version of Cylc'
@@ -531,7 +526,7 @@ def _version_check(
             ))
             if is_terminal():
                 # we are in interactive mode, ask the user if this is ok
-                return prompt(
+                options.upgrade = prompt(
                     cparse(
                         'Are you sure you want to upgrade from'
                         f' <yellow>{last_run_version}</yellow>'
@@ -540,6 +535,7 @@ def _version_check(
                     {'y': True, 'n': False},
                     process=str.lower,
                 )
+                return options.upgrade
             # we are in non-interactive mode, abort abort abort
             print('Use "--upgrade" to upgrade the workflow.', file=sys.stderr)
             return False
@@ -580,22 +576,23 @@ def _print_startup_message(options):
         LOG.warning(SUITERC_DEPR_MSG)
 
 
-def _distribute(host, workflow_id_raw, workflow_id, color):
+def _distribute(
+    workflow_id_raw: str, workflow_id: str, options: 'Values'
+) -> None:
     """Re-invoke this command on a different host if requested.
 
     Args:
-        host:
-            The remote host to re-invoke on.
         workflow_id_raw:
             The workflow ID as it appears in the CLI arguments.
         workflow_id:
             The workflow ID after it has gone through the CLI.
             This may be different (i.e. the run name may have been inferred).
+        options:
+            The CLI options.
 
     """
     # Check whether a run host is explicitly specified, else select one.
-    if not host:
-        host = select_workflow_host()[0]
+    host = options.host or select_workflow_host()[0]
     if is_remote_host(host):
         # Protect command args from second shell interpretation
         cmd = list(map(quote, sys.argv[1:]))
@@ -612,8 +609,12 @@ def _distribute(host, workflow_id_raw, workflow_id, color):
         # Prevent recursive host selection
         cmd.append("--host=localhost")
 
+        # Ensure interactive upgrade carries over:
+        if options.upgrade and '--upgrade' not in cmd:
+            cmd.append('--upgrade')
+
         # Preserve CLI colour
-        if is_terminal() and color != 'never':
+        if is_terminal() and options.color != 'never':
             # the detached process doesn't pass the is_terminal test
             # so we have to explicitly tell Cylc to use color
             cmd.append('--color=always')

tests/unit/test_scheduler_cli.py

@@ -14,16 +14,13 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from contextlib import contextmanager
 import sqlite3
+from contextlib import contextmanager
 
 import pytest
 
 from cylc.flow.exceptions import ServiceFileError
-from cylc.flow.scheduler_cli import (
-    _distribute,
-    _version_check,
-)
+from cylc.flow.scheduler_cli import RunOptions, _distribute, _version_check
 
 
 @pytest.fixture
@@ -184,7 +181,12 @@ def test_version_check_interactive(
     db_file = stopped_workflow_db(before)
     set_cylc_version(after)
     with answer(response):
-        assert _version_check(db_file, False, downgrade) is outcome
+        assert (
+            _version_check(
+                db_file, RunOptions(downgrade=downgrade)
+            )
+            is outcome
+        )
 
 
 def test_version_check_non_interactive(
@@ -200,29 +202,33 @@ def test_version_check_non_interactive(
     # upgrade
     db_file = stopped_workflow_db('8.0.0')
     set_cylc_version('8.1.0')
-    assert _version_check(db_file, False, False) is False
-    assert _version_check(db_file, True, False) is True  # CLI --upgrade
+    assert _version_check(db_file, RunOptions()) is False
+    assert (
+        _version_check(db_file, RunOptions(upgrade=True)) is True
+    )  # CLI --upgrade
 
     # downgrade
     db_file.unlink()
     db_file = stopped_workflow_db('8.1.0')
     set_cylc_version('8.0.0')
-    assert _version_check(db_file, False, False) is False
-    assert _version_check(db_file, False, True) is True  # CLI --downgrade
+    assert _version_check(db_file, RunOptions()) is False
+    assert (
+        _version_check(db_file, RunOptions(downgrade=True)) is True
+    )  # CLI --downgrade
 
 
 def test_version_check_incompat(tmp_path):
     """It should fail for a corrupted or invalid database file."""
     db_file = tmp_path / 'db'  # invalid DB file
     db_file.touch()
     with pytest.raises(ServiceFileError):
-        _version_check(db_file, False, False)
+        _version_check(db_file, RunOptions())
 
 
 def test_version_check_no_db(tmp_path):
     """It should pass if there is no DB file (e.g. on workflow first start)."""
     db_file = tmp_path / 'db'  # non-existent file
-    assert _version_check(db_file, False, False)
+    assert _version_check(db_file, RunOptions())
 
 
 @pytest.mark.parametrize(
@@ -257,5 +263,6 @@ def test_distribute_colour(
     _is_terminal = monkeymock('cylc.flow.scheduler_cli.is_terminal')
     _is_terminal.return_value = is_terminal
     _cylc_server_cmd = monkeymock('cylc.flow.scheduler_cli.cylc_server_cmd')
-    _distribute('myhost', 'foo', 'foo/run1', cli_colour)
+    opts = RunOptions(host='myhost', color=cli_colour)
+    _distribute('foo', 'foo/run1', opts)
     assert distribute_colour in _cylc_server_cmd.call_args[0][0]