Add back postgres auth backend support (#3753)

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"

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]):

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