How to reset / restart an SSH connection mid execution?

[pascal-cm]

Edit: see first comment.

One of my deploy scripts runs two imported deploys, step_1 which assigns a group to the current SSH user, and step_2, which runs a command that expects given SSH user to be in said group.

The first time this deploy runs, step_2 fails, as the newly assigned group isn't taken into account yet in the ongoing SSH session.

@deploy("Step 1")
def step_1():
    ... # Assign new group current SSH user

@deploy("Step 2")
def step_2():
    ... # Run a command that expects current SSH user to to be in group from step_1

step_1() # Succeeds
# Step missing here: connection reset
step_2() # Fails the first time, succeeds the second time

Note: the workaround I currently use in step_1 is utilizing sg as a wrapper, like so: sg newgroup -c 'command --flag'.

While this does work for step_1's functions, it would not be a feasible solution for step_2, as this is kept as an individual Python module that must not know about internals of step_1.

The other workarounds seem clumsy or didn't work.

• Doing a server.reboot() interrupts other services.
• Another attempt with interrupting the current session by issuing server.shell(commands=["kill -HUP $(echo $$)"]) doesn't work either, as pyinfra does not attempt a reconnect (Error: executed 0 commands).
• Finally, my happy go lucky approach to call the host object's disconnect() and connect() functions apparently didn't yield any logs or effect.

What's the suggested solution here? Could we have a server.reconnect() operation for this use case?

[pascal-cm]

The solution turned out to be simple – a custom operation modeled after server.py's reboot() function:

@operation()
def reconnect():
    def do_reconnect():
        host.disconnect()
        host.connect()

    yield FunctionCommand(do_reconnect, (), {})

step_1()
reconnect()
step_2()

(I do not understand what server.reboot() does with askpass; if it would be needed here, too.)

@Fizzadar: IMO this should be added as a stock operation.

[Fizzadar]

PRs welcome! This is the correct solution - the askpass bit is probably worth adding but it's only a just-in-case-of-failure cleanup, doesn't impact actual operation.