diff --git a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py index 961e37ba3fe40..b35ee7dfa1999 100644 --- a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py +++ b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py @@ -118,13 +118,9 @@ def __init__( self.credssp_disable_tlsv1_2 = credssp_disable_tlsv1_2 self.send_cbt = send_cbt - self.client = None self.winrm_protocol = None def get_conn(self): - if self.client: - return self.client - self.log.debug("Creating WinRM client for conn_id: %s", self.ssh_conn_id) if self.ssh_conn_id is not None: conn = self.get_connection(self.ssh_conn_id) @@ -213,15 +209,16 @@ def get_conn(self): send_cbt=self.send_cbt, ) - self.log.info("Establishing WinRM connection to host: %s", self.remote_host) - self.client = self.winrm_protocol.open_shell() - except Exception as error: - error_msg = f"Error connecting to host: {self.remote_host}, error: {error}" + error_msg = f"Error creating connection to host: {self.remote_host}, error: {error}" self.log.error(error_msg) raise AirflowException(error_msg) - return self.client + if not hasattr(self.winrm_protocol, "get_command_output_raw"): + # since pywinrm>=0.5 get_command_output_raw replace _raw_get_command_output + self.winrm_protocol.get_command_output_raw = self.winrm_protocol._raw_get_command_output + + return self.winrm_protocol def run( self, @@ -241,18 +238,25 @@ def run( :return: returns a tuple containing return_code, stdout and stderr in order. """ winrm_client = self.get_conn() + self.log.info("Establishing WinRM connection to host: %s", self.remote_host) + try: + shell_id = winrm_client.open_shell() + except Exception as error: + error_msg = f"Error connecting to host: {self.remote_host}, error: {error}" + self.log.error(error_msg) + raise AirflowException(error_msg) try: if ps_path is not None: self.log.info("Running command as powershell script: '%s'...", command) encoded_ps = b64encode(command.encode("utf_16_le")).decode("ascii") - command_id = self.winrm_protocol.run_command( # type: ignore[attr-defined] - winrm_client, f"{ps_path} -encodedcommand {encoded_ps}" + command_id = winrm_client.run_command( # type: ignore[attr-defined] + shell_id, f"{ps_path} -encodedcommand {encoded_ps}" ) else: self.log.info("Running command: '%s'...", command) - command_id = self.winrm_protocol.run_command( # type: ignore[attr-defined] - winrm_client, command + command_id = winrm_client.run_command( # type: ignore[attr-defined] + shell_id, command ) # See: https://github.com/diyan/pywinrm/blob/master/winrm/protocol.py @@ -267,8 +271,8 @@ def run( stderr, return_code, command_done, - ) = self.winrm_protocol._raw_get_command_output( # type: ignore[attr-defined] - winrm_client, command_id + ) = winrm_client.get_command_output_raw( # type: ignore[attr-defined] + shell_id, command_id ) # Only buffer stdout if we need to so that we minimize memory usage. @@ -281,12 +285,12 @@ def run( for line in stderr.decode(output_encoding).splitlines(): self.log.warning(line) - self.winrm_protocol.cleanup_command( # type: ignore[attr-defined] - winrm_client, command_id + winrm_client.cleanup_command( # type: ignore[attr-defined] + shell_id, command_id ) return return_code, stdout_buffer, stderr_buffer except Exception as e: raise AirflowException(f"WinRM operator error: {e}") finally: - self.winrm_protocol.close_shell(winrm_client) # type: ignore[attr-defined] + winrm_client.close_shell(shell_id) # type: ignore[attr-defined] diff --git a/providers/microsoft/winrm/tests/unit/microsoft/winrm/hooks/test_winrm.py b/providers/microsoft/winrm/tests/unit/microsoft/winrm/hooks/test_winrm.py index 7c2223cadc3f4..f2140e3a188bf 100644 --- a/providers/microsoft/winrm/tests/unit/microsoft/winrm/hooks/test_winrm.py +++ b/providers/microsoft/winrm/tests/unit/microsoft/winrm/hooks/test_winrm.py @@ -29,15 +29,6 @@ class TestWinRMHook: - @patch("airflow.providers.microsoft.winrm.hooks.winrm.Protocol") - def test_get_conn_exists(self, mock_protocol): - winrm_hook = WinRMHook() - winrm_hook.client = mock_protocol.return_value.open_shell.return_value - - conn = winrm_hook.get_conn() - - assert conn == winrm_hook.client - def test_get_conn_missing_remote_host(self): with pytest.raises(AirflowException): WinRMHook().get_conn() @@ -47,7 +38,7 @@ def test_get_conn_error(self, mock_protocol): mock_protocol.side_effect = Exception("Error") with pytest.raises(AirflowException): - WinRMHook(remote_host="host").get_conn() + WinRMHook(remote_host="host", password="pwd").get_conn() @patch("airflow.providers.microsoft.winrm.hooks.winrm.Protocol", autospec=True) @patch( @@ -151,7 +142,7 @@ def test_run_with_stdout(self, mock_get_connection, mock_protocol): winrm_hook = WinRMHook(ssh_conn_id="conn_id") mock_protocol.return_value.run_command = MagicMock(return_value="command_id") - mock_protocol.return_value._raw_get_command_output = MagicMock( + mock_protocol.return_value.get_command_output_raw = MagicMock( return_value=(b"stdout", b"stderr", 0, True) ) @@ -192,7 +183,7 @@ def test_run_without_stdout(self, mock_get_connection, mock_protocol): winrm_hook = WinRMHook(ssh_conn_id="conn_id") mock_protocol.return_value.run_command = MagicMock(return_value="command_id") - mock_protocol.return_value._raw_get_command_output = MagicMock( + mock_protocol.return_value.get_command_output_raw = MagicMock( return_value=(b"stdout", b"stderr", 0, True) )