diff --git a/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md b/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md index bf401dcbdfe8..b91b77eca1da 100644 --- a/sdk/postgresql/azure-postgresql-auth/CHANGELOG.md +++ b/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 diff --git a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/async_entra_connection.py b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/async_entra_connection.py index 252631b860d4..af5895f6a546 100644 --- a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/async_entra_connection.py +++ b/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. @@ -52,16 +52,19 @@ def provide_token( :type cargs: Any :param cparams: The keyword connection parameters. :type cparams: dict[str, Any] + :return: A DBAPI connection authenticated with Entra credentials. + :rtype: 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 +74,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) diff --git a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/entra_connection.py b/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/entra_connection.py index ad3d4c4d9048..bce39fc29d8b 100644 --- a/sdk/postgresql/azure-postgresql-auth/azure_postgresql_auth/sqlalchemy/entra_connection.py +++ b/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. @@ -49,16 +49,19 @@ def provide_token( :type cargs: Any :param cparams: The keyword connection parameters. :type cparams: dict[str, Any] + :return: A DBAPI connection authenticated with Entra credentials. + :rtype: 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 +71,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) diff --git a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy.py b/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy.py index e04b0faaf1b7..589218c92824 100644 --- a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy.py +++ b/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", diff --git a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py b/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py index 4e479e275c84..12a1475d2945 100644 --- a/sdk/postgresql/azure-postgresql-auth/tests/test_sqlalchemy_async.py +++ b/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",