Fix SQLAlchemy Entra authentication with connection pools

sdk/postgresql/azure-postgresql-auth/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## Unreleased
+
+### Bugs Fixed
+
+- Fixed SQLAlchemy connection pools failing to create subsequent connections when using Entra authentication.
+
 ## 1.0.2 (2026-04-28)
 
 ### Bugs Fixed

sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/async_entra_connection.py

@@ -41,7 +41,7 @@ def enable_entra_authentication_async(engine: AsyncEngine) -> None:
     @event.listens_for(engine.sync_engine, "do_connect")
     def provide_token(
         dialect: Dialect, conn_rec: Any, cargs: Any, cparams: dict[str, Any]  # pylint: disable=unused-argument
-    ) -> None:
+    ) -> Any:
         """Event handler that provides Entra credentials for each sync connection.
 
         :param dialect: The SQLAlchemy dialect being used.
@@ -53,15 +53,16 @@ def provide_token(
         :param cparams: The keyword connection parameters.
         :type cparams: dict[str, Any]
         """
-        credential = cparams.get("credential", None)
+        connection_params = cparams.copy()
+        credential = connection_params.pop("credential", None)
         if credential is None or not isinstance(credential, (TokenCredential)):
             raise CredentialValueError(
                 "credential is required and must be a TokenCredential. "
                 "Pass it via connect_args={'credential': DefaultAzureCredential()}"
             )
         # Check if credentials are already present
-        has_user = "user" in cparams
-        has_password = "password" in cparams
+        has_user = "user" in connection_params
+        has_password = "password" in connection_params
 
         # Only get Entra credentials if user or password is missing
         if not has_user or not has_password:
@@ -71,10 +72,8 @@ def provide_token(
                 raise EntraConnectionValueError("Could not retrieve Entra credentials") from e
             # Only update missing credentials
             if not has_user and "user" in entra_creds:
-                cparams["user"] = entra_creds["user"]
+                connection_params["user"] = entra_creds["user"]
             if not has_password and "password" in entra_creds:
-                cparams["password"] = entra_creds["password"]
+                connection_params["password"] = entra_creds["password"]
 
-        # Strip helper-only param before DBAPI connect to avoid 'invalid connection option'
-        if "credential" in cparams:
-            del cparams["credential"]
+        return dialect.connect(*cargs, **connection_params)

sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/entra_connection.py

@@ -38,7 +38,7 @@ def enable_entra_authentication(engine: Engine) -> None:
     @event.listens_for(engine, "do_connect")
     def provide_token(
         dialect: Dialect, conn_rec: Any, cargs: Any, cparams: dict[str, Any]  # pylint: disable=unused-argument
-    ) -> None:
+    ) -> Any:
         """Event handler that provides Entra credentials for each connection.
 
         :param dialect: The SQLAlchemy dialect being used.
@@ -50,15 +50,16 @@ def provide_token(
         :param cparams: The keyword connection parameters.
         :type cparams: dict[str, Any]
         """
-        credential = cparams.get("credential", None)
+        connection_params = cparams.copy()
+        credential = connection_params.pop("credential", None)
         if credential is None or not isinstance(credential, (TokenCredential)):
             raise CredentialValueError(
                 "credential is required and must be a TokenCredential. "
                 "Pass it via connect_args={'credential': DefaultAzureCredential()}"
             )
         # Check if credentials are already present
-        has_user = "user" in cparams
-        has_password = "password" in cparams
+        has_user = "user" in connection_params
+        has_password = "password" in connection_params
 
         # Only get Entra credentials if user or password is missing
         if not has_user or not has_password:
@@ -68,11 +69,8 @@ def provide_token(
                 raise EntraConnectionValueError("Could not retrieve Entra credentials") from e
             # Only update missing credentials
             if not has_user and "user" in entra_creds:
-                cparams["user"] = entra_creds["user"]
+                connection_params["user"] = entra_creds["user"]
             if not has_password and "password" in entra_creds:
-                cparams["password"] = entra_creds["password"]
+                connection_params["password"] = entra_creds["password"]
 
-        # Remove the helper-only parameter so the DBAPI (psycopg/psycopg2) doesn't see an
-        # unknown connection option and raise 'invalid connection option "credential"'.
-        if "credential" in cparams:
-            del cparams["credential"]
+        return dialect.connect(*cargs, **connection_params)

sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy.py

@@ -39,10 +39,13 @@ def test_provides_entra_credentials(self, mock_get_conninfo):
 
         credential = MockTokenCredential(token)
         cparams = {"credential": credential}
-        handler(MagicMock(), MagicMock(), [], cparams)
+        dialect = MagicMock()
+        dbapi_connection = MagicMock()
+        dialect.connect.return_value = dbapi_connection
+        assert handler(dialect, MagicMock(), [], cparams) is dbapi_connection
         mock_get_conninfo.assert_called_once_with(credential)
-        assert cparams["user"] == TEST_USERS["ENTRA_USER"]
-        assert cparams["password"] == token
+        dialect.connect.assert_called_once_with(user=TEST_USERS["ENTRA_USER"], password=token)
+        assert cparams == {"credential": credential}
 
     def test_missing_credential_raises_error(self):
         """Test that the event handler raises CredentialValueError when no credential."""
@@ -62,8 +65,8 @@ def test_invalid_credential_raises_error(self, credential_value):
             handler(MagicMock(), MagicMock(), [], {"credential": credential_value})
 
     @patch(f"{SYNC_MODULE}.get_entra_conninfo")
-    def test_credential_removed_from_cparams(self, mock_get_conninfo):
-        """Test that the credential parameter is removed before DBAPI connect."""
+    def test_credential_is_preserved_for_subsequent_connections(self, mock_get_conninfo):
+        """Test that pooled connections can reuse the credential configuration."""
         token = create_valid_jwt_token(TEST_USERS["ENTRA_USER"])
         mock_get_conninfo.return_value = {
             "user": TEST_USERS["ENTRA_USER"],
@@ -73,10 +76,13 @@ def test_credential_removed_from_cparams(self, mock_get_conninfo):
         cparams = {"credential": credential}
 
         handler, _, _ = capture_event_handler(enable_entra_authentication, SYNC_MODULE)
-        handler(MagicMock(), MagicMock(), [], cparams)
-        assert "credential" not in cparams
-        assert cparams["user"] == TEST_USERS["ENTRA_USER"]
-        assert cparams["password"] == token
+        dialect = MagicMock()
+        handler(dialect, MagicMock(), [], cparams)
+        handler(dialect, MagicMock(), [], cparams)
+
+        assert cparams == {"credential": credential}
+        assert mock_get_conninfo.call_count == 2
+        assert dialect.connect.call_count == 2
 
     @patch(f"{SYNC_MODULE}.get_entra_conninfo")
     def test_existing_credentials_preserved(self, mock_get_conninfo):
@@ -87,12 +93,12 @@ def test_existing_credentials_preserved(self, mock_get_conninfo):
         handler, _, _ = capture_event_handler(enable_entra_authentication, SYNC_MODULE)
 
         cparams = {"credential": credential, "user": "existing", "password": "secret"}
-        handler(MagicMock(), MagicMock(), [], cparams)
+        dialect = MagicMock()
+        handler(dialect, MagicMock(), [], cparams)
 
         mock_get_conninfo.assert_not_called()
-        assert cparams["user"] == "existing"
-        assert cparams["password"] == "secret"
-        assert "credential" not in cparams
+        dialect.connect.assert_called_once_with(user="existing", password="secret")
+        assert cparams == {"credential": credential, "user": "existing", "password": "secret"}
 
     @pytest.mark.parametrize(
         "cparams_extra,should_call",

sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py

@@ -41,10 +41,13 @@ def test_provides_entra_credentials_async(self, mock_get_conninfo):
 
         credential = MockTokenCredential(token)
         cparams = {"credential": credential}
-        handler(MagicMock(), MagicMock(), [], cparams)
+        dialect = MagicMock()
+        dbapi_connection = MagicMock()
+        dialect.connect.return_value = dbapi_connection
+        assert handler(dialect, MagicMock(), [], cparams) is dbapi_connection
         mock_get_conninfo.assert_called_once_with(credential)
-        assert cparams["user"] == TEST_USERS["ENTRA_USER"]
-        assert cparams["password"] == token
+        dialect.connect.assert_called_once_with(user=TEST_USERS["ENTRA_USER"], password=token)
+        assert cparams == {"credential": credential}
 
     def test_missing_credential_raises_error_async(self):
         """Test that the event handler raises CredentialValueError when no credential."""
@@ -64,8 +67,8 @@ def test_invalid_credential_raises_error_async(self, credential_value):
             handler(MagicMock(), MagicMock(), [], {"credential": credential_value})
 
     @patch(f"{ASYNC_MODULE}.get_entra_conninfo")
-    def test_credential_removed_from_cparams_async(self, mock_get_conninfo):
-        """Test that the credential parameter is removed before DBAPI connect (async)."""
+    def test_credential_is_preserved_for_subsequent_connections_async(self, mock_get_conninfo):
+        """Test that pooled connections can reuse the credential configuration."""
         token = create_valid_jwt_token(TEST_USERS["ENTRA_USER"])
         mock_get_conninfo.return_value = {
             "user": TEST_USERS["ENTRA_USER"],
@@ -75,10 +78,13 @@ def test_credential_removed_from_cparams_async(self, mock_get_conninfo):
         cparams = {"credential": credential}
 
         handler, _, _ = capture_event_handler(enable_entra_authentication_async, ASYNC_MODULE)
-        handler(MagicMock(), MagicMock(), [], cparams)
-        assert "credential" not in cparams
-        assert cparams["user"] == TEST_USERS["ENTRA_USER"]
-        assert cparams["password"] == token
+        dialect = MagicMock()
+        handler(dialect, MagicMock(), [], cparams)
+        handler(dialect, MagicMock(), [], cparams)
+
+        assert cparams == {"credential": credential}
+        assert mock_get_conninfo.call_count == 2
+        assert dialect.connect.call_count == 2
 
     @patch(f"{ASYNC_MODULE}.get_entra_conninfo")
     def test_existing_credentials_preserved_async(self, mock_get_conninfo):
@@ -89,12 +95,12 @@ def test_existing_credentials_preserved_async(self, mock_get_conninfo):
         handler, _, _ = capture_event_handler(enable_entra_authentication_async, ASYNC_MODULE)
 
         cparams = {"credential": credential, "user": "existing", "password": "secret"}
-        handler(MagicMock(), MagicMock(), [], cparams)
+        dialect = MagicMock()
+        handler(dialect, MagicMock(), [], cparams)
 
         mock_get_conninfo.assert_not_called()
-        assert cparams["user"] == "existing"
-        assert cparams["password"] == "secret"
-        assert "credential" not in cparams
+        dialect.connect.assert_called_once_with(user="existing", password="secret")
+        assert cparams == {"credential": credential, "user": "existing", "password": "secret"}
 
     @pytest.mark.parametrize(
         "cparams_extra,should_call",