Made WinRMOperator deferrable

providers/microsoft/winrm/provider.yaml

@@ -83,6 +83,11 @@ hooks:
     python-modules:
       - airflow.providers.microsoft.winrm.hooks.winrm
 
+triggers:
+  - integration-name: Windows Remote Management (WinRM)
+    python-modules:
+      - airflow.providers.microsoft.winrm.triggers.winrm
+
 connection-types:
   - hook-class-name: airflow.providers.microsoft.winrm.hooks.winrm.WinRMHook
     connection-type: winrm

providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/get_provider_info.py

@@ -47,6 +47,12 @@ def get_provider_info():
                 "python-modules": ["airflow.providers.microsoft.winrm.hooks.winrm"],
             }
         ],
+        "triggers": [
+            {
+                "integration-name": "Windows Remote Management (WinRM)",
+                "python-modules": ["airflow.providers.microsoft.winrm.triggers.winrm"],
+            }
+        ],
         "connection-types": [
             {
                 "hook-class-name": "airflow.providers.microsoft.winrm.hooks.winrm.WinRMHook",

providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py

@@ -19,16 +19,20 @@
 
 from __future__ import annotations
 
+import logging
 from base64 import b64encode
 from contextlib import suppress
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from winrm.exceptions import WinRMOperationTimeoutError
 from winrm.protocol import Protocol
 
+from airflow.providers.common.compat.connection import get_async_connection
 from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 from airflow.utils.platform import getuser
 
-# TODO: FIXME please - I have too complex implementation
+if TYPE_CHECKING:
+    from airflow.providers.common.compat.sdk import Connection
 
 
 class WinRMHook(BaseHook):
@@ -122,13 +126,10 @@ def __init__(
         self.credssp_disable_tlsv1_2 = credssp_disable_tlsv1_2
         self.send_cbt = send_cbt
 
-        self.winrm_protocol = None
-
-    def get_conn(self):
-        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)
+        self.winrm_protocol: Protocol | None = None
 
+    def create_protocol(self, conn: Connection | None) -> Protocol:
+        if conn:
             if self.username is None:
                 self.username = conn.login
             if self.password is None:
@@ -192,36 +193,57 @@ def get_conn(self):
             self.endpoint = f"http://{self.remote_host}:{self.remote_port}/wsman"
 
         try:
-            if self.password and self.password.strip():
-                self.winrm_protocol = Protocol(
-                    endpoint=self.endpoint,
-                    transport=self.transport,
-                    username=self.username,
-                    password=self.password,
-                    service=self.service,
-                    keytab=self.keytab,
-                    ca_trust_path=self.ca_trust_path,
-                    cert_pem=self.cert_pem,
-                    cert_key_pem=self.cert_key_pem,
-                    server_cert_validation=self.server_cert_validation,
-                    kerberos_delegation=self.kerberos_delegation,
-                    read_timeout_sec=self.read_timeout_sec,
-                    operation_timeout_sec=self.operation_timeout_sec,
-                    kerberos_hostname_override=self.kerberos_hostname_override,
-                    message_encryption=self.message_encryption,
-                    credssp_disable_tlsv1_2=self.credssp_disable_tlsv1_2,
-                    send_cbt=self.send_cbt,
+            winrm_protocol = Protocol(
+                endpoint=self.endpoint,
+                transport=cast(
+                    "Literal['auto', 'basic', 'certificate', 'ntlm', 'kerberos', 'credssp', 'plaintext', 'ssl']",
+                    self.transport,
+                ),
+                username=self.username,
+                password=self.password,
+                service=self.service,
+                keytab=cast("Any", self.keytab),
+                ca_trust_path=cast("str | Literal['legacy_requests']", self.ca_trust_path),
+                cert_pem=self.cert_pem,
+                cert_key_pem=self.cert_key_pem,
+                server_cert_validation=cast(
+                    "Literal['validate', 'ignore'] | None", self.server_cert_validation
+                ),
+                kerberos_delegation=self.kerberos_delegation,
+                read_timeout_sec=self.read_timeout_sec,
+                operation_timeout_sec=self.operation_timeout_sec,
+                kerberos_hostname_override=self.kerberos_hostname_override,
+                message_encryption=cast("Literal['auto', 'always', 'never']", self.message_encryption),
+                credssp_disable_tlsv1_2=self.credssp_disable_tlsv1_2,
+                send_cbt=self.send_cbt,
+            )
+
+            if not hasattr(winrm_protocol, "get_command_output_raw"):
+                # since pywinrm>=0.5 get_command_output_raw replace _raw_get_command_output
+                winrm_protocol.get_command_output_raw = (  # type: ignore[method-assign]
+                    winrm_protocol._raw_get_command_output
                 )
 
+            self.log.info("Establishing WinRM connection to host: %s", self.remote_host)
+
+            return winrm_protocol
         except Exception as error:
             error_msg = f"Error creating connection to host: {self.remote_host}, error: {error}"
             self.log.error(error_msg)
             raise AirflowException(error_msg)
 
-        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
+    def get_conn(self) -> Protocol:
+        if self.winrm_protocol is None:
+            self.winrm_protocol = self.create_protocol(
+                self.get_connection(self.ssh_conn_id) if self.ssh_conn_id else None
+            )
+        return self.winrm_protocol
 
+    async def get_async_conn(self) -> Protocol:
+        if self.winrm_protocol is None:
+            self.winrm_protocol = self.create_protocol(
+                await get_async_connection(self.ssh_conn_id) if self.ssh_conn_id else None
+            )
         return self.winrm_protocol
 
     def run(
@@ -231,7 +253,7 @@ def run(
         output_encoding: str = "utf-8",
         return_output: bool = True,
         working_directory: str | None = None,
-    ) -> tuple[int, list[bytes], list[bytes]]:
+    ) -> tuple[int | None, list[bytes], list[bytes]]:
         """
         Run a command.
 
@@ -243,55 +265,104 @@ def run(
         :param working_directory: specify working directory.
         :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(working_directory=working_directory)
-        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)
+        conn = self.get_conn()
+        shell_id, command_id = self._run_command(
+            conn=conn,
+            command=command,
+            ps_path=ps_path,
+            working_directory=working_directory,
+        )
 
         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 = winrm_client.run_command(shell_id, f"{ps_path} -encodedcommand {encoded_ps}")
-            else:
-                self.log.info("Running command: '%s'...", command)
-                command_id = winrm_client.run_command(shell_id, command)
-
-                # See: https://github.com/diyan/pywinrm/blob/master/winrm/protocol.py
+            command_done = False
             stdout_buffer = []
             stderr_buffer = []
-            command_done = False
+            return_code: int | None = None
+
+            # See: https://github.com/diyan/pywinrm/blob/master/winrm/protocol.py
             while not command_done:
-                # this is an expected error when waiting for a long-running process, just silently retry
-                with suppress(WinRMOperationTimeoutError):
-                    (
-                        stdout,
-                        stderr,
-                        return_code,
-                        command_done,
-                    ) = winrm_client.get_command_output_raw(shell_id, command_id)
-
-                    # Only buffer stdout if we need to so that we minimize memory usage.
-                    if return_output:
-                        stdout_buffer.append(stdout)
-                    stderr_buffer.append(stderr)
-
-                    for line in stdout.decode(output_encoding).splitlines():
-                        self.log.info(line)
-                    for line in stderr.decode(output_encoding).splitlines():
-                        self.log.warning(line)
-
-            winrm_client.cleanup_command(shell_id, command_id)
+                (
+                    stdout,
+                    stderr,
+                    return_code,
+                    command_done,
+                ) = self.get_command_output(conn, shell_id, command_id, output_encoding)
+
+                # Only buffer stdout if we need to so that we minimize memory usage.
+                if return_output:
+                    stdout_buffer.append(stdout)
+                stderr_buffer.append(stderr)
 
             return return_code, stdout_buffer, stderr_buffer
         except Exception as e:
             raise AirflowException(f"WinRM operator error: {e}")
         finally:
-            winrm_client.close_shell(shell_id)
+            conn.cleanup_command(shell_id, command_id)
+            conn.close_shell(shell_id)
+
+    def run_command(
+        self,
+        command: str,
+        ps_path: str | None = None,
+        working_directory: str | None = None,
+    ) -> tuple[str, str]:
+        return self._run_command(self.get_conn(), command, ps_path, working_directory)
+
+    def _run_command(
+        self,
+        conn: Protocol,
+        command: str,
+        ps_path: str | None = None,
+        working_directory: str | None = None,
+    ) -> tuple[str, str]:
+        if not command:
+            raise AirflowException("No command specified so nothing to execute here.")
+
+        try:
+            shell_id = conn.open_shell(working_directory=working_directory)
+
+            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 = conn.run_command(shell_id, f"{ps_path} -encodedcommand {encoded_ps}")
+            else:
+                self.log.info("Running command: '%s'...", command)
+                command_id = conn.run_command(shell_id, command)
+        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)
+
+        return shell_id, command_id
+
+    def get_command_output(
+        self, conn: Protocol, shell_id: str, command_id: str, output_encoding: str = "utf-8"
+    ) -> tuple[bytes, bytes, int | None, bool]:
+        with suppress(WinRMOperationTimeoutError):
+            (
+                stdout,
+                stderr,
+                return_code,
+                command_done,
+            ) = conn.get_command_output_raw(shell_id, command_id)
+
+            self.log.debug("return_code: %s", return_code)
+            self.log.debug("command_done: %s", command_done)
+            self.log_output(stdout, output_encoding=output_encoding)
+            self.log_output(stderr, level=logging.WARNING, output_encoding=output_encoding)
+
+            return stdout, stderr, return_code, command_done
+        return b"", b"", None, False
+
+    def log_output(
+        self,
+        output: bytes | None,
+        level: int = logging.INFO,
+        output_encoding: str = "utf-8",
+    ) -> None:
+        if output:
+            for line in output.decode(output_encoding).splitlines():
+                self.log.log(level, line)
 
     def test_connection(self):
         try: