[DBMON-5768] Avoid collision between multiple TokenProviders (DataDog#21560)

* Avoid collision between multiple TokenProviders

* Add changelog

Author: eric-weaver

postgres/changelog.d/21560.fixed

@@ -0,0 +1 @@
+Fixes a collision issue when token based authentication is configured for multiple Postgres instances

postgres/datadog_checks/postgres/connection_pool.py

@@ -97,15 +97,17 @@ class TokenAwareConnection(Connection):
     Connection that can be used for managed authentication.
     """
 
-    token_provider: Optional[TokenProvider] = None
-
     @classmethod
     def connect(cls, *args, **kwargs):
         """
         Override the connection method to pass a refreshable token as the connection password.
+
+        The token_provider can be passed via the 'token_provider' kwarg and will be used
+        to dynamically fetch authentication tokens.
         """
-        if cls.token_provider:
-            kwargs["password"] = cls.token_provider.get_token()
+        token_provider = kwargs.pop("token_provider", None)
+        if token_provider:
+            kwargs["password"] = token_provider.get_token()
         return super().connect(*args, **kwargs)
 
 
@@ -193,6 +195,7 @@ def __init__(
             pool_config (dict, optional): Additional ConnectionPool settings (min_size, max_size, etc).
             statement_timeout (int, optional): Statement timeout in milliseconds.
             sqlascii_encodings (list[str], optional): List of encodings to handle for SQLASCII text.
+            token_provider (TokenProvider, optional): Token provider for managed authentication.
         """
         self.max_db = max_db
         self.base_conn_args = base_conn_args
@@ -207,8 +210,6 @@ def __init__(
             "open": True,
         }
 
-        TokenAwareConnection.token_provider = self.token_provider
-
         self.lock = threading.Lock()
         self.pools: OrderedDict[str, Tuple[ConnectionPool, float, bool]] = OrderedDict()
         self._closed = False
@@ -240,6 +241,10 @@ def _create_pool(self, dbname: str) -> ConnectionPool:
         """
         kwargs = self.base_conn_args.as_kwargs(dbname=dbname)
 
+        # Pass the token_provider as a kwarg so it's available to TokenAwareConnection.connect()
+        if self.token_provider:
+            kwargs["token_provider"] = self.token_provider
+
         return ConnectionPool(
             kwargs=kwargs,
             configure=self._configure_connection,

postgres/datadog_checks/postgres/postgres.py

@@ -959,7 +959,13 @@ def _new_connection(self, dbname):
         # TODO: Keeping this main connection outside of the pool for now to keep existing behavior.
         # We should move this to the pool in the future.
         conn_args = self.build_connection_args()
-        conn = TokenAwareConnection.connect(**conn_args.as_kwargs(dbname=dbname))
+        kwargs = conn_args.as_kwargs(dbname=dbname)
+
+        # Pass the token_provider as a kwarg so it's available to TokenAwareConnection.connect()
+        if self.db_pool.token_provider:
+            kwargs["token_provider"] = self.db_pool.token_provider
+
+        conn = TokenAwareConnection.connect(**kwargs)
         self.db_pool._configure_connection(conn)
         return conn
 

postgres/tests/test_token_provider.py

@@ -3,9 +3,15 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
 import time
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
 
-from datadog_checks.postgres.connection_pool import AWSTokenProvider, AzureTokenProvider, TokenProvider
+from datadog_checks.postgres.connection_pool import (
+    AWSTokenProvider,
+    AzureTokenProvider,
+    LRUConnectionPoolManager,
+    PostgresConnectionArgs,
+    TokenProvider,
+)
 
 
 def test_get_token_first_call():
@@ -256,6 +262,109 @@ def test_azure_token_provider_integration():
         assert mock_credential.get_token.call_count == 1
 
 
+def test_multiple_connection_pools_no_token_collision():
+    """
+    Test that multiple LRUConnectionPoolManager instances with different token providers
+    don't have token collision issues.
+    """
+    # Create two different mock token providers that return different tokens
+    token_provider_1 = MockTokenProvider()
+    token_provider_1._fetch_token = Mock(return_value=("token_for_rds_1", time.time() + 3600))
+
+    token_provider_2 = MockTokenProvider()
+    token_provider_2._fetch_token = Mock(return_value=("token_for_rds_2", time.time() + 3600))
+
+    # Create connection args for two different RDS instances
+    conn_args_1 = PostgresConnectionArgs(
+        application_name="test_rds_1",
+        username="user1",
+        host="rds-instance-1.amazonaws.com",
+        port=5432,
+    )
+
+    conn_args_2 = PostgresConnectionArgs(
+        application_name="test_rds_2",
+        username="user2",
+        host="rds-instance-2.amazonaws.com",
+        port=5432,
+    )
+
+    # Create two connection pool managers with different token providers
+    pool_manager_1 = LRUConnectionPoolManager(
+        max_db=2,
+        base_conn_args=conn_args_1,
+        token_provider=token_provider_1,
+    )
+
+    pool_manager_2 = LRUConnectionPoolManager(
+        max_db=2,
+        base_conn_args=conn_args_2,
+        token_provider=token_provider_2,
+    )
+
+    # Mock the ConnectionPool to capture what connection_class is passed
+    # and simulate connections to verify the token provider behavior
+    captured_pools = []
+
+    def mock_ConnectionPool(*args, **pool_kwargs):
+        # Capture the kwargs and connection_class
+        connection_class = pool_kwargs.get('connection_class')
+        conn_kwargs = pool_kwargs.get('kwargs', {})
+
+        captured_pools.append(
+            {
+                'connection_class': connection_class,
+                'kwargs': conn_kwargs.copy() if conn_kwargs else {},
+            }
+        )
+
+        # Create a mock pool
+        mock_pool = MagicMock()
+        return mock_pool
+
+    with patch('datadog_checks.postgres.connection_pool.ConnectionPool', side_effect=mock_ConnectionPool):
+        # Create pools
+        pool_manager_1._create_pool("db1")
+        pool_manager_2._create_pool("db2")
+
+    # Verify that pools were created
+    assert len(captured_pools) == 2, "Should have created 2 connection pools"
+
+    # Get the pools for each RDS instance
+    pool_1 = captured_pools[0]
+    pool_2 = captured_pools[1]
+
+    # Simulate what happens when a connection is established by calling the connect method
+    # This will trigger the token provider logic
+    conn_class_1 = pool_1['connection_class']
+    conn_class_2 = pool_2['connection_class']
+
+    # Call connect to see what token each would use
+    # We need to mock the parent connect to capture the password
+    with patch('psycopg.Connection.connect') as mock_parent_connect:
+        mock_parent_connect.return_value = MagicMock()
+
+        # Call connect for pool 1 with its kwargs
+        conn_class_1.connect(**pool_1['kwargs'])
+        password_1 = mock_parent_connect.call_args[1].get('password')
+
+        # Call connect for pool 2 with its kwargs
+        conn_class_2.connect(**pool_2['kwargs'])
+        password_2 = mock_parent_connect.call_args[1].get('password')
+
+    # Verify no collision - each pool manager uses its own token provider
+    assert password_1 == "token_for_rds_1"
+    assert password_2 == "token_for_rds_2"
+
+    # Verify each token provider was called independently
+    assert token_provider_1._fetch_token.call_count >= 1, "Token provider 1 should be called at least once"
+    assert token_provider_2._fetch_token.call_count >= 1, "Token provider 2 should be called at least once"
+
+    # Clean up
+    pool_manager_1.close_all()
+    pool_manager_2.close_all()
+
+
 class MockTokenProvider(TokenProvider):
     """Mock implementation of TokenProvider for testing."""