From 4e70c665ebb901813268d61b1d4729d39d644711 Mon Sep 17 00:00:00 2001 From: Malthe Borch 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 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 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 --- 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 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")