Repository: apache/airflow
Pull request: Update AWS Base hook to use refreshable credentials (#16770)
Update AWS Base hook to use refreshable credentials when doing assume_role or similar (#16770)

TODO: Unit tests [DONE]
I've already tested in my proper Airflow environment and it works as expected.
I'm submitting this to CI so long so I can find issues more easily.

closes: #16770 


Problem statement:
AWS hook should automatically refresh credentials when using temporary credentials
**Apache Airflow version**: 1.10.8 (Patched with latest AWS Hook)

**Environment**:

- **Cloud provider or hardware configuration**: 4 VCPU 8GB RAM VM
- **OS** (e.g. from /etc/os-release): RHEL 7.7
- **Kernel** (e.g. `uname -a`): Linux 3.10.0-957.el7.x86_64
- **Install tools**:
- **Others**:
The AWS Hook functionality for AssumeRoleWithSAML is not available in this version, we manually added it via patching the hook file.

**What happened**:
We've been using this hook for a while now with this issue, basically sts.assume_role and sts.assume_role_with_saml will return temporary credentials that are only valid for eg 1 hour by default. Eventually with long running operators / hooks / sensors some of them fail because the credentials have expired.

Example error messages
An error occurred (ExpiredTokenException) when calling the AssumeRole operation: Response has expired
An error occurred (ExpiredTokenException) when calling the AssumeRoleWithSAML operation: Response has expired
botocore.exceptions.ClientError: An error occurred (ExpiredTokenException) when calling the <any operation here> operation: The security token included in the request is expired

**What you expected to happen**:

AWS hook should be updated to use boto3 RefreshableCredentials when temporary credentials are in use.

**How to reproduce it**:

Use any of the assume role methods with the AWS Hook, create a session, wait 1 hour (or whatever expiry period applies to your role), and try and use the hook again.

**Anything else we need to know**:

I have a solution, please self-assign this.

Hints from issue discussion:

Commit message: Update AWS Base hook to use refreshable credentials when doing assume_role or similar (#16770)

Patch to review:
diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py
index a1b7e7c2509ab..6589682ee82c4 100644
--- a/airflow/providers/amazon/aws/hooks/base_aws.py
+++ b/airflow/providers/amazon/aws/hooks/base_aws.py
@@ -36,7 +36,7 @@
 import requests
 import tenacity
 from botocore.config import Config
-from botocore.credentials import ReadOnlyCredentials
+from botocore.credentials import DeferredRefreshableCredentials, ReadOnlyCredentials, RefreshableCredentials
 
 try:
     from functools import cached_property
@@ -58,6 +58,8 @@ def __init__(self, conn: Connection, region_name: Optional[str], config: Config)
         self.region_name = region_name
         self.config = config
         self.extra_config = self.conn.extra_dejson
+        self.basic_session = None
+        self.role_arn = None
 
     def create_session(self) -> boto3.session.Session:
         """Create AWS session."""
@@ -68,13 +70,13 @@ def create_session(self) -> boto3.session.Session:
                 self.extra_config["session_kwargs"],
             )
             session_kwargs = self.extra_config["session_kwargs"]
-        session = self._create_basic_session(session_kwargs=session_kwargs)
-        role_arn = self._read_role_arn_from_extra_config()
+        self.basic_session = self._create_basic_session(session_kwargs=session_kwargs)
+        self.role_arn = self._read_role_arn_from_extra_config()
         # If role_arn was specified then STS + assume_role
-        if role_arn is None:
-            return session
+        if self.role_arn is None:
+            return self.basic_session
 
-        return self._impersonate_to_role(role_arn=role_arn, session=session, session_kwargs=session_kwargs)
+        return self._create_session_with_assume_role(session_kwargs=session_kwargs)
 
     def _create_basic_session(self, session_kwargs: Dict[str, Any]) -> boto3.session.Session:
         aws_access_key_id, aws_secret_access_key = self._read_credentials_from_connection()
@@ -97,57 +99,61 @@ def _create_basic_session(self, session_kwargs: Dict[str, Any]) -> boto3.session
             **session_kwargs,
         )
 
-    def _impersonate_to_role(
-        self, role_arn: str, session: boto3.session.Session, session_kwargs: Dict[str, Any]
-    ) -> boto3.session.Session:
-        assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
-        assume_role_method = self.extra_config.get('assume_role_method')
+    def _create_session_with_assume_role(self, session_kwargs: Dict[str, Any]) -> boto3.session.Session:
+        assume_role_method = self.extra_config.get('assume_role_method', 'assume_role')
         self.log.info("assume_role_method=%s", assume_role_method)
-        if not assume_role_method or assume_role_method == 'assume_role':
-            sts_client = session.client("sts", config=self.config)
-            sts_response = self._assume_role(
-                sts_client=sts_client, role_arn=role_arn, assume_role_kwargs=assume_role_kwargs
-            )
-        elif assume_role_method == 'assume_role_with_saml':
-            sts_client = session.client("sts", config=self.config)
-            sts_response = self._assume_role_with_saml(
-                sts_client=sts_client, role_arn=role_arn, assume_role_kwargs=assume_role_kwargs
-            )
-        elif assume_role_method == 'assume_role_with_web_identity':
-            botocore_session = self._assume_role_with_web_identity(
-                role_arn=role_arn,
-                assume_role_kwargs=assume_role_kwargs,
-                base_session=session._session,
-            )
-            return boto3.session.Session(
-                region_name=session.region_name,
-                botocore_session=botocore_session,
-                **session_kwargs,
-            )
-        else:
+        supported_methods = ['assume_role', 'assume_role_with_saml', 'assume_role_with_web_identity']
+        if assume_role_method not in supported_methods:
             raise NotImplementedError(
                 f'assume_role_method={assume_role_method} in Connection {self.conn.conn_id} Extra.'
-                'Currently "assume_role" or "assume_role_with_saml" are supported.'
+                f'Currently {supported_methods} are supported.'
                 '(Exclude this setting will default to "assume_role").'
             )
-        # Use credentials retrieved from STS
-        credentials = sts_response["Credentials"]
-        aws_access_key_id = credentials["AccessKeyId"]
-        aws_secret_access_key = credentials["SecretAccessKey"]
-        aws_session_token = credentials["SessionToken"]
-        self.log.info(
-            "Creating session with aws_access_key_id=%s region_name=%s",
-            aws_access_key_id,
-            session.region_name,
-        )
-
-        return boto3.session.Session(
-            aws_access_key_id=aws_access_key_id,
-            aws_secret_access_key=aws_secret_access_key,
-            region_name=session.region_name,
-            aws_session_token=aws_session_token,
-            **session_kwargs,
-        )
+        if assume_role_method == 'assume_role_with_web_identity':
+            # Deferred credentials have no initial credentials
+            credential_fetcher = self._get_web_identity_credential_fetcher()
+            credentials = DeferredRefreshableCredentials(
+                method='assume-role-with-web-identity',
+                refresh_using=credential_fetcher.fetch_credentials,
+                time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
+            )
+        else:
+            # Refreshable credentials do have initial credentials
+            credentials = RefreshableCredentials.create_from_metadata(
+                metadata=self._refresh_credentials(),
+                refresh_using=self._refresh_credentials,
+                method="sts-assume-role",
+            )
+        session = botocore.session.get_session()
+        session._credentials = credentials  # pylint: disable=protected-access
+        region_name = self.basic_session.region_name
+        session.set_config_variable("region", region_name)
+        return boto3.session.Session(botocore_session=session, **session_kwargs)
+
+    def _refresh_credentials(self) -> boto3.session.Session:
+        self.log.info('Refreshing credentials')
+        assume_role_method = self.extra_config.get('assume_role_method')
+        sts_session = self.basic_session
+        if assume_role_method == 'assume_role':
+            sts_client = sts_session.client("sts", config=self.config)
+            sts_response = self._assume_role(sts_client=sts_client)
+        elif assume_role_method == 'assume_role_with_saml':
+            sts_client = sts_session.client("sts", config=self.config)
+            sts_response = self._assume_role_with_saml(sts_client=sts_client)
+        else:
+            raise NotImplementedError(f'assume_role_method={assume_role_method} not expected')
+        if not sts_response['ResponseMetadata']['HTTPStatusCode'] == 200:
+            raise Exception('An HTTP error occurred')
+        credentials = sts_response['Credentials']
+        expiry_time = credentials.get('Expiration').isoformat()
+        self.log.info(f'New credentials expiry_time:{credentials["expiry_time"]}')
+        credentials = {
+            "access_key": credentials.get("AccessKeyId"),
+            "secret_key": credentials.get("SecretAccessKey"),
+            "token": credentials.get("SessionToken"),
+            "expiry_time": expiry_time,
+        }
+        return credentials
 
     def _read_role_arn_from_extra_config(self) -> Optional[str]:
         aws_account_id = self.extra_config.get("aws_account_id")
@@ -181,24 +187,21 @@ def _read_credentials_from_connection(self) -> Tuple[Optional[str], Optional[str
             self.log.info("No credentials retrieved from Connection")
         return aws_access_key_id, aws_secret_access_key
 
-    def _assume_role(
-        self, sts_client: boto3.client, role_arn: str, assume_role_kwargs: Dict[str, Any]
-    ) -> Dict:
+    def _assume_role(self, sts_client: boto3.client) -> Dict:
+        assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
         if "external_id" in self.extra_config:  # Backwards compatibility
             assume_role_kwargs["ExternalId"] = self.extra_config.get("external_id")
         role_session_name = f"Airflow_{self.conn.conn_id}"
         self.log.info(
             "Doing sts_client.assume_role to role_arn=%s (role_session_name=%s)",
-            role_arn,
+            self.role_arn,
             role_session_name,
         )
         return sts_client.assume_role(
-            RoleArn=role_arn, RoleSessionName=role_session_name, **assume_role_kwargs
+            RoleArn=self.role_arn, RoleSessionName=role_session_name, **assume_role_kwargs
         )
 
-    def _assume_role_with_saml(
-        self, sts_client: boto3.client, role_arn: str, assume_role_kwargs: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    def _assume_role_with_saml(self, sts_client: boto3.client) -> Dict[str, Any]:
         saml_config = self.extra_config['assume_role_with_saml']
         principal_arn = saml_config['principal_arn']
 
@@ -211,9 +214,13 @@ def _assume_role_with_saml(
                 'Currently only "http_spegno_auth" is supported, and must be specified.'
             )
 
-        self.log.info("Doing sts_client.assume_role_with_saml to role_arn=%s", role_arn)
+        self.log.info("Doing sts_client.assume_role_with_saml to role_arn=%s", self.role_arn)
+        assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
         return sts_client.assume_role_with_saml(
-            RoleArn=role_arn, PrincipalArn=principal_arn, SAMLAssertion=saml_assertion, **assume_role_kwargs
+            RoleArn=self.role_arn,
+            PrincipalArn=principal_arn,
+            SAMLAssertion=saml_assertion,
+            **assume_role_kwargs,
         )
 
     def _get_idp_response(
@@ -289,7 +296,8 @@ def _fetch_saml_assertion_using_http_spegno_auth(self, saml_config: Dict[str, An
             raise ValueError('Invalid SAML Assertion')
         return saml_assertion
 
-    def _assume_role_with_web_identity(self, role_arn, assume_role_kwargs, base_session):
+    def _get_web_identity_credential_fetcher(self):
+        base_session = self.basic_session._session  # pylint: disable=protected-access
         base_session = base_session or botocore.session.get_session()
         client_creator = base_session.create_client
         federation = self.extra_config.get('assume_role_with_web_identity_federation')
@@ -299,20 +307,13 @@ def _assume_role_with_web_identity(self, role_arn, assume_role_kwargs, base_sess
             raise AirflowException(
                 f'Unsupported federation: {federation}. Currently "google" only are supported.'
             )
-        fetcher = botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher(
+        assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
+        return botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher(
             client_creator=client_creator,
             web_identity_token_loader=web_identity_token_loader,
-            role_arn=role_arn,
-            extra_args=assume_role_kwargs or {},
-        )
-        aws_creds = botocore.credentials.DeferredRefreshableCredentials(
-            method='assume-role-with-web-identity',
-            refresh_using=fetcher.fetch_credentials,
-            time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
+            role_arn=self.role_arn,
+            extra_args=assume_role_kwargs,
         )
-        botocore_session = botocore.session.Session()
-        botocore_session._credentials = aws_creds
-        return botocore_session
 
     def _get_google_identity_token_loader(self):
         from google.auth.transport import requests as requests_transport


Code review comments:
diff_hunk: @@ -97,57 +99,61 @@ def _create_basic_session(self, session_kwargs: Dict[str, Any]) -> boto3.session
             **session_kwargs,
         )
 
-    def _impersonate_to_role(
-        self, role_arn: str, session: boto3.session.Session, session_kwargs: Dict[str, Any]
-    ) -> boto3.session.Session:
-        assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
-        assume_role_method = self.extra_config.get('assume_role_method')
+    def _create_session_with_assume_role(self, session_kwargs: Dict[str, Any]) -> boto3.session.Session:
+        assume_role_method = self.extra_config.get('assume_role_method', 'assume_role')
         self.log.info("assume_role_method=%s", assume_role_method)
-        if not assume_role_method or assume_role_method == 'assume_role':
-            sts_client = session.client("sts", config=self.config)
-            sts_response = self._assume_role(
-                sts_client=sts_client, role_arn=role_arn, assume_role_kwargs=assume_role_kwargs
-            )
-        elif assume_role_method == 'assume_role_with_saml':
-            sts_client = session.client("sts", config=self.config)
-            sts_response = self._assume_role_with_saml(
-                sts_client=sts_client, role_arn=role_arn, assume_role_kwargs=assume_role_kwargs
-            )
-        elif assume_role_method == 'assume_role_with_web_identity':
-            botocore_session = self._assume_role_with_web_identity(
-                role_arn=role_arn,
-                assume_role_kwargs=assume_role_kwargs,
-                base_session=session._session,
-            )
-            return boto3.session.Session(
-                region_name=session.region_name,
-                botocore_session=botocore_session,
-                **session_kwargs,
-            )
-        else:
+        supported_methods = ['assume_role', 'assume_role_with_saml', 'assume_role_with_web_identity']
+        if assume_role_method not in supported_methods:
             raise NotImplementedError(
                 f'assume_role_method={assume_role_method} in Connection {self.conn.conn_id} Extra.'
-                'Currently "assume_role" or "assume_role_with_saml" are supported.'
+                f'Currently {supported_methods} are supported.'
                 '(Exclude this setting will default to "assume_role").'
             )
-        # Use credentials retrieved from STS
-        credentials = sts_response["Credentials"]
-        aws_access_key_id = credentials["AccessKeyId"]
-        aws_secret_access_key = credentials["SecretAccessKey"]
-        aws_session_token = credentials["SessionToken"]
-        self.log.info(
-            "Creating session with aws_access_key_id=%s region_name=%s",
-            aws_access_key_id,
-            session.region_name,
-        )
-
-        return boto3.session.Session(
-            aws_access_key_id=aws_access_key_id,
-            aws_secret_access_key=aws_secret_access_key,
-            region_name=session.region_name,
-            aws_session_token=aws_session_token,
-            **session_kwargs,
-        )
+        if assume_role_method == 'assume_role_with_web_identity':
+            # Deferred credentials have no initial credentials
+            credential_fetcher = self._get_web_identity_credential_fetcher()
+            credentials = DeferredRefreshableCredentials(
+                method='assume-role-with-web-identity',
+                refresh_using=credential_fetcher.fetch_credentials,
+                time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
+            )
+        else:
+            # Refreshable credentials do have initial credentials
+            credentials = RefreshableCredentials.create_from_metadata(
+                metadata=self._refresh_credentials(),
+                refresh_using=self._refresh_credentials,
+                method="sts-assume-role",
+            )
+        session = botocore.session.get_session()
+        session._credentials = credentials  # pylint: disable=protected-access
+        region_name = self.basic_session.region_name
+        session.set_config_variable("region", region_name)
+        return boto3.session.Session(botocore_session=session, **session_kwargs)
+
+    def _refresh_credentials(self) -> boto3.session.Session:
path: airflow/providers/amazon/aws/hooks/base_aws.py
text: @user1:
I have a feeling this function returns a dictionary not `boto3.session.Session`. Am I right?

@author:
Fixed up, thank you for the catch