From 69d2ed65cb7c9384d309ae5e499d5798c2c3ac96 Mon Sep 17 00:00:00 2001 From: Malthe Borch Date: Sun, 8 Aug 2021 14:02:59 +0200 Subject: [PATCH] Add Microsoft PSRP provider (#17361) --- CONTRIBUTING.rst | 9 +- INSTALL | 9 +- .../providers/microsoft/psrp/CHANGELOG.rst | 25 ++++ airflow/providers/microsoft/psrp/__init__.py | 17 +++ .../microsoft/psrp/hooks/__init__.py | 17 +++ .../providers/microsoft/psrp/hooks/psrp.py | 118 ++++++++++++++++++ .../microsoft/psrp/operators/__init__.py | 17 +++ .../microsoft/psrp/operators/psrp.py | 66 ++++++++++ .../providers/microsoft/psrp/provider.yaml | 45 +++++++ .../commits.rst | 20 +++ .../index.rst | 44 +++++++ docs/apache-airflow/extra-packages-ref.rst | 2 + docs/spelling_wordlist.txt | 3 + setup.py | 5 + tests/providers/microsoft/psrp/__init__.py | 17 +++ .../microsoft/psrp/hooks/__init__.py | 17 +++ .../microsoft/psrp/hooks/test_psrp.py | 69 ++++++++++ .../microsoft/psrp/operators/__init__.py | 17 +++ .../microsoft/psrp/operators/test_psrp.py | 55 ++++++++ 19 files changed, 564 insertions(+), 8 deletions(-) create mode 100644 airflow/providers/microsoft/psrp/CHANGELOG.rst create mode 100644 airflow/providers/microsoft/psrp/__init__.py create mode 100644 airflow/providers/microsoft/psrp/hooks/__init__.py create mode 100644 airflow/providers/microsoft/psrp/hooks/psrp.py create mode 100644 airflow/providers/microsoft/psrp/operators/__init__.py create mode 100644 airflow/providers/microsoft/psrp/operators/psrp.py create mode 100644 airflow/providers/microsoft/psrp/provider.yaml create mode 100644 docs/apache-airflow-providers-microsoft-psrp/commits.rst create mode 100644 docs/apache-airflow-providers-microsoft-psrp/index.rst create mode 100644 tests/providers/microsoft/psrp/__init__.py create mode 100644 tests/providers/microsoft/psrp/hooks/__init__.py create mode 100644 tests/providers/microsoft/psrp/hooks/test_psrp.py create mode 100644 tests/providers/microsoft/psrp/operators/__init__.py create mode 100644 tests/providers/microsoft/psrp/operators/test_psrp.py diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 5222edab2264b..d8e2c431b7e69 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -592,10 +592,11 @@ cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, deprecate devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, jdbc, jenkins, jira, kerberos, kubernetes, ldap, leveldb, microsoft.azure, -microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, opsgenie, oracle, -pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, rabbitmq, redis, s3, -salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, snowflake, spark, sqlite, -ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, winrm, yandex, zendesk +microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, +opsgenie, oracle, pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, +rabbitmq, redis, s3, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, +snowflake, spark, sqlite, ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, +winrm, yandex, zendesk .. END EXTRAS HERE diff --git a/INSTALL b/INSTALL index 616005cb587b1..47f48c3ece76c 100644 --- a/INSTALL +++ b/INSTALL @@ -96,10 +96,11 @@ cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, deprecate devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, jdbc, jenkins, jira, kerberos, kubernetes, ldap, leveldb, microsoft.azure, -microsoft.mssql, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, opsgenie, oracle, -pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, rabbitmq, redis, s3, -salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, snowflake, spark, sqlite, -ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, winrm, yandex, zendesk +microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, +opsgenie, oracle, pagerduty, papermill, password, pinot, plexus, postgres, presto, qds, qubole, +rabbitmq, redis, s3, salesforce, samba, segment, sendgrid, sentry, sftp, singularity, slack, +snowflake, spark, sqlite, ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, +winrm, yandex, zendesk # END EXTRAS HERE diff --git a/airflow/providers/microsoft/psrp/CHANGELOG.rst b/airflow/providers/microsoft/psrp/CHANGELOG.rst new file mode 100644 index 0000000000000..cef7dda80708a --- /dev/null +++ b/airflow/providers/microsoft/psrp/CHANGELOG.rst @@ -0,0 +1,25 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + +Changelog +--------- + +1.0.0 +..... + +Initial version of the provider. diff --git a/airflow/providers/microsoft/psrp/__init__.py b/airflow/providers/microsoft/psrp/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/providers/microsoft/psrp/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/providers/microsoft/psrp/hooks/__init__.py b/airflow/providers/microsoft/psrp/hooks/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/providers/microsoft/psrp/hooks/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py new file mode 100644 index 0000000000000..8cb809646e76e --- /dev/null +++ b/airflow/providers/microsoft/psrp/hooks/psrp.py @@ -0,0 +1,118 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from time import sleep + +from pypsrp.messages import ErrorRecord, InformationRecord, ProgressRecord +from pypsrp.powershell import PowerShell, PSInvocationState, RunspacePool +from pypsrp.wsman import WSMan + +from airflow.exceptions import AirflowException +from airflow.hooks.base import BaseHook + + +class PSRPHook(BaseHook): + """ + Hook for PowerShell Remoting Protocol execution. + + The hook must be used as a context manager. + """ + + _client = None + _poll_interval = 1 + + def __init__(self, psrp_conn_id: str): + self.conn_id = psrp_conn_id + + def __enter__(self): + conn = self.get_connection(self.conn_id) + + self.log.info("Establishing WinRM connection %s to host: %s", self.conn_id, conn.host) + self._client = WSMan( + conn.host, + ssl=True, + auth="ntlm", + encryption="never", + username=conn.login, + password=conn.password, + cert_validation=False, + ) + self._client.__enter__() + return self + + def __exit__(self, exc_type, exc_value, traceback): + try: + self._client.__exit__() + finally: + self._client = None + + def invoke_powershell(self, script: str) -> PowerShell: + with RunspacePool(self._client) as pool: + ps = PowerShell(pool) + ps.add_script(script) + ps.begin_invoke() + streams = [ + (ps.output, self._log_output), + (ps.streams.debug, self._log_record), + (ps.streams.information, self._log_record), + (ps.streams.error, self._log_record), + ] + offsets = [0 for _ in streams] + + # We're using polling to make sure output and streams are + # handled while the process is running. + while ps.state == PSInvocationState.RUNNING: + sleep(self._poll_interval) + ps.poll_invoke() + + for (i, (stream, handler)) in enumerate(streams): + offset = offsets[i] + while len(stream) > offset: + handler(stream[offset]) + offset += 1 + offsets[i] = offset + + # For good measure, we'll make sure the process has + # stopped running. + ps.end_invoke() + + if ps.streams.error: + raise AirflowException("Process had one or more errors") + + self.log.info("Invocation state: %s", str(PSInvocationState(ps.state))) + return ps + + def _log_output(self, message: str): + self.log.info("%s", message) + + def _log_record(self, record): + # TODO: Consider translating some or all of these records into + # normal logging levels, using `log(level, msg, *args)`. + if isinstance(record, ErrorRecord): + self.log.info("Error: %s", record) + return + + if isinstance(record, InformationRecord): + self.log.info("Information: %s", record.message_data) + return + + if isinstance(record, ProgressRecord): + self.log.info("Progress: %s (%s)", record.activity, record.description) + return + + self.log.info("Unsupported record type: %s", type(record).__name__) diff --git a/airflow/providers/microsoft/psrp/operators/__init__.py b/airflow/providers/microsoft/psrp/operators/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/providers/microsoft/psrp/operators/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/providers/microsoft/psrp/operators/psrp.py b/airflow/providers/microsoft/psrp/operators/psrp.py new file mode 100644 index 0000000000000..7b29de5d4b9eb --- /dev/null +++ b/airflow/providers/microsoft/psrp/operators/psrp.py @@ -0,0 +1,66 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from typing import List, Optional + +from airflow.exceptions import AirflowException +from airflow.models import BaseOperator +from airflow.providers.microsoft.psrp.hooks.psrp import PSRPHook + + +class PSRPOperator(BaseOperator): + """PowerShell Remoting Protocol operator. + + :param psrp_conn_id: connection id + :type psrp_conn_id: str + :param command: command to execute on remote host. (templated) + :type command: str + :param powershell: powershell to execute on remote host. (templated) + :type powershell: str + """ + + template_fields = ( + "command", + "powershell", + ) + template_fields_renderers = {"command": "powershell", "powershell": "powershell"} + ui_color = "#901dd2" + + def __init__( + self, + *, + psrp_conn_id: str, + command: Optional[str] = None, + powershell: Optional[str] = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if not (command or powershell): + raise ValueError("Must provide either 'command' or 'powershell'") + self.conn_id = psrp_conn_id + self.command = command + self.powershell = powershell + + def execute(self, context: dict) -> List[str]: + with PSRPHook(self.conn_id) as hook: + ps = hook.invoke_powershell( + f"cmd.exe /c @'\n{self.command}\n'@" if self.command else self.powershell + ) + if ps.had_errors: + raise AirflowException("Process failed") + return ps.output diff --git a/airflow/providers/microsoft/psrp/provider.yaml b/airflow/providers/microsoft/psrp/provider.yaml new file mode 100644 index 0000000000000..024f8acbe4028 --- /dev/null +++ b/airflow/providers/microsoft/psrp/provider.yaml @@ -0,0 +1,45 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +--- +package-name: apache-airflow-providers-microsoft-psrp +name: PowerShell Remoting Protocol (PSRP) +description: | + `PowerShell Remoting Protocol (PSRP) + `__ + +versions: + - 1.0.0 + +additional-dependencies: + - pypsrp>=0.5.0 + +integrations: + - integration-name: Windows Remote Management (WinRM) + external-doc-url: https://docs.microsoft.com/en-us/windows/win32/winrm/portal + logo: /integration-logos/winrm/WinRM.png + tags: [protocol] + +operators: + - integration-name: Windows Remote Management (WinRM) + python-modules: + - airflow.providers.microsoft.winrm.operators.winrm + +hooks: + - integration-name: Windows Remote Management (WinRM) + python-modules: + - airflow.providers.microsoft.winrm.hooks.winrm diff --git a/docs/apache-airflow-providers-microsoft-psrp/commits.rst b/docs/apache-airflow-providers-microsoft-psrp/commits.rst new file mode 100644 index 0000000000000..a8128bcb267df --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-psrp/commits.rst @@ -0,0 +1,20 @@ + + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Package apache-airflow-providers-microsoft-psrp +----------------------------------------------- diff --git a/docs/apache-airflow-providers-microsoft-psrp/index.rst b/docs/apache-airflow-providers-microsoft-psrp/index.rst new file mode 100644 index 0000000000000..9e30528763770 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-psrp/index.rst @@ -0,0 +1,44 @@ + + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +``apache-airflow-providers-microsoft-psrp`` +=========================================== + +Content +------- + +.. toctree:: + :maxdepth: 1 + :caption: References + + Python API <_api/airflow/providers/microsoft/psrp/index> + +.. toctree:: + :maxdepth: 1 + :caption: Resources + + PyPI Repository + +.. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN AT RELEASE TIME! + + +.. toctree:: + :maxdepth: 1 + :caption: Commits + + Detailed list of commits diff --git a/docs/apache-airflow/extra-packages-ref.rst b/docs/apache-airflow/extra-packages-ref.rst index 2af6fda6abd8c..ae8872883d274 100644 --- a/docs/apache-airflow/extra-packages-ref.rst +++ b/docs/apache-airflow/extra-packages-ref.rst @@ -276,6 +276,8 @@ Those are extras that provide support for integration with external systems via +---------------------+-----------------------------------------------------+--------------------------------------+--------------+ | ssh | ``pip install 'apache-airflow[ssh]'`` | SSH hooks and operators | | +---------------------+-----------------------------------------------------+--------------------------------------+--------------+ +| microsoft.psrp | ``pip install 'apache-airflow[microsoft.psrp]'`` | PSRP hooks and operators | | ++---------------------+-----------------------------------------------------+--------------------------------------+--------------+ | microsoft.winrm | ``pip install 'apache-airflow[microsoft.winrm]'`` | WinRM hooks and operators | | +---------------------+-----------------------------------------------------+--------------------------------------+--------------+ diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 53f886a628ec7..47c82baf461bb 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -293,6 +293,7 @@ Reddit Redhat ReidentifyContentResponse Reinitialising +Remoting ResourceRequirements Roadmap Robinhood @@ -1077,6 +1078,7 @@ proto protobuf provisioner psql +psrp psycopg pty pubsub @@ -1092,6 +1094,7 @@ pymssql pymysql pyodbc pypa +pypsrp pytest pythonic pythonpath diff --git a/setup.py b/setup.py index 9441fab0687d0..2eaaa477df1bf 100644 --- a/setup.py +++ b/setup.py @@ -415,6 +415,9 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'psycopg2-binary>=2.7.4', ] presto = ['presto-python-client>=0.7.0,<0.8'] +psrp = [ + 'pypsrp~=0.5', +] qubole = [ 'qds-sdk>=1.10.4', ] @@ -513,6 +516,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'paramiko', 'pipdeptree', 'pre-commit', + 'pypsrp', 'pygithub', 'pysftp', 'pytest~=6.0', @@ -573,6 +577,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'jira': jira, 'microsoft.azure': azure, 'microsoft.mssql': mssql, + 'microsoft.psrp': psrp, 'microsoft.winrm': winrm, 'mongo': mongo, 'mysql': mysql, diff --git a/tests/providers/microsoft/psrp/__init__.py b/tests/providers/microsoft/psrp/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/providers/microsoft/psrp/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/providers/microsoft/psrp/hooks/__init__.py b/tests/providers/microsoft/psrp/hooks/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/providers/microsoft/psrp/hooks/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py new file mode 100644 index 0000000000000..84e04e88d264b --- /dev/null +++ b/tests/providers/microsoft/psrp/hooks/test_psrp.py @@ -0,0 +1,69 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +import unittest +from unittest.mock import MagicMock, call, patch + +from pypsrp.messages import InformationRecord +from pypsrp.powershell import PSInvocationState + +from airflow.models import Connection +from airflow.providers.microsoft.psrp.hooks.psrp import PSRPHook + +CONNECTION_ID = "conn_id" + + +class TestPSRPHook(unittest.TestCase): + @patch( + f"{PSRPHook.__module__}.{PSRPHook.__name__}.get_connection", + return_value=Connection( + login='username', + password='password', + host='remote_host', + ), + ) + @patch(f"{PSRPHook.__module__}.WSMan") + @patch(f"{PSRPHook.__module__}.PowerShell") + @patch(f"{PSRPHook.__module__}.RunspacePool") + @patch("logging.Logger.info") + def test_invoke_powershell(self, log_info, runspace_pool, powershell, ws_man, get_connection): + with PSRPHook(CONNECTION_ID) as hook: + ps = powershell.return_value = MagicMock() + ps.state = PSInvocationState.RUNNING + ps.output = [] + ps.streams.debug = [] + ps.streams.information = [] + ps.streams.error = [] + + def poll_invoke(): + ps.output.append("") + ps.streams.debug.append(MagicMock(spec=InformationRecord, message_data="")) + ps.state = PSInvocationState.COMPLETED + + def end_invoke(): + ps.streams.error = [] + + ps.poll_invoke.side_effect = poll_invoke + ps.end_invoke.side_effect = end_invoke + + hook.invoke_powershell("foo") + + assert call('%s', '') in log_info.mock_calls + assert call('Information: %s', '') in log_info.mock_calls + assert call('Invocation state: %s', 'Completed') in log_info.mock_calls diff --git a/tests/providers/microsoft/psrp/operators/__init__.py b/tests/providers/microsoft/psrp/operators/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/tests/providers/microsoft/psrp/operators/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/providers/microsoft/psrp/operators/test_psrp.py b/tests/providers/microsoft/psrp/operators/test_psrp.py new file mode 100644 index 0000000000000..df4b6fe54c3fb --- /dev/null +++ b/tests/providers/microsoft/psrp/operators/test_psrp.py @@ -0,0 +1,55 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import unittest +from unittest.mock import patch + +import pytest +from parameterized import parameterized + +from airflow.exceptions import AirflowException +from airflow.providers.microsoft.psrp.operators.psrp import PSRPOperator + +CONNECTION_ID = "conn_id" + + +class TestPSRPOperator(unittest.TestCase): + def test_no_command_or_powershell(self): + exception_msg = "Must provide either 'command' or 'powershell'" + with pytest.raises(ValueError, match=exception_msg): + PSRPOperator(task_id='test_task_id', psrp_conn_id=CONNECTION_ID) + + @parameterized.expand( + [ + (False,), + (True,), + ] + ) + @patch(f"{PSRPOperator.__module__}.PSRPHook") + def test_execute(self, had_errors, hook): + op = PSRPOperator(task_id='test_task_id', psrp_conn_id=CONNECTION_ID, command='dummy') + ps = hook.return_value.__enter__.return_value.invoke_powershell.return_value + ps.output = [""] + ps.had_errors = had_errors + if had_errors: + exception_msg = "Process failed" + with pytest.raises(AirflowException, match=exception_msg): + op.execute(None) + else: + output = op.execute(None) + assert output == ps.output