fix: only request a token when near expiration

[strantalis]

Proposed Changes

Checklist

• I have added or updated unit tests
• I have added or updated integration tests (if appropriate)
• I have added or updated documentation

Testing Instructions

[gemini-code-assist]

Hello @strantalis, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello team,

Gemini here to provide a summary of this pull request. The primary goal of this PR, as indicated by the title, is to optimize the Keycloak integration within the Entity Resolution Service (ERS). Previously, the service would perform a client login to Keycloak to obtain an access token on potentially every request that required interaction with Keycloak. This PR changes that behavior to only request a new token when the currently cached token is nil, expired, or nearing its expiration.

This is achieved by adding state to the ERS service struct to hold the Keycloak connector and its expiration time, protected by a mutex for concurrent access. A new internal method getConnector is introduced to manage the token lifecycle, fetching a new token only when necessary. The existing handler functions (ResolveEntities, CreateEntityChainFromJwt, CreateEntityChainsFromTokens) and helper functions now utilize this cached connector instance.

Highlights

• Keycloak Token Caching: The Keycloak connector and its access token are now cached within the ERS service instance. A mutex is used to ensure thread-safe access to the connector.
• Token Refresh Logic: A new getConnector method handles retrieving the connector. It checks if the current token is valid and not nearing expiration (with a 60-second buffer). If the token is missing or expiring, it performs a client login to fetch a new one.
• Refactored Keycloak Interaction: The functions that interact with Keycloak (EntityResolution, CreateEntityChainFromJwt, CreateEntityChainsFromTokens, getEntitiesFromToken, getServiceAccountClient) have been updated to accept the cached connector instance as a parameter, removing the need for them to repeatedly call the old getKCClient function.
• Removed Redundant Login Logic: The getKCClient function, which previously performed a Keycloak client login on every call, has been removed as its functionality is now integrated into the getConnector method.
• Test Updates: Unit tests have been updated to reflect the changes in function signatures, passing a mock or initialized connector where required. The test package name was also changed from keycloak_test to keycloak to allow access to internal types for testing.

Changelog

Click here to see the changelog

• service/entityresolution/keycloak/entity_resolution.go
  • Added sync and time imports (lines 11-12).
  • Added connector and connectorMu fields to KeycloakEntityResolutionService struct (lines 49-50).
  • Changed receiver for ResolveEntities and CreateEntityChainFromJwt methods to pointer (*KeycloakEntityResolutionService) (lines 73, 86).
  • Added connector field to KeyCloakConnector struct (line 124).
  • Added connector parameter to CreateEntityChainFromJwt function (line 131).
  • Added connector parameter to EntityResolution function (line 148).
  • Removed the getKCClient function (lines 337-364).
  • Added connector parameter to getEntitiesFromToken function (line 373).
  • Added connector parameter to getServiceAccountClient function (line 441).
  • Added getConnector method to KeycloakEntityResolutionService for managing the cached Keycloak connection and token refresh (lines 478-513).
• service/entityresolution/keycloak/entity_resolution_test.go
  • Changed package name from keycloak_test to keycloak (line 1).
  • Removed explicit import alias for keycloak package (line 15).
  • Updated return types of testKeycloakConfig and testKeycloakConfigInferID to use local types (KeycloakConfig, InferredIdentityConfig, EntityImpliedFrom) (lines 78, 89, 96).
  • Initialized KeyCloakConnector instances in various test functions (Test_KCEntityResolutionByClientId, Test_KCEntityResolutionByEmail, Test_KCEntityResolutionByUsername, Test_KCEntityResolutionByGroupEmail, Test_KCEntityResolutionNotFoundError, Test_JwtClientAndUsernameClientCredentials, Test_JwtClientAndUsernamePasswordPub, Test_JwtClientAndUsernamePasswordPriv, Test_JwtClientAndUsernameAuthPub, Test_JwtClientAndUsernameAuthPriv, Test_JwtClientAndUsernameImplicitPub, Test_JwtClientAndUsernameImplicitPriv, Test_JwtClientAndClientTokenExchange, Test_JwtMultiple, Test_KCEntityResolutionNotFoundInferEmail, Test_KCEntityResolutionNotFoundInferClientId, Test_KCEntityResolutionNotFoundNotInferUsername) (e.g., lines 159-162).
  • Passed the initialized connector to calls to EntityResolution and CreateEntityChainFromJwt (e.g., lines 164, 341).
  • Updated error message comparison in Test_KCEntityResolutionNotFoundError and Test_KCEntityResolutionNotFoundNotInferUsername to use local ErrGetRetrievalFailed (lines 321, 322, 652, 653).
• service/entityresolution/keycloak/v2/entity_resolution.go
  • Added sync and time imports (lines 11-12).
  • Added connector and connectorMu fields to EntityResolutionServiceV2 struct (lines 48-49).
  • Changed receiver for ResolveEntities and CreateEntityChainsFromTokens methods to pointer (*EntityResolutionServiceV2) (lines 72, 83).
  • Added expiresAt field to Connector struct (line 120).
  • Added connector parameter to CreateEntityChainsFromTokens function (line 127).
  • Added connector parameter to EntityResolution function (line 144).
  • Removed the getKCClient function (lines 337-364).
  • Added connector parameter to getEntitiesFromToken function (line 369).
  • Added connector parameter to getServiceAccountClient function (line 437).
  • Added getConnector method to EntityResolutionServiceV2 for managing the cached Keycloak connection and token refresh (lines 474-509).
• service/entityresolution/keycloak/v2/entity_resolution_test.go
  • Changed package name from keycloak_test to keycloak (line 1).
  • Removed explicit import alias for keycloak package (line 15).
  • Updated return types of testConfig and testConfigInferID to use local types (Config, InferredIdentityConfig, EntityImpliedFrom) (lines 78, 89, 96).
  • Initialized Connector instances in various test functions (Test_KCEntityResolutionByClientId, Test_KCEntityResolutionByEmail, Test_KCEntityResolutionByUsername, Test_KCEntityResolutionByGroupEmail, Test_KCEntityResolutionNotFoundError, Test_JwtClientAndUsernameClientCredentials, Test_JwtClientAndUsernamePasswordPub, Test_JwtClientAndUsernamePasswordPriv, Test_JwtClientAndUsernameAuthPub, Test_JwtClientAndUsernameAuthPriv, Test_JwtClientAndUsernameImplicitPub, Test_JwtClientAndUsernameImplicitPriv, Test_JwtClientAndClientTokenExchange, Test_JwtMultiple, Test_KCEntityResolutionNotFoundInferEmail, Test_KCEntityResolutionNotFoundInferClientId, Test_KCEntityResolutionNotFoundNotInferUsername) (e.g., lines 160-163).
  • Passed the initialized connector to calls to EntityResolution and CreateEntityChainsFromTokens (e.g., lines 164, 340).
  • Updated error message comparison in Test_KCEntityResolutionNotFoundError and Test_KCEntityResolutionNotFoundNotInferUsername to use local ErrGetRetrievalFailed (lines 321, 322, 643, 644).

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

A token in the cache,
Saves a network dash.
Check time, then renew,
Efficiency is true.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces an important optimization for the Keycloak Entity Resolution Service by caching the authentication token and only refreshing it when it's nil, expired, or nearing expiration. This should significantly reduce the number of login requests to Keycloak, improving performance and reducing load.

The core logic in the new getConnector method appears sound, incorporating a mutex for thread-safe access and a refresh buffer for proactive token renewal. The changes are consistently applied across both the v1 and v2 ERS implementations, and method receivers have been correctly updated to pointers to manage the shared connector state.

While the existing tests have been adapted to the new function signatures, the PR description is quite minimal. Providing more details on the changes and testing strategy would be helpful for reviewers.

Overall, this is a valuable improvement. My main suggestion revolves around enhancing test coverage for the new token management logic, as detailed in the specific comment.

Summary of Findings

• Token Caching and Refresh Logic: The PR successfully implements token caching and refresh logic in the getConnector method, only requesting new tokens when necessary. This includes mutex protection for concurrent access and a refresh buffer.
• Method Receiver Updates: Service methods (ResolveEntities, CreateEntityChainFromJwt, etc.) were correctly changed to use pointer receivers (*KeycloakEntityResolutionService) to allow modification of the internal connector state.
• Test Adaptations: Existing unit tests were updated to pass the new KeyCloakConnector / Connector instance to helper functions, accommodating the refactoring.
• Test Coverage for getConnector: The new getConnector method, which contains the core token management logic (caching, expiry checks, refresh), is not unit-tested in isolation. Adding dedicated tests for this method would improve robustness and maintainability. (Commented with 'medium' severity)
• Test Package Renaming (Not Commented - Low Severity): Test packages were renamed (e.g., keycloak_test to keycloak), shifting from a black-box to a white-box testing style. This is an acceptable change, likely to simplify testing of internal components.

Merge Readiness

The pull request introduces a significant and valuable optimization for token handling. The core implementation is well done. However, to ensure the robustness of the new token management logic in getConnector, I recommend adding dedicated unit tests for this critical component as suggested in the review comments. Addressing this medium severity testing concern would increase confidence in the change. Therefore, I suggest these changes be considered before merging. As an AI reviewer, I am not authorized to approve pull requests.

[jakedoublev]

Not worth blocking this, but you could rely on the s.idpConfig.TokenBuffer state within s.getConnector instead of passing it as an argument to each call.

[strantalis]

/backport

[opentdf-automation]

Successfully created backport PR for release/service/v0.6:

• fix: only request a token when near expiration [backport to release/service/v0.6] #2408