baolsen commented on a change in pull request #7619: [AIRFLOW-6975] Base
AWSHook AssumeRoleWithSAML
URL: https://github.com/apache/airflow/pull/7619#discussion_r391410752
##########
File path: airflow/providers/amazon/aws/hooks/base_aws.py
##########
@@ -206,6 +214,106 @@ def _get_credentials(self, region_name):
endpoint_url,
)
+ def _assume_role(
+ self,
+ sts_client: boto3.client,
+ extra_config: dict,
+ role_arn: str,
+ assume_role_kwargs: dict):
+ if "external_id" in extra_config: # Backwards compatibility
+ assume_role_kwargs["ExternalId"] = extra_config.get(
+ "external_id"
+ )
+ role_session_name = "Airflow_" + self.aws_conn_id
+ self.log.info(
+ "Doing sts_client.assume_role to role_arn=%s
(role_session_name=%s)",
+ role_arn,
+ role_session_name,
+ )
+ return sts_client.assume_role(
+ RoleArn=role_arn,
+ RoleSessionName=role_session_name,
+ **assume_role_kwargs
+ )
+
+ def _assume_role_with_saml(
+ self,
+ sts_client: boto3.client,
+ extra_config: dict,
+ role_arn: str,
+ assume_role_kwargs: dict):
+
+ saml_config = extra_config['assume_role_with_saml']
+ principal_arn = saml_config['principal_arn']
+
+ idp_url = saml_config["idp_url"]
+ self.log.info("idp_url= %s", idp_url)
+
+ idp_request_kwargs = saml_config["idp_request_kwargs"]
+
+ idp_auth_method = saml_config['idp_auth_method']
+ if idp_auth_method == 'http_spegno_auth':
+ # requests_gssapi will need paramiko > 2.6 since you'll need
+ # 'gssapi' not 'python-gssapi' from PyPi.
+ # https://github.com/paramiko/paramiko/pull/1311
Review comment:
Added to the PR description; will quickly add it to the how-to documentation
as well :)
----------------------------------------------------------------
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]
With regards,
Apache Git Services
