From 87e73c68c9890b965e359e4cf4d3f1bf4439873a Mon Sep 17 00:00:00 2001
From: Nick Mills-Barrett <nick@fizzadar.com>
Date: Sun, 25 Feb 2024 19:41:50 +0000
Subject: [PATCH] Handle errors in nested operations like unnested operations

This correctly fails hosts when they fail executing a nested operation,
just like unnested operations.
---
 pyinfra/api/exceptions.py | 6 ++++++
 pyinfra/api/operation.py  | 4 +++-
 pyinfra/api/operations.py | 5 +++--
 pyinfra/api/state.py      | 6 +++---
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/pyinfra/api/exceptions.py b/pyinfra/api/exceptions.py
index 4a3e1410b..cf3454774 100644
--- a/pyinfra/api/exceptions.py
+++ b/pyinfra/api/exceptions.py
@@ -4,6 +4,12 @@ class PyinfraError(Exception):
     """
 
 
+class NoMoreHostsError(PyinfraError):
+    """
+    Exception raised when pyinfra runs out of hosts (they all failed).
+    """
+
+
 class ConnectError(PyinfraError):
     """
     Exception raised when connecting fails.
diff --git a/pyinfra/api/operation.py b/pyinfra/api/operation.py
index e4ce11882..962f36b1d 100644
--- a/pyinfra/api/operation.py
+++ b/pyinfra/api/operation.py
@@ -343,7 +343,9 @@ def _execute_immediately(state, host, op_data, op_meta, op_hash):
     )
     op_data["parent_op_hash"] = host.executing_op_hash
     log_operation_start(op_meta, op_types=["nested"], prefix="")
-    run_host_op(state, host, op_hash)
+    status = run_host_op(state, host, op_hash)
+    if status is False:
+        state.fail_hosts({host})
 
 
 def _attach_args(op_meta, args, kwargs):
diff --git a/pyinfra/api/operations.py b/pyinfra/api/operations.py
index ce137cee4..b91d971fd 100644
--- a/pyinfra/api/operations.py
+++ b/pyinfra/api/operations.py
@@ -14,7 +14,7 @@
 
 from .arguments import get_executor_kwarg_keys
 from .command import FunctionCommand, PyinfraCommand, StringCommand
-from .exceptions import PyinfraError
+from .exceptions import NoMoreHostsError, PyinfraError
 from .util import (
     format_exception,
     log_error_or_warning,
@@ -115,7 +115,6 @@ def run_condition(condition_name: str) -> bool:
     all_combined_output_lines = []
 
     for i, command in enumerate(op_data["commands"]):
-
         status = False
 
         executor_kwargs = base_executor_kwargs.copy()
@@ -130,6 +129,8 @@ def run_condition(condition_name: str) -> bool:
         if isinstance(command, FunctionCommand):
             try:
                 status = command.execute(state, host, executor_kwargs)
+            except NoMoreHostsError:
+                status = False
             except Exception as e:  # Custom functions could do anything, so expect anything!
                 _formatted_exc = format_exception(e)
                 _error_msg = "Unexpected error in Python callback: {0}".format(_formatted_exc)
diff --git a/pyinfra/api/state.py b/pyinfra/api/state.py
index e639d645d..378ce8b82 100644
--- a/pyinfra/api/state.py
+++ b/pyinfra/api/state.py
@@ -10,7 +10,7 @@
 from pyinfra import logger
 
 from .config import Config
-from .exceptions import PyinfraError
+from .exceptions import NoMoreHostsError, PyinfraError
 from .util import sha1_hash
 
 if TYPE_CHECKING:
@@ -330,13 +330,13 @@ def fail_hosts(self, hosts_to_fail, activated_count=None):
 
         # No hosts left!
         if not active_hosts:
-            raise PyinfraError("No hosts remaining!")
+            raise NoMoreHostsError("No hosts remaining!")
 
         if self.config.FAIL_PERCENT is not None:
             percent_failed = (1 - len(active_hosts) / activated_count) * 100
 
             if percent_failed > self.config.FAIL_PERCENT:
-                raise PyinfraError(
+                raise NoMoreHostsError(
                     "Over {0}% of hosts failed ({1}%)".format(
                         self.config.FAIL_PERCENT,
                         int(round(percent_failed)),