nuclearpinguin commented on a change in pull request #7619: [AIRFLOW-6975] Base
AWSHook AssumeRoleWithSAML
URL: https://github.com/apache/airflow/pull/7619#discussion_r387532488
##########
File path: airflow/providers/amazon/aws/hooks/base_aws.py
##########
@@ -206,6 +212,104 @@ def _get_credentials(self, region_name):
endpoint_url,
)
+ def _assume_role(
+ self,
+ sts_client,
+ extra_config,
+ role_arn,
+ assume_role_kwargs):
+ 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,
+ extra_config,
+ role_arn,
+ assume_role_kwargs):
+
+ 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
+ import requests_gssapi
+ auth = requests_gssapi.HTTPSPNEGOAuth()
+ if 'mutual_authentication' in saml_config:
+ mutual_auth = saml_config['mutual_authentication']
+ if mutual_auth == 'REQUIRED':
+ auth = requests_gssapi.HTTPSPNEGOAuth(
+ requests_gssapi.REQUIRED)
Review comment:
Is there any reason why we cannot keep this in one line? 🤔
----------------------------------------------------------------
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
