in aws_advanced_python_wrapper/iam_plugin.py [0:0]
def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callable) -> Connection:
user = WrapperProperties.USER.get(props)
if not user:
raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.IsNoneOrEmpty", WrapperProperties.USER.name))
host = IamAuthUtils.get_iam_host(props, host_info)
region = self._region_utils.get_region(props, WrapperProperties.IAM_REGION.name, host, self._session)
if not region:
error_message = "RdsUtils.UnsupportedHostname"
logger.debug(error_message, host)
raise AwsWrapperError(Messages.get_formatted(error_message, host))
port = IamAuthUtils.get_port(props, host_info, self._plugin_service.database_dialect.default_port)
token_expiration_sec: int = WrapperProperties.IAM_EXPIRATION.get_int(props)
cache_key: str = IamAuthUtils.get_cache_key(
user,
host,
port,
region
)
token_info = IamAuthPlugin._token_cache.get(cache_key)
if token_info is not None and not token_info.is_expired():
logger.debug("IamAuthPlugin.UseCachedIamToken", token_info.token)
self._plugin_service.driver_dialect.set_password(props, token_info.token)
else:
token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
self._fetch_token_counter.inc()
token: str = IamAuthUtils.generate_authentication_token(self._plugin_service, user, host, port, region, client_session=self._session)
self._plugin_service.driver_dialect.set_password(props, token)
IamAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)
try:
return connect_func()
except Exception as e:
logger.debug("IamAuthPlugin.ConnectException", e)
is_cached_token = (token_info is not None and not token_info.is_expired())
if not self._plugin_service.is_login_exception(error=e) or not is_cached_token:
raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
# Login unsuccessful with cached token
# Try to generate a new token and try to connect again
token_expiry = datetime.now() + timedelta(seconds=token_expiration_sec)
self._fetch_token_counter.inc()
token = IamAuthUtils.generate_authentication_token(self._plugin_service, user, host, port, region, client_session=self._session)
self._plugin_service.driver_dialect.set_password(props, token)
IamAuthPlugin._token_cache[cache_key] = TokenInfo(token, token_expiry)
try:
return connect_func()
except Exception as e:
raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.UnhandledException", e)) from e