[airflow] branch main updated: Add a new SSM hook and use it in the System Test context builder (#28755)

[on...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

onikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 870ecd477a Add a new SSM hook and use it in the System Test context builder (#28755)
870ecd477a is described below

commit 870ecd477af3774546bd82bb71921a03914a2b64
Author: D. Ferruzzi <fe...@amazon.com>
AuthorDate: Wed Jan 11 22:49:44 2023 -0800

    Add a new SSM hook and use it in the System Test context builder (#28755)
    
    * Add a new SSM hook and use it in the System Test context builder
    
    Includes unit tests
    
    Co-authored-by: Andrey Anshin <An...@taragol.is>
---
 airflow/providers/amazon/aws/hooks/ssm.py          |  53 ++++++++++++++++
 airflow/providers/amazon/provider.yaml             |   7 +++
 .../aws/AWS-Systems-Manager_light-bg@4x.png        | Bin 0 -> 75092 bytes
 tests/providers/amazon/aws/hooks/test_ssm.py       |  67 +++++++++++++++++++++
 .../system/providers/amazon/aws/utils/__init__.py  |   7 ++-
 5 files changed, 131 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/ssm.py b/airflow/providers/amazon/aws/hooks/ssm.py
new file mode 100644
index 0000000000..25a7f01f90
--- /dev/null
+++ b/airflow/providers/amazon/aws/hooks/ssm.py
@@ -0,0 +1,53 @@
+# 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 __future__ import annotations
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.utils.types import NOTSET, ArgNotSet
+
+
+class SsmHook(AwsBaseHook):
+    """
+    Interact with Amazon Systems Manager (SSM) using the boto3 library.
+    All API calls available through the Boto API are also available here.
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ssm.html#client
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = "ssm"
+        super().__init__(*args, **kwargs)
+
+    def get_parameter_value(self, parameter: str, default: str | ArgNotSet = NOTSET) -> str:
+        """
+        Returns the value of the provided Parameter or an optional default.
+
+        :param parameter: The SSM Parameter name to return the value for.
+        :param default: Optional default value to return if none is found.
+        """
+        try:
+            return self.conn.get_parameter(Name=parameter)["Parameter"]["Value"]
+        except self.conn.exceptions.ParameterNotFound:
+            if isinstance(default, ArgNotSet):
+                raise
+            return default
diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
index dc401b8884..b1a8aa0f89 100644
--- a/airflow/providers/amazon/provider.yaml
+++ b/airflow/providers/amazon/provider.yaml
@@ -195,6 +195,10 @@ integrations:
     how-to-guide:
       - /docs/apache-airflow-providers-amazon/operators/s3.rst
     tags: [aws]
+  - integration-name: Amazon Systems Manager (SSM)
+    external-doc-url: https://aws.amazon.com/systems-manager/
+    logo: /integration-logos/aws/AWS-Systems-Manager_light-bg@4x.png
+    tags: [aws]
   - integration-name: Amazon Web Services
     external-doc-url: https://aws.amazon.com/
     logo: /integration-logos/aws/AWS-Cloud-alt_light-bg@4x.png
@@ -461,6 +465,9 @@ hooks:
   - integration-name: Amazon Simple Email Service (SES)
     python-modules:
       - airflow.providers.amazon.aws.hooks.ses
+  - integration-name: Amazon Systems Manager (SSM)
+    python-modules:
+      - airflow.providers.amazon.aws.hooks.ssm
   - integration-name: Amazon SecretsManager
     python-modules:
       - airflow.providers.amazon.aws.hooks.secrets_manager
diff --git a/docs/integration-logos/aws/AWS-Systems-Manager_light-bg@4x.png b/docs/integration-logos/aws/AWS-Systems-Manager_light-bg@4x.png
new file mode 100644
index 0000000000..0aae71d6cb
Binary files /dev/null and b/docs/integration-logos/aws/AWS-Systems-Manager_light-bg@4x.png differ
diff --git a/tests/providers/amazon/aws/hooks/test_ssm.py b/tests/providers/amazon/aws/hooks/test_ssm.py
new file mode 100644
index 0000000000..20eb390310
--- /dev/null
+++ b/tests/providers/amazon/aws/hooks/test_ssm.py
@@ -0,0 +1,67 @@
+# 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 __future__ import annotations
+
+import botocore.exceptions
+import pytest
+from moto import mock_ssm
+from pytest import param
+
+from airflow.providers.amazon.aws.hooks.ssm import SsmHook
+
+DEFAULT_CONN_ID: str = "aws_default"
+REGION: str = "us-east-1"
+
+EXISTING_PARAM_NAME = "parameter"
+BAD_PARAM_NAME = "parameter_does_not_exist"
+PARAM_VALUE = "value"
+DEFAULT_VALUE = "default"
+
+
+@mock_ssm
+class TestSsmHooks:
+    def setup(self):
+        self.hook = SsmHook(region_name=REGION)
+        self.hook.conn.put_parameter(Name=EXISTING_PARAM_NAME, Value=PARAM_VALUE, Overwrite=True)
+
+    def test_hook(self) -> None:
+        assert self.hook.conn is not None
+        assert self.hook.aws_conn_id == DEFAULT_CONN_ID
+        assert self.hook.region_name == REGION
+
+    @pytest.mark.parametrize(
+        "param_name, default_value, expected_result",
+        [
+            param(EXISTING_PARAM_NAME, None, PARAM_VALUE, id="param_exists_no_default_provided"),
+            param(EXISTING_PARAM_NAME, DEFAULT_VALUE, PARAM_VALUE, id="param_exists_with_default"),
+            param(BAD_PARAM_NAME, DEFAULT_VALUE, DEFAULT_VALUE, id="param_does_not_exist_uses_default"),
+        ],
+    )
+    def test_get_parameter_value_happy_cases(self, param_name, default_value, expected_result) -> None:
+        if default_value:
+            assert self.hook.get_parameter_value(param_name, default=default_value) == expected_result
+        else:
+            assert self.hook.get_parameter_value(param_name) == expected_result
+
+    def test_get_parameter_value_param_does_not_exist_no_default_provided(self) -> None:
+        with pytest.raises(botocore.exceptions.ClientError) as raised_exception:
+            self.hook.get_parameter_value(BAD_PARAM_NAME)
+
+        error = raised_exception.value.response["Error"]
+        assert error["Code"] == "ParameterNotFound"
+        assert BAD_PARAM_NAME in error["Message"]
diff --git a/tests/system/providers/amazon/aws/utils/__init__.py b/tests/system/providers/amazon/aws/utils/__init__.py
index eea520c085..11d88074a6 100644
--- a/tests/system/providers/amazon/aws/utils/__init__.py
+++ b/tests/system/providers/amazon/aws/utils/__init__.py
@@ -29,6 +29,7 @@ from botocore.client import BaseClient
 from botocore.exceptions import ClientError, NoCredentialsError
 
 from airflow.decorators import task
+from airflow.providers.amazon.aws.hooks.ssm import SsmHook
 
 ENV_ID_ENVIRON_KEY: str = "SYSTEM_TESTS_ENV_ID"
 ENV_ID_KEY: str = "ENV_ID"
@@ -92,15 +93,15 @@ def _fetch_from_ssm(key: str, test_name: str | None = None) -> str:
     :return: The value of the provided key from SSM
     """
     _test_name: str = test_name if test_name else _get_test_name()
-    ssm_client: BaseClient = boto3.client("ssm")
+    hook = SsmHook(aws_conn_id=None)
     value: str = ""
 
     try:
-        value = json.loads(ssm_client.get_parameter(Name=_test_name)["Parameter"]["Value"])[key]
+        value = json.loads(hook.get_parameter_value(_test_name))[key]
     # Since a default value after the SSM check is allowed, these exceptions should not stop execution.
     except NoCredentialsError as e:
         log.info("No boto credentials found: %s", e)
-    except ssm_client.exceptions.ParameterNotFound as e:
+    except hook.conn.exceptions.ParameterNotFound as e:
         log.info("SSM does not contain any parameter for this test: %s", e)
     except KeyError as e:
         log.info("SSM contains one parameter for this test, but not the requested value: %s", e)