chore: add type verification for custom handlers

Author: karenc-bq

aws_advanced_python_wrapper/aws_credentials_manager.py

@@ -18,7 +18,10 @@
 from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from boto3 import Session
+
+from aws_advanced_python_wrapper.utils.messages import Messages
 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
@@ -32,6 +35,8 @@ class AwsCredentialsManager:
 
     @staticmethod
     def set_custom_handler(custom_handler: Callable[[HostInfo, Properties], Optional[Session]]) -> None:
+        if not callable(custom_handler):
+            raise TypeError("custom_handler must be callable")
         with AwsCredentialsManager._lock:
             AwsCredentialsManager._handler = custom_handler
 
@@ -52,27 +57,30 @@ def get_session(host_info: HostInfo, props: Properties, region: str) -> Session:
 
         # Initialize session outside of lock.
         session = handler(host_info, props) if handler else None
-        
+
+        if session is not None and not isinstance(session, Session):
+            raise TypeError(Messages.get_formatted("AwsCredentialsManager.InvalidHandler", type(session).__name__))
+
         if session is None:
             profile_name = WrapperProperties.AWS_PROFILE.get(props)
             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: str, region: str):
+    def get_client(service_name: str, session: Session, host: Optional[str], region: Optional[str]):
         key = f'{host}{region}{service_name}'
-        
+
         with AwsCredentialsManager._lock:
             if key in AwsCredentialsManager._clients:
                 return AwsCredentialsManager._clients[key]
 
         # Initialize client outside of lock.
-        client = session.client(service_name)
-        
+        client = session.client(service_name)  # type: ignore[call-overload]
+
         with AwsCredentialsManager._lock:
             if key not in AwsCredentialsManager._clients:
                 AwsCredentialsManager._clients[key] = client

aws_advanced_python_wrapper/cleanup.py

@@ -12,7 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+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 \

aws_advanced_python_wrapper/credentials_provider_factory.py

@@ -16,14 +16,14 @@
 
 from typing import TYPE_CHECKING, Dict, Optional, Protocol
 
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
-from aws_advanced_python_wrapper.hostinfo import HostInfo
-
 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
 
 

aws_advanced_python_wrapper/federated_plugin.py

@@ -19,9 +19,8 @@
 from typing import TYPE_CHECKING, List
 from urllib.parse import urlencode
 
-import boto3
-
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+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
@@ -166,6 +165,7 @@ def release_resources() -> None:
         AwsCredentialsManager.release_resources()
         return None
 
+
 class FederatedAuthPluginFactory(PluginFactory):
     @staticmethod
     def get_instance(plugin_service: PluginService, props: Properties) -> Plugin:

aws_advanced_python_wrapper/iam_plugin.py

@@ -16,9 +16,8 @@
 
 from typing import TYPE_CHECKING
 
-import boto3
-
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+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
 

aws_advanced_python_wrapper/okta_plugin.py

@@ -19,7 +19,8 @@
 from re import search
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Set
 
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+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

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -26,6 +26,8 @@ AdfsCredentialsProviderFactory.SignOnPagePostActionUrl=[AdfsCredentialsProviderF
 AdfsCredentialsProviderFactory.SignOnPagePostActionRequestFailed=[AdfsCredentialsProviderFactory] ADFS SignOn Page POST action failed with HTTP status '{}', reason phrase '{}', and response '{}'
 AdfsCredentialsProviderFactory.SignOnPageUrl=[AdfsCredentialsProviderFactory] ADFS SignOn URL: '{}'
 
+AwsCredentialsManager.InvalidHandler=[AwsCredentialsManager] Custom credentials provider set via AwsCredentialsManager.set_custom_handler must return a boto3.Session or None, got '{}'.
+
 AwsSdk.UnsupportedRegion=[AwsSdk] Unsupported AWS region {}. For supported regions please read https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Concepts.RegionsAndAvailabilityZones.html
 
 AwsSecretsManagerPlugin.ConnectException=[AwsSecretsManagerPlugin] Error occurred while opening a connection: {}

aws_advanced_python_wrapper/utils/iam_utils.py

@@ -17,9 +17,8 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, Dict, Optional
 
-import boto3
-
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.errors import AwsWrapperError
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.messages import Messages

docs/examples/PGAwsCredentialsManager.py

@@ -18,7 +18,8 @@
 import psycopg
 
 from aws_advanced_python_wrapper import AwsWrapperConnection, release_resources
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 
 
 def custom_credentials_handler(host_info, props):

tests/unit/test_aws_credentials_manager.py

@@ -19,7 +19,8 @@
 import pytest
 from boto3 import Session
 
-from aws_advanced_python_wrapper.aws_credentials_manager import AwsCredentialsManager
+from aws_advanced_python_wrapper.aws_credentials_manager import \
+    AwsCredentialsManager
 from aws_advanced_python_wrapper.hostinfo import HostInfo
 from aws_advanced_python_wrapper.utils.atomic import AtomicInt
 from aws_advanced_python_wrapper.utils.properties import Properties
@@ -104,7 +105,7 @@ def test_get_session_different_regions(self, host_info, props, mocker):
         mock_session1.region_name = "us-east-1"
         mock_session2 = mocker.MagicMock(spec=Session)
         mock_session2.region_name = "us-west-2"
-        
+
         mock_session_class = mocker.patch('aws_advanced_python_wrapper.aws_credentials_manager.Session')
         mock_session_class.side_effect = [mock_session1, mock_session2]
 
@@ -121,7 +122,7 @@ def test_get_session_different_hosts(self, props, region, mocker):
 
         mock_session1 = mocker.MagicMock(spec=Session)
         mock_session2 = mocker.MagicMock(spec=Session)
-        
+
         mock_session_class = mocker.patch('aws_advanced_python_wrapper.aws_credentials_manager.Session')
         mock_session_class.side_effect = [mock_session1, mock_session2]
 
@@ -144,7 +145,7 @@ def test_get_session_with_custom_handler(self, mock_session, host_info, props, r
     def test_reset_custom_handler(self, host_info, props, region, mocker):
         custom_session = mocker.MagicMock(spec=Session)
         custom_handler = mocker.MagicMock(return_value=custom_session)
-        
+
         mock_default_session = mocker.MagicMock(spec=Session)
         mock_default_session.region_name = region
         mocker.patch('aws_advanced_python_wrapper.aws_credentials_manager.Session', return_value=mock_default_session)
@@ -200,10 +201,10 @@ def test_release_resources_clears_caches(self, host_info, props, region, mocker)
         mock_session1 = mocker.MagicMock(spec=Session)
         mock_session2 = mocker.MagicMock(spec=Session)
         mock_client = mocker.MagicMock()
-        
+
         mock_session_class = mocker.patch('aws_advanced_python_wrapper.aws_credentials_manager.Session')
         mock_session_class.side_effect = [mock_session1, mock_session2]
-        
+
         mock_session1.client.return_value = mock_client
 
         session = AwsCredentialsManager.get_session(host_info, props, region)
@@ -218,7 +219,7 @@ def test_release_resources_clears_caches(self, host_info, props, region, mocker)
     def test_concurrent_get_session_same_host(self, mock_session, host_info, props, region, counter, concurrent_counter, num_threads):
         barrier = Barrier(num_threads)
         sessions = []
-        
+
         def get_session_thread():
             barrier.wait()
             val = counter.get_and_increment()
@@ -275,13 +276,13 @@ def get_session_thread(thread_id):
     def test_concurrent_get_session_different_regions(self, num_threads, host_info, props, counter, concurrent_counter, regions, mocker):
         barrier = Barrier(num_threads)
         results = []
-        
+
         # One session per region
         mock_sessions = {region: mocker.MagicMock(spec=Session) for region in regions}
-        
+
         def session_factory(region_name=None, **kwargs):
             return mock_sessions[region_name]
-        
+
         mocker.patch('aws_advanced_python_wrapper.aws_credentials_manager.Session', side_effect=session_factory)
 
         def get_session_thread(thread_id):