Merge pull request #43 from passren/0.5.4

Fix bug in _assume_role_with_web_identity
passren · Sep 8, 2023 · 664d05c · 664d05c
2 parents f70ae7f + 03df313
commit 664d05c
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 10 deletions.
diff --git a/pydynamodb/connection.py b/pydynamodb/connection.py
@@ -204,6 +204,10 @@ def _assume_role_with_saml(
         client = session.client(
             "sts", region_name=region_name, config=self.config, **self._client_kwargs
         )
+
+        # Resolve the issue which plus sign will be replaced to space by url.parse_qsl
+        saml_assertion = saml_assertion.replace(" ", "+")
+
         request = {
             "RoleArn": role_arn,
             "PrincipalArn": principal_arn,
@@ -236,9 +240,15 @@ def _assume_role_with_web_identity(
             "RoleSessionName": role_session_name,
             "DurationSeconds": duration_seconds,
             "WebIdentityToken": web_identity_token,
-            "ProviderId": provider_id if provider_id else "",
         }
 
+        if provider_id:
+            request.update(
+                {
+                    "ProviderId": provider_id,
+                }
+            )
+
         response = client.assume_role_with_web_identity(**request)
         creds: Dict[str, Any] = response["Credentials"]
         return creds

diff --git a/tests/test_connection.py b/tests/test_connection.py
@@ -239,6 +239,15 @@ def test_conn_with_credentials(self):
         assert conn._session_kwargs["aws_session_token"] is not None
 
         # assume_role_with_web_identity
+        conn = connect(
+            region_name=TestConnectionWithSTS.TEST_REGION_NAME,
+            role_arn=TestConnectionWithSTS.TEST_ROLE_ARN,
+            web_identity_token=TestConnectionWithSTS.TEST_WEB_IDENTITY_TOKEN,
+        )
+        assert conn._session_kwargs["aws_access_key_id"] is not None
+        assert conn._session_kwargs["aws_secret_access_key"] is not None
+        assert conn._session_kwargs["aws_session_token"] is not None
+
         conn = connect(
             region_name=TestConnectionWithSTS.TEST_REGION_NAME,
             role_arn=TestConnectionWithSTS.TEST_ROLE_ARN,
@@ -278,16 +287,29 @@ def test_conn_with_sqlalchemy(self):
         with contextlib.closing(engine.connect()) as conn:
             assert conn is not None
 
-        # TODO: Some particular characters in connection string will be converted by sqlalchemy.
-        # Need to figure out a solution for this.
+        conn_str = (
+            CONN_STR_PREFIX
+            + CONN_STR_URL
+            + CONN_STR_PARAM
+            + "&principal_arn="
+            + TestConnectionWithSTS.TEST_PRINCIPAL_ARN
+            + "&saml_assertion="
+            + TestConnectionWithSTS.TEST_SAML_ASSERTION
+        )
+        engine = sqlalchemy.engine.create_engine(conn_str, echo=True)
+        with contextlib.closing(engine.connect()) as conn:
+            assert conn is not None
 
-        # conn_str = (CONN_STR_PREFIX + CONN_STR_URL + CONN_STR_PARAM
-        #             + "&principal_arn=" + TestConnectionWithSTS.TEST_PRINCIPAL_ARN
-        #             + "&saml_assertion=" + TestConnectionWithSTS.TEST_SAML_ASSERTION
-        # )
-        # engine = sqlalchemy.engine.create_engine(conn_str, echo=True)
-        # with contextlib.closing(engine.connect()) as conn:
-        #     assert conn is not None
+        conn_str = (
+            CONN_STR_PREFIX
+            + CONN_STR_URL
+            + CONN_STR_PARAM
+            + "&web_identity_token="
+            + TestConnectionWithSTS.TEST_WEB_IDENTITY_TOKEN
+        )
+        engine = sqlalchemy.engine.create_engine(conn_str, echo=True)
+        with contextlib.closing(engine.connect()) as conn:
+            assert conn is not None
 
         conn_str = (
             CONN_STR_PREFIX