fix: enhance Google Drive OAuth error handling and re-authentication logic#1751
fix: enhance Google Drive OAuth error handling and re-authentication logic#1751ricofurtado wants to merge 3 commits into
Conversation
WalkthroughThis PR reworks authentication handling: Google Drive OAuth refresh failures are classified and token files removed only for permanent errors; orphan reconciliation now re-authenticates each connector and aborts (returns None) on failed authentication; tests mock the new async authenticate method. ChangesAuthentication Failure Handling
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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 |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/api/connectors.py (1)
188-206:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon’t short-circuit before
authenticate()has a chance to recover the connector.This still returns
Nonewheneverget_connector()yields an instance with a falsey cachedis_authenticatedflag, so the new re-auth path never runs in that case. Elsewhere in this module (connector_sync()/sync_all_connectors()), readiness is determined byawait connector.authenticate()directly, so a valid persisted connection can be treated as unavailable here even though sync would succeed.Suggested fix
connector = await connector_service.get_connector(conn.connection_id) - if not connector or not connector.is_authenticated: + if not connector: logger.info( - "Skipping orphan compute — connection unauthenticated", + "Skipping orphan compute — connector unavailable", connector_type=connector_type, connection_id=conn.connection_id, ) return None # Re-authenticate to refresh any stale cached credentials before # making API calls. get_connector() may return a cached connector # whose access token has since expired (Google tokens last 1 hour). # This mirrors the pattern used in connector_sync. if not await connector.authenticate(): logger.info( "Skipping orphan compute — re-authentication failed",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/api/connectors.py` around lines 188 - 206, The early return that checks connector.is_authenticated is preventing the re-authentication path from running; remove the short-circuit so that after obtaining connector via get_connector() you always call await connector.authenticate() (and only log/return None if that call fails). Update the block that currently inspects connector.is_authenticated to instead proceed to call connector.authenticate() unconditionally, keeping the same logger.info messages (with connector_type and conn.connection_id) when authenticate() returns false.src/connectors/google_drive/oauth.py (1)
135-145:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUse a transient-specific refresh error message.
Non-permanent failures now keep the stored token, but this branch still raises an error that says the token was revoked/expired and tells the user to re-authenticate. Callers like
connector_token()can surface that misleading guidance during a temporary Google or network outage.Suggested fix
error_str = str(e).lower() is_permanent = any( k in error_str for k in ("invalid_grant", "revoked", "unauthorized_client") ) if is_permanent: self._remove_token_file() - raise ValueError( - f"Failed to refresh Google Drive credentials. " - f"The refresh token may have expired or been revoked. " - f"Please re-authenticate: {str(e)}" - ) from e + if is_permanent: + raise ValueError( + "Failed to refresh Google Drive credentials. " + "The refresh token may have expired or been revoked. " + f"Please re-authenticate: {e}" + ) from e + raise ValueError( + "Failed to refresh Google Drive credentials due to a transient error. " + f"Please retry later: {e}" + ) from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/connectors/google_drive/oauth.py` around lines 135 - 145, The error handling in the refresh block uses is_permanent to decide whether to remove the stored token, but always raises a ValueError that tells users to re-authenticate even for transient failures; update the branch so that when is_permanent is True you keep the existing behavior (call self._remove_token_file() and raise a ValueError indicating the refresh token expired/revoked and re-authentication is required), and when is_permanent is False raise a different ValueError that clearly signals a transient refresh failure (network/Google outage) and advises retrying later rather than re-authenticating; adjust the messages raised by the refresh logic (the ValueError raised in this block) so callers like connector_token() will not surface misleading guidance.
🧹 Nitpick comments (1)
tests/unit/api/test_reconcile_orphans_for_connector_type.py (1)
30-34: ⚡ Quick winLet the helper model stale auth state separately from re-auth success.
Right now
_make_connector()forcesis_authenticatedandauthenticate()to agree, so this suite still can’t express the important case where a cached connector starts unauthenticated butauthenticate()refreshes it successfully. Adding that shape would catch the gating bug incompute_orphans_for_connector_type().Suggested fix
-def _make_connector(remote_file_ids, *, authenticated=True, raise_on_list=False): +def _make_connector( + remote_file_ids, + *, + authenticated=True, + authenticate_result=None, + raise_on_list=False, +): connector = MagicMock() connector.is_authenticated = authenticated - connector.authenticate = AsyncMock(return_value=authenticated) + connector.authenticate = AsyncMock( + return_value=authenticated if authenticate_result is None else authenticate_result + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/api/test_reconcile_orphans_for_connector_type.py` around lines 30 - 34, The helper _make_connector forces the initial cached auth state (connector.is_authenticated) and the result of connector.authenticate() to be the same, preventing tests from simulating a stale/expired cached state that is successfully refreshed; change _make_connector to accept separate parameters (e.g., initial_authenticated and authenticate_succeeds or reauth_success) and set connector.is_authenticated = initial_authenticated while making connector.authenticate = AsyncMock(return_value=authenticate_succeeds); update any calls in this test to use the new args when you need an unauthenticated cached connector that re-authenticates successfully and ensure compute_orphans_for_connector_type() tests exercise that path.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/api/connectors.py`:
- Around line 188-206: The early return that checks connector.is_authenticated
is preventing the re-authentication path from running; remove the short-circuit
so that after obtaining connector via get_connector() you always call await
connector.authenticate() (and only log/return None if that call fails). Update
the block that currently inspects connector.is_authenticated to instead proceed
to call connector.authenticate() unconditionally, keeping the same logger.info
messages (with connector_type and conn.connection_id) when authenticate()
returns false.
In `@src/connectors/google_drive/oauth.py`:
- Around line 135-145: The error handling in the refresh block uses is_permanent
to decide whether to remove the stored token, but always raises a ValueError
that tells users to re-authenticate even for transient failures; update the
branch so that when is_permanent is True you keep the existing behavior (call
self._remove_token_file() and raise a ValueError indicating the refresh token
expired/revoked and re-authentication is required), and when is_permanent is
False raise a different ValueError that clearly signals a transient refresh
failure (network/Google outage) and advises retrying later rather than
re-authenticating; adjust the messages raised by the refresh logic (the
ValueError raised in this block) so callers like connector_token() will not
surface misleading guidance.
---
Nitpick comments:
In `@tests/unit/api/test_reconcile_orphans_for_connector_type.py`:
- Around line 30-34: The helper _make_connector forces the initial cached auth
state (connector.is_authenticated) and the result of connector.authenticate() to
be the same, preventing tests from simulating a stale/expired cached state that
is successfully refreshed; change _make_connector to accept separate parameters
(e.g., initial_authenticated and authenticate_succeeds or reauth_success) and
set connector.is_authenticated = initial_authenticated while making
connector.authenticate = AsyncMock(return_value=authenticate_succeeds); update
any calls in this test to use the new args when you need an unauthenticated
cached connector that re-authenticates successfully and ensure
compute_orphans_for_connector_type() tests exercise that path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: c42ddb8d-77fc-4df9-bfe2-5405c0c843bd
📒 Files selected for processing (3)
src/api/connectors.pysrc/connectors/google_drive/oauth.pytests/unit/api/test_reconcile_orphans_for_connector_type.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/api/connectors.py (1)
243-250:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winBroaden this warning beyond “listing failed.”
This
exceptnow also catchesget_connector()andauthenticate()failures, so a revoked/expired OAuth refresh can be logged as a listing problem. That makes the new auth path harder to debug. Please either split auth vs. list failures or rename the message to something connector-access related.Suggested tweak
except Exception as e: logger.warning( - "Skipping orphan compute — listing failed", + "Skipping orphan compute — connector access failed", connector_type=connector_type, connection_id=conn.connection_id, error=str(e), error_type=type(e).__name__, ) return None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/api/connectors.py` around lines 243 - 250, The current except around the connector access flow logs "Skipping orphan compute — listing failed" but also catches failures from get_connector() and authenticate(), obscuring auth errors; update the error handling in the connector access sequence (surrounding get_connector(), authenticate(), and the listing call) to either: 1) split into separate try/except blocks and log distinct messages (e.g., "connector authentication failed" for authenticate() and "resource listing failed" for the listing call), or 2) keep a single except but change the log message to a generic "connector access failed" and include a new context field (e.g., stage="get_connector"|"authenticate"|"list") so callers can see which step failed; modify the logger.warning invocation (the block that currently logs connector_type, connection_id, error, error_type) to include this stage/context or create the separate except handlers around get_connector() and authenticate() accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/api/connectors.py`:
- Around line 243-250: The current except around the connector access flow logs
"Skipping orphan compute — listing failed" but also catches failures from
get_connector() and authenticate(), obscuring auth errors; update the error
handling in the connector access sequence (surrounding get_connector(),
authenticate(), and the listing call) to either: 1) split into separate
try/except blocks and log distinct messages (e.g., "connector authentication
failed" for authenticate() and "resource listing failed" for the listing call),
or 2) keep a single except but change the log message to a generic "connector
access failed" and include a new context field (e.g.,
stage="get_connector"|"authenticate"|"list") so callers can see which step
failed; modify the logger.warning invocation (the block that currently logs
connector_type, connection_id, error, error_type) to include this stage/context
or create the separate except handlers around get_connector() and authenticate()
accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ca0702a4-2566-42a5-b68a-98520a26c0c5
📒 Files selected for processing (1)
src/api/connectors.py
This pull request improves the reliability of connector authentication and error handling, particularly for Google Drive connectors, and enhances test coverage to reflect these changes. The updates ensure that expired credentials are refreshed before performing operations and that error reporting is more detailed. The most important changes are:
Authentication and Credential Handling:
compute_orphans_for_connector_type, ensuring that any cached but expired credentials (e.g., Google tokens) are refreshed, and gracefully skipping processing if re-authentication fails.GoogleDriveOAuth2.load_credentialsto only remove the token file on permanent authentication failures (like revoked or invalid grants), avoiding unnecessary user re-authentication on transient errors.Error Reporting:
compute_orphans_for_connector_typeto include the exception type, making debugging easier.Test Improvements:
authenticatemethod, aligning test behavior with the new authentication flow. [1] [2]lSummary by CodeRabbit
Bug Fixes
Tests