Support OAuth2 Workload Identity (federated credentials) for Entra ID #9536
Conversation
WalkthroughAdds OIDC support and ID-token-first logic, implements workload identity (AKS) client-assertion flow, tightens PKCE/public/confidential validation and registration behavior, guards auth-source mutation to explicit POST button actions, and converts OAuth2 tests to a data-driven suite covering many scenarios. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant Browser as Browser (Login Page)
participant Server as pgAdmin Server
participant Provider as OAuth2/OIDC Provider
participant UserInfo as UserInfo Endpoint
Note over Browser,Server: Explicit login POST required to mutate auth sources
User->>Browser: Click "Login with OAuth2"
Browser->>Server: POST /authenticate (oauth2_button)
Server->>Provider: Authorization Request (PKCE if configured)
Provider-->>Browser: Redirect back with code
Browser->>Server: Callback with code
alt OIDC with sufficient ID token claims
Server->>Provider: Token Request
Provider-->>Server: ID Token + Access Token
Server->>Server: Extract ID token claims, resolve username
else Insufficient ID token claims or non-OIDC
Server->>Provider: Token Request
Provider-->>Server: Access/ID Tokens
Server->>UserInfo: GET /userinfo
UserInfo-->>Server: Profile JSON
Server->>Server: Resolve username from profile or claims
end
Server-->>Browser: Session established
sequenceDiagram
participant Server as pgAdmin Server
participant FS as Workload Identity Token File
participant Provider as OAuth2 Token Endpoint
Note over Server: Workload Identity (AKS) client_assertion flow
Server->>FS: Read projected service account JWT (token file)
FS-->>Server: JWT assertion
Server->>Provider: Token Request (client_assertion & type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer)
Provider-->>Server: Access/ID Tokens
Server->>Server: Extract claims / resolve identity
sequenceDiagram
actor User
participant Browser as Browser (Login Page)
participant Server as pgAdmin Server
Note over Server: update_auth_sources only on POST with explicit button
User->>Browser: Click specific login button
Browser->>Server: POST with `internal_button` or `oauth2_button`
alt internal_button
Server->>Server: Remove KERBEROS, OAUTH2 from auth_sources
else oauth2_button
Server->>Server: Remove INTERNAL, LDAP from auth_sources
end
Server-->>Browser: Continue selected auth flow
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧹 Recent nitpick comments
📜 Recent review detailsConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (1)
🧰 Additional context used🧬 Code graph analysis (2)web/pgadmin/authenticate/__init__.py (1)
web/pgadmin/browser/tests/test_oauth2_with_mocking.py (1)
🪛 Ruff (0.14.11)web/pgadmin/authenticate/oauth2.py157-163: Avoid specifying long messages outside the exception class (TRY003) 166-171: Avoid specifying long messages outside the exception class (TRY003) 177-182: Avoid specifying long messages outside the exception class (TRY003) 186-192: Avoid specifying long messages outside the exception class (TRY003) 199-205: Avoid specifying long messages outside the exception class (TRY003) 230-230: Possible hardcoded password assigned to: "token_endpoint_auth_method" (S105) 239-243: Avoid specifying long messages outside the exception class (TRY003) 247-252: Avoid specifying long messages outside the exception class (TRY003) 258-261: Avoid specifying long messages outside the exception class (TRY003) 295-295: Unused method argument: (ARG002) 438-438: Unused method argument: (ARG002) web/pgadmin/browser/tests/test_oauth2_with_mocking.py140-142: Avoid specifying long messages outside the exception class (TRY003) 320-320: Do not call Replace (B010) 356-356: Do not catch blind exception: (BLE001) 373-373: Unused function argument: (ARG001) 408-408: Unused function argument: (ARG001) 🔇 Additional comments (17)
✏️ Tip: You can disable this entire section by setting Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
web/pgadmin/authenticate/oauth2.py (2)
141-146: Consider simplifying PKCE detection.The
any()with a list works but could be simplified:♻️ Optional simplification
- pkce_is_configured = any( - [ - pkce_method is not None, - pkce_response_type is not None - ] - ) + pkce_is_configured = ( + pkce_method is not None or + pkce_response_type is not None + )
645-656: Verify handling when userinfo endpoint is missing for non-OIDC providers.When
OAUTH2_USERINFO_ENDPOINTis not configured for a non-OIDC provider, the method returns an empty dict{}. This will cause the login to fail downstream with "No profile data found." which is appropriate, but the error message at line 654-655 could be more actionable.Consider enhancing the error message to guide users:
♻️ Optional improvement
else: current_app.logger.error( - 'OAUTH2_USERINFO_ENDPOINT not configured for ' - 'non-OIDC provider') + 'OAUTH2_USERINFO_ENDPOINT is required for non-OIDC ' + f'provider "{self.oauth2_current_client}". ' + 'Please configure the userinfo endpoint.')web/pgadmin/browser/tests/test_oauth2_with_mocking.py (1)
356-360: Consider catching a more specific exception type.The broad
Exceptioncatch is intentional here to verify external redirect behavior, but catchingRuntimeErroror the specific werkzeug exception would be more precise.♻️ Optional improvement
- except Exception as e: + except RuntimeError as e: self.assertEqual( 'Following external redirects is not supported.', str(e) )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
docs/en_US/oauth2.rstweb/pgadmin/authenticate/__init__.pyweb/pgadmin/authenticate/oauth2.pyweb/pgadmin/browser/tests/__init__.pyweb/pgadmin/browser/tests/test_oauth2_with_mocking.py
🧰 Additional context used
🪛 Ruff (0.14.10)
web/pgadmin/browser/tests/test_oauth2_with_mocking.py
140-142: Avoid specifying long messages outside the exception class
(TRY003)
320-320: Do not call setattr with a constant attribute value. It is not any safer than normal property access.
Replace setattr with assignment
(B010)
356-356: Do not catch blind exception: Exception
(BLE001)
373-373: Unused function argument: self
(ARG001)
408-408: Unused function argument: self
(ARG001)
web/pgadmin/authenticate/oauth2.py
157-163: Avoid specifying long messages outside the exception class
(TRY003)
166-171: Avoid specifying long messages outside the exception class
(TRY003)
177-182: Avoid specifying long messages outside the exception class
(TRY003)
186-192: Avoid specifying long messages outside the exception class
(TRY003)
199-205: Avoid specifying long messages outside the exception class
(TRY003)
230-230: Possible hardcoded password assigned to: "token_endpoint_auth_method"
(S105)
239-243: Avoid specifying long messages outside the exception class
(TRY003)
247-252: Avoid specifying long messages outside the exception class
(TRY003)
258-261: Avoid specifying long messages outside the exception class
(TRY003)
295-295: Unused method argument: form
(ARG002)
438-438: Unused method argument: form
(ARG002)
🔇 Additional comments (13)
web/pgadmin/authenticate/__init__.py (1)
229-250: Well-structured guard for authentication source mutation.The POST-check guard prevents unintended side-effects during non-form-button requests. The early returns after handling each button type are clean and prevent fall-through.
One minor observation: the comment block (lines 230-235) clearly documents the rationale, which is helpful for maintainability.
web/pgadmin/browser/tests/__init__.py (1)
14-22: LGTM - placeholder test lifecycle methods.These no-op
setUpandtearDownmethods serve as smoke/placeholder hooks for test discovery. Consider usingpassinstead ofreturnfor slightly more idiomatic Python in methods that don't need to return anything, but this is a minor style preference.docs/en_US/oauth2.rst (2)
129-216: Comprehensive workload identity documentation.The Azure Entra ID Workload Identity section is thorough, covering:
- Configuration parameters
- App registration setup
- Federated credential configuration
- AKS ServiceAccount example
This provides clear guidance for users deploying pgAdmin on AKS with federated credentials.
273-302: Clear username resolution documentation.The username resolution order for OIDC providers is well-documented:
- Custom
OAUTH2_USERNAME_CLAIM(if configured)preferred_usernameclaimsubclaim (last resort)This matches the implementation in
oauth2.pyand provides clear guidance for users.web/pgadmin/authenticate/oauth2.py (5)
107-109: LGTM - Standard JWT bearer assertion type constant.The constant follows the OAuth2 RFC 7523 specification for JWT client assertions.
236-263: Secure token file handling with proper validation.The method correctly:
- Validates the path configuration
- Expands environment variables and user paths
- Checks file existence before reading
- Validates the token content is non-empty
Reading the token fresh on each call is appropriate since Kubernetes projected tokens rotate.
350-436: Well-implemented OIDC-aware username resolution.The resolution logic correctly implements the documented priority:
- Custom
OAUTH2_USERNAME_CLAIM(ID token → userinfo)- OIDC hierarchy:
preferred_username→sub- Non-OIDC: email only
The debug logging at each decision point aids troubleshooting.
438-442: Good defensive guard for missing provider.The early return with a clear error message prevents confusing downstream errors if
oauth2_current_clientis not set.
666-671: Good defensive form access.Using
request.form.get('oauth2_button')instead of direct key access preventsBadRequestKeyErrorwhen the OAuth2 button wasn't clicked, providing a cleaner error path.web/pgadmin/browser/tests/test_oauth2_with_mocking.py (4)
309-325: Good test isolation via state reset.The
_reset_oauth2_statemethod properly clears:
AuthSourceRegistry._objects- Per-app
_pgadmin_auth_sourcescacheOAuth2Authenticationclass-level cachesThis prevents test pollution between scenarios.
Regarding the static analysis hint about
setattrat line 320: this is necessary because the attribute may not exist, and the current approach is correct for initializing a potentially missing attribute.
362-397: Test helper methods are correctly structured.The fake
_fake_authenticateand_fake_get_user_profilemethods correctly simulate OAuth2 flows. Theselfparameter is required since these replace instance methods, so the static analysis warnings are false positives.The test properly simulates the OIDC token response by populating
session['oauth2_token']['userinfo']with ID token claims.
617-664: Comprehensive workload identity client assertion test.The test correctly verifies that:
- The token file is read via
mock_open- The whitespace is stripped from the token
authorize_access_tokenis called with the correctclient_assertion_typeandclient_assertionparameters
25-127: Excellent test scenario coverage.The data-driven scenarios comprehensively cover:
- OAuth2 external redirects and login flows
- PKCE for confidential and public clients
- Workload identity registration and token exchange
- OIDC claim extraction from ID tokens
- Additional claims authorization (accept/reject)
- Userinfo endpoint skip/call decisions
This provides good confidence in the OAuth2/OIDC implementation.
|
@Pyodin |
Pyodin
force-pushed
the
entra-workload-identity
branch
from
ae88afa to
f80c276
Compare
|
@yogeshmahajan-1903 Ok you're right. I squashed the commit from the two previous PR #9525 and #9534 and force pushed them into this PR. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@docs/en_US/oauth2.rst`:
- Around line 190-212: The YAML example block is indented too far causing
copy-paste errors; remove the leading spaces so the YAML content (starting from
"apiVersion: v1") is flush left (column 0) inside the code block, keeping the
surrounding ".. code-block:: yaml" marker and maintaining the same lines for
ServiceAccount, metadata (name/namespace/annotations), and the
Deployment/spec/template/metadata/labels/serviceAccountName entries so those
keys (e.g., "apiVersion", "kind", "metadata", "annotations",
"azure.workload.identity/client-id", "Deployment", "labels",
"azure.workload.identity/use", "serviceAccountName") begin at column 0.
🧹 Nitpick comments (2)
web/pgadmin/authenticate/oauth2.py (1)
139-213: Inconsistent handling of partial PKCE configuration.
pkce_is_configuredisTrueif eitherOAUTH2_CHALLENGE_METHODorOAUTH2_RESPONSE_TYPEis set (line 141-146). However:
- The public client validation (line 197) correctly requires both to be set (
pkce_method and pkce_response_type == 'code')- But the kwargs population (lines 209-213) only checks
pkce_is_configured, so partial PKCE config will add incomplete kwargs (e.g.,code_challenge_method='S256'withresponse_type=None)This could lead to confusing behavior where the OAuth provider receives malformed PKCE parameters.
♻️ Suggested fix: Only add PKCE kwargs when both fields are properly configured
- if pkce_is_configured: + if pkce_method and pkce_response_type: client_kwargs.update({ 'code_challenge_method': pkce_method, 'response_type': pkce_response_type })web/pgadmin/browser/tests/test_oauth2_with_mocking.py (1)
349-360: Broad exception catch is intentional but could be more robust.The test catches
Exceptionbecause Flask's test client raises a generic exception for external redirects. Consider usingself.assertRaisesRegexfor cleaner assertion if the test framework supports it, or add a fallbackself.fail()to catch unexpected exceptions.♻️ Alternative using assertRaisesRegex
- try: - self.tester.login( - email=None, password=None, - _follow_redirects=True, - headers=None, - extra_form_data=dict(oauth2_button=provider) - ) - except Exception as e: - self.assertEqual( - 'Following external redirects is not supported.', - str(e) - ) + with self.assertRaisesRegex( + Exception, + r'Following external redirects is not supported\.' + ): + self.tester.login( + email=None, password=None, + _follow_redirects=True, + headers=None, + extra_form_data=dict(oauth2_button=provider) + )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
docs/en_US/oauth2.rstweb/pgadmin/authenticate/__init__.pyweb/pgadmin/authenticate/oauth2.pyweb/pgadmin/browser/tests/__init__.pyweb/pgadmin/browser/tests/test_oauth2_with_mocking.py
🧰 Additional context used
🧬 Code graph analysis (2)
web/pgadmin/authenticate/__init__.py (1)
web/pgadmin/tools/user_management/__init__.py (1)
auth_sources(499-507)
web/pgadmin/authenticate/oauth2.py (2)
web/pgadmin/authenticate/internal.py (5)
get_friendly_name(40-41)get_friendly_name(95-96)validate(47-62)validate(98-105)login(64-83)web/pgadmin/authenticate/__init__.py (3)
validate(273-283)login(94-100)login(309-323)
🪛 Ruff (0.14.11)
web/pgadmin/authenticate/oauth2.py
157-163: Avoid specifying long messages outside the exception class
(TRY003)
166-171: Avoid specifying long messages outside the exception class
(TRY003)
177-182: Avoid specifying long messages outside the exception class
(TRY003)
186-192: Avoid specifying long messages outside the exception class
(TRY003)
199-205: Avoid specifying long messages outside the exception class
(TRY003)
230-230: Possible hardcoded password assigned to: "token_endpoint_auth_method"
(S105)
239-243: Avoid specifying long messages outside the exception class
(TRY003)
247-252: Avoid specifying long messages outside the exception class
(TRY003)
258-261: Avoid specifying long messages outside the exception class
(TRY003)
295-295: Unused method argument: form
(ARG002)
438-438: Unused method argument: form
(ARG002)
web/pgadmin/browser/tests/test_oauth2_with_mocking.py
140-142: Avoid specifying long messages outside the exception class
(TRY003)
320-320: Do not call setattr with a constant attribute value. It is not any safer than normal property access.
Replace setattr with assignment
(B010)
356-356: Do not catch blind exception: Exception
(BLE001)
373-373: Unused function argument: self
(ARG001)
408-408: Unused function argument: self
(ARG001)
🔇 Additional comments (12)
web/pgadmin/browser/tests/__init__.py (1)
14-24: LGTM!The added
setUpandtearDownlifecycle hooks are appropriate placeholders for test discovery compatibility. The docstring clearly explains the intent.docs/en_US/oauth2.rst (1)
218-345: Documentation is comprehensive and well-aligned with the implementation.The OIDC configuration examples, username resolution order, and legacy OAuth2 fallback behavior are clearly documented and match the code behavior in
oauth2.py.web/pgadmin/authenticate/oauth2.py (5)
236-284: Workload identity implementation is solid.The helper methods properly validate the token file and handle the client assertion flow. The early existence check in
__init__combined with runtime validation provides good defense-in-depth.One consideration: if the token file becomes unreadable (permissions change) after startup, the
open()call at line 254 will raisePermissionError. This is acceptable since it would be an operational issue requiring admin intervention.
350-436: OIDC username resolution logic is well-implemented.The fallback chain (custom claim → email → preferred_username → sub) matches the documentation and handles edge cases appropriately. The debug logging will help troubleshoot authentication issues in production.
438-578: Login flow properly handles OIDC and non-OIDC providers.The early guard for missing
oauth2_current_clientprevents confusing errors downstream. The asymmetric profile requirements (OIDC can use ID token OR userinfo; non-OIDC requires userinfo) are correctly implemented. Audit logging includes claim names without values, which is appropriate for security.
598-664: OIDC userinfo optimization is well-reasoned.The logic correctly determines when the userinfo endpoint can be skipped based on ID token claim sufficiency. This reduces unnecessary network calls while ensuring required claims are available.
666-678: Safer form access prevents BadRequestKeyError.Using
request.form.get()with a fallback error message is a good defensive improvement over direct dictionary access.web/pgadmin/browser/tests/test_oauth2_with_mocking.py (4)
25-128: Comprehensive test scenario coverage.The 17 scenarios cover the key OAuth2/OIDC flows including edge cases like rejected additional claims and missing token files. The kind-based dispatch pattern improves test organization and readability.
309-326: Test isolation logic is necessary but tightly coupled.The state reset correctly clears singleton caches for test isolation. The direct access to
AuthSourceRegistry._objects(line 312) couples the test to internal implementation details. This is acceptable for test code but may require updates if the registry implementation changes.
362-432: Login success and failure test helpers are well-structured.The mock functions correctly simulate OAuth2 authentication flows. The
selfparameter is required for method signature compatibility withpatch.object, so the static analysis warning about unused arguments is a false positive.
579-694: Workload identity tests provide excellent coverage.The tests verify:
- Registration without client secret or PKCE (lines 579-616)
- Client assertion flow with correct JWT bearer assertion (lines 617-664)
- Fast-fail on missing token file (lines 666-694)
The
mock_openusage correctly simulates the token file read, and the assertion value verification confirms whitespace is properly trimmed.web/pgadmin/authenticate/__init__.py (1)
229-250: Guards for auth source mutation look correct.The logic properly ensures auth sources are only filtered when an explicit button is clicked. The early return on non-POST requests prevents unintended mutations.
Note:
kerberos_buttondoes not exist in the codebase. While KERBEROS is a valid auth source, onlyinternal_buttonandoauth2_buttonare provided in the login form, so the current handling is complete.Likely an incorrect or invalid review comment.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
Pyodin
force-pushed
the
entra-workload-identity
branch
2 times, most recently
from
0e63dba to
557823c
Compare
…andling and configuration options Change logging level from exception to error for OIDC profile data issues Refactor debug logging in OAuth2 authentication to improve clarity and consistency Add error handling for missing OAuth2 provider and enhance claims processing logic Enhance OIDC ID token handling by implementing JWT parsing and updating tests to mock claims extraction Refactor ID token claims extraction for OIDC providers and update tests to mock userinfo handling Refactor OAuth2 configuration to use get method for optional URLs Enhance OAuth2 documentation and implement PKCE support for public clients in authentication logic Fix typo in OAUTH2 authentication documentation Implement Azure Entra ID Workload Identity authentication support and add corresponding tests
Pyodin
force-pushed
the
entra-workload-identity
branch
from
557823c to
4c106e2
Compare
2 participants
Related to issue: #9292
Depends on #9525 and #9534
This PR adds support for OAuth2 client authentication using federated Workload Identity, allowing pgAdmin to authenticate with Microsoft Entra ID without a client secret when running on Azure (e.g. AKS).
When OAUTH2_CLIENT_AUTH_METHOD is set to workload_identity, pgAdmin authenticates to the token endpoint using a JWT client assertion read from a projected service account token, as supported by Entra ID federated credentials.
The implementation is standards-based (client_assertion, JWT bearer), backward-compatible with existing OAuth2 configurations, and requires no Azure-specific SDK.
Summary by CodeRabbit
New Features
Documentation
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.