fix: credential caching improvements (#431)

* fix: credential caching improvements

Signed-off-by: Stephen Crowe <[email protected]>

src/deadline/client/api/_loginout.py

@@ -13,12 +13,12 @@
 
 
 from ._session import (
-    invalidate_boto3_session_cache,
     get_credentials_source,
     check_authentication_status,
     AwsCredentialsSource,
     AwsAuthenticationStatus,
 )
+from . import _session
 from ..config import get_setting
 from ..exceptions import DeadlineOperationError
 import time
@@ -143,7 +143,7 @@ def logout(config: Optional[ConfigParser] = None) -> str:
             )
 
         # Force a refresh of the cached boto3 Session
-        invalidate_boto3_session_cache()
+        _session.invalidate_boto3_session_cache()
         return output.decode("utf8")
     raise UnsupportedProfileTypeForLoginLogout(
         "Logging out is only supported for AWS Profiles created by Deadline Cloud monitor."

src/deadline/client/api/_session.py

@@ -5,6 +5,7 @@
 of the Deadline-configured IAM credentials.
 """
 from __future__ import annotations
+from functools import lru_cache
 import logging
 from configparser import ConfigParser
 from contextlib import contextmanager
@@ -23,12 +24,6 @@
 from ..config import get_setting
 from ..exceptions import DeadlineOperationError
 
-__cached_boto3_session = None
-__cached_boto3_session_profile_name = None
-__cached_boto3_queue_session = None
-__cached_farm_id_for_queue_session = None
-__cached_queue_id_for_queue_session = None
-
 
 class AwsCredentialsSource(Enum):
     NOT_VALID = 0
@@ -57,8 +52,6 @@ def get_boto3_session(
         force_refresh (bool, optional): If set to True, forces a cache refresh.
         config (ConfigParser, optional): If provided, the AWS Deadline Cloud config to use.
     """
-    global __cached_boto3_session
-    global __cached_boto3_session_profile_name
 
     profile_name: Optional[str] = get_setting("defaults.aws_profile_name", config)
 
@@ -67,36 +60,40 @@ def get_boto3_session(
     if profile_name in ("(default)", "default", ""):
         profile_name = None
 
-    # If a config was provided, don't use the Session caching mechanism.
-    if config:
-        return boto3.Session(profile_name=profile_name)
-
     if force_refresh:
         invalidate_boto3_session_cache()
 
-    # If this is the first call or the profile name has changed, make a fresh Session
-    if not __cached_boto3_session or __cached_boto3_session_profile_name != profile_name:
-        __cached_boto3_session = boto3.Session(profile_name=profile_name)
-        __cached_boto3_session_profile_name = profile_name
+    return _get_boto3_session_for_profile(profile_name)
 
-    return __cached_boto3_session
 
+@lru_cache
+def _get_boto3_session_for_profile(profile_name: str):
+    session = boto3.Session(profile_name=profile_name)
 
-def invalidate_boto3_session_cache() -> None:
-    """
-    Invalidates the cached boto3 session and boto3 queue session.
-    """
-    global __cached_boto3_session
-    global __cached_boto3_session_profile_name
-    global __cached_boto3_queue_session
-    global __cached_farm_id_for_queue_session
-    global __cached_queue_id_for_queue_session
+    # By default, DCM returns creds that expire after 15 minutes, and boto3's RefreshableCredentials
+    # class refreshes creds that are within 15 minutes of expiring, so credentials would never be reused.
+    # Also DCM credentials currently take several seconds to refresh. Lower the refresh timeouts so creds
+    # are reused between API calls to save time.
+    # See https://github.com/boto/botocore/blob/develop/botocore/credentials.py#L342-L362
 
-    __cached_boto3_session = None
-    __cached_boto3_session_profile_name = None
-    __cached_boto3_queue_session = None
-    __cached_farm_id_for_queue_session = None
-    __cached_queue_id_for_queue_session = None
+    try:
+        credentials = session.get_credentials()
+        if (
+            isinstance(credentials, RefreshableCredentials)
+            and credentials.method == "custom-process"
+        ):
+            credentials._advisory_refresh_timeout = 5 * 60  # 5 minutes
+            credentials._mandatory_refresh_timeout = 2.5 * 60  # 2.5 minutes
+    except:  # noqa: E722
+        # Attempt to patch the timeouts but ignore any errors. These patched proeprties are internal and could change
+        # without notice. Creds are functional without patching timeouts.
+        pass
+    return session
+
+
+def invalidate_boto3_session_cache() -> None:
+    _get_boto3_session_for_profile.cache_clear()
+    _get_queue_user_boto3_session.cache_clear()
 
 
 def get_boto3_client(service_name: str, config: Optional[ConfigParser] = None) -> BaseClient:
@@ -182,39 +179,19 @@ def get_queue_user_boto3_session(
         force_refresh (bool, optional): If True, forces a cache refresh.
     """
 
-    global __cached_boto3_queue_session
-    global __cached_farm_id_for_queue_session
-    global __cached_queue_id_for_queue_session
-
     base_session = get_boto3_session(config=config, force_refresh=force_refresh)
 
     if farm_id is None:
         farm_id = get_setting("defaults.farm_id")
     if queue_id is None:
         queue_id = get_setting("defaults.queue_id")
 
-    # If a config was provided, don't use the Session caching mechanism.
-    if config:
-        return _get_queue_user_boto3_session(
-            deadline, base_session, farm_id, queue_id, queue_display_name
-        )
-
-    # If this is the first call or the farm ID/queue ID has changed, make a fresh Session and cache it
-    if (
-        not __cached_boto3_queue_session
-        or __cached_farm_id_for_queue_session != farm_id
-        or __cached_queue_id_for_queue_session != queue_id
-    ):
-        __cached_boto3_queue_session = _get_queue_user_boto3_session(
-            deadline, base_session, farm_id, queue_id, queue_display_name
-        )
-
-        __cached_farm_id_for_queue_session = farm_id
-        __cached_queue_id_for_queue_session = queue_id
-
-    return __cached_boto3_queue_session
+    return _get_queue_user_boto3_session(
+        deadline, base_session, farm_id, queue_id, queue_display_name
+    )
 
 
+@lru_cache
 def _get_queue_user_boto3_session(
     deadline: BaseClient,
     base_session: boto3.Session,

test/unit/conftest.py

@@ -14,8 +14,15 @@
 def fresh_deadline_config():
     """
     Fixture to start with a blank AWS Deadline Cloud config file.
+
     """
 
+    # Clear the session cache. Importing the cache invalidator at runtime is necessary
+    # to make sure the import order doesn't bypass moto mocking in other areas.
+    from deadline.client.api._session import invalidate_boto3_session_cache
+
+    invalidate_boto3_session_cache()
+
     try:
         # Create an empty temp file to set as the AWS Deadline Cloud config
         temp_dir = tempfile.TemporaryDirectory()

test/unit/deadline_client/api/test_api_session.py

@@ -4,6 +4,7 @@
 tests the deadline.client.api functions relating to boto3.Client
 """
 
+from typing import Optional
 from unittest.mock import call, patch, MagicMock, ANY
 
 import boto3  # type: ignore[import]
@@ -14,12 +15,13 @@ def test_get_boto3_session(fresh_deadline_config):
     """Confirm that api.get_boto3_session gets a session for the configured profile"""
     config.set_setting("defaults.aws_profile_name", "SomeRandomProfileName")
 
-    with patch.object(boto3, "Session", return_value="SomeReturnValue") as boto3_session:
+    mock_session = MagicMock()
+    with patch.object(boto3, "Session", return_value=mock_session) as boto3_session:
         # Testing this function
         result = api.get_boto3_session()
 
         # Confirm it returned the mocked value, and was called with the correct args
-        assert result == "SomeReturnValue"
+        assert result == mock_session
         boto3_session.assert_called_once_with(profile_name="SomeRandomProfileName")
 
 
@@ -28,31 +30,30 @@ def test_get_boto3_session_caching_behavior(fresh_deadline_config):
     Confirm that api.get_boto3_session caches the session, and refreshes if
     the configured profile name changes
     """
+
     # mock boto3.Session to return a fresh object based on the input profile name
-    with patch.object(
-        boto3, "Session", side_effect=lambda profile_name: f"session for {profile_name}"
-    ) as boto3_session:
+    def mock_create_session(profile_name: Optional[str]):
+        session = MagicMock()
+        session._profile_name = profile_name
+        return session
+
+    with patch.object(boto3, "Session", side_effect=mock_create_session) as boto3_session:
         # This is a session with the default profile name
         session0 = api.get_boto3_session()
 
-        assert session0 == "session for None"  # type: ignore
+        assert session0._profile_name is None
 
         # This should return the cached object, and not call boto3.Session
         session1 = api.get_boto3_session()
 
         assert session1 is session0
 
-        # Overriding the config object means to not use the cached object
-        session1_override_config = api.get_boto3_session(config=config.config_file.read_config())
-
-        assert session1_override_config is not session0
-
         # Configuring a new session name should result in a new Session object
         config.set_setting("defaults.aws_profile_name", "SomeRandomProfileName")
         session2 = api.get_boto3_session()
 
         assert session2 is not session0
-        assert session2 == "session for SomeRandomProfileName"  # type: ignore
+        assert session2._profile_name == "SomeRandomProfileName"
 
         # This should return the cached object, and not call boto3.Session
         session3 = api.get_boto3_session()
@@ -93,39 +94,6 @@ def test_get_check_authentication_status_configuration_error(fresh_deadline_conf
         assert api.check_authentication_status() == api.AwsAuthenticationStatus.CONFIGURATION_ERROR
 
 
-def test_get_queue_user_boto3_session_cache(fresh_deadline_config):
-    session_mock = MagicMock()
-    session_mock.profile_name = "test_profile"
-    session_mock.region_name = "us-west-2"
-    deadline_mock = MagicMock()
-    mock_botocore_session = MagicMock()
-    mock_botocore_session.get_config_variable = lambda name: (
-        "test_profile" if name == "profile" else None
-    )
-
-    with patch.object(api._session, "get_boto3_session", return_value=session_mock), patch(
-        "botocore.session.Session", return_value=mock_botocore_session
-    ), patch.object(
-        api._session, "_get_queue_user_boto3_session"
-    ) as _get_queue_user_boto3_session_mock:
-        _ = api.get_queue_user_boto3_session(
-            deadline_mock, farm_id="farm-1234", queue_id="queue-1234", queue_display_name="queue"
-        )
-        # Same farm ID and queue ID, returns cached session
-        _ = api.get_queue_user_boto3_session(
-            deadline_mock, farm_id="farm-1234", queue_id="queue-1234", queue_display_name="queue"
-        )
-        # Different queue ID, makes a fresh session
-        _ = api.get_queue_user_boto3_session(
-            deadline_mock, farm_id="farm-1234", queue_id="queue-5678", queue_display_name="queue"
-        )
-        # Different queue ID, makes a fresh session
-        _ = api.get_queue_user_boto3_session(
-            deadline_mock, farm_id="farm-5678", queue_id="queue-1234", queue_display_name="queue"
-        )
-        assert _get_queue_user_boto3_session_mock.call_count == 3
-
-
 def test_get_queue_user_boto3_session_no_profile(fresh_deadline_config):
     """Make sure that boto3.Session gets called with profile_name=None for the default profile."""
     session_mock = MagicMock()

test/unit/deadline_client/cli/test_cli_auth.py

@@ -31,8 +31,14 @@ def test_cli_deadline_cloud_monitor_login_and_logout(fresh_deadline_config):
     config.set_setting("defaults.aws_profile_name", profile_name)
 
     with patch.object(api._session, "get_boto3_session") as session_mock, patch.object(
+        api._session._get_boto3_session_for_profile, "cache_clear"
+    ) as mock_profile_session_cache_clear, patch.object(
+        api._session._get_queue_user_boto3_session, "cache_clear"
+    ) as mock_queue_session_cache_clear, patch.object(
         api, "get_boto3_session", new=session_mock
-    ), patch.object(subprocess, "Popen") as popen_mock, patch.object(
+    ), patch.object(
+        subprocess, "Popen"
+    ) as popen_mock, patch.object(
         subprocess, "check_output"
     ) as check_output_mock:
         # The profile name
@@ -79,9 +85,8 @@ def test_cli_deadline_cloud_monitor_login_and_logout(fresh_deadline_config):
         )
 
         assert "Successfully logged out" in result.output
-
-        # Verify that the logout call resets the cached session to None
-        assert api._session.__cached_boto3_session is None
+        mock_profile_session_cache_clear.assert_called()
+        mock_queue_session_cache_clear.assert_called()
 
 
 def test_cli_auth_status(fresh_deadline_config):