addressing code review comments

leszekamz · leszekamz · commit cf805a822554 · 2025-07-24T11:03:03.000-07:00
diff --git a/aws_advanced_python_wrapper/federated_plugin.py b/aws_advanced_python_wrapper/federated_plugin.py
@@ -57,7 +57,8 @@ class FederatedAuthPlugin(Plugin):
     _rds_utils: RdsUtils = RdsUtils()
     _token_cache: Dict[str, TokenInfo] = {}
 
-    def __init__(self, plugin_service: PluginService,
+    def __init__(self,
+                 plugin_service: PluginService,
                  credentials_provider_factory: CredentialsProviderFactory,
                  token_utils: TokenUtils,
                  session: Optional[Session] = None):
diff --git a/aws_advanced_python_wrapper/okta_plugin.py b/aws_advanced_python_wrapper/okta_plugin.py
@@ -53,7 +53,8 @@ class OktaAuthPlugin(Plugin):
     _rds_utils: RdsUtils = RdsUtils()
     _token_cache: Dict[str, TokenInfo] = {}
 
-    def __init__(self, plugin_service: PluginService,
+    def __init__(self,
+                 plugin_service: PluginService,
                  credentials_provider_factory: CredentialsProviderFactory,
                  token_utils: TokenUtils,
                  session: Optional[Session] = None):
diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties
@@ -151,7 +151,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.
@@ -316,6 +315,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.
 
diff --git a/aws_advanced_python_wrapper/utils/dsql_token_utils.py b/aws_advanced_python_wrapper/utils/dsql_token_utils.py
@@ -41,7 +41,7 @@ def generate_authentication_token(
             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 authentication token", TelemetryTraceLevel.NESTED)
+        context = telemetry_factory.open_telemetry_context("fetch DSQL authentication token", TelemetryTraceLevel.NESTED)
 
         try:
             session = client_session if client_session else boto3.Session()
@@ -64,7 +64,8 @@ def generate_authentication_token(
             else:
                 token = client.generate_db_connect_auth_token(host_name, region)
 
-            logger.debug("IamAuthUtils.GeneratedNewAuthToken", token)
+            logger.debug("TokenUtils.GeneratedNewAuthTokenLength", len(token) if token else 0)
+            client.close()
             return token
         except Exception as ex:
             context.set_success(False)
diff --git a/aws_advanced_python_wrapper/utils/rds_token_utils.py b/aws_advanced_python_wrapper/utils/rds_token_utils.py
@@ -69,7 +69,7 @@ def generate_authentication_token(
 
             client.close()
 
-            logger.debug("IamAuthUtils.GeneratedNewAuthToken", token)
+            logger.debug("TokenUtils.GeneratedNewAuthTokenLength", len(token) if token else 0)
             return token
         except Exception as ex:
             context.set_success(False)
diff --git a/docs/using-the-python-driver/using-plugins/UsingTheDSQLIamAuthenticationPlugin.md b/docs/using-the-python-driver/using-plugins/UsingTheDSQLIamAuthenticationPlugin.md
@@ -9,7 +9,7 @@ AWS Identity and Access Management (IAM) grants users access control across all
 > [!WARNING]\
 > To preserve compatibility with customers using the community driver, IAM Authentication requires the AWS SDK for Python; [Boto3](https://pypi.org/project/boto3/). Boto3 is a runtime dependency and must be resolved. It can be installed via pip like so: `pip install boto3`.
 
-The IAM Authentication plugin requires authentication via AWS Credentials. These credentials can be defined in `~/.aws/credentials` or set as environment variables. All users must set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`. Users who are using temporary security credentials will also need to additionally set `AWS_SESSION_TOKEN`.
+The DSQL IAM Authentication plugin requires authentication via AWS Credentials. These credentials can be defined in `~/.aws/credentials` or set as environment variables. All users must set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`. Users who are using temporary security credentials will also need to additionally set `AWS_SESSION_TOKEN`.
 
 To enable the AWS Aurora DSQL IAM Authentication Plugin, add the plugin code `iam_dsql` to the [`plugins`](../UsingThePythonDriver.md#connection-plugin-manager-parameters) parameter.
 
@@ -20,6 +20,8 @@ To enable the AWS Aurora DSQL IAM Authentication Plugin, add the plugin code `ia
 The AWS Python Driver supports Amazon AWS Identity and Access Management (IAM) authentication. When using AWS IAM database authentication, the host URL must be a valid AWS Aurora DSQL endpoint, and not a custom domain or an IP address.
 <br>i.e. `cluster-identifier.dsql.us-east-1.on.aws`
 
+Connections established by the `iamDsql` plugin are beholden to the [Cluster quotas and database limits in Amazon Aurora DSQL](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/CHAP_quotas.html). In particular, applications need to consider the maximum transaction duration, and maximum connection duration limits. Ensure connections are returned to the pool regularly, and not retained for long periods.
+
 
 ## How do I use IAM with the AWS Python Driver?
 1. Configure IAM roles for the cluster according to [Using database roles and IAM authentication](https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-database-and-iam-roles.html).
@@ -28,7 +30,7 @@ The AWS Python Driver supports Amazon AWS Identity and Access Management (IAM) a
 | Parameter          |  Value  | Required | Description                                                                                                                                                                                                                                                                                            | Example Value                                       |
 |--------------------|:-------:|:--------:|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------|
 | `iam_host`         | String  |    No    | This property will override the default hostname that is used to generate the IAM token. The default hostname is derived from the connection string. This parameter is required when users are connecting with custom endpoints.                                                                       | `cluster-identifier.dsql.us-east-1.on.aws` |
-| `iam_region`       | String  |    No    | This property will override the default region that is used to generate the IAM token. The default region is parsed from the connection string.                                                                                                                                                        | `us-east-2`                                         |
+| `iam_region`       | String  |    No    | This property will override the default region that is used to generate the IAM token. The default region is parsed from the connection string where possible. Some connection string formats may not be supported, and the `iam_region` must be provided in these cases.                                                                                                                                                     | `us-east-2`                                         |
 | `iam_expiration`   | Integer |    No    | This property determines how long an IAM token is kept in the driver cache before a new one is generated. The default expiration time is set to 14 minutes and 30 seconds. Note that IAM database authentication tokens have a lifetime of 15 minutes.                                                 | `600`                                               |
 
 ## Sample code
diff --git a/docs/using-the-python-driver/using-plugins/UsingTheIamAuthenticationPlugin.md b/docs/using-the-python-driver/using-plugins/UsingTheIamAuthenticationPlugin.md
@@ -11,6 +11,9 @@ The IAM Authentication plugin requires authentication via AWS Credentials. These
 
 To enable the IAM Authentication Connection Plugin, add the plugin code `iam` to the [`plugins`](../UsingThePythonDriver.md#connection-plugin-manager-parameters) parameter.
 
+> [!WARNING]\
+> The `iam` plugin must NOT be specified when using the `iam_dsql` plugin.
+
 ## AWS IAM Database Authentication
 The AWS Python Driver supports Amazon AWS Identity and Access Management (IAM) authentication. When using AWS IAM database authentication, the host URL must be a valid Amazon endpoint, and not a custom domain or an IP address.
 <br>i.e. `db-identifier.cluster-XYZ.us-east-2.rds.amazonaws.com`
