feat: add AWS IAM Database Authentication support for RDS/Aurora PosgreSQL #9579
+431
−35
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…greSQL This adds support for AWS IAM Database Authentication, allowing users to connect to Amazon RDS and Aurora PostgreSQL databases using IAM credentials instead of traditional passwords. Features: - New "AWS IAM authentication" toggle in the Server Connection tab - AWS Profile and Region configuration fields - Automatic IAM token generation using botocore - Automatic SSL enforcement (required by IAM auth) - Token refresh retry on authentication failures - No password prompt when IAM auth is enabled Implementation: - Database migration for new server fields (use_iam_auth, aws_profile, aws_region, aws_role_arn) - AWS IAM token generator module (web/pgadmin/utils/aws_iam.py) - Integration with ServerManager connection string generation - Connection retry logic for expired tokens - UI fields in server dialog - RST documentation Prerequisites for users: - AWS credentials configured (profile or environment variables) - RDS/Aurora instance with IAM auth enabled - Database user with rds_iam role granted - IAM policy allowing rds-db:connect action
WalkthroughThis pull request adds AWS IAM Database Authentication support to pgAdmin. Changes include documentation, database schema migration for new IAM-related columns, model and UI updates to expose IAM configuration fields, token generation utility, server manager integration for IAM-based connections with SSL enforcement, and connection retry logic for handling token expiration. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User/UI
participant ServerMgr as Server Manager
participant AWSUtil as AWS IAM Util
participant Boto as botocore
participant RDS as RDS Instance
participant DB as PostgreSQL DB
User->>ServerMgr: Configure connection with use_iam_auth=true
ServerMgr->>AWSUtil: generate_rds_auth_token(host, port, user, region, profile, role_arn)
AWSUtil->>Boto: Create session with optional profile
Boto->>Boto: Resolve AWS credentials
Boto->>RDS: generate_db_auth_token()
RDS-->>Boto: Auth token (~15 min validity)
Boto-->>AWSUtil: Token
AWSUtil-->>ServerMgr: Token (or error)
ServerMgr->>ServerMgr: Inject token as password, enforce SSL
ServerMgr->>DB: Connect with token as password
DB-->>ServerMgr: Success or auth failure (SQLSTATE 28P01)
alt Auth Success
ServerMgr-->>User: Connected
else Token Expired/Invalid (SQLSTATE 28P01)
ServerMgr->>ServerMgr: Retry: regenerate token
ServerMgr->>AWSUtil: generate_rds_auth_token() [retry]
AWSUtil-->>ServerMgr: Fresh token
ServerMgr->>DB: Connect with new token
DB-->>ServerMgr: Success
ServerMgr-->>User: Connected
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~30 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 🧪 Unit Test Generation v2 is now available!We have significantly improved our unit test generation capabilities. To enable: Add this to your reviews:
finishing_touches:
unit_tests:
enabled: trueTry it out by using the Have feedback? Share your thoughts on our Discord thread! Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Fix all issues with AI agents
In `@docs/en_US/server_dialog.rst`:
- Around line 73-79: Add a new bullet describing the optional "AWS Role ARN"
field in the same list as "AWS Profile" and "AWS Region": explain that "AWS Role
ARN" is the IAM role Amazon Resource Name to assume when using AWS IAM Database
Authentication, when to provide it (e.g., to assume a role different from the
profile’s default), and mention it is optional and only used when "AWS IAM
authentication?" is True; edit the paragraph around the Connection tab bullets
to include this new bullet so it matches the UI.
In `@requirements.txt`:
- Around line 21-22: Update the botocore version constraint in requirements.txt
to match boto3 1.42.x's internal requirement so transitive resolution cannot
pick an incompatible botocore; replace the loose "botocore>=1.31.0" entry with a
tight constraint "botocore>=1.42.0,<1.43.0" to align with the installed
boto3==1.42.* dependency.
In `@web/pgadmin/browser/server_groups/servers/__init__.py`:
- Around line 1236-1239: When use_iam_auth is set (use_iam_auth key truthy)
validate required IAM-related fields before constructing the server entry:
ensure aws_region is provided and non-empty and, if your flow expects a service
file or uses host/port for token generation, ensure host and port are present as
well; raise a clear error (or return a validation response) referencing
use_iam_auth so callers see "aws_region (and host/port when using service files)
required when use_iam_auth=true". Update the place building the server from data
(where use_iam_auth, aws_profile, aws_region, aws_role_arn are read) to perform
these checks and return a user-friendly validation message rather than allowing
token generation to fail later.
In `@web/pgadmin/browser/server_groups/servers/static/js/server.ui.js`:
- Around line 385-409: Add a validation rule in the validate() function to
require aws_region when use_iam_auth is true: when state.use_iam_auth is truthy,
ensure state.aws_region is non-empty and push a validation error for the field
id 'aws_region' (or return validation failure) so the form cannot be saved
without an AWS region; reference the existing form field ids 'use_iam_auth' and
'aws_region' and respect obj.isConnected/read-only handling already present.
In `@web/pgadmin/utils/aws_iam.py`:
- Around line 30-110: The function generate_rds_auth_token currently has an
unused parameter role_arn which triggers lint warnings; add a no-op debug log
that references role_arn (e.g. logging.debug or logging.info) near the start of
generate_rds_auth_token to acknowledge the parameter and indicate STS
assume-role is not yet implemented, preserving the public signature and
satisfying Ruff ARG checks while keeping the existing TODO about role
assumption.
In `@web/pgadmin/utils/driver/psycopg3/connection.py`:
- Around line 345-369: Remove the unused last_error variable and eliminate the
nested async function connectdbserver to avoid loop-variable closure; inside the
IAM retry loop (around ConnectionLocker and manager.create_connection_string)
call asyncio.run(...) directly with psycopg.AsyncConnection.connect(...) using
the freshly built connection_string, cursor_factory=AsyncDictCursor, autocommit
(or kwargs['auto_commit']), and prepare_threshold=manager.prepare_threshold, and
delete the connectdbserver definition so no async helper captures mutable loop
variables like attempt/connection_string.
In `@web/pgadmin/utils/driver/psycopg3/server_manager.py`:
- Around line 674-692: The create_connection_string() path is prematurely
generating IAM tokens (via generate_rds_auth_token) when ServerManager.update()
only needs a display string; change this so IAM token generation is deferred to
the actual connection flow. Remove or guard the token-generation block in
create_connection_string() behind an explicit flag/parameter (e.g.,
generate_token=False) or move that logic into the real connection code path (the
method that performs connections/auth, e.g., connect() or the routine that
builds live connection params), and ensure use_iam_auth, aws_region,
aws_profile, aws_role_arn and connection_params are still considered for display
without resolving credentials or calling generate_rds_auth_token; keep
generate_rds_auth_token only in the code path that attempts real authentication
so update() and UI rendering no longer trigger AWS credential resolution.
🧹 Nitpick comments (2)
web/pgadmin/browser/server_groups/servers/__init__.py (2)
800-803: Normalizeuse_iam_authto 0/1 in updates for SQLite consistency.
create()storesuse_iam_authas 1/0, but update() doesn’t coerce it like other boolean flags. Consider adding it to the boolean conversion list in_set_valid_attr_value.🔁 Suggested normalization
- if arg in ('sslcompression', 'use_ssh_tunnel', - 'tunnel_authentication', 'kerberos_conn', 'shared'): + if arg in ('sslcompression', 'use_ssh_tunnel', + 'tunnel_authentication', 'kerberos_conn', 'shared', + 'use_iam_auth'): value = 1 if value else 0
1470-1472: Use info‑level logging for normal connection events.These are expected operations; error‑level logging will inflate error noise.
📝 Downgrade to info
- current_app.logger.error( + current_app.logger.info( 'Connection Request for server#{0}'.format(sid) ) ... - current_app.logger.error('Connection Established for server: \ + current_app.logger.info('Connection Established for server: \ %s - %s' % (server.id, server.name))Also applies to: 1673-1674
Comment on lines
+73
to
+79
| * When *AWS IAM authentication?* is set to *True*, pgAdmin will use AWS IAM | ||
| Database Authentication to connect to Amazon RDS or Aurora PostgreSQL. | ||
| For more information, see :ref:`AWS IAM Database Authentication <aws_iam_authentication>`. | ||
|
|
||
| * *AWS Profile*: (Optional) The AWS profile name from your credentials file. | ||
| * *AWS Region*: The AWS region where your RDS/Aurora instance is located. | ||
|
|
There was a problem hiding this comment.
Document the AWS Role ARN field to match the UI.
The Connection tab now has an optional AWS Role ARN field, but it isn’t documented here. Please add a bullet so users know how/when to use it.
Doc snippet
* *AWS Profile*: (Optional) The AWS profile name from your credentials file.
* *AWS Region*: The AWS region where your RDS/Aurora instance is located.
+ * *AWS Role ARN*: (Optional) IAM role ARN for cross-account or assumed-role access.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| * When *AWS IAM authentication?* is set to *True*, pgAdmin will use AWS IAM | |
| Database Authentication to connect to Amazon RDS or Aurora PostgreSQL. | |
| For more information, see :ref:`AWS IAM Database Authentication <aws_iam_authentication>`. | |
| * *AWS Profile*: (Optional) The AWS profile name from your credentials file. | |
| * *AWS Region*: The AWS region where your RDS/Aurora instance is located. | |
| * When *AWS IAM authentication?* is set to *True*, pgAdmin will use AWS IAM | |
| Database Authentication to connect to Amazon RDS or Aurora PostgreSQL. | |
| For more information, see :ref:`AWS IAM Database Authentication <aws_iam_authentication>`. | |
| * *AWS Profile*: (Optional) The AWS profile name from your credentials file. | |
| * *AWS Region*: The AWS region where your RDS/Aurora instance is located. | |
| * *AWS Role ARN*: (Optional) IAM role ARN for cross-account or assumed-role access. | |
🤖 Prompt for AI Agents
In `@docs/en_US/server_dialog.rst` around lines 73 - 79, Add a new bullet
describing the optional "AWS Role ARN" field in the same list as "AWS Profile"
and "AWS Region": explain that "AWS Role ARN" is the IAM role Amazon Resource
Name to assume when using AWS IAM Database Authentication, when to provide it
(e.g., to assume a role different from the profile’s default), and mention it is
optional and only used when "AWS IAM authentication?" is True; edit the
paragraph around the Connection tab bullets to include this new bullet so it
matches the UI.
Comment on lines
21
to
+22
| boto3==1.42.* | ||
| botocore>=1.31.0 |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
boto3 1.42.x botocore version requirement dependency
💡 Result:
For boto3 1.42.x, the botocore dependency is pinned to the same minor line:
- botocore >= 1.42.* and < 1.43.0 (upper-bound prevents pulling the next minor series). (secure.software)
Example (from 1.42.3 metadata as reported by a dependency scanner):
botocore >= 1.42.3, < 1.43.0(secure.software)
Citations:
- 1: https://secure.software/pypi/packages/boto3/dependencies/1.42.3?utm_source=openai
- 2: https://secure.software/pypi/packages/boto3/dependencies/1.42.3?utm_source=openai
Tighten botocore constraint to match boto3's internal requirement.
boto3 1.42.x internally requires botocore >= 1.42.0, < 1.43.0. The current botocore>=1.31.0 is unnecessarily loose and could allow incompatible versions in environments where package resolvers don't strictly enforce transitive constraints. Align it with boto3's minor series:
Suggested change
-botocore>=1.31.0
+botocore==1.42.*🤖 Prompt for AI Agents
In `@requirements.txt` around lines 21 - 22, Update the botocore version
constraint in requirements.txt to match boto3 1.42.x's internal requirement so
transitive resolution cannot pick an incompatible botocore; replace the loose
"botocore>=1.31.0" entry with a tight constraint "botocore>=1.42.0,<1.43.0" to
align with the installed boto3==1.42.* dependency.
Comment on lines
+1236
to
+1239
| use_iam_auth=1 if data.get('use_iam_auth', False) else 0, | ||
| aws_profile=data.get('aws_profile', None), | ||
| aws_region=data.get('aws_region', None), | ||
| aws_role_arn=data.get('aws_role_arn', None), |
There was a problem hiding this comment.
Validate required IAM fields when use_iam_auth is enabled.
If aws_region (and potentially host/port when using service files) is missing, token generation fails later with a generic error. A small server‑side check gives clearer feedback.
✅ Add minimal validation
for arg in required_args:
if arg not in data:
return make_json_response(
status=410,
success=0,
errormsg=gettext(
"Could not find the required parameter ({})."
).format(arg)
)
+
+ if data.get('use_iam_auth') and not data.get('aws_region'):
+ return make_json_response(
+ status=400,
+ success=0,
+ errormsg=gettext("AWS region is required when IAM auth is enabled.")
+ )🤖 Prompt for AI Agents
In `@web/pgadmin/browser/server_groups/servers/__init__.py` around lines 1236 -
1239, When use_iam_auth is set (use_iam_auth key truthy) validate required
IAM-related fields before constructing the server entry: ensure aws_region is
provided and non-empty and, if your flow expects a service file or uses
host/port for token generation, ensure host and port are present as well; raise
a clear error (or return a validation response) referencing use_iam_auth so
callers see "aws_region (and host/port when using service files) required when
use_iam_auth=true". Update the place building the server from data (where
use_iam_auth, aws_profile, aws_region, aws_role_arn are read) to perform these
checks and return a user-friendly validation message rather than allowing token
generation to fail later.
Comment on lines
+385
to
+409
| id: 'use_iam_auth', label: gettext('AWS IAM authentication?'), type: 'switch', | ||
| group: gettext('Connection'), mode: ['create', 'edit'], | ||
| disabled: obj.isShared, | ||
| helpMessage: gettext('Use AWS IAM authentication tokens for RDS/Aurora PostgreSQL databases') | ||
| },{ | ||
| id: 'aws_profile', label: gettext('AWS Profile'), type: 'text', | ||
| group: gettext('Connection'), mode: ['create', 'edit'], | ||
| deps: ['use_iam_auth'], | ||
| disabled: function(state) { return !state.use_iam_auth; }, | ||
| readonly: obj.isConnected, | ||
| helpMessage: gettext('AWS profile name for credentials (leave empty for default)') | ||
| },{ | ||
| id: 'aws_region', label: gettext('AWS Region'), type: 'text', | ||
| group: gettext('Connection'), mode: ['create', 'edit'], | ||
| deps: ['use_iam_auth'], | ||
| disabled: function(state) { return !state.use_iam_auth; }, | ||
| readonly: obj.isConnected, | ||
| helpMessage: gettext('AWS region where the database is located (e.g., us-east-1)') | ||
| },{ | ||
| id: 'aws_role_arn', label: gettext('AWS Role ARN (Optional)'), type: 'text', | ||
| group: gettext('Connection'), mode: ['create', 'edit'], | ||
| deps: ['use_iam_auth'], | ||
| disabled: function(state) { return !state.use_iam_auth; }, | ||
| readonly: obj.isConnected, | ||
| helpMessage: gettext('IAM role ARN for cross-account or assumed role access') |
There was a problem hiding this comment.
Validate AWS Region when IAM auth is enabled.
Users can currently save a server with IAM enabled but no region, which will likely cause token generation failures. Consider validating aws_region when use_iam_auth is true (or explicitly marking it optional if the backend derives it).
Possible validation addition (in validate())
+ if (state.use_iam_auth) {
+ if (isEmptyString(state.aws_region)) {
+ errmsg = gettext('AWS Region must be specified when IAM authentication is enabled.');
+ setError('aws_region', errmsg);
+ return true;
+ } else {
+ setError('aws_region', null);
+ }
+ }🤖 Prompt for AI Agents
In `@web/pgadmin/browser/server_groups/servers/static/js/server.ui.js` around
lines 385 - 409, Add a validation rule in the validate() function to require
aws_region when use_iam_auth is true: when state.use_iam_auth is truthy, ensure
state.aws_region is non-empty and push a validation error for the field id
'aws_region' (or return validation failure) so the form cannot be saved without
an AWS region; reference the existing form field ids 'use_iam_auth' and
'aws_region' and respect obj.isConnected/read-only handling already present.
Comment on lines
+30
to
+110
| def generate_rds_auth_token(host, port, username, region, profile=None, | ||
| role_arn=None): | ||
| """ | ||
| Generate an AWS RDS IAM authentication token. | ||
|
|
||
| Args: | ||
| host (str): The hostname of the database server | ||
| port (int): The port number of the database server | ||
| username (str): The database username to authenticate as | ||
| region (str): The AWS region where the database is located | ||
| profile (str, optional): The AWS profile name to use for credentials | ||
| role_arn (str, optional): The ARN of an IAM role to assume (future) | ||
|
|
||
| Returns: | ||
| str: The generated authentication token (valid for 15 minutes) | ||
|
|
||
| Raises: | ||
| Exception: If botocore is not available or token generation fails | ||
| """ | ||
| if not BOTOCORE_AVAILABLE: | ||
| raise Exception( | ||
| _("AWS IAM authentication requires the 'botocore' package. " | ||
| "Please install it with: pip install botocore") | ||
| ) | ||
|
|
||
| try: | ||
| # Create a botocore session with the specified profile (if any) | ||
| session = botocore.session.Session(profile=profile) | ||
|
|
||
| # TODO: Future enhancement - support role assumption via STS | ||
| # if role_arn: | ||
| # sts_client = session.create_client('sts', region_name=region) | ||
| # assumed_role = sts_client.assume_role( | ||
| # RoleArn=role_arn, | ||
| # RoleSessionName='pgAdmin4-IAM-Session' | ||
| # ) | ||
| # # Use temporary credentials from assumed role | ||
| # ... | ||
|
|
||
| # Create RDS client | ||
| rds_client = session.create_client('rds', region_name=region) | ||
|
|
||
| # Generate the authentication token | ||
| token = rds_client.generate_db_auth_token( | ||
| DBHostname=host, | ||
| Port=port, | ||
| DBUsername=username, | ||
| Region=region | ||
| ) | ||
|
|
||
| logging.info( | ||
| f"Successfully generated IAM auth token for {username}@{host}" | ||
| ) | ||
|
|
||
| return token | ||
|
|
||
| except NoCredentialsError: | ||
| raise Exception( | ||
| _("AWS credentials not found. Please configure your AWS " | ||
| "credentials using 'aws configure' or ensure your " | ||
| "environment has valid credentials.") | ||
| ) | ||
| except PartialCredentialsError: | ||
| raise Exception( | ||
| _("Incomplete AWS credentials found. Please check your " | ||
| "AWS configuration.") | ||
| ) | ||
| except (BotoCoreError, ClientError) as e: | ||
| error_msg = str(e) | ||
| logging.error(f"AWS IAM token generation failed: {error_msg}") | ||
| raise Exception( | ||
| _("AWS IAM token generation failed: {error}").format( | ||
| error=error_msg | ||
| ) | ||
| ) | ||
| except Exception as e: | ||
| error_msg = str(e) | ||
| logging.error(f"Unexpected error during IAM auth: {error_msg}") | ||
| raise Exception( | ||
| _("IAM authentication error: {error}").format(error=error_msg) | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's find and inspect the file
find . -name "aws_iam.py" -type fRepository: pgadmin-org/pgadmin4
Length of output: 94
🏁 Script executed:
# Check Ruff configuration in the repo
find . -name "ruff.toml" -o -name "pyproject.toml" -o -name ".ruff.toml" | head -5Repository: pgadmin-org/pgadmin4
Length of output: 75
🏁 Script executed:
# Search for role_arn usage in the codebase
rg "role_arn" --type py -B 2 -A 2Repository: pgadmin-org/pgadmin4
Length of output: 5719
🏁 Script executed:
# Let's check the Ruff configuration
cat ruff.tomlRepository: pgadmin-org/pgadmin4
Length of output: 970
🏁 Script executed:
# Let's see the full generate_rds_auth_token function implementation
cat -n web/pgadmin/utils/aws_iam.py | sed -n '30,110p'Repository: pgadmin-org/pgadmin4
Length of output: 3553
🏁 Script executed:
# Check how the function is called to understand if role_arn should stay in the signature
rg "generate_rds_auth_token" -B 2 -A 3 --type pyRepository: pgadmin-org/pgadmin4
Length of output: 1207
Add debug logging for unused role_arn parameter.
role_arn is currently unused and Ruff's ARG rule will flag it as an unused argument. Adding a no-op debug statement preserves the function signature (which is already part of the public API and called with this parameter) while satisfying lint and documenting that STS assume-role is not yet implemented.
Fix
def generate_rds_auth_token(host, port, username, region, profile=None,
role_arn=None):
"""
Generate an AWS RDS IAM authentication token.
"""
+ if role_arn:
+ logging.debug("role_arn provided; STS assume-role not yet implemented")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def generate_rds_auth_token(host, port, username, region, profile=None, | |
| role_arn=None): | |
| """ | |
| Generate an AWS RDS IAM authentication token. | |
| Args: | |
| host (str): The hostname of the database server | |
| port (int): The port number of the database server | |
| username (str): The database username to authenticate as | |
| region (str): The AWS region where the database is located | |
| profile (str, optional): The AWS profile name to use for credentials | |
| role_arn (str, optional): The ARN of an IAM role to assume (future) | |
| Returns: | |
| str: The generated authentication token (valid for 15 minutes) | |
| Raises: | |
| Exception: If botocore is not available or token generation fails | |
| """ | |
| if not BOTOCORE_AVAILABLE: | |
| raise Exception( | |
| _("AWS IAM authentication requires the 'botocore' package. " | |
| "Please install it with: pip install botocore") | |
| ) | |
| try: | |
| # Create a botocore session with the specified profile (if any) | |
| session = botocore.session.Session(profile=profile) | |
| # TODO: Future enhancement - support role assumption via STS | |
| # if role_arn: | |
| # sts_client = session.create_client('sts', region_name=region) | |
| # assumed_role = sts_client.assume_role( | |
| # RoleArn=role_arn, | |
| # RoleSessionName='pgAdmin4-IAM-Session' | |
| # ) | |
| # # Use temporary credentials from assumed role | |
| # ... | |
| # Create RDS client | |
| rds_client = session.create_client('rds', region_name=region) | |
| # Generate the authentication token | |
| token = rds_client.generate_db_auth_token( | |
| DBHostname=host, | |
| Port=port, | |
| DBUsername=username, | |
| Region=region | |
| ) | |
| logging.info( | |
| f"Successfully generated IAM auth token for {username}@{host}" | |
| ) | |
| return token | |
| except NoCredentialsError: | |
| raise Exception( | |
| _("AWS credentials not found. Please configure your AWS " | |
| "credentials using 'aws configure' or ensure your " | |
| "environment has valid credentials.") | |
| ) | |
| except PartialCredentialsError: | |
| raise Exception( | |
| _("Incomplete AWS credentials found. Please check your " | |
| "AWS configuration.") | |
| ) | |
| except (BotoCoreError, ClientError) as e: | |
| error_msg = str(e) | |
| logging.error(f"AWS IAM token generation failed: {error_msg}") | |
| raise Exception( | |
| _("AWS IAM token generation failed: {error}").format( | |
| error=error_msg | |
| ) | |
| ) | |
| except Exception as e: | |
| error_msg = str(e) | |
| logging.error(f"Unexpected error during IAM auth: {error_msg}") | |
| raise Exception( | |
| _("IAM authentication error: {error}").format(error=error_msg) | |
| ) | |
| def generate_rds_auth_token(host, port, username, region, profile=None, | |
| role_arn=None): | |
| """ | |
| Generate an AWS RDS IAM authentication token. | |
| Args: | |
| host (str): The hostname of the database server | |
| port (int): The port number of the database server | |
| username (str): The database username to authenticate as | |
| region (str): The AWS region where the database is located | |
| profile (str, optional): The AWS profile name to use for credentials | |
| role_arn (str, optional): The ARN of an IAM role to assume (future) | |
| Returns: | |
| str: The generated authentication token (valid for 15 minutes) | |
| Raises: | |
| Exception: If botocore is not available or token generation fails | |
| """ | |
| if role_arn: | |
| logging.debug("role_arn provided; STS assume-role not yet implemented") | |
| if not BOTOCORE_AVAILABLE: | |
| raise Exception( | |
| _("AWS IAM authentication requires the 'botocore' package. " | |
| "Please install it with: pip install botocore") | |
| ) | |
| try: | |
| # Create a botocore session with the specified profile (if any) | |
| session = botocore.session.Session(profile=profile) | |
| # TODO: Future enhancement - support role assumption via STS | |
| # if role_arn: | |
| # sts_client = session.create_client('sts', region_name=region) | |
| # assumed_role = sts_client.assume_role( | |
| # RoleArn=role_arn, | |
| # RoleSessionName='pgAdmin4-IAM-Session' | |
| # ) | |
| # # Use temporary credentials from assumed role | |
| # ... | |
| # Create RDS client | |
| rds_client = session.create_client('rds', region_name=region) | |
| # Generate the authentication token | |
| token = rds_client.generate_db_auth_token( | |
| DBHostname=host, | |
| Port=port, | |
| DBUsername=username, | |
| Region=region | |
| ) | |
| logging.info( | |
| f"Successfully generated IAM auth token for {username}@{host}" | |
| ) | |
| return token | |
| except NoCredentialsError: | |
| raise Exception( | |
| _("AWS credentials not found. Please configure your AWS " | |
| "credentials using 'aws configure' or ensure your " | |
| "environment has valid credentials.") | |
| ) | |
| except PartialCredentialsError: | |
| raise Exception( | |
| _("Incomplete AWS credentials found. Please check your " | |
| "AWS configuration.") | |
| ) | |
| except (BotoCoreError, ClientError) as e: | |
| error_msg = str(e) | |
| logging.error(f"AWS IAM token generation failed: {error_msg}") | |
| raise Exception( | |
| _("AWS IAM token generation failed: {error}").format( | |
| error=error_msg | |
| ) | |
| ) | |
| except Exception as e: | |
| error_msg = str(e) | |
| logging.error(f"Unexpected error during IAM auth: {error_msg}") | |
| raise Exception( | |
| _("IAM authentication error: {error}").format(error=error_msg) | |
| ) |
🧰 Tools
🪛 Ruff (0.14.14)
[warning] 31-31: Unused function argument: role_arn
(ARG001)
[warning] 50-53: Create your own exception
(TRY002)
[warning] 84-84: Consider moving this statement to an else block
(TRY300)
[warning] 87-91: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 87-91: Create your own exception
(TRY002)
[warning] 93-96: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 93-96: Create your own exception
(TRY002)
[warning] 99-99: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
[warning] 100-104: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 100-104: Create your own exception
(TRY002)
[warning] 105-105: Do not catch blind exception: Exception
(BLE001)
[warning] 107-107: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
[warning] 108-110: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 108-110: Create your own exception
(TRY002)
🤖 Prompt for AI Agents
In `@web/pgadmin/utils/aws_iam.py` around lines 30 - 110, The function
generate_rds_auth_token currently has an unused parameter role_arn which
triggers lint warnings; add a no-op debug log that references role_arn (e.g.
logging.debug or logging.info) near the start of generate_rds_auth_token to
acknowledge the parameter and indicate STS assume-role is not yet implemented,
preserving the public signature and satisfying Ruff ARG checks while keeping the
existing TODO about role assumption.
Comment on lines
+345
to
+369
| # Retry logic for IAM token refresh | ||
| max_retries = 2 if (hasattr(manager, 'use_iam_auth') and | ||
| manager.use_iam_auth) else 1 | ||
| last_error = None | ||
|
|
||
| for attempt in range(max_retries): | ||
| try: | ||
| with ConnectionLocker(manager.kerberos_conn): | ||
| # Create/regenerate connection string (gets fresh IAM token) | ||
| connection_string = manager.create_connection_string( | ||
| database, user, password) | ||
|
|
||
| if self.async_: | ||
| autocommit = True | ||
| if 'auto_commit' in kwargs: | ||
| autocommit = kwargs['auto_commit'] | ||
|
|
||
| async def connectdbserver(): | ||
| return await psycopg.AsyncConnection.connect( | ||
| connection_string, | ||
| cursor_factory=AsyncDictCursor, | ||
| autocommit=autocommit, | ||
| prepare_threshold=manager.prepare_threshold | ||
| ) | ||
| pg_conn = asyncio.run(connectdbserver()) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's see the complete code block with error handling
head -400 web/pgadmin/utils/driver/psycopg3/connection.py | tail -80Repository: pgadmin-org/pgadmin4
Length of output: 3659
🏁 Script executed:
# Search for last_error in the file to confirm it's not used elsewhere
rg -n "last_error" web/pgadmin/utils/driver/psycopg3/connection.pyRepository: pgadmin-org/pgadmin4
Length of output: 136
Resolve unused variable and loop variable capture in the IAM retry loop.
last_error is assigned but never read (F841), and the async helper closes over mutable loop variables (B023). Both can be fixed without changing behavior: remove last_error and inline the coroutine directly into asyncio.run().
🔧 Suggested cleanup
- last_error = None
-
for attempt in range(max_retries):
try:
with ConnectionLocker(manager.kerberos_conn):
# Create/regenerate connection string (gets fresh IAM token)
connection_string = manager.create_connection_string(
database, user, password)
if self.async_:
autocommit = True
if 'auto_commit' in kwargs:
autocommit = kwargs['auto_commit']
- async def connectdbserver():
- return await psycopg.AsyncConnection.connect(
- connection_string,
- cursor_factory=AsyncDictCursor,
- autocommit=autocommit,
- prepare_threshold=manager.prepare_threshold
- )
- pg_conn = asyncio.run(connectdbserver())
+ pg_conn = asyncio.run(
+ psycopg.AsyncConnection.connect(
+ connection_string,
+ cursor_factory=AsyncDictCursor,
+ autocommit=autocommit,
+ prepare_threshold=manager.prepare_threshold
+ )
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # Retry logic for IAM token refresh | |
| max_retries = 2 if (hasattr(manager, 'use_iam_auth') and | |
| manager.use_iam_auth) else 1 | |
| last_error = None | |
| for attempt in range(max_retries): | |
| try: | |
| with ConnectionLocker(manager.kerberos_conn): | |
| # Create/regenerate connection string (gets fresh IAM token) | |
| connection_string = manager.create_connection_string( | |
| database, user, password) | |
| if self.async_: | |
| autocommit = True | |
| if 'auto_commit' in kwargs: | |
| autocommit = kwargs['auto_commit'] | |
| async def connectdbserver(): | |
| return await psycopg.AsyncConnection.connect( | |
| connection_string, | |
| cursor_factory=AsyncDictCursor, | |
| autocommit=autocommit, | |
| prepare_threshold=manager.prepare_threshold | |
| ) | |
| pg_conn = asyncio.run(connectdbserver()) | |
| # Retry logic for IAM token refresh | |
| max_retries = 2 if (hasattr(manager, 'use_iam_auth') and | |
| manager.use_iam_auth) else 1 | |
| for attempt in range(max_retries): | |
| try: | |
| with ConnectionLocker(manager.kerberos_conn): | |
| # Create/regenerate connection string (gets fresh IAM token) | |
| connection_string = manager.create_connection_string( | |
| database, user, password) | |
| if self.async_: | |
| autocommit = True | |
| if 'auto_commit' in kwargs: | |
| autocommit = kwargs['auto_commit'] | |
| pg_conn = asyncio.run( | |
| psycopg.AsyncConnection.connect( | |
| connection_string, | |
| cursor_factory=AsyncDictCursor, | |
| autocommit=autocommit, | |
| prepare_threshold=manager.prepare_threshold | |
| ) | |
| ) |
🧰 Tools
🪛 Ruff (0.14.14)
[warning] 364-364: Function definition does not bind loop variable connection_string
(B023)
[warning] 366-366: Function definition does not bind loop variable autocommit
(B023)
🤖 Prompt for AI Agents
In `@web/pgadmin/utils/driver/psycopg3/connection.py` around lines 345 - 369,
Remove the unused last_error variable and eliminate the nested async function
connectdbserver to avoid loop-variable closure; inside the IAM retry loop
(around ConnectionLocker and manager.create_connection_string) call
asyncio.run(...) directly with psycopg.AsyncConnection.connect(...) using the
freshly built connection_string, cursor_factory=AsyncDictCursor, autocommit (or
kwargs['auto_commit']), and prepare_threshold=manager.prepare_threshold, and
delete the connectdbserver definition so no async helper captures mutable loop
variables like attempt/connection_string.
Comment on lines
+674
to
+692
| # AWS IAM Authentication: Generate token and force SSL | ||
| if hasattr(self, 'use_iam_auth') and self.use_iam_auth: | ||
| from pgadmin.utils.aws_iam import generate_rds_auth_token | ||
| try: | ||
| password = generate_rds_auth_token( | ||
| host=self.host, | ||
| port=self.port, | ||
| username=user, | ||
| region=self.aws_region, | ||
| profile=self.aws_profile if self.aws_profile else None, | ||
| role_arn=self.aws_role_arn if self.aws_role_arn else None | ||
| ) | ||
| # IAM authentication requires SSL | ||
| if not self.connection_params: | ||
| self.connection_params = {} | ||
| if 'sslmode' not in self.connection_params: | ||
| self.connection_params['sslmode'] = 'require' | ||
| except Exception as e: | ||
| raise Exception(f"IAM authentication failed: {str(e)}") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's find and examine the ServerManager class
fd server_manager.pyRepository: pgadmin-org/pgadmin4
Length of output: 115
🏁 Script executed:
# Let's look at the create_connection_string method and the update method
cd web/pgadmin/utils/driver/psycopg3 && wc -l server_manager.pyRepository: pgadmin-org/pgadmin4
Length of output: 85
🏁 Script executed:
# Read the file to understand the structure
cat -n web/pgadmin/utils/driver/psycopg3/server_manager.py | head -100Repository: pgadmin-org/pgadmin4
Length of output: 4090
🏁 Script executed:
# Search for the update method and create_connection_string calls
rg "def update|create_connection_string" web/pgadmin/utils/driver/psycopg3/server_manager.py -A 2Repository: pgadmin-org/pgadmin4
Length of output: 664
🏁 Script executed:
# Let's see the context around the create_connection_string call in update()
cat -n web/pgadmin/utils/driver/psycopg3/server_manager.py | sed -n '60,180p'Repository: pgadmin-org/pgadmin4
Length of output: 5469
🏁 Script executed:
# Now let's find all calls to create_connection_string in the file
rg "create_connection_string\(" web/pgadmin/utils/driver/psycopg3/server_manager.py -nRepository: pgadmin-org/pgadmin4
Length of output: 200
🏁 Script executed:
# Let's also search in the entire web/pgadmin directory for calls to create_connection_string
rg "create_connection_string\(" web/pgadmin/ -n --type py | head -30Repository: pgadmin-org/pgadmin4
Length of output: 793
🏁 Script executed:
# Let's examine the create_connection_string method implementation
cat -n web/pgadmin/utils/driver/psycopg3/server_manager.py | sed -n '658,730p'Repository: pgadmin-org/pgadmin4
Length of output: 3790
🏁 Script executed:
# Let's search for usage of display_connection_string
rg "display_connection_string" web/pgadmin/ -n --type pyRepository: pgadmin-org/pgadmin4
Length of output: 423
🏁 Script executed:
# Let's check the context around the calls in connection.py
cat -n web/pgadmin/utils/driver/psycopg3/connection.py | sed -n '350,360p'
cat -n web/pgadmin/utils/driver/psycopg3/connection.py | sed -n '1470,1480p'Repository: pgadmin-org/pgadmin4
Length of output: 1224
🏁 Script executed:
# Let's verify if display_connection_string should ever contain real IAM tokens
# by looking at how it's used more carefully
cat -n web/pgadmin/browser/server_groups/servers/__init__.py | sed -n '1115,1130p'Repository: pgadmin-org/pgadmin4
Length of output: 809
🏁 Script executed:
# Let's also check what conninfo_to_dict does
rg "def conninfo_to_dict" web/pgadmin/ -A 3 --type pyRepository: pgadmin-org/pgadmin4
Length of output: 46
🏁 Script executed:
# Let's verify there are no other places where display_connection_string is used that might need real IAM tokens
rg "display_connection_string" web/pgadmin/ -B 2 -A 2 --type pyRepository: pgadmin-org/pgadmin4
Length of output: 1387
Gate IAM token generation to actual connection attempts only.
ServerManager.update() calls create_connection_string() solely to generate a display string for the UI; with IAM enabled, this unnecessarily triggers token generation and AWS credential resolution on every page load, slowing UI operations and causing failures when credentials aren't available. Tokens will be regenerated anyway during real connection attempts, so generate them only when needed.
- def create_connection_string(self, database, user, password=None):
+ def create_connection_string(self, database, user, password=None,
+ generate_iam_token=True):
"""
This function is used to create connection string based on the
parameters.
"""
...
- if hasattr(self, 'use_iam_auth') and self.use_iam_auth:
+ if generate_iam_token and hasattr(self, 'use_iam_auth') and self.use_iam_auth:
from pgadmin.utils.aws_iam import generate_rds_auth_token
...- self.create_connection_string(self.db, self.user)
+ self.create_connection_string(self.db, self.user, generate_iam_token=False)🧰 Tools
🪛 Ruff (0.14.14)
[warning] 691-691: Do not catch blind exception: Exception
(BLE001)
[warning] 692-692: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 692-692: Create your own exception
(TRY002)
[warning] 692-692: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 692-692: Use explicit conversion flag
Replace with conversion flag
(RUF010)
🤖 Prompt for AI Agents
In `@web/pgadmin/utils/driver/psycopg3/server_manager.py` around lines 674 - 692,
The create_connection_string() path is prematurely generating IAM tokens (via
generate_rds_auth_token) when ServerManager.update() only needs a display
string; change this so IAM token generation is deferred to the actual connection
flow. Remove or guard the token-generation block in create_connection_string()
behind an explicit flag/parameter (e.g., generate_token=False) or move that
logic into the real connection code path (the method that performs
connections/auth, e.g., connect() or the routine that builds live connection
params), and ensure use_iam_auth, aws_region, aws_profile, aws_role_arn and
connection_params are still considered for display without resolving credentials
or calling generate_rds_auth_token; keep generate_rds_auth_token only in the
code path that attempts real authentication so update() and UI rendering no
longer trigger AWS credential resolution.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This adds support for AWS IAM Database Authentication, allowing users to connect to Amazon RDS and Aurora PostgreSQL databases using IAM credentials instead of traditional passwords.
Features:
Implementation:
Prerequisites for users:
Summary by CodeRabbit
Release Notes
New Features
Documentation
✏️ Tip: You can customize this high-level summary in your review settings.