chore: performance optimization for auth plugins by caching clients and sessions (#1084)

Author: karenc-bq

.github/workflows/main.yml

@@ -23,8 +23,8 @@ jobs:
   markdown-link-check:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@master
-      - uses: gaurav-nelson/github-action-markdown-link-check@v1
+      - uses: actions/checkout@v4
+      - uses: tcort/github-action-markdown-link-check@v1
         with:
           use-quiet-mode: 'yes'
           folder-path: 'docs'

aws_advanced_python_wrapper/aws_credentials_manager.py

@@ -0,0 +1,77 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from __future__ import annotations
+
+from threading import Lock
+from typing import TYPE_CHECKING, Any, Optional
+
+from boto3 import Session
+
+from aws_advanced_python_wrapper.utils.properties import WrapperProperties
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.hostinfo import HostInfo
+    from aws_advanced_python_wrapper.utils.properties import Properties
+
+
+class AwsCredentialsManager:
+    _lock = Lock()
+    _sessions: dict[str, Session] = {}
+    _clients: dict[str, Any] = {}
+
+    @staticmethod
+    def get_session(host_info: HostInfo, props: Properties, region: str) -> Session:
+        profile_name = WrapperProperties.AWS_PROFILE.get(props)
+        host_key = f'{host_info.as_alias()}{region}{profile_name}'
+
+        with AwsCredentialsManager._lock:
+            if host_key in AwsCredentialsManager._sessions:
+                return AwsCredentialsManager._sessions[host_key]
+
+        # Initialize session outside of lock.
+        session = Session(profile_name=profile_name, region_name=region) if profile_name else Session(region_name=region)
+
+        with AwsCredentialsManager._lock:
+            if host_key not in AwsCredentialsManager._sessions:
+                AwsCredentialsManager._sessions[host_key] = session
+            return AwsCredentialsManager._sessions[host_key]
+
+    @staticmethod
+    def get_client(service_name: str, session: Session, host: Optional[str], region: Optional[str], endpoint_url: Optional[str] = None):
+        key = f'{host}{region}{service_name}{endpoint_url}'
+
+        with AwsCredentialsManager._lock:
+            if key in AwsCredentialsManager._clients:
+                return AwsCredentialsManager._clients[key]
+
+        # Initialize client outside of lock.
+        if endpoint_url:
+            client = session.client(service_name=service_name, endpoint_url=endpoint_url)  # type: ignore[call-overload]
+        else:
+            client = session.client(service_name=service_name)  # type: ignore[call-overload]
+
+        with AwsCredentialsManager._lock:
+            if key not in AwsCredentialsManager._clients:
+                AwsCredentialsManager._clients[key] = client
+            return AwsCredentialsManager._clients[key]
+
+    @staticmethod
+    def release_resources() -> None:
+        with AwsCredentialsManager._lock:
+            for key, client in AwsCredentialsManager._clients.items():
+                client.close()
+            AwsCredentialsManager._clients.clear()
+            AwsCredentialsManager._sessions.clear()
+        return None

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -19,9 +19,10 @@
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Callable, Optional, Set, Tuple
 
-import boto3
 from botocore.exceptions import ClientError, EndpointConnectionError
 
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.utils.cache_map import CacheMap
 
 if TYPE_CHECKING:
@@ -86,7 +87,7 @@ def connect(
             props: Properties,
             is_initial_connection: bool,
             connect_func: Callable) -> Connection:
-        return self._connect(props, connect_func)
+        return self._connect(host_info, props, connect_func)
 
     def force_connect(
             self,
@@ -96,16 +97,16 @@ def force_connect(
             props: Properties,
             is_initial_connection: bool,
             force_connect_func: Callable) -> Connection:
-        return self._connect(props, force_connect_func)
+        return self._connect(host_info, props, force_connect_func)
 
-    def _connect(self, props: Properties, connect_func: Callable) -> Connection:
+    def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callable) -> Connection:
         token_expiration_sec: int = WrapperProperties.SECRETS_MANAGER_EXPIRATION.get_int(props)
         # if value is less than 0, default to one year
         if token_expiration_sec < 0:
             token_expiration_sec = AwsSecretsManagerPlugin._ONE_YEAR_IN_SECONDS
         token_expiration_ns = token_expiration_sec * 1_000_000_000
 
-        secret_fetched: bool = self._update_secret(token_expiration_ns=token_expiration_ns)
+        secret_fetched: bool = self._update_secret(host_info, props, token_expiration_ns=token_expiration_ns)
 
         try:
             self._apply_secret_to_properties(props)
@@ -116,7 +117,7 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
                 raise AwsWrapperError(
                     Messages.get_formatted("AwsSecretsManagerPlugin.ConnectException", e)) from e
 
-            secret_fetched = self._update_secret(token_expiration_ns=token_expiration_ns, force_refetch=True)
+            secret_fetched = self._update_secret(host_info, props, token_expiration_ns=token_expiration_ns, force_refetch=True)
 
             if secret_fetched:
                 try:
@@ -128,7 +129,7 @@ def _connect(self, props: Properties, connect_func: Callable) -> Connection:
                                                unhandled_error)) from unhandled_error
             raise AwsWrapperError(Messages.get_formatted("AwsSecretsManagerPlugin.FailedLogin", e)) from e
 
-    def _update_secret(self, token_expiration_ns: int, force_refetch: bool = False) -> bool:
+    def _update_secret(self, host_info: HostInfo, props: Properties, token_expiration_ns: int, force_refetch: bool = False) -> bool:
         """
         Called to update credentials from the cache, or from the AWS Secrets Manager service.
         :param token_expiration_ns: Expiration time in nanoseconds for secret stored in cache.
@@ -146,7 +147,7 @@ def _update_secret(self, token_expiration_ns: int, force_refetch: bool = False)
             endpoint = self._secret_key[2]
             if not self._secret or force_refetch:
                 try:
-                    self._secret = self._fetch_latest_credentials()
+                    self._secret = self._fetch_latest_credentials(host_info, props)
                     if self._secret:
                         AwsSecretsManagerPlugin._secrets_cache.put(self._secret_key, self._secret, token_expiration_ns)
                         fetched = True
@@ -177,26 +178,19 @@ def _update_secret(self, token_expiration_ns: int, force_refetch: bool = False)
             if context is not None:
                 context.close_context()
 
-    def _fetch_latest_credentials(self):
+    def _fetch_latest_credentials(self, host_info: HostInfo, props: Properties):
         """
         Fetches the current credentials from AWS Secrets Manager service.
 
         :return: a Secret object containing the credentials fetched from the AWS Secrets Manager service.
         """
-        session = self._session if self._session else boto3.Session()
-
-        client = session.client(
-            'secretsmanager',
-            region_name=self._secret_key[1],
-            endpoint_url=self._secret_key[2],
-        )
+        session = AwsCredentialsManager.get_session(host_info, props, self._secret_key[1])
+        client = AwsCredentialsManager.get_client("secretsmanager", session, host_info.host, self._secret_key[1], self._secret_key[2])
 
         secret = client.get_secret_value(
             SecretId=self._secret_key[0],
         )
 
-        client.close()
-
         return loads(secret.get("SecretString"), object_hook=lambda d: SimpleNamespace(**d))
 
     def _apply_secret_to_properties(self, properties: Properties):

aws_advanced_python_wrapper/cleanup.py

@@ -14,6 +14,8 @@
 
 from aws_advanced_python_wrapper.aurora_connection_tracker_plugin import \
     OpenedConnectionTracker
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.host_monitoring_plugin import \
     MonitoringThreadContainer
 from aws_advanced_python_wrapper.thread_pool_container import \
@@ -26,5 +28,6 @@ def release_resources() -> None:
     """Release all global resources used by the wrapper."""
     MonitoringThreadContainer.clean_up()
     ThreadPoolContainer.release_resources()
+    AwsCredentialsManager.release_resources()
     OpenedConnectionTracker.release_resources()
     SlidingExpirationCacheContainer.release_resources()

aws_advanced_python_wrapper/credentials_provider_factory.py

@@ -16,32 +16,29 @@
 
 from typing import TYPE_CHECKING, Dict, Optional, Protocol
 
-import boto3
-
 if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.utils.properties import Properties
 
 from abc import abstractmethod
 
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.utils.properties import WrapperProperties
 
 
 class CredentialsProviderFactory(Protocol):
     @abstractmethod
-    def get_aws_credentials(self, region: str, props: Properties) -> Optional[Dict[str, str]]:
+    def get_aws_credentials(self, region: str, props: Properties, host_info: HostInfo) -> Optional[Dict[str, str]]:
         ...
 
 
 class SamlCredentialsProviderFactory(CredentialsProviderFactory):
 
-    def get_aws_credentials(self, region: str, props: Properties) -> Optional[Dict[str, str]]:
+    def get_aws_credentials(self, region: str, props: Properties, host_info: HostInfo) -> Optional[Dict[str, str]]:
         saml_assertion: str = self.get_saml_assertion(props)
-        session = boto3.Session()
-
-        sts_client = session.client(
-            'sts',
-            region_name=region
-        )
+        session = AwsCredentialsManager.get_session(host_info, props, region)
+        sts_client = AwsCredentialsManager.get_client("sts", session, host_info.host, region)
 
         response: Dict[str, Dict[str, str]] = sts_client.assume_role_with_saml(
             RoleArn=WrapperProperties.IAM_ROLE_ARN.get(props),

aws_advanced_python_wrapper/federated_plugin.py

@@ -14,19 +14,21 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
 from html import unescape
 from re import DOTALL, findall, search
 from typing import TYPE_CHECKING, List
 from urllib.parse import urlencode
 
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.credentials_provider_factory import (
     CredentialsProviderFactory, SamlCredentialsProviderFactory)
 from aws_advanced_python_wrapper.utils.iam_utils import IamAuthUtils, TokenInfo
 from aws_advanced_python_wrapper.utils.region_utils import RegionUtils
 from aws_advanced_python_wrapper.utils.saml_utils import SamlUtils
 
 if TYPE_CHECKING:
-    from boto3 import Session
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
@@ -55,10 +57,9 @@ class FederatedAuthPlugin(Plugin):
     _rds_utils: RdsUtils = RdsUtils()
     _token_cache: Dict[str, TokenInfo] = {}
 
-    def __init__(self, plugin_service: PluginService, credentials_provider_factory: CredentialsProviderFactory, session: Optional[Session] = None):
+    def __init__(self, plugin_service: PluginService, credentials_provider_factory: CredentialsProviderFactory):
         self._plugin_service = plugin_service
         self._credentials_provider_factory = credentials_provider_factory
-        self._session = session
 
         self._region_utils = RegionUtils()
         telemetry_factory = self._plugin_service.get_telemetry_factory()
@@ -100,11 +101,13 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
 
         token_info: Optional[TokenInfo] = FederatedAuthPlugin._token_cache.get(cache_key)
 
+        token_host_info = deepcopy(host_info)
+        token_host_info.host = host
         if token_info is not None and not token_info.is_expired():
             logger.debug("FederatedAuthPlugin.UseCachedToken", token_info.token)
             self._plugin_service.driver_dialect.set_password(props, token_info.token)
         else:
-            self._update_authentication_token(host_info, props, user, region, cache_key)
+            self._update_authentication_token(token_host_info, props, user, region, cache_key)
 
         WrapperProperties.USER.set(props, WrapperProperties.DB_USER.get(props))
 
@@ -114,7 +117,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             if token_info is None or token_info.is_expired() or not self._plugin_service.is_login_exception(e):
                 raise e
 
-            self._update_authentication_token(host_info, props, user, region, cache_key)
+            self._update_authentication_token(token_host_info, props, user, region, cache_key)
 
             try:
                 return connect_func()
@@ -142,18 +145,19 @@ def _update_authentication_token(self,
         token_expiration_sec: int = WrapperProperties.IAM_TOKEN_EXPIRATION.get_int(props)
         token_expiry: datetime = datetime.now() + timedelta(seconds=token_expiration_sec)
         port: int = IamAuthUtils.get_port(props, host_info, self._plugin_service.database_dialect.default_port)
-        credentials: Optional[Dict[str, str]] = self._credentials_provider_factory.get_aws_credentials(region, props)
+        credentials: Optional[Dict[str, str]] = self._credentials_provider_factory.get_aws_credentials(region, props, host_info)
 
         if self._fetch_token_counter is not None:
             self._fetch_token_counter.inc()
+        session = AwsCredentialsManager.get_session(host_info, props, region)
         token: str = IamAuthUtils.generate_authentication_token(
             self._plugin_service,
             user,
             host_info.host,
             port,
             region,
-            credentials,
-            self._session)
+            session,
+            credentials)
         WrapperProperties.PASSWORD.set(props, token)
         FederatedAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)
 

aws_advanced_python_wrapper/iam_plugin.py

@@ -14,20 +14,22 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
 from typing import TYPE_CHECKING
 
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.utils.iam_utils import IamAuthUtils, TokenInfo
 from aws_advanced_python_wrapper.utils.region_utils import RegionUtils
 
 if TYPE_CHECKING:
-    from boto3 import Session
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
     from aws_advanced_python_wrapper.hostinfo import HostInfo
     from aws_advanced_python_wrapper.pep249 import Connection
     from aws_advanced_python_wrapper.plugin_service import PluginService
 
 from datetime import datetime, timedelta
-from typing import Callable, Dict, Optional, Set
+from typing import Callable, Dict, Set
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.pep249_methods import DbApiMethod
@@ -49,9 +51,8 @@ class IamAuthPlugin(Plugin):
     _rds_utils: RdsUtils = RdsUtils()
     _token_cache: Dict[str, TokenInfo] = {}
 
-    def __init__(self, plugin_service: PluginService, session: Optional[Session] = None):
+    def __init__(self, plugin_service: PluginService):
         self._plugin_service = plugin_service
-        self._session = session
 
         self._region_utils = RegionUtils()
         telemetry_factory = self._plugin_service.get_telemetry_factory()
@@ -104,7 +105,17 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
             if self._fetch_token_counter is not None:
                 self._fetch_token_counter.inc()
-            token: str = IamAuthUtils.generate_authentication_token(self._plugin_service, user, host, port, region, client_session=self._session)
+
+            session_host_info = deepcopy(host_info)
+            session_host_info.host = host
+            session = AwsCredentialsManager.get_session(host_info, props, region)
+            token: str = IamAuthUtils.generate_authentication_token(
+                self._plugin_service,
+                user,
+                host,
+                port,
+                region,
+                session)
             self._plugin_service.driver_dialect.set_password(props, token)
             IamAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)
 
@@ -123,7 +134,11 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
             if self._fetch_token_counter is not None:
                 self._fetch_token_counter.inc()
-            token = IamAuthUtils.generate_authentication_token(self._plugin_service, user, host, port, region, client_session=self._session)
+
+            session_host_info = deepcopy(host_info)
+            session_host_info.host = host
+            session = AwsCredentialsManager.get_session(session_host_info, props, region)
+            token = IamAuthUtils.generate_authentication_token(self._plugin_service, user, host, port, region, session)
             self._plugin_service.driver_dialect.set_password(props, token)
             IamAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)