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:
[email protected]
