diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py index d5f70b7867892..005f1e215d74d 100644 --- a/airflow/providers/microsoft/psrp/hooks/psrp.py +++ b/airflow/providers/microsoft/psrp/hooks/psrp.py @@ -22,6 +22,7 @@ from typing import Any, Callable, Dict, Iterator, Optional from weakref import WeakKeyDictionary +from pypsrp.host import PSHost from pypsrp.messages import MessageType from pypsrp.powershell import PowerShell, PSInvocationState, RunspacePool from pypsrp.wsman import WSMan @@ -64,6 +65,9 @@ class PsrpHook(BaseHook): :param exchange_keys: If true (default), automatically initiate a session key exchange when the hook is used as a context manager. + :param host: + Optional PowerShell host instance. If this is not set, the default + implementation will be used. You can provide an alternative `configuration_name` using either `runspace_options` or by setting this key as the extra fields of your connection. @@ -82,6 +86,7 @@ def __init__( wsman_options: Optional[Dict[str, Any]] = None, on_output_callback: Optional[OutputCallback] = None, exchange_keys: bool = True, + host: Optional[PSHost] = None, ): self.conn_id = psrp_conn_id self._logging_level = logging_level @@ -90,6 +95,7 @@ def __init__( self._wsman_options = wsman_options or {} self._on_output_callback = on_output_callback self._exchange_keys = exchange_keys + self._host = host or PSHost(None, None, False, type(self).__name__, None, None, "1.0") def __enter__(self): conn = self.get_conn() @@ -144,7 +150,7 @@ def apply_extra(d, keys): if extra: raise AirflowException(f"Unexpected extra configuration keys: {', '.join(sorted(extra))}") - pool = RunspacePool(wsman, **runspace_options) + pool = RunspacePool(wsman, host=self._host, **runspace_options) self._wsman_ref[pool] = wsman return pool diff --git a/airflow/providers/microsoft/psrp/operators/psrp.py b/airflow/providers/microsoft/psrp/operators/psrp.py index dc63881109e37..5ec70f292f536 100644 --- a/airflow/providers/microsoft/psrp/operators/psrp.py +++ b/airflow/providers/microsoft/psrp/operators/psrp.py @@ -146,6 +146,10 @@ def execute(self, context: "Context") -> Optional[List[Any]]: if ps.had_errors: raise AirflowException("Process failed") + rc = ps.runspace_pool.host.rc + if rc: + raise AirflowException(f"Process exited with non-zero status code: {rc}") + if not self.do_xcom_push: return None diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py index 67bf63164dc1d..8826adf66bea7 100644 --- a/tests/providers/microsoft/psrp/hooks/test_psrp.py +++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py @@ -22,6 +22,7 @@ from unittest.mock import MagicMock, Mock, call, patch from parameterized import parameterized +from pypsrp.host import PSHost from pypsrp.messages import MessageType from pypsrp.powershell import PSInvocationState from pytest import raises @@ -180,8 +181,10 @@ def assert_log(level, *args): assert_log(INFO, DUMMY_STACKTRACE[1]) assert call('Invocation state: %s', 'Completed') in logger.info.mock_calls - - assert runspace_pool.call_args == call(ws_man.return_value, connection_name='foo') + args, kwargs = runspace_pool.call_args + assert args == (ws_man.return_value,) + assert kwargs["connection_name"] == "foo" + assert isinstance(kwargs["host"], PSHost) def test_invoke_cmdlet(self, *mocks): with PsrpHook(CONNECTION_ID) as hook: diff --git a/tests/providers/microsoft/psrp/operators/test_psrp.py b/tests/providers/microsoft/psrp/operators/test_psrp.py index 288d630e8db11..6fb77b5f75909 100644 --- a/tests/providers/microsoft/psrp/operators/test_psrp.py +++ b/tests/providers/microsoft/psrp/operators/test_psrp.py @@ -60,13 +60,20 @@ def test_cmdlet_task_id_default(self): ExecuteParameter("powershell", call.add_script("foo"), None), ExecuteParameter("cmdlet", call.add_cmdlet("foo"), {"bar": "baz"}), ], - [False, True], + [ + (False, 0), + (False, None), + (True, None), + (False, 1), + (True, 1), + ], [False, True], ) ) ) @patch(f"{PsrpOperator.__module__}.PsrpHook") - def test_execute(self, parameter, had_errors, do_xcom_push, hook_impl): + def test_execute(self, parameter, result, do_xcom_push, hook_impl): + had_errors, rc = result kwargs = {parameter.name: "foo"} if parameter.expected_parameters: kwargs["parameters"] = parameter.expected_parameters @@ -78,12 +85,19 @@ def test_execute(self, parameter, had_errors, do_xcom_push, hook_impl): do_xcom_push=do_xcom_push, **kwargs, ) - ps = Mock(spec=PowerShell, output=[json.dumps("")], had_errors=had_errors) + runspace_pool = Mock() + runspace_pool.host.rc = rc + ps = Mock( + spec=PowerShell, + output=[json.dumps("")], + had_errors=had_errors, + runspace_pool=runspace_pool, + ) hook_impl.configure_mock( **{"return_value.__enter__.return_value.invoke.return_value.__enter__.return_value": ps} ) - if had_errors: - exception_msg = "Process failed" + if had_errors or rc: + exception_msg = "Process failed" if had_errors else "Process exited with non-zero status code: 1" with pytest.raises(AirflowException, match=exception_msg): op.execute(None) else: