diff --git a/backend/onyx/auth/schemas.py b/backend/onyx/auth/schemas.py index 526264c2f..88a8f6111 100644 --- a/backend/onyx/auth/schemas.py +++ b/backend/onyx/auth/schemas.py @@ -47,3 +47,8 @@ class UserUpdate(schemas.BaseUserUpdate): Role updates are not allowed through the user update endpoint for security reasons Role changes should be handled through a separate, admin-only process """ + + +class AuthBackend(str, Enum): + REDIS = "redis" + POSTGRES = "postgres" diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index aa299d57a..db79ed29a 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -33,6 +33,8 @@ from fastapi_users.authentication import AuthenticationBackend from fastapi_users.authentication import CookieTransport from fastapi_users.authentication import RedisStrategy from fastapi_users.authentication import Strategy +from fastapi_users.authentication.strategy.db import AccessTokenDatabase +from fastapi_users.authentication.strategy.db import DatabaseStrategy from fastapi_users.exceptions import UserAlreadyExists from fastapi_users.jwt import decode_jwt from fastapi_users.jwt import generate_jwt @@ -52,14 +54,15 @@ from onyx.auth.api_key import get_hashed_api_key_from_request from onyx.auth.email_utils import send_forgot_password_email from onyx.auth.email_utils import send_user_verification_email from onyx.auth.invited_users import get_invited_users +from onyx.auth.schemas import AuthBackend from onyx.auth.schemas import UserCreate from onyx.auth.schemas import UserRole from onyx.auth.schemas import UserUpdate +from onyx.configs.app_configs import AUTH_BACKEND from onyx.configs.app_configs import AUTH_COOKIE_EXPIRE_TIME_SECONDS from onyx.configs.app_configs import AUTH_TYPE from onyx.configs.app_configs import DISABLE_AUTH from onyx.configs.app_configs import EMAIL_CONFIGURED -from onyx.configs.app_configs import REDIS_AUTH_EXPIRE_TIME_SECONDS from onyx.configs.app_configs import REDIS_AUTH_KEY_PREFIX from onyx.configs.app_configs import REQUIRE_EMAIL_VERIFICATION from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS @@ -75,6 +78,7 @@ from onyx.configs.constants import OnyxRedisLocks from onyx.configs.constants import PASSWORD_SPECIAL_CHARS from onyx.configs.constants import UNNAMED_KEY_PLACEHOLDER from onyx.db.api_key import fetch_user_for_api_key +from onyx.db.auth import get_access_token_db from onyx.db.auth import get_default_admin_user_emails from onyx.db.auth import get_user_count from onyx.db.auth import get_user_db @@ -83,6 +87,7 @@ from onyx.db.engine import get_async_session from onyx.db.engine import get_async_session_with_tenant from onyx.db.engine import get_current_tenant_id from onyx.db.engine import get_session_with_tenant +from onyx.db.models import AccessToken from onyx.db.models import OAuthAccount from onyx.db.models import User from onyx.db.users import get_user_by_email @@ -582,6 +587,14 @@ def get_redis_strategy() -> RedisStrategy: return TenantAwareRedisStrategy() +def get_database_strategy( + access_token_db: AccessTokenDatabase[AccessToken] = Depends(get_access_token_db), +) -> DatabaseStrategy: + return DatabaseStrategy( + access_token_db, lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS + ) + + class TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]): """ A custom strategy that fetches the actual async Redis connection inside each method. @@ -590,7 +603,7 @@ class TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]): def __init__( self, - lifetime_seconds: Optional[int] = REDIS_AUTH_EXPIRE_TIME_SECONDS, + lifetime_seconds: Optional[int] = SESSION_EXPIRE_TIME_SECONDS, key_prefix: str = REDIS_AUTH_KEY_PREFIX, ): self.lifetime_seconds = lifetime_seconds @@ -639,9 +652,16 @@ class TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]): await redis.delete(f"{self.key_prefix}{token}") -auth_backend = AuthenticationBackend( - name="redis", transport=cookie_transport, get_strategy=get_redis_strategy -) +if AUTH_BACKEND == AuthBackend.REDIS: + auth_backend = AuthenticationBackend( + name="redis", transport=cookie_transport, get_strategy=get_redis_strategy + ) +elif AUTH_BACKEND == AuthBackend.POSTGRES: + auth_backend = AuthenticationBackend( + name="postgres", transport=cookie_transport, get_strategy=get_database_strategy + ) +else: + raise ValueError(f"Invalid auth backend: {AUTH_BACKEND}") class FastAPIUserWithLogoutRouter(FastAPIUsers[models.UP, models.ID]): diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py index 8afdd6875..ad6b53c60 100644 --- a/backend/onyx/configs/app_configs.py +++ b/backend/onyx/configs/app_configs.py @@ -3,6 +3,7 @@ import os import urllib.parse from typing import cast +from onyx.auth.schemas import AuthBackend from onyx.configs.constants import AuthType from onyx.configs.constants import DocumentIndexType from onyx.file_processing.enums import HtmlBasedConnectorTransformLinksStrategy @@ -55,12 +56,12 @@ MASK_CREDENTIAL_PREFIX = ( os.environ.get("MASK_CREDENTIAL_PREFIX", "True").lower() != "false" ) -REDIS_AUTH_EXPIRE_TIME_SECONDS = int( - os.environ.get("REDIS_AUTH_EXPIRE_TIME_SECONDS") or 86400 * 7 -) # 7 days +AUTH_BACKEND = AuthBackend(os.environ.get("AUTH_BACKEND") or AuthBackend.REDIS.value) SESSION_EXPIRE_TIME_SECONDS = int( - os.environ.get("SESSION_EXPIRE_TIME_SECONDS") or 86400 * 7 + os.environ.get("SESSION_EXPIRE_TIME_SECONDS") + or os.environ.get("REDIS_AUTH_EXPIRE_TIME_SECONDS") + or 86400 * 7 ) # 7 days # Default request timeout, mostly used by connectors