You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2023/01/05 21:05:00 UTC

[GitHub] [airflow] ferruzzi opened a new pull request, #28755: Add a new SSM hook and use it in the System Test context builder

ferruzzi opened a new pull request, #28755:
URL: https://github.com/apache/airflow/pull/28755

   Adds a new hook for AWS Systems Manager (SSM), including unit tests, and incorporates it into existing system tests.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on PR #28755:
URL: https://github.com/apache/airflow/pull/28755#issuecomment-1372875398

   I see some tests are failing, I've got a lot on my plate today, I'll try to get to them tomorrow.  Sorry about that.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062928452


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   I don't think we should create this in Hook constructor because it immediately call Airflow connection



##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound
+
+    def get_parameter_value(self, parameter: str, **kwargs) -> 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.ParameterNotFoundException as ex:
+            try:
+                # None may be a valid default value, so pulling it from kwargs
+                # if available instead of setting it as an optional parameter.
+                return kwargs["default"]

Review Comment:
   As another option we could use sentinel  NOTSET from airflow.utils.types, this is available from Airflow 2.3 and we use in some places.
   
   https://github.com/apache/airflow/blob/38f08f6c37a125cc2d80f117be682110af877350/airflow/utils/types.py#L28-L44



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067251104


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   I had two solutions I was playing with and that was not one of them.  I'll try that out, if it works it would be a much less intrusive fix than the one I decided to implement.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1063363611


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   BTW, any concern why we can't use?
   
   ```python
   try:
       return self.conn.get_parameter(Name=parameter)["Parameter"]["Value"]
   except self.conn.exceptions.ParameterNotFound as ex:
       ...
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1063781365


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   Done.  Required a small tweak in the unit tests to match.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on PR #28755:
URL: https://github.com/apache/airflow/pull/28755#issuecomment-1377824253

   Reverting the change [here](https://github.com/apache/airflow/pull/28755/files#diff-9a10c331ad3eee9de5dde9e28f6f26b708be356b390b6e8dd74e4b2d715efca0R96) does pass tests (locally at least) but surely there must e a way to use the hook there.  If I can't figure it out by tomorrow I'll just revert that one line so the rest of this can move forward.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062952380


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound
+
+    def get_parameter_value(self, parameter: str, **kwargs) -> 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.ParameterNotFoundException as ex:
+            try:
+                # None may be a valid default value, so pulling it from kwargs
+                # if available instead of setting it as an optional parameter.
+                return kwargs["default"]

Review Comment:
   Interesting.   I had not seen that before, thanks.  I'll give it some thought.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062974964


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   Yeah, I did it this way specifically to avoid unwrapping the ClientError but that's an option too.  
   
   I'll do it as a property, but I'd like to keep it CamelCase since exception names are... well... exceptions to that rule and are (always?) CamelCase in python.  Would you object to leaving it that way?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] o-nikolas merged pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
o-nikolas merged PR #28755:
URL: https://github.com/apache/airflow/pull/28755


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1063363611


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   BTW, any concern why we can just use
   
   ```python
   try:
       return self.conn.get_parameter(Name=parameter)["Parameter"]["Value"]
   except self.conn.exceptions.ParameterNotFound as ex:
       ...
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062971436


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   Yep, i think better avoid camel case
   
   Another case is catch `botocore.exception.ClientError` and check is attribute `response.get('Error', {}).get('Code', 'Unknown') == 'ParameterNotFound'` - it is looks dumb but it is generic way in case if error happen during client initialisation
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067316123


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   > Thinking this through, that would mean that the connection itself is not being tested.
   
   We tests connections for AwsBaseHook and all other helpers. Actually there is 3 cases and all of them tested:
   1. aws_conn_id = None
   2. aws_conn_id = not-existed-conn, we know only after we tried connect to DB (one in the future we would remove this one)
   3. aws_conn_id = exists-conn
   
   So we don't need test for all Hooks which  based on AwsBaseHook and AwsGenericHook and how they obtain connection from Airflow.
   
   The actual error here happen in this test
   
   https://github.com/apache/airflow/blob/877189916900c930b2c4b92101d46d2f98eb9077/tests/always/test_example_dags.py#L62-L68
   
   I'm not familiar with this test, but seems like it test that example DAGs do not perform queries to Airflow Database during import example dags: https://github.com/apache/airflow/pull/7419 
   
   I don't know maybe this test less relevant nowadays 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067316123


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   > Thinking this through, that would mean that the connection itself is not being tested.
   
   We tests connections for AwsBaseHook and all other helpers. Actually there is 3 cases and all of them tested:
   1. aws_conn_id = None
   2. aws_conn_id = not-existed-conn, we know only after we tried connect to DB (one in the future we would remove this one)
   3. aws_conn_id = exists-conn
   
   So we don't need test for all Hooks which  based on AwsBaseHook and AwsGenericHook and how they obtain connection from Airflow.
   
   The actual error happen in this test
   
   https://github.com/apache/airflow/blob/877189916900c930b2c4b92101d46d2f98eb9077/tests/always/test_example_dags.py#L62-L68
   
   I'm not familiar with this test, but seems like it test that example DAGs do not perform queries to Airflow Database during import example dags: https://github.com/apache/airflow/pull/7419 
   
   I don't know maybe this test less relevant nowadays 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] uranusjr commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
uranusjr commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1063067666


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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)
+
+    @property
+    def ParameterNotFoundException(self):
+        return self.conn.exceptions.ParameterNotFound
+
+    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.ParameterNotFoundException as ex:
+            if isinstance(default, ArgNotSet):
+                raise ex
+            return default

Review Comment:
   ```suggestion
           except self.ParameterNotFoundException:
               if isinstance(default, ArgNotSet):
                   raise
               return default
   ```
   
   This produces a better traceback.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062960492


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound

Review Comment:
   How do you feel about moving it to a @property:
   
   ```
   class SsmHook(AwsBaseHook):
       def __init__(self, *args, **kwargs) -> None:
           kwargs["client_type"] = "ssm"
           super().__init__(*args, **kwargs)
       
       @property
       def ParameterNotFoundException(self):
           return self.conn.exceptions.ParameterNotFound
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1062984408


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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
+
+
+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)
+
+        self.ParameterNotFoundException = self.conn.exceptions.ParameterNotFound
+
+    def get_parameter_value(self, parameter: str, **kwargs) -> 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.ParameterNotFoundException as ex:
+            try:
+                # None may be a valid default value, so pulling it from kwargs
+                # if available instead of setting it as an optional parameter.
+                return kwargs["default"]

Review Comment:
   I'll have to tweak the example a little and do `if isinstance(default, ArgNotSet)` instead of `if default == NOTSET` otherwise MyPy will make me cast the return to a string, but this is much more elegant.   I'll make the change.  



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on PR #28755:
URL: https://github.com/apache/airflow/pull/28755#issuecomment-1374168649

   Causing `test_should_not_do_database_queries` to fail on a whole slew of system tests with the following message:
   
   ```
   E           AssertionError: The expected number of db queries is 0 with extra margin: 0. The current number is 1.
   E           
   E           Recorded query locations:
   E           	cached_property.py:__get__:36>base_aws.py:conn_config:513>base.py:get_connection:72>connection.py:get_connection_from_secrets:425>metastore.py:get_connection:39:	1
   
   tests/test_utils/asserts.py:115: AssertionError
   ```
   
   I'll try to figure it out.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067293206


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   Thinking this through, that would mean that the connection itself is not being tested.  Is that something we want to accept in a system test?
   
   I think it's alright in this case.  The system tests are supposed to be testing the specific services, not whether the connection to SSM is working.... but that does mean it's a less-complete system test.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1063781931


##########
airflow/providers/amazon/aws/hooks/ssm.py:
##########
@@ -0,0 +1,57 @@
+# 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)
+
+    @property
+    def ParameterNotFoundException(self):
+        return self.conn.exceptions.ParameterNotFound
+
+    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.ParameterNotFoundException as ex:
+            if isinstance(default, ArgNotSet):
+                raise ex
+            return default

Review Comment:
   Thanks for the suggestion.  Made the change.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] o-nikolas commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
o-nikolas commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067757109


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   I don't love changing the "prod" code in this way to make the tests pass (I think the expectations in that unit test should be updated for AWS examples). But on the other hand the "prod" code here is technically also test code, so it's not so bad and it's a clean change :+1: Happy to merge!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067015096


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   Hey @ferruzzi I guess the issues with tests here.
   Due to the nature of AwsBaseHook by default it use `aws_conn_id="aws_default"` as result when you try to create SSM Client it Hook tried to obtain connection first.
   
   I think if you set `aws_conn_id` to `None` it should not use Airflow Connection
   
   ```suggestion
       hook = SsmHook(aws_conn_id=None)
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067015096


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   Hey @ferruzzi I guess the issues with tests here.
   Due to the nature of AwsBaseHook by default it use `aws_conn_id="aws_default"` as result when you try to create SSM Client than hook obtain connection first.
   
   https://github.com/apache/airflow/blob/352d492c66e69e816fb1547e46fc1e3b7ba32066/airflow/providers/amazon/aws/hooks/base_aws.py#L507-L526
   
   I think if you set `aws_conn_id` to `None` it should not use Airflow Connection
   
   ```suggestion
       hook = SsmHook(aws_conn_id=None)
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] Taragolis commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067015096


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   Hey @ferruzzi I guess the issues with tests here.
   Due to the nature of AwsBaseHook by default it use `aws_conn_id=aws_default` as result when you try to create SSM Client it Hook tried to obtain connection first.
   
   I think if you set `aws_conn_id` to `None` it should not use Airflow Connection
   
   ```suggestion
       hook = SsmHook(aws_conn_id=None)
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


[GitHub] [airflow] ferruzzi commented on a diff in pull request #28755: Add a new SSM hook and use it in the System Test context builder

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #28755:
URL: https://github.com/apache/airflow/pull/28755#discussion_r1067357578


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -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()

Review Comment:
   Yeah, I was digging in that code yesterday and one option I worked on was adding AWS system tests as an exception and to assert that AWS tests make exactly one query, but I don't really know how important that restriction is. so I was trying to look into the restriction and why it's there.  It looks like your change does fix the test (this PR is passing now) so maybe let's run with this.
   
   @o-nikolas, can you have a look at this when you get a few minutes and let me know what you think?  I think I'm inclined to agree with Taragolis that this is the right solution here.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org