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 2022/07/25 19:14:11 UTC
[GitHub] [airflow] ferruzzi commented on a diff in pull request #25256: Resolve and validate AWS Connection parameters in wrapper
ferruzzi commented on code in PR #25256:
URL: https://github.com/apache/airflow/pull/25256#discussion_r929204741
##########
airflow/providers/amazon/aws/hooks/base_aws.py:
##########
@@ -194,70 +171,13 @@ def _refresh_credentials(self) -> Dict[str, Any]:
}
return credentials
- def _read_role_arn_from_extra_config(self) -> Optional[str]:
- aws_account_id = self.extra_config.get("aws_account_id")
- aws_iam_role = self.extra_config.get("aws_iam_role")
- role_arn = self.extra_config.get("role_arn")
- if role_arn is None and aws_account_id is not None and aws_iam_role is not None:
- self.log.info("Constructing role_arn from aws_account_id and aws_iam_role")
- warnings.warn(
- "Constructing 'role_arn' from 'aws_account_id' and 'aws_iam_role' is deprecated and "
- "will be removed in a future releases. Please set 'role_arn' in extra config.",
- DeprecationWarning,
- stacklevel=3,
- )
- role_arn = f"arn:aws:iam::{aws_account_id}:role/{aws_iam_role}"
- self.log.debug("role_arn is %s", role_arn)
- return role_arn
-
- def _read_credentials_from_connection(self) -> Tuple[Optional[str], Optional[str]]:
- aws_access_key_id = None
- aws_secret_access_key = None
- if self.conn.login:
- aws_access_key_id = self.conn.login
- aws_secret_access_key = self.conn.password
- self.log.info("Credentials retrieved from login")
- elif "aws_access_key_id" in self.extra_config and "aws_secret_access_key" in self.extra_config:
- aws_access_key_id = self.extra_config["aws_access_key_id"]
- aws_secret_access_key = self.extra_config["aws_secret_access_key"]
- self.log.info("Credentials retrieved from extra_config")
- elif "s3_config_file" in self.extra_config:
- warnings.warn(
- "Use local credentials file is never documented and well tested. "
- "Obtain credentials by this way deprecated and will be removed in a future releases.",
- DeprecationWarning,
- stacklevel=3,
- )
- aws_access_key_id, aws_secret_access_key = _parse_s3_config(
- self.extra_config["s3_config_file"],
- self.extra_config.get("s3_config_format"),
- self.extra_config.get("profile"),
- )
- self.log.info("Credentials retrieved from extra_config['s3_config_file']")
- return aws_access_key_id, aws_secret_access_key
-
- def _strip_invalid_session_name_characters(self, role_session_name: str) -> str:
- return slugify(role_session_name, regex_pattern=r'[^\w+=,.@-]+')
-
def _assume_role(self, sts_client: boto3.client) -> Dict:
- assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {})
- if "ExternalId" not in assume_role_kwargs and "external_id" in self.extra_config:
- warnings.warn(
- "'external_id' in extra config is deprecated and will be removed in a future releases. "
- "Set 'ExternalId' in 'assume_role_kwargs' in extra config.",
- DeprecationWarning,
- stacklevel=3,
- )
- assume_role_kwargs["ExternalId"] = self.extra_config.get("external_id")
- role_session_name = self._strip_invalid_session_name_characters(f"Airflow_{self.conn.conn_id}")
- self.log.debug(
- "Doing sts_client.assume_role to role_arn=%s (role_session_name=%s)",
- self.role_arn,
- role_session_name,
- )
- return sts_client.assume_role(
- RoleArn=self.role_arn, RoleSessionName=role_session_name, **assume_role_kwargs
- )
+ kw = {
+ "RoleSessionName": self._strip_invalid_session_name_characters(f"Airflow_{self.conn.conn_id}"),
+ **self.conn.assume_role_kwargs,
+ "RoleArn": self.role_arn,
+ }
+ return sts_client.assume_role(**kw)
Review Comment:
Beautiful 😙👌
##########
airflow/providers/amazon/aws/hooks/base_aws.py:
##########
@@ -151,38 +135,31 @@ def _create_session_with_assume_role(self, session_kwargs: Dict[str, Any]) -> bo
refresh_using=self._refresh_credentials,
method="sts-assume-role",
)
+
session = botocore.session.get_session()
session._credentials = credentials
-
- if self.basic_session is None:
- raise RuntimeError("The basic session should be created here!")
-
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) -> Dict[str, Any]:
self.log.debug('Refreshing credentials')
- assume_role_method = self.extra_config.get('assume_role_method', 'assume_role')
- sts_session = self.basic_session
-
- if sts_session is None:
- raise RuntimeError(
- "Session should be initialized when refresh credentials with assume_role is used!"
- )
+ assume_role_method = self.conn.assume_role_method
+ if assume_role_method not in ('assume_role', 'assume_role_with_saml'):
+ raise NotImplementedError(f'assume_role_method={assume_role_method} not expected')
- sts_client = sts_session.client("sts", config=self.config)
+ sts_client = self.basic_session.client("sts", config=self.config)
if assume_role_method == 'assume_role':
sts_response = self._assume_role(sts_client=sts_client)
- elif assume_role_method == 'assume_role_with_saml':
- sts_response = self._assume_role_with_saml(sts_client=sts_client)
else:
- raise NotImplementedError(f'assume_role_method={assume_role_method} not expected')
Review Comment:
Is there a reason you removed this branch? It seems like a reasonable catch to me. Or maybe enforcing the values by using an enum for the possible/supported assume_role_method values? But that sounds like a lot more work.
--
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