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/23 19:16:00 UTC

[GitHub] [airflow] Taragolis opened a new pull request, #25256: Resolve and validate AWS Connection parameters in wrapper

Taragolis opened a new pull request, #25256:
URL: https://github.com/apache/airflow/pull/25256

   Separate class which help resolve and validate AWS connection parameters.
   Currently, might for historical reason, connections properties validate/read/resolve/construct in `BaseSessionFactory` and `AwsGenericHook` in different methods.
   
   This changes would help to make changes/deprecate in AWS Connection settings easier than now - need to make changes in one place, rather tried to found in multiple places.
   
   Also it should make easier tests Connections configuration relates.
   
   WDYT? @o-nikolas @ferruzzi 


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


[GitHub] [airflow] potiuk merged pull request #25256: Resolve and validate AWS Connection parameters in wrapper

Posted by GitBox <gi...@apache.org>.
potiuk merged PR #25256:
URL: https://github.com/apache/airflow/pull/25256


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


[GitHub] [airflow] ferruzzi commented on a diff in pull request #25256: Resolve and validate AWS Connection parameters in wrapper

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #25256:
URL: https://github.com/apache/airflow/pull/25256#discussion_r929241072


##########
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:
   Gotcha, sorry I missed that.  :+1:   Feel free to resolve this thread



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


[GitHub] [airflow] ferruzzi commented on a diff in pull request #25256: Resolve and validate AWS Connection parameters in wrapper

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on code in PR #25256:
URL: https://github.com/apache/airflow/pull/25256#discussion_r929241072


##########
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:
   Gotcha, sorry I missed that.  :+1: 
   



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


[GitHub] [airflow] ferruzzi commented on a diff in pull request #25256: Resolve and validate AWS Connection parameters in wrapper

Posted by GitBox <gi...@apache.org>.
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


[GitHub] [airflow] Taragolis commented on a diff in pull request #25256: Resolve and validate AWS Connection parameters in wrapper

Posted by GitBox <gi...@apache.org>.
Taragolis commented on code in PR #25256:
URL: https://github.com/apache/airflow/pull/25256#discussion_r929238315


##########
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:
   I've just move this check before obtain the session. 
   
   https://github.com/apache/airflow/blob/79279cf38337c4e3250bbb0a8f6bd31328b5f43d/airflow/providers/amazon/aws/hooks/base_aws.py#L149-L157
   
   Seems like it not possible that this exception raised. Only if user tried to use this private method inside their own SessionFactory.
   



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