[GitHub] [airflow] feluelle commented on a change in pull request #9878: Refactor AwsBaseHook._get_credentials

[GitBox <gi...@apache.org>]

feluelle commented on a change in pull request #9878:
URL: https://github.com/apache/airflow/pull/9878#discussion_r457228408



##########
File path: airflow/providers/amazon/aws/hooks/base_aws.py
##########
@@ -26,14 +26,215 @@
 
 import configparser
 import logging
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import boto3
 from botocore.config import Config
 from cached_property import cached_property
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base_hook import BaseHook
+from airflow.models.connection import Connection
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+
+class _SessionFactory(LoggingMixin):
+    def __init__(self, conn: Connection, region_name: str, config: Config):
+        super().__init__()
+        self.conn = conn
+        self.region_name = region_name
+        self.config = config
+        self.extra_config = self.conn.extra_dejson
+
+    def create_session(self) -> boto3.session.Session:
+        """Create AWS session."""
+        session_kwargs = {}
+        if "session_kwargs" in self.extra_config:
+            self.log.info(
+                "Retrieving session_kwargs from Connection.extra_config['session_kwargs']: %s",
+                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()
+        # If role_arn was specified then STS + assume_role
+        if role_arn is None:
+            return session
+
+        return self._impersonate_to_role(role_arn=role_arn, session=session, 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()
+        aws_session_token = self.extra_config.get("aws_session_token")
+        region_name = self.region_name
+        if self.region_name is None and 'region_name' in self.extra_config:
+            self.log.info("Retrieving region_name from Connection.extra_config['region_name']")
+            region_name = self.extra_config.get("region_name")

Review comment:
       having `and 'region_name' in self.extra_config:` you can be sure that it exists.
   ```suggestion
               region_name = self.extra_config["region_name"]
   ```




----------------------------------------------------------------
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.

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