Add support for DSQL iam authentication

.github/workflows/integration_tests.yml

@@ -19,7 +19,7 @@ jobs:
       matrix:
         python-version: [ "3.8", "3.11" ]
         engine-version: [ "lts", "latest"]
-        environment: ["mysql", "pg"]
+        environment: ["mysql", "pg", "dsql"]
 
     steps:
       - name: 'Clone repository'

README.md

@@ -26,6 +26,12 @@ Since a database failover is usually identified by reaching a network or a conne
 
 Enhanced Failure Monitoring (EFM) is a feature available from the [Host Monitoring Connection Plugin](./docs/using-the-python-driver/using-plugins/UsingTheHostMonitoringPlugin.md#enhanced-failure-monitoring) that periodically checks the connected database host's health and availability. If a database host is determined to be unhealthy, the connection is aborted (and potentially routed to another healthy host in the cluster).
 
+### Using the AWS Advanced Python Driver with AWS Aurora DSQL
+The AWS Advanced Python Driver is able to handle IAM authentication when working with AWS Aurora DSQL clusters.
+
+Please visit [this page](./docs/using-the-python-driver/using-plugins/UsingTheDSQLIamAuthenticationPlugin.md) for more information.
+
+
 ### Using the AWS Advanced Python Driver with plain RDS databases
 
 The AWS Advanced Python Driver also works with RDS provided databases that are not Aurora.

aws_advanced_python_wrapper/dsql_iam_auth_plugin_factory.py

@@ -0,0 +1,30 @@
+#  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 typing import TYPE_CHECKING
+
+from aws_advanced_python_wrapper.iam_plugin import IamAuthPlugin
+from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory
+from aws_advanced_python_wrapper.utils.dsql_token_utils import DSQLTokenUtils
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from aws_advanced_python_wrapper.utils.properties import Properties
+
+
+class DsqlIamAuthPluginFactory(PluginFactory):
+    def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
+        return IamAuthPlugin(plugin_service, DSQLTokenUtils())

aws_advanced_python_wrapper/federated_plugin.py

@@ -31,6 +31,7 @@
     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 aws_advanced_python_wrapper.utils.token_utils import TokenUtils
 
 from datetime import datetime, timedelta
 from typing import Callable, Dict, Optional, Set
@@ -43,6 +44,7 @@
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
+from aws_advanced_python_wrapper.utils.rds_token_utils import RDSTokenUtils
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 
 logger = Logger(__name__)
@@ -55,12 +57,17 @@ 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 +152,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 +166,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/iam_plugin.py

@@ -17,6 +17,8 @@
 from typing import TYPE_CHECKING
 
 from aws_advanced_python_wrapper.utils.iam_utils import IamAuthUtils, TokenInfo
+from aws_advanced_python_wrapper.utils.rds_token_utils import RDSTokenUtils
+from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 from aws_advanced_python_wrapper.utils.region_utils import RegionUtils
 
 if TYPE_CHECKING:
@@ -25,6 +27,7 @@
     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 aws_advanced_python_wrapper.utils.token_utils import TokenUtils
 
 from datetime import datetime, timedelta
 from typing import Callable, Dict, Optional, Set
@@ -35,7 +38,6 @@
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
-from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 
 logger = Logger(__name__)
 
@@ -48,11 +50,12 @@ 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, token_utils: TokenUtils, session: Optional[Session] = None):
         self._plugin_service = plugin_service
         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("iam.fetch_token.count")
         self._cache_size_gauge = telemetry_factory.create_gauge(
@@ -102,7 +105,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
         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)
+            token: str = self._token_utils.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)
 
@@ -120,7 +123,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
             # 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)
+            token = self._token_utils.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)
 
@@ -142,4 +145,4 @@ def force_connect(
 
 class IamAuthPluginFactory(PluginFactory):
     def get_instance(self, plugin_service: PluginService, props: Properties) -> Plugin:
-        return IamAuthPlugin(plugin_service)
+        return IamAuthPlugin(plugin_service, RDSTokenUtils())

aws_advanced_python_wrapper/okta_plugin.py

@@ -31,6 +31,7 @@
     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 aws_advanced_python_wrapper.utils.token_utils import TokenUtils
 
 import requests
 
@@ -40,6 +41,7 @@
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
+from aws_advanced_python_wrapper.utils.rds_token_utils import RDSTokenUtils
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
 
 logger = Logger(__name__)
@@ -51,12 +53,17 @@ 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 +147,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 +235,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)

aws_advanced_python_wrapper/plugin_service.py

@@ -61,6 +61,8 @@
 from aws_advanced_python_wrapper.developer_plugin import DeveloperPluginFactory
 from aws_advanced_python_wrapper.driver_configuration_profiles import \
     DriverConfigurationProfiles
+from aws_advanced_python_wrapper.dsql_iam_auth_plugin_factory import \
+    DsqlIamAuthPluginFactory
 from aws_advanced_python_wrapper.errors import (AwsWrapperError,
                                                 QueryTimeoutError,
                                                 UnsupportedOperationError)
@@ -752,6 +754,7 @@ class PluginManager(CanReleaseResources):
 
     PLUGIN_FACTORIES: Dict[str, Type[PluginFactory]] = {
         "iam": IamAuthPluginFactory,
+        "iam_dsql": DsqlIamAuthPluginFactory,
         "aws_secrets_manager": AwsSecretsManagerPluginFactory,
         "aurora_connection_tracker": AuroraConnectionTrackerPluginFactory,
         "host_monitoring": HostMonitoringPluginFactory,
@@ -786,6 +789,7 @@ class PluginManager(CanReleaseResources):
         BlueGreenPluginFactory: 550,
         FastestResponseStrategyPluginFactory: 600,
         IamAuthPluginFactory: 700,
+        DsqlIamAuthPluginFactory: 710,
         AwsSecretsManagerPluginFactory: 800,
         FederatedAuthPluginFactory: 900,
         LimitlessPluginFactory: 950,

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -186,7 +186,6 @@ IamAuthPlugin.UnhandledException=[IamAuthPlugin] Unhandled exception: {}
 IamAuthPlugin.UseCachedIamToken=[IamAuthPlugin] Used cached IAM token = {}
 IamAuthPlugin.InvalidHost=[IamAuthPlugin] Invalid IAM host {}. The IAM host must be a valid RDS or Aurora endpoint.
 IamAuthPlugin.IsNoneOrEmpty=[IamAuthPlugin] Property "{}" is None or empty.
-IamAuthUtils.GeneratedNewAuthToken=Generated new authentication token = {}
 
 LimitlessPlugin.FailedToConnectToHost=[LimitlessPlugin] Failed to connect to host {}.
 LimitlessPlugin.UnsupportedDialectOrDatabase=[LimitlessPlugin] Unsupported dialect '{}' encountered. Please ensure the connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin.
@@ -357,6 +356,8 @@ RoundRobinHostSelector.ClusterInfoNone=[RoundRobinHostSelector] The round robin
 RoundRobinHostSelector.RoundRobinInvalidDefaultWeight=[RoundRobinHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
 
+TokenUtils.GeneratedNewAuthTokenLength=Generated new authentication token length = {}
+
 WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs= [WeightedRandomHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
 WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight=[WeightedRandomHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 

aws_advanced_python_wrapper/utils/dsql_token_utils.py

@@ -0,0 +1,75 @@
+#  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 typing import TYPE_CHECKING, Dict, Optional
+
+from aws_advanced_python_wrapper.utils.log import Logger
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
+    TelemetryTraceLevel
+from aws_advanced_python_wrapper.utils.token_utils import TokenUtils
+
+if TYPE_CHECKING:
+    from aws_advanced_python_wrapper.plugin_service import PluginService
+    from boto3 import Session
+
+import boto3
+
+logger = Logger(__name__)
+
+
+class DSQLTokenUtils(TokenUtils):
+    def generate_authentication_token(
+            self,
+            plugin_service: PluginService,
+            user: Optional[str],
+            host_name: Optional[str],
+            port: Optional[int],
+            region: Optional[str],
+            credentials: Optional[Dict[str, str]] = None,
+            client_session: Optional[Session] = None) -> str:
+        telemetry_factory = plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context("fetch DSQL authentication token", TelemetryTraceLevel.NESTED)
+
+        try:
+            session = client_session if client_session else boto3.Session()
+            if credentials is not None:
+                client = session.client(
+                    'dsql',
+                    region_name=region,
+                    aws_access_key_id=credentials.get('AccessKeyId'),
+                    aws_secret_access_key=credentials.get('SecretAccessKey'),
+                    aws_session_token=credentials.get('SessionToken')
+                )
+            else:
+                client = session.client(
+                    'dsql',
+                    region_name=region
+                )
+
+            if user == "admin":
+                token = client.generate_db_connect_admin_auth_token(host_name, region)
+            else:
+                token = client.generate_db_connect_auth_token(host_name, region)
+
+            logger.debug("TokenUtils.GeneratedNewAuthTokenLength", len(token) if token else 0)
+            client.close()
+            return token
+        except Exception as ex:
+            context.set_success(False)
+            context.set_exception(ex)
+            raise ex
+        finally:
+            context.close_context()