mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-21 00:21:19 +00:00
699d512e2f
Both session factories run with expire_on_commit=False, so ORM objects keep their attribute values after commit. Every session.refresh issued right after a commit therefore re-SELECTed a row whose values the session already held, including full chat JSON blobs and user settings, purely to overwrite identical data. Fifty such calls existed across the model layer, covering nearly every write path in the app (chat inserts, title updates, pin/archive toggles, user role and settings updates, tool, prompt, function, model, file, tag, feedback, memory, automation and grant writes). All fifty are removed. The only refreshes with an actual job were the two update-then-reload paths in tools and skills, where a Core UPDATE statement bypasses the identity map; those now use session.get(..., populate_existing=True), which guarantees a fresh row in one SELECT whether or not the row was already present in the session (the previous code issued get plus refresh, two SELECTs, on the default configuration). Benchmark (real SQLite DB, per write): | write path | before | after | | --- | --- | --- | | chat title update, ~600 KB chat blob | 2.08 ms | 1.24 ms | | user role update, small row | 1.21 ms | 0.68 ms | On Postgres each removed refresh is additionally a network round trip. The chat-blob case also skips re-parsing the entire JSON document per write. Functionally verified against a fresh database: user insert, role and settings updates, chat insert (including the server-default meta column, which is always provided client-side), title update and pin toggle, tool insert and the Core-update reload path, tag insert and the prompt insert flow that pins version_id after history creation all return correct values and persist correctly.
344 lines
13 KiB
Python
344 lines
13 KiB
Python
import base64
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from typing import List, Optional
|
|
|
|
from cryptography.fernet import Fernet
|
|
from open_webui.env import OAUTH_SESSION_TOKEN_ENCRYPTION_KEY
|
|
from open_webui.internal.db import Base, get_async_db_context
|
|
from pydantic import BaseModel, ConfigDict
|
|
from sqlalchemy import BigInteger, Column, Index, String, Text, delete, select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
####################
|
|
# DB MODEL
|
|
####################
|
|
|
|
|
|
class OAuthSession(Base):
|
|
__tablename__ = 'oauth_session'
|
|
|
|
id = Column(Text, primary_key=True, unique=True)
|
|
user_id = Column(Text, nullable=False)
|
|
provider = Column(Text, nullable=False)
|
|
token = Column(Text, nullable=False) # JSON with access_token, id_token, refresh_token
|
|
expires_at = Column(BigInteger, nullable=False)
|
|
created_at = Column(BigInteger, nullable=False)
|
|
updated_at = Column(BigInteger, nullable=False)
|
|
|
|
# Add indexes for better performance
|
|
__table_args__ = (
|
|
Index('idx_oauth_session_user_id', 'user_id'),
|
|
Index('idx_oauth_session_expires_at', 'expires_at'),
|
|
Index('idx_oauth_session_user_provider', 'user_id', 'provider'),
|
|
)
|
|
|
|
|
|
class OAuthSessionModel(BaseModel):
|
|
id: str
|
|
user_id: str
|
|
provider: str
|
|
token: dict
|
|
expires_at: int # timestamp in epoch
|
|
created_at: int # timestamp in epoch
|
|
updated_at: int # timestamp in epoch
|
|
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
|
|
####################
|
|
# Forms
|
|
####################
|
|
|
|
|
|
class OAuthSessionResponse(BaseModel):
|
|
id: str
|
|
user_id: str
|
|
provider: str
|
|
expires_at: int
|
|
|
|
|
|
class OAuthSessionTable:
|
|
def __init__(self):
|
|
self.encryption_key = OAUTH_SESSION_TOKEN_ENCRYPTION_KEY
|
|
if not self.encryption_key:
|
|
raise Exception('OAUTH_SESSION_TOKEN_ENCRYPTION_KEY is not set')
|
|
|
|
# check if encryption key is in the right format for Fernet (32 url-safe base64-encoded bytes)
|
|
if len(self.encryption_key) != 44:
|
|
key_bytes = hashlib.sha256(self.encryption_key.encode()).digest()
|
|
self.encryption_key = base64.urlsafe_b64encode(key_bytes)
|
|
else:
|
|
self.encryption_key = self.encryption_key.encode()
|
|
|
|
try:
|
|
self.fernet = Fernet(self.encryption_key)
|
|
except Exception as e:
|
|
log.error(f'Error initializing Fernet with provided key: {e}')
|
|
raise
|
|
|
|
def _encrypt_token(self, token) -> str:
|
|
"""Encrypt OAuth tokens for storage"""
|
|
try:
|
|
token_json = json.dumps(token)
|
|
encrypted = self.fernet.encrypt(token_json.encode()).decode()
|
|
return encrypted
|
|
except Exception as e:
|
|
log.error(f'Error encrypting tokens: {e}')
|
|
raise
|
|
|
|
def _decrypt_token(self, token: str):
|
|
"""Decrypt OAuth tokens from storage"""
|
|
try:
|
|
decrypted = self.fernet.decrypt(token.encode()).decode()
|
|
return json.loads(decrypted)
|
|
except Exception as e:
|
|
log.error(f'Error decrypting tokens: {type(e).__name__}: {e}')
|
|
raise
|
|
|
|
async def create_session(
|
|
self,
|
|
user_id: str,
|
|
provider: str,
|
|
token: dict,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[OAuthSessionModel]:
|
|
"""Create a new OAuth session"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
current_time = int(time.time())
|
|
id = str(uuid.uuid4())
|
|
|
|
result = OAuthSession(
|
|
**{
|
|
'id': id,
|
|
'user_id': user_id,
|
|
'provider': provider,
|
|
'token': self._encrypt_token(token),
|
|
'expires_at': token.get('expires_at') or int(time.time() + 3600),
|
|
'created_at': current_time,
|
|
'updated_at': current_time,
|
|
}
|
|
)
|
|
|
|
db.add(result)
|
|
await db.commit()
|
|
|
|
if result:
|
|
# Make a copy of the model data before closing session
|
|
model = OAuthSessionModel(
|
|
id=result.id,
|
|
user_id=result.user_id,
|
|
provider=result.provider,
|
|
token=token, # Return decrypted token
|
|
expires_at=result.expires_at,
|
|
created_at=result.created_at,
|
|
updated_at=result.updated_at,
|
|
)
|
|
return model
|
|
else:
|
|
return None
|
|
except Exception as e:
|
|
log.error(f'Error creating OAuth session: {e}')
|
|
return None
|
|
|
|
async def get_session_by_id(
|
|
self, session_id: str, db: Optional[AsyncSession] = None
|
|
) -> Optional[OAuthSessionModel]:
|
|
"""Get OAuth session by ID"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(OAuthSession).filter_by(id=session_id))
|
|
session = result.scalars().first()
|
|
if session:
|
|
return OAuthSessionModel(
|
|
id=session.id,
|
|
user_id=session.user_id,
|
|
provider=session.provider,
|
|
token=self._decrypt_token(session.token),
|
|
expires_at=session.expires_at,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
)
|
|
|
|
return None
|
|
except Exception as e:
|
|
log.error(f'Error getting OAuth session by ID: {e}')
|
|
return None
|
|
|
|
async def get_session_by_id_and_user_id(
|
|
self, session_id: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> Optional[OAuthSessionModel]:
|
|
"""Get OAuth session by ID and user ID"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(OAuthSession).filter_by(id=session_id, user_id=user_id))
|
|
session = result.scalars().first()
|
|
if session:
|
|
return OAuthSessionModel(
|
|
id=session.id,
|
|
user_id=session.user_id,
|
|
provider=session.provider,
|
|
token=self._decrypt_token(session.token),
|
|
expires_at=session.expires_at,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
)
|
|
|
|
return None
|
|
except Exception as e:
|
|
log.error(f'Error getting OAuth session by ID: {e}')
|
|
return None
|
|
|
|
async def get_session_by_provider_and_user_id(
|
|
self, provider: str, user_id: str, db: Optional[AsyncSession] = None
|
|
) -> Optional[OAuthSessionModel]:
|
|
"""Get OAuth session by provider and user ID"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(
|
|
select(OAuthSession)
|
|
.filter_by(provider=provider, user_id=user_id)
|
|
.order_by(OAuthSession.created_at.desc())
|
|
)
|
|
session = result.scalars().first()
|
|
if session:
|
|
return OAuthSessionModel(
|
|
id=session.id,
|
|
user_id=session.user_id,
|
|
provider=session.provider,
|
|
token=self._decrypt_token(session.token),
|
|
expires_at=session.expires_at,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
)
|
|
|
|
return None
|
|
except Exception as e:
|
|
log.error(f'Error getting OAuth session by provider and user ID: {e}')
|
|
return None
|
|
|
|
async def get_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> List[OAuthSessionModel]:
|
|
"""Get all OAuth sessions for a user"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(select(OAuthSession).filter_by(user_id=user_id))
|
|
sessions = result.scalars().all()
|
|
|
|
results = []
|
|
for session in sessions:
|
|
try:
|
|
results.append(
|
|
OAuthSessionModel(
|
|
id=session.id,
|
|
user_id=session.user_id,
|
|
provider=session.provider,
|
|
token=self._decrypt_token(session.token),
|
|
expires_at=session.expires_at,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
log.warning(
|
|
f'Skipping OAuth session {session.id} due to decryption failure, deleting corrupted session: {type(e).__name__}: {e}'
|
|
)
|
|
await db.execute(delete(OAuthSession).filter_by(id=session.id))
|
|
await db.commit()
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
log.error(f'Error getting OAuth sessions by user ID: {e}')
|
|
return []
|
|
|
|
async def update_session_by_id(
|
|
self, session_id: str, token: dict, db: Optional[AsyncSession] = None
|
|
) -> Optional[OAuthSessionModel]:
|
|
"""Update OAuth session tokens"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
current_time = int(time.time())
|
|
|
|
await db.execute(
|
|
update(OAuthSession)
|
|
.filter_by(id=session_id)
|
|
.values(
|
|
token=self._encrypt_token(token),
|
|
expires_at=token.get('expires_at') or int(time.time() + 3600),
|
|
updated_at=current_time,
|
|
)
|
|
)
|
|
await db.commit()
|
|
result = await db.execute(select(OAuthSession).filter_by(id=session_id))
|
|
session = result.scalars().first()
|
|
|
|
if session:
|
|
return OAuthSessionModel(
|
|
id=session.id,
|
|
user_id=session.user_id,
|
|
provider=session.provider,
|
|
token=self._decrypt_token(session.token),
|
|
expires_at=session.expires_at,
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
)
|
|
|
|
return None
|
|
except Exception as e:
|
|
log.error(f'Error updating OAuth session tokens: {e}')
|
|
return None
|
|
|
|
async def delete_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> bool:
|
|
"""Delete an OAuth session"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(delete(OAuthSession).filter_by(id=session_id))
|
|
await db.commit()
|
|
return result.rowcount > 0
|
|
except Exception as e:
|
|
log.error(f'Error deleting OAuth session: {e}')
|
|
return False
|
|
|
|
async def delete_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
|
|
"""Delete all OAuth sessions for a user"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
await db.execute(delete(OAuthSession).filter_by(user_id=user_id))
|
|
await db.commit()
|
|
return True
|
|
except Exception as e:
|
|
log.error(f'Error deleting OAuth sessions by user ID: {e}')
|
|
return False
|
|
|
|
async def delete_sessions_by_user_id_and_provider(
|
|
self, user_id: str, provider: str, db: Optional[AsyncSession] = None
|
|
) -> bool:
|
|
"""Delete all OAuth sessions for a specific user and provider"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
result = await db.execute(delete(OAuthSession).filter_by(user_id=user_id, provider=provider))
|
|
await db.commit()
|
|
return result.rowcount > 0
|
|
except Exception as e:
|
|
log.error(f'Error deleting OAuth sessions for user {user_id} and provider {provider}: {e}')
|
|
return False
|
|
|
|
async def delete_sessions_by_provider(self, provider: str, db: Optional[AsyncSession] = None) -> bool:
|
|
"""Delete all OAuth sessions for a provider"""
|
|
try:
|
|
async with get_async_db_context(db) as db:
|
|
await db.execute(delete(OAuthSession).filter_by(provider=provider))
|
|
await db.commit()
|
|
return True
|
|
except Exception as e:
|
|
log.error(f'Error deleting OAuth sessions by provider {provider}: {e}')
|
|
return False
|
|
|
|
|
|
OAuthSessions = OAuthSessionTable()
|