refactor(be): replace HTTPException with OnyxError in OAuth files

backend/ee/onyx/server/oauth/api.py

@@ -2,7 +2,6 @@ import base64
 import uuid
 
 from fastapi import Depends
-from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 
 from ee.onyx.server.oauth.api_router import router
@@ -13,6 +12,8 @@ from onyx.auth.users import current_admin_user
 from onyx.configs.app_configs import DEV_MODE
 from onyx.configs.constants import DocumentSource
 from onyx.db.models import User
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.redis.redis_pool import get_redis_client
 from onyx.utils.logger import setup_logger
 from shared_configs.contextvars import get_current_tenant_id
@@ -71,15 +72,15 @@ def prepare_authorization_request(
         oauth_url = None
 
     if not oauth_url:
-        raise HTTPException(
-            status_code=404,
-            detail=f"The document source type {connector} does not have OAuth implemented",
+        raise OnyxError(
+            OnyxErrorCode.NOT_FOUND,
+            f"The document source type {connector} does not have OAuth implemented",
         )
 
     if not session:
-        raise HTTPException(
-            status_code=500,
-            detail=f"The document source type {connector} failed to generate an OAuth session.",
+        raise OnyxError(
+            OnyxErrorCode.INTERNAL_ERROR,
+            f"The document source type {connector} failed to generate an OAuth session.",
         )
 
     r = get_redis_client(tenant_id=tenant_id)

backend/ee/onyx/server/oauth/confluence_cloud.py

@@ -8,7 +8,6 @@ from typing import cast
 
 import requests
 from fastapi import Depends
-from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from pydantic import ValidationError
@@ -27,6 +26,8 @@ from onyx.db.credentials import fetch_credential_by_id_for_user
 from onyx.db.credentials import update_credential_json
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.models import User
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.redis.redis_pool import get_redis_client
 from onyx.server.documents.models import CredentialBase
 from onyx.utils.logger import setup_logger
@@ -154,9 +155,9 @@ def confluence_oauth_callback(
     after visiting the oauth authorization url."""
 
     if not ConfluenceCloudOAuth.CLIENT_ID or not ConfluenceCloudOAuth.CLIENT_SECRET:
-        raise HTTPException(
-            status_code=500,
-            detail="Confluence Cloud client ID or client secret is not configured.",
+        raise OnyxError(
+            OnyxErrorCode.INTERNAL_ERROR,
+            "Confluence Cloud client ID or client secret is not configured.",
         )
 
     r = get_redis_client(tenant_id=tenant_id)
@@ -177,9 +178,9 @@ def confluence_oauth_callback(
 
     session_json_bytes = cast(bytes, r.get(r_key))
     if not session_json_bytes:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Confluence Cloud OAuth failed - OAuth state key not found: key={r_key}",
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR,
+            f"Confluence Cloud OAuth failed - OAuth state key not found: key={r_key}",
         )
 
     session_json = session_json_bytes.decode("utf-8")
@@ -268,7 +269,10 @@ def confluence_oauth_accessible_resources(
 
     credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
     if not credential:
-        raise HTTPException(400, f"Credential {credential_id} not found.")
+        raise OnyxError(
+            OnyxErrorCode.CREDENTIAL_NOT_FOUND,
+            f"Credential {credential_id} not found.",
+        )
 
     credential_dict = (
         credential.credential_json.get_value(apply_mask=False)
@@ -336,9 +340,9 @@ def confluence_oauth_finalize(
 
     credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
     if not credential:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Confluence Cloud OAuth failed - credential {credential_id} not found.",
+        raise OnyxError(
+            OnyxErrorCode.CREDENTIAL_NOT_FOUND,
+            f"Confluence Cloud OAuth failed - credential {credential_id} not found.",
         )
 
     existing_credential_json = (

backend/ee/onyx/server/oauth/google_drive.py

@@ -6,7 +6,6 @@ from typing import cast
 
 import requests
 from fastapi import Depends
-from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
@@ -35,6 +34,8 @@ from onyx.connectors.google_utils.shared_constants import (
 from onyx.db.credentials import create_credential
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.models import User
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.redis.redis_pool import get_redis_client
 from onyx.server.documents.models import CredentialBase
 from shared_configs.contextvars import get_current_tenant_id
@@ -119,9 +120,9 @@ def handle_google_drive_oauth_callback(
     tenant_id: str | None = Depends(get_current_tenant_id),
 ) -> JSONResponse:
     if not GoogleDriveOAuth.CLIENT_ID or not GoogleDriveOAuth.CLIENT_SECRET:
-        raise HTTPException(
-            status_code=500,
-            detail="Google Drive client ID or client secret is not configured.",
+        raise OnyxError(
+            OnyxErrorCode.INTERNAL_ERROR,
+            "Google Drive client ID or client secret is not configured.",
         )
 
     r = get_redis_client(tenant_id=tenant_id)
@@ -142,9 +143,9 @@ def handle_google_drive_oauth_callback(
 
     session_json_bytes = cast(bytes, r.get(r_key))
     if not session_json_bytes:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Google Drive OAuth failed - OAuth state key not found: key={r_key}",
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR,
+            f"Google Drive OAuth failed - OAuth state key not found: key={r_key}",
         )
 
     session_json = session_json_bytes.decode("utf-8")

backend/ee/onyx/server/oauth/slack.py

@@ -4,7 +4,6 @@ from typing import cast
 
 import requests
 from fastapi import Depends
-from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
@@ -19,6 +18,8 @@ from onyx.configs.constants import DocumentSource
 from onyx.db.credentials import create_credential
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.models import User
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.redis.redis_pool import get_redis_client
 from onyx.server.documents.models import CredentialBase
 from shared_configs.contextvars import get_current_tenant_id
@@ -103,9 +104,9 @@ def handle_slack_oauth_callback(
     tenant_id: str | None = Depends(get_current_tenant_id),
 ) -> JSONResponse:
     if not SlackOAuth.CLIENT_ID or not SlackOAuth.CLIENT_SECRET:
-        raise HTTPException(
-            status_code=500,
-            detail="Slack client ID or client secret is not configured.",
+        raise OnyxError(
+            OnyxErrorCode.INTERNAL_ERROR,
+            "Slack client ID or client secret is not configured.",
         )
 
     r = get_redis_client(tenant_id=tenant_id)
@@ -126,9 +127,9 @@ def handle_slack_oauth_callback(
 
     session_json_bytes = cast(bytes, r.get(r_key))
     if not session_json_bytes:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Slack OAuth failed - OAuth state key not found: key={r_key}",
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR,
+            f"Slack OAuth failed - OAuth state key not found: key={r_key}",
         )
 
     session_json = session_json_bytes.decode("utf-8")
@@ -155,9 +156,9 @@ def handle_slack_oauth_callback(
         response_data = response.json()
 
         if not response_data.get("ok"):
-            raise HTTPException(
-                status_code=400,
-                detail=f"Slack OAuth failed: {response_data.get('error')}",
+            raise OnyxError(
+                OnyxErrorCode.VALIDATION_ERROR,
+                f"Slack OAuth failed: {response_data.get('error')}",
             )
 
         # Extract token and team information
@@ -173,6 +174,8 @@ def handle_slack_oauth_callback(
         )
 
         create_credential(credential_info, user, db_session)
+    except OnyxError:
+        raise
     except Exception as e:
         return JSONResponse(
             status_code=500,

backend/onyx/server/documents/standard_oauth.py

@@ -5,7 +5,6 @@ from typing import cast
 
 from fastapi import APIRouter
 from fastapi import Depends
-from fastapi import HTTPException
 from fastapi import Query
 from fastapi import Request
 from pydantic import BaseModel
@@ -19,6 +18,8 @@ from onyx.connectors.interfaces import OAuthConnector
 from onyx.db.credentials import create_credential
 from onyx.db.engine.sql_engine import get_session
 from onyx.db.models import User
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
 from onyx.redis.redis_pool import get_redis_client
 from onyx.server.documents.models import CredentialBase
 from onyx.utils.logger import setup_logger
@@ -69,12 +70,10 @@ def _get_additional_kwargs(
         # validate
         connector_cls.AdditionalOauthKwargs(**additional_kwargs_dict)
     except ValidationError:
-        raise HTTPException(
-            status_code=400,
-            detail=(
-                f"Invalid additional kwargs. Got {additional_kwargs_dict}, expected "
-                f"{connector_cls.AdditionalOauthKwargs.model_json_schema()}"
-            ),
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR,
+            f"Invalid additional kwargs. Got {additional_kwargs_dict}, expected "
+            f"{connector_cls.AdditionalOauthKwargs.model_json_schema()}",
         )
 
     return additional_kwargs_dict
@@ -97,7 +96,9 @@ def oauth_authorize(
     oauth_connectors = _discover_oauth_connectors()
 
     if source not in oauth_connectors:
-        raise HTTPException(status_code=400, detail=f"Unknown OAuth source: {source}")
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR, f"Unknown OAuth source: {source}"
+        )
 
     connector_cls = oauth_connectors[source]
     base_url = WEB_DOMAIN
@@ -147,7 +148,9 @@ def oauth_callback(
     oauth_connectors = _discover_oauth_connectors()
 
     if source not in oauth_connectors:
-        raise HTTPException(status_code=400, detail=f"Unknown OAuth source: {source}")
+        raise OnyxError(
+            OnyxErrorCode.VALIDATION_ERROR, f"Unknown OAuth source: {source}"
+        )
 
     connector_cls = oauth_connectors[source]
 
@@ -157,7 +160,7 @@ def oauth_callback(
         bytes, redis_client.get(_OAUTH_STATE_KEY_FMT.format(state=state))
     )
     if not oauth_state_bytes:
-        raise HTTPException(status_code=400, detail="Invalid OAuth state")
+        raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, "Invalid OAuth state")
     oauth_state = json.loads(oauth_state_bytes.decode("utf-8"))
 
     desired_return_url = cast(str, oauth_state[_DESIRED_RETURN_URL_KEY])