Compare commits

...

1 Commits

Author SHA1 Message Date
Nik
dbf3e111ff refactor: replace HTTPException with OnyxError in remaining EE server files
Convert 9 EE server files from HTTPException to OnyxError:
- query_history/api.py
- reporting/usage_export_api.py
- user_group/api.py
- documents/cc_pair.py
- enterprise_settings/api.py
- middleware/tenant_tracking.py
- query_and_chat/query_backend.py
- query_and_chat/search_backend.py
- query_and_chat/token_limit.py
2026-03-09 15:49:38 -07:00
9 changed files with 115 additions and 105 deletions

View File

@@ -1,9 +1,7 @@
from datetime import datetime
from http import HTTPStatus
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from sqlalchemy.orm import Session
from ee.onyx.background.celery.tasks.doc_permission_syncing.tasks import (
@@ -19,6 +17,8 @@ from onyx.db.connector_credential_pair import (
)
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_connector import RedisConnector
from onyx.redis.redis_pool import get_redis_client
from onyx.server.models import StatusResponse
@@ -42,9 +42,9 @@ def get_cc_pair_latest_sync(
get_editable=False,
)
if not cc_pair:
raise HTTPException(
status_code=400,
detail="cc_pair not found for current user's permissions",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"cc_pair not found for current user's permissions",
)
return cc_pair.last_time_perm_sync
@@ -66,18 +66,18 @@ def sync_cc_pair(
get_editable=False,
)
if not cc_pair:
raise HTTPException(
status_code=400,
detail="Connection not found for current user's permissions",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"Connection not found for current user's permissions",
)
r = get_redis_client()
redis_connector = RedisConnector(tenant_id, cc_pair_id)
if redis_connector.permissions.fenced:
raise HTTPException(
status_code=HTTPStatus.CONFLICT,
detail="Permissions sync task already in progress.",
raise OnyxError(
OnyxErrorCode.CONFLICT,
"Permissions sync task already in progress.",
)
logger.info(
@@ -90,9 +90,9 @@ def sync_cc_pair(
client_app, cc_pair_id, r, tenant_id
)
if not payload_id:
raise HTTPException(
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
detail="Permissions sync task creation failed.",
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
"Permissions sync task creation failed.",
)
logger.info(f"Permissions sync queued: cc_pair={cc_pair_id} id={payload_id}")
@@ -116,9 +116,9 @@ def get_cc_pair_latest_group_sync(
get_editable=False,
)
if not cc_pair:
raise HTTPException(
status_code=400,
detail="cc_pair not found for current user's permissions",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"cc_pair not found for current user's permissions",
)
return cc_pair.last_time_external_group_sync
@@ -140,18 +140,18 @@ def sync_cc_pair_groups(
get_editable=False,
)
if not cc_pair:
raise HTTPException(
status_code=400,
detail="Connection not found for current user's permissions",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"Connection not found for current user's permissions",
)
r = get_redis_client()
redis_connector = RedisConnector(tenant_id, cc_pair_id)
if redis_connector.external_group_sync.fenced:
raise HTTPException(
status_code=HTTPStatus.CONFLICT,
detail="External group sync task already in progress.",
raise OnyxError(
OnyxErrorCode.CONFLICT,
"External group sync task already in progress.",
)
logger.info(
@@ -164,9 +164,9 @@ def sync_cc_pair_groups(
client_app, cc_pair_id, r, tenant_id
)
if not payload_id:
raise HTTPException(
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
detail="External group sync task creation failed.",
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
"External group sync task creation failed.",
)
logger.info(f"External group sync queued: cc_pair={cc_pair_id} id={payload_id}")

View File

@@ -5,9 +5,7 @@ from typing import Any
import httpx
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from fastapi import Response
from fastapi import status
from fastapi import UploadFile
from pydantic import BaseModel
from pydantic import Field
@@ -33,6 +31,8 @@ from onyx.auth.users import get_user_manager
from onyx.auth.users import UserManager
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.file_store.file_store import get_default_file_store
from onyx.server.utils import BasicAuthenticationError
from onyx.utils.logger import setup_logger
@@ -97,24 +97,24 @@ async def refresh_access_token(
except httpx.HTTPStatusError as e:
if e.response.status_code == 401:
logger.warning(f"Full authentication required for user {user.id}")
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Full authentication required",
raise OnyxError(
OnyxErrorCode.UNAUTHENTICATED,
"Full authentication required",
)
logger.error(
f"HTTP error occurred while refreshing token for user {user.id}: {str(e)}"
)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to refresh token",
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
"Failed to refresh token",
)
except Exception as e:
logger.error(
f"Unexpected error occurred while refreshing token for user {user.id}: {str(e)}"
)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="An unexpected error occurred",
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
"An unexpected error occurred",
)
@@ -152,9 +152,9 @@ def fetch_logo_helper(db_session: Session) -> Response: # noqa: ARG001
raise ValueError("get_onyx_file returned None!")
except Exception:
logger.exception("Faield to fetch logo file")
raise HTTPException(
status_code=404,
detail="No logo file found",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"No logo file found",
)
else:
return Response(content=onyx_file.data, media_type=onyx_file.mime_type)
@@ -167,9 +167,9 @@ def fetch_logotype_helper(db_session: Session) -> Response: # noqa: ARG001
if not onyx_file:
raise ValueError("get_onyx_file returned None!")
except Exception:
raise HTTPException(
status_code=404,
detail="No logotype file found",
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
"No logotype file found",
)
else:
return Response(content=onyx_file.data, media_type=onyx_file.mime_type)
@@ -197,7 +197,7 @@ def upload_custom_analytics_script(
try:
store_analytics_script(script_upload)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, str(e))
@basic_router.get("/custom-analytics-script")
@@ -222,7 +222,7 @@ def get_active_scim_token(
"""Return the currently active SCIM token's metadata, or 404 if none."""
token = dal.get_active_token()
if not token:
raise HTTPException(status_code=404, detail="No active SCIM token")
raise OnyxError(OnyxErrorCode.NOT_FOUND, "No active SCIM token")
# Derive the IdP domain from the first synced user as a heuristic.
idp_domain: str | None = None

View File

@@ -3,7 +3,6 @@ from collections.abc import Awaitable
from collections.abc import Callable
from fastapi import FastAPI
from fastapi import HTTPException
from fastapi import Request
from fastapi import Response
@@ -12,6 +11,8 @@ from onyx.auth.utils import extract_tenant_from_auth_header
from onyx.configs.constants import ANONYMOUS_USER_COOKIE_NAME
from onyx.configs.constants import TENANT_ID_COOKIE_NAME
from onyx.db.engine.sql_engine import is_valid_schema_name
from onyx.error_handling.error_codes import OnyxErrorCode
from onyx.error_handling.exceptions import OnyxError
from onyx.redis.redis_pool import retrieve_auth_token_data_from_redis
from shared_configs.configs import MULTI_TENANT
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
@@ -76,7 +77,9 @@ async def _get_tenant_id_from_request(
)
if tenant_id and not is_valid_schema_name(tenant_id):
raise HTTPException(status_code=400, detail="Invalid tenant ID format")
raise OnyxError(
OnyxErrorCode.VALIDATION_ERROR, "Invalid tenant ID format"
)
# Check for anonymous user cookie
anonymous_user_cookie = request.cookies.get(ANONYMOUS_USER_COOKIE_NAME)
@@ -90,8 +93,8 @@ async def _get_tenant_id_from_request(
)
if not tenant_id or not is_valid_schema_name(tenant_id):
raise HTTPException(
status_code=400, detail="Invalid tenant ID format"
raise OnyxError(
OnyxErrorCode.VALIDATION_ERROR, "Invalid tenant ID format"
)
return tenant_id
@@ -111,7 +114,7 @@ async def _get_tenant_id_from_request(
except Exception as e:
logger.error(f"Unexpected error in _get_tenant_id_from_request: {str(e)}")
raise HTTPException(status_code=500, detail="Internal server error")
raise OnyxError(OnyxErrorCode.INTERNAL_ERROR, "Internal server error")
finally:
if tenant_id:

View File

@@ -1,6 +1,5 @@
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from sqlalchemy.orm import Session
from ee.onyx.onyxbot.slack.handlers.handle_standard_answers import (
@@ -11,6 +10,8 @@ from ee.onyx.server.query_and_chat.models import StandardAnswerResponse
from onyx.auth.users import current_user
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.utils.logger import setup_logger
logger = setup_logger()
@@ -33,4 +34,6 @@ def get_standard_answer(
return StandardAnswerResponse(standard_answers=standard_answers)
except Exception as e:
logger.error(f"Error in get_standard_answer: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail="An internal server error occurred")
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR, "An internal server error occurred"
)

View File

@@ -2,7 +2,6 @@ from collections.abc import Generator
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
@@ -23,6 +22,8 @@ from onyx.auth.users import current_user
from onyx.db.engine.sql_engine import get_session
from onyx.db.engine.sql_engine import get_session_with_current_tenant
from onyx.db.models import User
from onyx.error_handling.error_codes import OnyxErrorCode
from onyx.error_handling.exceptions import OnyxError
from onyx.llm.factory import get_default_llm
from onyx.server.usage_limits import check_llm_cost_limit_for_provider
from onyx.server.utils import get_json_line
@@ -107,7 +108,7 @@ def handle_send_search_message(
yield get_json_line(packet.model_dump())
except NotImplementedError as e:
yield get_json_line(SearchErrorPacket(error=str(e)).model_dump())
except HTTPException:
except OnyxError:
raise
except Exception as e:
logger.exception("Error in search streaming")
@@ -135,21 +136,21 @@ def get_search_history(
"""
# Validate limit
if limit <= 0:
raise HTTPException(
status_code=400,
detail="limit must be greater than 0",
raise OnyxError(
OnyxErrorCode.INVALID_INPUT,
"limit must be greater than 0",
)
if limit > 1000:
raise HTTPException(
status_code=400,
detail="limit must be at most 1000",
raise OnyxError(
OnyxErrorCode.INVALID_INPUT,
"limit must be at most 1000",
)
# Validate filter_days
if filter_days is not None and filter_days <= 0:
raise HTTPException(
status_code=400,
detail="filter_days must be greater than 0",
raise OnyxError(
OnyxErrorCode.INVALID_INPUT,
"filter_days must be greater than 0",
)
search_queries = fetch_search_queries_for_user(

View File

@@ -7,7 +7,6 @@ from typing import List
from typing import Tuple
from uuid import UUID
from fastapi import HTTPException
from sqlalchemy import func
from sqlalchemy import select
from sqlalchemy.orm import Session
@@ -22,6 +21,8 @@ from onyx.db.models import User
from onyx.db.models import User__UserGroup
from onyx.db.models import UserGroup
from onyx.db.token_limit import fetch_all_user_token_rate_limits
from onyx.error_handling.error_codes import OnyxErrorCode
from onyx.error_handling.exceptions import OnyxError
from onyx.server.query_and_chat.token_limit import _get_cutoff_time
from onyx.server.query_and_chat.token_limit import _is_rate_limited
from onyx.server.query_and_chat.token_limit import _user_is_rate_limited_by_global
@@ -63,9 +64,9 @@ def _user_is_rate_limited(user_id: UUID) -> None:
user_usage = _fetch_user_usage(user_id, user_cutoff_time, db_session)
if _is_rate_limited(user_rate_limits, user_usage):
raise HTTPException(
status_code=429,
detail="Token budget exceeded for user. Try again later.",
raise OnyxError(
OnyxErrorCode.RATE_LIMITED,
"Token budget exceeded for user. Try again later.",
)
@@ -119,9 +120,9 @@ def _user_is_rate_limited_by_group(user_id: UUID) -> None:
break
if not has_at_least_one_untriggered_limit:
raise HTTPException(
status_code=429,
detail="Token budget exceeded for user's groups. Try again later.",
raise OnyxError(
OnyxErrorCode.RATE_LIMITED,
"Token budget exceeded for user's groups. Try again later.",
)

View File

@@ -2,13 +2,12 @@ import uuid
from collections.abc import Generator
from datetime import datetime
from datetime import timezone
from http import HTTPStatus
from uuid import UUID
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from fastapi import Query
from fastapi.responses import JSONResponse
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
@@ -45,6 +44,8 @@ from onyx.db.models import ChatSession
from onyx.db.models import User
from onyx.db.tasks import get_task_with_id
from onyx.db.tasks import register_task
from onyx.error_handling.error_codes import OnyxErrorCode
from onyx.error_handling.exceptions import OnyxError
from onyx.file_store.file_store import get_default_file_store
from onyx.server.documents.models import PaginatedReturn
from onyx.server.query_and_chat.models import ChatSessionDetails
@@ -61,9 +62,9 @@ def ensure_query_history_is_enabled(
disallowed: list[QueryHistoryType],
) -> None:
if ONYX_QUERY_HISTORY_TYPE in disallowed:
raise HTTPException(
status_code=HTTPStatus.FORBIDDEN,
detail="Query history has been disabled by the administrator.",
raise OnyxError(
OnyxErrorCode.INSUFFICIENT_PERMISSIONS,
"Query history has been disabled by the administrator.",
)
@@ -247,8 +248,8 @@ def get_chat_session_admin(
include_deleted=True,
)
except ValueError:
raise HTTPException(
HTTPStatus.BAD_REQUEST,
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
f"Chat session with id '{chat_session_id}' does not exist.",
)
snapshot = snapshot_from_chat_session(
@@ -256,8 +257,8 @@ def get_chat_session_admin(
)
if snapshot is None:
raise HTTPException(
HTTPStatus.BAD_REQUEST,
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
f"Could not create snapshot for chat session with id '{chat_session_id}'",
)
@@ -290,9 +291,7 @@ def list_all_query_history_exports(
return merged
except Exception as e:
raise HTTPException(
HTTPStatus.INTERNAL_SERVER_ERROR, f"Failed to get all tasks: {e}"
)
raise OnyxError(OnyxErrorCode.INTERNAL_ERROR, f"Failed to get all tasks: {e}")
@router.post("/admin/query-history/start-export", tags=PUBLIC_API_TAGS)
@@ -308,8 +307,8 @@ def start_query_history_export(
end = end or datetime.now(tz=timezone.utc)
if start >= end:
raise HTTPException(
HTTPStatus.BAD_REQUEST,
raise OnyxError(
OnyxErrorCode.INVALID_INPUT,
f"Start time must come before end time, but instead got the start time coming after; {start=} {end=}",
)
@@ -367,8 +366,8 @@ def get_query_history_export_status(
)
if not has_file:
raise HTTPException(
HTTPStatus.NOT_FOUND,
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
f"No task with {request_id=} was found",
)
@@ -395,8 +394,8 @@ def download_query_history_csv(
try:
csv_stream = file_store.read_file(report_name)
except Exception as e:
raise HTTPException(
HTTPStatus.INTERNAL_SERVER_ERROR,
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
f"Failed to read query history file: {str(e)}",
)
csv_stream.seek(0)
@@ -410,19 +409,20 @@ def download_query_history_csv(
# Therefore, we check the task queue to determine its status, if there is any.
task = get_task_with_id(db_session=db_session, task_id=request_id)
if not task:
raise HTTPException(
HTTPStatus.NOT_FOUND,
raise OnyxError(
OnyxErrorCode.NOT_FOUND,
f"No task with {request_id=} was found",
)
if task.status in [TaskStatus.STARTED, TaskStatus.PENDING]:
raise HTTPException(
HTTPStatus.ACCEPTED, f"Task with {request_id=} is still being worked on"
return JSONResponse(
status_code=202,
content={"message": f"Task with {request_id=} is still being worked on"},
)
elif task.status == TaskStatus.FAILURE:
raise HTTPException(
HTTPStatus.INTERNAL_SERVER_ERROR,
raise OnyxError(
OnyxErrorCode.INTERNAL_ERROR,
f"Task with {request_id=} failed to be processed",
)
else:

View File

@@ -3,7 +3,6 @@ from datetime import datetime
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from fastapi import Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
@@ -17,6 +16,8 @@ from onyx.background.celery.versioned_apps.client import app as client_app
from onyx.configs.constants import OnyxCeleryTask
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.file_store.constants import STANDARD_CHUNK_SIZE
from shared_configs.contextvars import get_current_tenant_id
@@ -39,7 +40,7 @@ def generate_report(
datetime.fromisoformat(params.period_from)
datetime.fromisoformat(params.period_to)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e))
raise OnyxError(OnyxErrorCode.VALIDATION_ERROR, str(e))
tenant_id = get_current_tenant_id()
client_app.send_task(
@@ -64,7 +65,7 @@ def read_usage_report(
try:
file = get_usage_report_data(report_name)
except (ValueError, RuntimeError) as e:
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))
def iterfile() -> Generator[bytes, None, None]:
while True:
@@ -88,4 +89,4 @@ def fetch_usage_reports(
try:
return get_all_usage_reports(db_session)
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))

View File

@@ -1,6 +1,5 @@
from fastapi import APIRouter
from fastapi import Depends
from fastapi import HTTPException
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
@@ -27,6 +26,8 @@ from onyx.configs.constants import PUBLIC_API_TAGS
from onyx.db.engine.sql_engine import get_session
from onyx.db.models import User
from onyx.db.models import UserRole
from onyx.error_handling.error_codes import OnyxErrorCode
from onyx.error_handling.exceptions import OnyxError
from onyx.utils.logger import setup_logger
logger = setup_logger()
@@ -79,8 +80,8 @@ def create_user_group(
try:
db_user_group = insert_user_group(db_session, user_group)
except IntegrityError:
raise HTTPException(
400,
raise OnyxError(
OnyxErrorCode.DUPLICATE_RESOURCE,
f"User group with name '{user_group.name}' already exists. Please "
+ "choose a different name.",
)
@@ -104,7 +105,7 @@ def patch_user_group(
)
)
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))
@router.post("/admin/user-group/{user_group_id}/add-users")
@@ -124,7 +125,7 @@ def add_users(
)
)
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))
@router.post("/admin/user-group/{user_group_id}/set-curator")
@@ -143,7 +144,7 @@ def set_user_curator(
)
except ValueError as e:
logger.error(f"Error setting user curator: {e}")
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))
@router.delete("/admin/user-group/{user_group_id}")
@@ -155,7 +156,7 @@ def delete_user_group(
try:
prepare_user_group_for_deletion(db_session, user_group_id)
except ValueError as e:
raise HTTPException(status_code=404, detail=str(e))
raise OnyxError(OnyxErrorCode.NOT_FOUND, str(e))
if DISABLE_VECTOR_DB:
user_group = fetch_user_group(db_session, user_group_id)