From 4e70c665ebb901813268d61b1d4729d39d644711 Mon Sep 17 00:00:00 2001
From: Malthe Borch <mborch@gmail.com>
Date: Tue, 15 Nov 2022 13:43:53 +0100
Subject: [PATCH 1/4] Add option to add arguments to PSRP hook and operator

This also adds 'parameters' as an optional keyword argument to
'invoke_cmdlet', preferred over passing **kwargs (now deprecated).

The change to 'parameters' is motivated by the fact that most
parameters to cmdlets in Windows are capitalized, which is at odds
with the common Python parameter style.
---
 .../providers/microsoft/psrp/hooks/psrp.py    | 27 +++++++++++++++++--
 .../microsoft/psrp/operators/psrp.py          | 12 ++++++++-
 .../microsoft/psrp/hooks/test_psrp.py         |  9 +++++++
 .../microsoft/psrp/operators/test_psrp.py     | 11 +++++---
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py
index 00cdae17a28a3..132c286bdd3b3 100644
--- a/airflow/providers/microsoft/psrp/hooks/psrp.py
+++ b/airflow/providers/microsoft/psrp/hooks/psrp.py
@@ -21,6 +21,7 @@
 from copy import copy
 from logging import DEBUG, ERROR, INFO, WARNING
 from typing import Any, Callable, Generator
+from warnings import warn
 from weakref import WeakKeyDictionary
 
 from pypsrp.host import PSHost
@@ -215,11 +216,33 @@ def invoke(self) -> Generator[PowerShell, None, None]:
             if local_context:
                 self.__exit__(None, None, None)
 
-    def invoke_cmdlet(self, name: str, use_local_scope=None, **parameters: dict[str, str]) -> PowerShell:
+    def invoke_cmdlet(
+        self,
+        name: str,
+        use_local_scope=None,
+        arguments: list[str] | None = None,
+        parameters: dict[str, str] | None = None,
+        **kwargs: str,
+    ) -> PowerShell:
         """Invoke a PowerShell cmdlet and return session."""
+        if kwargs:
+            if parameters:
+                raise ValueError("**kwargs not allowed when 'parameters' is used at the same time.")
+            warn(
+                "Passing **kwargs to 'invoke_cmdlet' is deprecated "
+                "and will be removed in a future release. Please use 'parameters' "
+                "instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            parameters = kwargs
+
         with self.invoke() as ps:
             ps.add_cmdlet(name, use_local_scope=use_local_scope)
-            ps.add_parameters(parameters)
+            for argument in arguments or ():
+                ps.add_argument(argument)
+            if parameters:
+                ps.add_parameters(parameters)
         return ps
 
     def invoke_powershell(self, script: str) -> PowerShell:
diff --git a/airflow/providers/microsoft/psrp/operators/psrp.py b/airflow/providers/microsoft/psrp/operators/psrp.py
index 733b8cb29fd3a..1f74225de3be6 100644
--- a/airflow/providers/microsoft/psrp/operators/psrp.py
+++ b/airflow/providers/microsoft/psrp/operators/psrp.py
@@ -56,8 +56,11 @@ class PsrpOperator(BaseOperator):
     :param cmdlet:
         cmdlet to execute on remote host (templated). Also used as the default
         value for `task_id`.
+    :param arguments:
+        When using the `cmdlet` or `powershell` option, use `arguments` to
+        provide arguments (templated).
     :param parameters:
-        When using the `cmdlet` or `powershell` arguments, use this parameter to
+        When using the `cmdlet` or `powershell` option, use `parameters` to
         provide parameters (templated). Note that a parameter with a value of `None`
         becomes an *argument* (i.e., switch).
     :param logging_level:
@@ -79,6 +82,7 @@ class PsrpOperator(BaseOperator):
     template_fields: Sequence[str] = (
         "cmdlet",
         "command",
+        "arguments",
         "parameters",
         "powershell",
     )
@@ -92,6 +96,7 @@ def __init__(
         command: str | None = None,
         powershell: str | None = None,
         cmdlet: str | None = None,
+        arguments: list[str] | None = None,
         parameters: dict[str, str] | None = None,
         logging_level: int = DEBUG,
         runspace_options: dict[str, Any] | None = None,
@@ -102,6 +107,8 @@ def __init__(
         args = {command, powershell, cmdlet}
         if not exactly_one(*args):
             raise ValueError("Must provide exactly one of 'command', 'powershell', or 'cmdlet'")
+        if arguments and not cmdlet:
+            raise ValueError("Arguments only allowed with 'cmdlet'")
         if parameters and not cmdlet:
             raise ValueError("Parameters only allowed with 'cmdlet'")
         if cmdlet:
@@ -111,6 +118,7 @@ def __init__(
         self.command = command
         self.powershell = powershell
         self.cmdlet = cmdlet
+        self.arguments = arguments
         self.parameters = parameters
         self.logging_level = logging_level
         self.runspace_options = runspace_options
@@ -134,6 +142,8 @@ def execute(self, context: Context) -> list[Any] | None:
                     ps.add_cmdlet(self.cmdlet)
                 else:
                     ps.add_script(self.powershell)
+                for argument in self.arguments or ():
+                    ps.add_argument(argument)
                 if self.parameters:
                     ps.add_parameters(self.parameters)
                 if self.do_xcom_push:
diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py
index bc1b06aa5f7d1..07c66fb639640 100644
--- a/tests/providers/microsoft/psrp/hooks/test_psrp.py
+++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py
@@ -188,6 +188,15 @@ def assert_log(level, *args):
         assert isinstance(kwargs["host"], PSHost)
 
     def test_invoke_cmdlet(self, *mocks):
+        arguments = ("a", "b", "c")
+        parameters = {"bar": 1, "baz": "2"}
+        with PsrpHook(CONNECTION_ID) as hook:
+            ps = hook.invoke_cmdlet("foo", arguments=arguments, parameters=parameters)
+            assert [call("foo", use_local_scope=None)] == ps.add_cmdlet.mock_calls
+            assert [call({"bar": "1", "baz": "2"})] == ps.add_parameters.mock_calls
+            assert [call(arg) for arg in arguments] == ps.add_argument.mock_calls
+
+    def test_invoke_cmdlet_deprecated_kwargs(self, *mocks):
         with PsrpHook(CONNECTION_ID) as hook:
             ps = hook.invoke_cmdlet("foo", bar="1", baz="2")
             assert [call("foo", use_local_scope=None)] == ps.add_cmdlet.mock_calls
diff --git a/tests/providers/microsoft/psrp/operators/test_psrp.py b/tests/providers/microsoft/psrp/operators/test_psrp.py
index 980c7efc8d48e..bad5f02f5685b 100644
--- a/tests/providers/microsoft/psrp/operators/test_psrp.py
+++ b/tests/providers/microsoft/psrp/operators/test_psrp.py
@@ -35,6 +35,7 @@
 class ExecuteParameter(NamedTuple):
     name: str
     expected_method: str
+    expected_arguments: list[str] | None
     expected_parameters: dict[str, Any] | None
 
 
@@ -57,15 +58,17 @@ def test_cmdlet_task_id_default(self):
         [
             # These tuples map the command parameter to an execution method and parameter set.
             pytest.param(
-                ExecuteParameter("command", call.add_script("cmd.exe /c @'\nfoo\n'@"), None), id="command"
+                ExecuteParameter("command", call.add_script("cmd.exe /c @'\nfoo\n'@"), None, None), id="command"
             ),
-            pytest.param(ExecuteParameter("powershell", call.add_script("foo"), None), id="powershell"),
-            pytest.param(ExecuteParameter("cmdlet", call.add_cmdlet("foo"), {"bar": "baz"}), id="cmdlet"),
+            pytest.param(ExecuteParameter("powershell", call.add_script("foo"), None, None), id="powershell"),
+            pytest.param(ExecuteParameter("cmdlet", call.add_cmdlet("foo"), ["abc"], {"bar": "baz"}), id="cmdlet"),
         ],
     )
     @patch(f"{PsrpOperator.__module__}.PsrpHook")
     def test_execute(self, hook_impl, parameter, had_errors, rc, do_xcom_push):
         kwargs = {parameter.name: "foo"}
+        if parameter.expected_arguments:
+            kwargs["arguments"] = parameter.expected_arguments
         if parameter.expected_parameters:
             kwargs["parameters"] = parameter.expected_parameters
         psrp_session_init = Mock(spec=Command)
@@ -100,6 +103,8 @@ def test_execute(self, hook_impl, parameter, had_errors, rc, do_xcom_push):
             call.add_command(psrp_session_init),
             parameter.expected_method,
         ]
+        if parameter.expected_arguments:
+            expected_ps_calls.append(call.add_argument("abc"))
         if parameter.expected_parameters:
             expected_ps_calls.extend([call.add_parameters({"bar": "baz"})])
         if parameter.name in ("cmdlet", "powershell") and do_xcom_push:

From 82db35d35936372f4d66ae26b35ac06d1c33e41c Mon Sep 17 00:00:00 2001
From: Malthe Borch <mborch@gmail.com>
Date: Wed, 16 Nov 2022 07:56:08 +0100
Subject: [PATCH 2/4] Fix test

---
 tests/providers/microsoft/psrp/hooks/test_psrp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py
index 07c66fb639640..563375e95842d 100644
--- a/tests/providers/microsoft/psrp/hooks/test_psrp.py
+++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py
@@ -189,7 +189,7 @@ def assert_log(level, *args):
 
     def test_invoke_cmdlet(self, *mocks):
         arguments = ("a", "b", "c")
-        parameters = {"bar": 1, "baz": "2"}
+        parameters = {"bar": "1", "baz": "2"}
         with PsrpHook(CONNECTION_ID) as hook:
             ps = hook.invoke_cmdlet("foo", arguments=arguments, parameters=parameters)
             assert [call("foo", use_local_scope=None)] == ps.add_cmdlet.mock_calls

From ee84ad3761518b17ec9f9fcbb66135f73ae67c81 Mon Sep 17 00:00:00 2001
From: Malthe Borch <mborch@gmail.com>
Date: Wed, 16 Nov 2022 07:09:19 +0000
Subject: [PATCH 3/4] Add type to 'use_local_scope' parameter

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/providers/microsoft/psrp/hooks/psrp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py
index 132c286bdd3b3..4e4518609f6ac 100644
--- a/airflow/providers/microsoft/psrp/hooks/psrp.py
+++ b/airflow/providers/microsoft/psrp/hooks/psrp.py
@@ -219,7 +219,7 @@ def invoke(self) -> Generator[PowerShell, None, None]:
     def invoke_cmdlet(
         self,
         name: str,
-        use_local_scope=None,
+        use_local_scope: bool | None = None,
         arguments: list[str] | None = None,
         parameters: dict[str, str] | None = None,
         **kwargs: str,

From 9037f10ae63f42e5e0e4fbad1c7aeac2f4127766 Mon Sep 17 00:00:00 2001
From: Malthe Borch <mborch@gmail.com>
Date: Sun, 1 Jan 2023 23:27:27 +0100
Subject: [PATCH 4/4] Fix syntax formatting

---
 tests/providers/microsoft/psrp/operators/test_psrp.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/providers/microsoft/psrp/operators/test_psrp.py b/tests/providers/microsoft/psrp/operators/test_psrp.py
index bad5f02f5685b..6e16e4ab47237 100644
--- a/tests/providers/microsoft/psrp/operators/test_psrp.py
+++ b/tests/providers/microsoft/psrp/operators/test_psrp.py
@@ -58,10 +58,13 @@ def test_cmdlet_task_id_default(self):
         [
             # These tuples map the command parameter to an execution method and parameter set.
             pytest.param(
-                ExecuteParameter("command", call.add_script("cmd.exe /c @'\nfoo\n'@"), None, None), id="command"
+                ExecuteParameter("command", call.add_script("cmd.exe /c @'\nfoo\n'@"), None, None),
+                id="command",
             ),
             pytest.param(ExecuteParameter("powershell", call.add_script("foo"), None, None), id="powershell"),
-            pytest.param(ExecuteParameter("cmdlet", call.add_cmdlet("foo"), ["abc"], {"bar": "baz"}), id="cmdlet"),
+            pytest.param(
+                ExecuteParameter("cmdlet", call.add_cmdlet("foo"), ["abc"], {"bar": "baz"}), id="cmdlet"
+            ),
         ],
     )
     @patch(f"{PsrpOperator.__module__}.PsrpHook")
