Fixed errors with updated interfaces for federated and octa plugins.

Author: leszekamz

aws_advanced_python_wrapper/federated_plugin.py

@@ -44,6 +44,8 @@
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+from aws_advanced_python_wrapper.utils.token_utils import TokenUtils
+from aws_advanced_python_wrapper.utils.rds_token_utils import RDSTokenUtils
 
 logger = Logger(__name__)
 
@@ -55,12 +57,13 @@ 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, token_utils: TokenUtils, session: Optional[Session] = None):
         self._plugin_service = plugin_service
         self._credentials_provider_factory = credentials_provider_factory
         self._session = session
 
         self._region_utils = RegionUtils()
+        self._token_utils = token_utils
         telemetry_factory = self._plugin_service.get_telemetry_factory()
         self._fetch_token_counter = telemetry_factory.create_counter("federated.fetch_token.count")
         self._cache_size_gauge = telemetry_factory.create_gauge("federated.token_cache.size", lambda: len(FederatedAuthPlugin._token_cache))
@@ -145,7 +148,7 @@ def _update_authentication_token(self,
         credentials: Optional[Dict[str, str]] = self._credentials_provider_factory.get_aws_credentials(region, props)
 
         self._fetch_token_counter.inc()
-        token: str = IamAuthUtils.generate_authentication_token(
+        token: str = self._token_utils.generate_authentication_token(
             self._plugin_service,
             user,
             host_info.host,
@@ -159,7 +162,7 @@ def _update_authentication_token(self,
 
 class FederatedAuthPluginFactory(PluginFactory):
     def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
-        return FederatedAuthPlugin(plugin_service, self.get_credentials_provider_factory(plugin_service, props))
+        return FederatedAuthPlugin(plugin_service, self.get_credentials_provider_factory(plugin_service, props), RDSTokenUtils())
 
     def get_credentials_provider_factory(self, plugin_service: PluginService, props: Properties) -> AdfsCredentialsProviderFactory:
         idp_name = WrapperProperties.IDP_NAME.get(props)

aws_advanced_python_wrapper/okta_plugin.py

@@ -41,6 +41,8 @@
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+from aws_advanced_python_wrapper.utils.token_utils import TokenUtils
+from aws_advanced_python_wrapper.utils.rds_token_utils import RDSTokenUtils
 
 logger = Logger(__name__)
 
@@ -51,12 +53,13 @@ class OktaAuthPlugin(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, token_utils: TokenUtils, session: Optional[Session] = None):
         self._plugin_service = plugin_service
         self._credentials_provider_factory = credentials_provider_factory
         self._session = session
 
         self._region_utils = RegionUtils()
+        self._token_utils = token_utils
         telemetry_factory = self._plugin_service.get_telemetry_factory()
         self._fetch_token_counter = telemetry_factory.create_counter("okta.fetch_token.count")
         self._cache_size_gauge = telemetry_factory.create_gauge("okta.token_cache.size", lambda: len(OktaAuthPlugin._token_cache))
@@ -140,7 +143,7 @@ def _update_authentication_token(self,
         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)
 
-        token: str = IamAuthUtils.generate_authentication_token(
+        token: str = self._token_utils.generate_authentication_token(
             self._plugin_service,
             user,
             host_info.host,
@@ -228,7 +231,7 @@ def get_saml_assertion(self, props: Properties):
 
 class OktaAuthPluginFactory(PluginFactory):
     def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
-        return OktaAuthPlugin(plugin_service, self.get_credentials_provider_factory(plugin_service, props))
+        return OktaAuthPlugin(plugin_service, self.get_credentials_provider_factory(plugin_service, props), RDSTokenUtils())
 
     def get_credentials_provider_factory(self, plugin_service: PluginService, props: Properties) -> OktaCredentialsProviderFactory:
         return OktaCredentialsProviderFactory(plugin_service, props)