You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by fo...@apache.org on 2018/02/28 15:17:52 UTC
incubator-airflow git commit: [AIRFLOW-2151] Allow getting the
session from AwsHook
Repository: incubator-airflow
Updated Branches:
refs/heads/master 02292c5d9 -> 68679ae52
[AIRFLOW-2151] Allow getting the session from AwsHook
Add new tests for get_credentials and get_session
methods of AwsHook Implement get_credentials and
get_session methods in AwsHook
Closes #3079 from inytar/aws-session
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/68679ae5
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/68679ae5
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/68679ae5
Branch: refs/heads/master
Commit: 68679ae5263b8c01ea348624ae1c74a8e57a6c8d
Parents: 02292c5
Author: inytar <pi...@fastmail.net>
Authored: Wed Feb 28 16:17:12 2018 +0100
Committer: Fokko Driesprong <fo...@godatadriven.com>
Committed: Wed Feb 28 16:17:35 2018 +0100
----------------------------------------------------------------------
airflow/contrib/hooks/aws_hook.py | 16 ++++++
tests/contrib/hooks/test_aws_hook.py | 84 ++++++++++++++++++++++++++++++-
2 files changed, 99 insertions(+), 1 deletion(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/68679ae5/airflow/contrib/hooks/aws_hook.py
----------------------------------------------------------------------
diff --git a/airflow/contrib/hooks/aws_hook.py b/airflow/contrib/hooks/aws_hook.py
index 7e1b024..2a8fa5f 100644
--- a/airflow/contrib/hooks/aws_hook.py
+++ b/airflow/contrib/hooks/aws_hook.py
@@ -151,3 +151,19 @@ class AwsHook(BaseHook):
session, endpoint_url = self._get_credentials(region_name)
return session.resource(resource_type, endpoint_url=endpoint_url)
+
+ def get_session(self, region_name=None):
+ """Get the underlying boto3.session."""
+ session, _ = self._get_credentials(region_name)
+ return session
+
+ def get_credentials(self, region_name=None):
+ """Get the underlying `botocore.Credentials` object.
+
+ This contains the attributes: access_key, secret_key and token.
+ """
+ session, _ = self._get_credentials(region_name)
+ # Credentials are refreshable, so accessing your access key / secret key
+ # separately can lead to a race condition.
+ # See https://stackoverflow.com/a/36291428/8283373
+ return session.get_credentials().get_frozen_credentials()
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/68679ae5/tests/contrib/hooks/test_aws_hook.py
----------------------------------------------------------------------
diff --git a/tests/contrib/hooks/test_aws_hook.py b/tests/contrib/hooks/test_aws_hook.py
index aa246f0..086e486 100644
--- a/tests/contrib/hooks/test_aws_hook.py
+++ b/tests/contrib/hooks/test_aws_hook.py
@@ -17,14 +17,23 @@ import unittest
import boto3
from airflow import configuration
+from airflow.models import Connection
from airflow.contrib.hooks.aws_hook import AwsHook
+try:
+ from unittest import mock
+except ImportError:
+ try:
+ import mock
+ except ImportError:
+ mock = None
try:
- from moto import mock_emr, mock_dynamodb2
+ from moto import mock_emr, mock_dynamodb2, mock_sts
except ImportError:
mock_emr = None
mock_dynamodb2 = None
+ mock_sts = None
class TestAwsHook(unittest.TestCase):
@@ -77,6 +86,79 @@ class TestAwsHook(unittest.TestCase):
self.assertEqual(table.item_count, 0)
+ @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
+ @mock_dynamodb2
+ def test_get_session_returns_a_boto3_session(self):
+ hook = AwsHook(aws_conn_id='aws_default')
+ session_from_hook = hook.get_session()
+ resource_from_session = session_from_hook.resource('dynamodb')
+ table = resource_from_session.create_table(
+ TableName='test_airflow',
+ KeySchema=[
+ {
+ 'AttributeName': 'id',
+ 'KeyType': 'HASH'
+ },
+ ],
+ AttributeDefinitions=[
+ {
+ 'AttributeName': 'name',
+ 'AttributeType': 'S'
+ }
+ ],
+ ProvisionedThroughput={
+ 'ReadCapacityUnits': 10,
+ 'WriteCapacityUnits': 10
+ }
+ )
+
+ table.meta.client.get_waiter(
+ 'table_exists').wait(TableName='test_airflow')
+
+ self.assertEqual(table.item_count, 0)
+
+ @mock.patch.object(AwsHook, 'get_connection')
+ def test_get_credentials_from_login(self, mock_get_connection):
+ mock_connection = Connection(login='aws_access_key_id',
+ password='aws_secret_access_key')
+ mock_get_connection.return_value = mock_connection
+ hook = AwsHook()
+ credentials_from_hook = hook.get_credentials()
+ self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
+ self.assertEqual(credentials_from_hook.secret_key, 'aws_secret_access_key')
+ self.assertIsNone(credentials_from_hook.token)
+
+ @mock.patch.object(AwsHook, 'get_connection')
+ def test_get_credentials_from_extra(self, mock_get_connection):
+ mock_connection = Connection(
+ extra='{"aws_access_key_id": "aws_access_key_id",'
+ '"aws_secret_access_key": "aws_secret_access_key"}'
+ )
+ mock_get_connection.return_value = mock_connection
+ hook = AwsHook()
+ credentials_from_hook = hook.get_credentials()
+ self.assertEqual(credentials_from_hook.access_key, 'aws_access_key_id')
+ self.assertEqual(credentials_from_hook.secret_key, 'aws_secret_access_key')
+ self.assertIsNone(credentials_from_hook.token)
+
+ @unittest.skipIf(mock_sts is None, 'mock_sts package not present')
+ @mock.patch.object(AwsHook, 'get_connection')
+ @mock_sts
+ def test_get_credentials_from_role_arn(self, mock_get_connection):
+ mock_connection = Connection(
+ extra='{"role_arn":"arn:aws:iam::123456:role/role_arn"}')
+ mock_get_connection.return_value = mock_connection
+ hook = AwsHook()
+ credentials_from_hook = hook.get_credentials()
+ self.assertEqual(credentials_from_hook.access_key, 'AKIAIOSFODNN7EXAMPLE')
+ self.assertEqual(credentials_from_hook.secret_key,
+ 'aJalrXUtnFEMI/K7MDENG/bPxRfiCYzEXAMPLEKEY')
+ self.assertEqual(credentials_from_hook.token,
+ 'BQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh'
+ '3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4I'
+ 'gRmpRV3zrkuWJOgQs8IZZaIv2BXIa2R4OlgkBN9bkUDNCJiBeb/AXlzBBko7b15'
+ 'fjrBs2+cTQtpZ3CYWFXG8C5zqx37wnOE49mRl/+OtkIKGO7fAE')
+
if __name__ == '__main__':
unittest.main()