Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[FEATURE REQ] azure-identity-extensions for PostgreSQL is inneficient #39393

Open
paulaustin-automutatio opened this issue Mar 25, 2024 · 13 comments · Fixed by #43659
Open

[FEATURE REQ] azure-identity-extensions for PostgreSQL is inneficient #39393

paulaustin-automutatio opened this issue Mar 25, 2024 · 13 comments · Fixed by #43659
Assignees
@moarychan
Labels
azure-spring All azure-spring related issues Client This issue points to a problem in the data-plane of the library. customer-reported Issues that are reported by GitHub users external to the Azure organization. needs-team-attention Workflow: This issue needs attention from Azure service team or SDK team
Milestone
2025-01

Comments

@paulaustin-automutatio
Copy link

https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/identity/azure-identity-extensions

I just wanted to let you know that the way the above is implemented for PostgreSQL is inefficient.

This is due to the way that the postgreSQL driver creates an instance of the authenticator plugin for each connection request.

This means that a new token is requested for each connection request.

I'm just reporting this as a courtesy as I'm not going to be using this code
@github-actions github-actions bot added Azure.Identity Client This issue points to a problem in the data-plane of the library. customer-reported Issues that are reported by GitHub users external to the Azure organization. needs-team-attention Workflow: This issue needs attention from Azure service team or SDK team question The issue doesn't require a change to the product in order to be resolved. Most issues start as that labels Mar 25, 2024
@github-actions GitHub Actions
Copy link

@billwert @g2vinay

@github-actions GitHub Actions
Copy link

Thank you for your feedback. Tagging and routing to the team member best able to assist.

@g2vinay
Copy link
Member

g2vinay commented Jul 29, 2024

@paulaustin-automutatio
Thank you for reporting this, would be great if you could point out to the code as well.

Assigning to @saragluna to take a look.

@g2vinay g2vinay assigned saragluna and unassigned g2vinay Jul 29, 2024
@paulaustin-automutatio
Copy link
Author

Have a look at com.azure.identity.extensions.jdbc.postgresql.AzurePostgresqlAuthenticationPlugin

Put a breakpoint in the constructor and you will see that it is created for every new database connection. That means it has to get a new token every time as that instance can't be cached. PostgreSQL driver provides no way to cache it.

You'd have to do some kind of caching using a static variable.

@saragluna saragluna added this to the Backlog milestone Aug 19, 2024
@saragluna saragluna added azure-spring All azure-spring related issues and removed question The issue doesn't require a change to the product in order to be resolved. Most issues start as that labels Aug 19, 2024
@g2vinay g2vinay moved this from Untriaged to Backlog in Azure Identity SDK Improvements Oct 16, 2024
@pdefreitas
Copy link

This behavior makes it very unusable in scenarios where we want to use pgBouncer. A new token is requested all the time whenever a new connection is made by an application.

@saragluna
Copy link
Member

We will look into it.

@saragluna saragluna modified the milestones: Backlog, 2024-12 Nov 15, 2024
@pdefreitas pdefreitas mentioned this issue Nov 18, 2024
Open
2 tasks
@saragluna saragluna modified the milestones: 2024-12, 2025-01 Nov 29, 2024
@saragluna saragluna assigned moarychan and unassigned saragluna Nov 29, 2024
@moarychan moarychan moved this from Todo to In Progress in Spring Cloud Azure Dec 27, 2024
@moarychan
Copy link
Member

Hi @paulaustin-automutatio and @pdefreitas , the token is cached in the layer msal4j, below is their dependency relationship:

azure-identity-extensions -> azure-identity -> msal4j

the azure-identity’s cache capability is also reflected in layer msal4j.

For your reference, I pasted the logs for MySQL token acquisition, first time was from internal AAD endpoint and second time was from the cache:

First request responded from internal AAD endpoint:

2024-12-26T06:11:30.059Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : ForceRefresh set to false. Attempting cache lookup
2024-12-26T06:11:30.061Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.m.AadInstanceDiscoveryProvider     : Instance discovery set to false, caching a default entry.
2024-12-26T06:11:30.068Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : Cache lookup failed: Token not found in the cache
2024-12-26T06:11:30.068Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Unable to find AZURE_POD_IDENTITY_AUTHORITY_HOST environment variable for IMDS, using the default endpoint.
2024-12-26T06:11:30.069Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Creating IMDS managed identity source. Endpoint URI: http://169.254.169.254/metadata/identity/oauth2/token
2024-12-26T06:11:30.069Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : [Managed Identity] Managed Identity source and ID type identified and set successfully, request will use Managed Identity for IMDS
2024-12-26T06:11:30.069Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Adding user assigned ID to the request for IMDS Managed Identity.
2024-12-26T06:11:30.069Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.aad.msal4j.ManagedIdentityRequest    : [Managed Identity] Adding user assigned client id to the request.
2024-12-26T06:11:30.438Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] com.microsoft.aad.msal4j.HttpHelper      : [Correlation ID: null] Sent (null) Correlation Id is not same as received (null).
2024-12-26T06:11:30.438Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.m.AbstractManagedIdentitySource    : [Managed Identity] Successful response received.
2024-12-26T06:11:30.521Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.ManagedIdentityApplication  : [Correlation ID: 47e2c671-59b9-4a18-8907-65082e41bda4] Access Token was returned
2024-12-26T06:11:30.522Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => Managed Identity environment: Managed Identity
2024-12-26T06:11:30.522Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => getToken() result for scopes [https://ossrdbms-aad.database.windows.net/.default]: SUCCESS

Second request responded from the cache:

2024-12-26T06:13:23.514Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-3] .m.AcquireTokenByManagedIdentitySupplier : ForceRefresh set to false. Attempting cache lookup
2024-12-26T06:13:23.514Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-3] c.m.a.msal4j.AcquireTokenSilentSupplier  : Returning token from cache
2024-12-26T06:13:23.515Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-3] c.m.a.msal4j.ManagedIdentityApplication  : [Correlation ID: 83f51f18-19a9-43d2-b224-52f1fc83d316] Access Token was returned
2024-12-26T06:13:23.515Z  INFO 13736 --- [ForkJoinPool.commonPool-worker-3] c.a.identity.ManagedIdentityCredential   : Azure Identity => Managed Identity environment: Managed Identity
2024-12-26T06:13:23.515Z DEBUG 13736 --- [ForkJoinPool.commonPool-worker-3] c.a.identity.ManagedIdentityCredential   : Azure Identity => getToken() result for scopes [https://ossrdbms-aad.database.windows.net/.default]: SUCCESS

@pdefreitas
Copy link

@moarychan what you're stating is correct for MySQL it indeed caches the token, but for PostgreSQL it doesn't have the same behavior.
You can compare this in the authentication plugin for MySQL:
https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/jdbc/mysql/AzureMysqlAuthenticationPlugin.java#L75

During our tests we've noticed the token is not cached for PostgreSQL, the implementation shows signs it doesn't have the same caching logic because the MySQL implementation uses the same instance. From our tests, pgjdbc creates an instance of AzurePostgresqlAuthenticationPlugin for each connection request.
login.microsoftonline.com is spammed with /token requests depending how the amount of queries that are ran. To make it unusable, in scenarios, where you using connection bouncers like pgbouncer a new token will be generated every time the application opens a connection to the bouncer.

https://github.com/Azure/azure-sdk-for-java/blob/main/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/jdbc/postgresql/AzurePostgresqlAuthenticationPlugin.java#L54

@moarychan
Copy link
Member

Hi @pdefreitas thanks for your update!

As far as I know, the cache under msal4j should have the same behavior for each authentication plugin implementation. Not familiar with pgbouncer, could you help provide a minimal project to reproduce your issue?

@moarychan moarychan mentioned this issue Jan 3, 2025
Merged
6 tasks
@pdefreitas
Copy link

@moarychan Azure provides pgbouncer on PostgreSQL - Flexible Server Link.
The workload should not be using local connection pooling like c3p0 so that it can freely open connections without reusing them.

@github-project-automation github-project-automation bot moved this from In Progress to Done in Spring Cloud Azure Jan 8, 2025
@moarychan moarychan reopened this Jan 8, 2025
@github-project-automation github-project-automation bot moved this from Done to In Progress in Spring Cloud Azure Jan 8, 2025
@moarychan
Copy link
Member

Hi @pdefreitas , I tried with PGBouncer enabled, which results using jdbc url with another port 6432, not the 5432.
We can see the token acquisition only occurs one time ([Managed Identity] Successful response received.), then cached it, subsequent token acquisition requests return cached token (AcquireTokenSilentSupplier : Returning token from cache).

Below are the logs and removed some unwanted info:

2025-01-07T07:00:38.990Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initializing AzureAuthenticationTemplate.
2025-01-07T07:00:39.087Z DEBUG 5612 --- [main] c.a.identity.ManagedIdentityCredential   : Azure Identity => Found the following environment variables: 
2025-01-07T07:00:39.092Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initialized AzureAuthenticationTemplate.
2025-01-07T07:00:39.522Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : ForceRefresh set to false. Attempting cache lookup
2025-01-07T07:00:39.525Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.m.AadInstanceDiscoveryProvider     : Instance discovery set to false, caching a default entry.
2025-01-07T07:00:39.537Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : Cache lookup failed: Token not found in the cache
2025-01-07T07:00:39.538Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Unable to find AZURE_POD_IDENTITY_AUTHORITY_HOST environment variable for IMDS, using the default endpoint.
2025-01-07T07:00:39.538Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Creating IMDS managed identity source. Endpoint URI: http://169.254.169.254/metadata/identity/oauth2/token
2025-01-07T07:00:39.539Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : [Managed Identity] Managed Identity source and ID type identified and set successfully, request will use Managed Identity for IMDS
2025-01-07T07:00:39.539Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.IMDSManagedIdentitySource   : [Managed Identity] Adding user assigned ID to the request for IMDS Managed Identity.
2025-01-07T07:00:39.540Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.aad.msal4j.ManagedIdentityRequest    : [Managed Identity] Adding user assigned client id to the request.
2025-01-07T07:00:39.956Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] com.microsoft.aad.msal4j.HttpHelper      : [Correlation ID: null] Sent (null) Correlation Id is not same as received (null).
2025-01-07T07:00:39.957Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.m.AbstractManagedIdentitySource    : [Managed Identity] Successful response received.
2025-01-07T07:00:40.032Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.ManagedIdentityApplication  : [Correlation ID: 666ce3b4-6d67-41e1-82b2-e1ba85b3d1d0] Access Token was returned
2025-01-07T07:00:40.033Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => Managed Identity environment: Managed Identity
2025-01-07T07:00:40.034Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => getToken() result for scopes [https://ossrdbms-aad.database.windows.net/.default]: SUCCESS

2025-01-07T07:00:40.083Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [0 - Connection via Default credential provider]  - 0 Connection successful!
2025-01-07T07:00:40.101Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [0 - Connection via Default credential provider] Current number from PostgreSQL: 2025-01-07 07:00:40.094219+00

2025-01-07T07:00:40.127Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initializing AzureAuthenticationTemplate.
2025-01-07T07:00:40.128Z DEBUG 5612 --- [main] c.a.identity.ManagedIdentityCredential   : Azure Identity => Found the following environment variables: 
2025-01-07T07:00:40.128Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initialized AzureAuthenticationTemplate.
2025-01-07T07:00:40.129Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : ForceRefresh set to false. Attempting cache lookup
2025-01-07T07:00:40.130Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.AcquireTokenSilentSupplier  : Returning token from cache
2025-01-07T07:00:40.130Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.ManagedIdentityApplication  : [Correlation ID: a0d20c65-a2e9-4ac5-a278-fe6cf61a3602] Access Token was returned
2025-01-07T07:00:40.130Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => Managed Identity environment: Managed Identity
2025-01-07T07:00:40.131Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => getToken() result for scopes [https://ossrdbms-aad.database.windows.net/.default]: SUCCESS
2025-01-07T07:00:40.135Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [1 - Connection via Default credential provider]  - 1 Connection successful!
2025-01-07T07:00:40.137Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [1 - Connection via Default credential provider] Current number from PostgreSQL: 2025-01-07 07:00:40.137537+00

2025-01-07T07:00:40.160Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initializing AzureAuthenticationTemplate.
2025-01-07T07:00:40.161Z DEBUG 5612 --- [main] c.a.identity.ManagedIdentityCredential   : Azure Identity => Found the following environment variables: 
2025-01-07T07:00:40.162Z DEBUG 5612 --- [main] c.a.i.e.i.t.AzureAuthenticationTemplate  : Initialized AzureAuthenticationTemplate.
2025-01-07T07:00:40.163Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] .m.AcquireTokenByManagedIdentitySupplier : ForceRefresh set to false. Attempting cache lookup
2025-01-07T07:00:40.164Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.AcquireTokenSilentSupplier  : Returning token from cache
2025-01-07T07:00:40.165Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.m.a.msal4j.ManagedIdentityApplication  : [Correlation ID: a5034f8c-b499-4b64-b8a7-73686c79432b] Access Token was returned
2025-01-07T07:00:40.165Z  INFO 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => Managed Identity environment: Managed Identity
2025-01-07T07:00:40.166Z DEBUG 5612 --- [ForkJoinPool.commonPool-worker-1] c.a.identity.ManagedIdentityCredential   : Azure Identity => getToken() result for scopes [https://ossrdbms-aad.database.windows.net/.default]: SUCCESS
2025-01-07T07:00:40.170Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [2 - Connection via Default credential provider]  - 2 Connection successful!
2025-01-07T07:00:40.171Z  INFO 5612 --- [main] o.e.PostgresqlWithAuthenticationPlugin   : [2 - Connection via Default credential provider] Current number from PostgreSQL: 2025-01-07 07:00:40.171459+00

@github-project-automation github-project-automation bot moved this from In Progress to Done in Spring Cloud Azure Jan 8, 2025
@moarychan moarychan reopened this Jan 8, 2025
@github-project-automation github-project-automation bot moved this from Done to In Progress in Spring Cloud Azure Jan 8, 2025
@moarychan
Copy link
Member

Hi @pdefreitas, I want to double confirm which authentication method you used, no managed identity credential is used, right?

@pdefreitas
Copy link

@moarychan we're using workload identities through Microsoft Entra Workload ID. I guess we will have to tweak logging levels to logging and instrumentation, we can try it out when #43659 is available.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@moarychan moarychan

Labels
azure-spring All azure-spring related issues Client This issue points to a problem in the data-plane of the library. customer-reported Issues that are reported by GitHub users external to the Azure organization. needs-team-attention Workflow: This issue needs attention from Azure service team or SDK team
Projects
Spring Cloud Azure
Status: In Progress
Milestone
2025-01
Development

Successfully merging a pull request may close this issue.

Support token credential cache for azure-identity-extension moarychan/azure-sdk-for-java
6 participants
@joshfree @g2vinay @pdefreitas @moarychan @saragluna @paulaustin-automutatio