[Test] Add token refresh support for Postgres connectors #5
Conversation
…gres integration (DataDog#21503) * Update connection pool to handle cloud provider token refreshing * run format * Fix kwarg usernname * fixup * Add token provider tests * Lint * Add changelog --------- Co-authored-by: Seth Samuel <[email protected]>
|
/describe |
![mesa-dot-dev[bot]](https://avatars.githubusercontent.com/in/1050077?s=60&v=4){kind=small}
There was a problem hiding this comment.
Performed full review of 32e0fdf...63d2d29
Analysis
-
The TokenAwareConnection implementation uses a class-level
token_providerattribute, creating global state that will cause conflicts in multi-instance deployments where the last-configured provider overwrites others. -
The Azure token provider incorrectly expects an
expires_atfield instead ofexpires_onon Azure's AccessToken objects, which will cause runtime errors in production as the tests aren't aligned with real SDK behavior. -
While the token provider abstraction is architecturally sound with good separation of concerns, these implementation issues compromise the reliability of the authentication mechanism across different database configurations.
Tip
Help
Slash Commands:
/review- Request a full code review/review latest- Review only changes since the last review/describe- Generate PR description. This will update the PR body or issue comment depending on your configuration/help- Get help with Mesa commands and configuration options
5 files reviewed | 3 comments | Edit Agent Settings • Read Docs
| """Test Azure token fetching with custom scope.""" | ||
| mock_token = Mock() | ||
| mock_token.token = "azure_token_123" | ||
| mock_token.expires_at = 1900.0 |
There was a problem hiding this comment.
This test mock uses expires_at attribute which doesn't exist on Azure's AccessToken. The actual AccessToken object has an expires_on attribute. This incorrect mock masks the bug in the implementation and should be changed to expires_on to match the real Azure SDK behavior.
Agent: 🤖 Integrations reviewer •
Prompt for Agent
Task: Address review feedback left on GitHub.
Repository: mesa-dot-dev/integrations-core#5
File: postgres/tests/test_token_provider.py#L198
Action: Open this file location in your editor, inspect the highlighted code, and resolve the issue described below.
Feedback:
This test mock uses `expires_at` attribute which doesn't exist on Azure's AccessToken. The actual AccessToken object has an `expires_on` attribute. This incorrect mock masks the bug in the implementation and should be changed to `expires_on` to match the real Azure SDK behavior.
| from .azure import generate_managed_identity_token | ||
|
|
||
| token = generate_managed_identity_token(client_id=self.client_id, identity_scope=self.identity_scope) | ||
| return token.token, float(token.expires_at) |
There was a problem hiding this comment.
Azure's AccessToken object uses expires_on attribute, not expires_at. This will cause an AttributeError at runtime when the AzureTokenProvider attempts to access token.expires_at. The correct attribute is expires_on which contains the expiration time as seconds since epoch.
Agent: 🤖 Integrations reviewer •
Prompt for Agent
Task: Address review feedback left on GitHub.
Repository: mesa-dot-dev/integrations-core#5
File: postgres/datadog_checks/postgres/connection_pool.py#L92
Action: Open this file location in your editor, inspect the highlighted code, and resolve the issue described below.
Feedback:
Azure's AccessToken object uses `expires_on` attribute, not `expires_at`. This will cause an AttributeError at runtime when the AzureTokenProvider attempts to access `token.expires_at`. The correct attribute is `expires_on` which contains the expiration time as seconds since epoch.
| "open": True, | ||
| } | ||
|
|
||
| TokenAwareConnection.token_provider = self.token_provider |
There was a problem hiding this comment.
This stores the token provider as a class-level attribute, which creates a global state shared across all PostgreSQL check instances in the Agent process. If multiple PostgreSQL checks are configured (e.g., different hosts, or some with managed auth and others without), the last initialized check will overwrite this attribute, causing incorrect authentication behavior for other checks. The token provider should be scoped to the specific connection pool instance rather than the class.
Agent: 🤖 Integrations reviewer •
Prompt for Agent
Task: Address review feedback left on GitHub.
Repository: mesa-dot-dev/integrations-core#5
File: postgres/datadog_checks/postgres/connection_pool.py#L210
Action: Open this file location in your editor, inspect the highlighted code, and resolve the issue described below.
Feedback:
This stores the token provider as a class-level attribute, which creates a global state shared across all PostgreSQL check instances in the Agent process. If multiple PostgreSQL checks are configured (e.g., different hosts, or some with managed auth and others without), the last initialized check will overwrite this attribute, causing incorrect authentication behavior for other checks. The token provider should be scoped to the specific connection pool instance rather than the class.
This PR refactors the token-based authentication for the Postgres integration to support automatic, thread-safe token refreshing for AWS IAM and Azure Managed Identity. This change fixes a bug where connections would fail upon token expiration by proactively managing the token lifecycle within the connection pool.
The core of the change is a new
TokenProviderabstract class, which centralizes token fetching, caching, and refreshing logic. Concrete implementations forAWSTokenProviderandAzureTokenProviderencapsulate the platform-specific details.A new
TokenAwareConnectionclass is introduced, which utilizes aTokenProviderto retrieve a valid token for the password just before establishing a connection. TheConnectionPoolManagerhas been updated to manage these providers, ensuring that connections remain resilient to token expiration.Finally, comprehensive unit and integration tests have been added for the new token provider mechanism, validating its caching, refresh logic, and thread safety.
Description generated by Mesa. Update settings