refac: async db
This commit is contained in:
Timothy Jaeryang Baek
@@ -8,8 +8,9 @@ import json
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
from open_webui.internal.db import Base, get_db, get_db_context
|
||||
from sqlalchemy import select, delete, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from open_webui.internal.db import Base, get_async_db_context
|
||||
from open_webui.env import OAUTH_SESSION_TOKEN_ENCRYPTION_KEY
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
@@ -103,16 +104,16 @@ class OAuthSessionTable:
|
||||
log.error(f'Error decrypting tokens: {type(e).__name__}: {e}')
|
||||
raise
|
||||
|
||||
def create_session(
|
||||
async def create_session(
|
||||
self,
|
||||
user_id: str,
|
||||
provider: str,
|
||||
token: dict,
|
||||
db: Optional[Session] = None,
|
||||
db: Optional[AsyncSession] = None,
|
||||
) -> Optional[OAuthSessionModel]:
|
||||
"""Create a new OAuth session"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
async with get_async_db_context(db) as db:
|
||||
current_time = int(time.time())
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
@@ -129,91 +130,126 @@ class OAuthSessionTable:
|
||||
)
|
||||
|
||||
db.add(result)
|
||||
db.commit()
|
||||
db.refresh(result)
|
||||
await db.commit()
|
||||
await db.refresh(result)
|
||||
|
||||
if result:
|
||||
db.expunge(result) # Detach so dict swap is never flushed
|
||||
result.token = token # Return decrypted token
|
||||
return OAuthSessionModel.model_validate(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
|
||||
|
||||
def get_session_by_id(self, session_id: str, db: Optional[Session] = None) -> Optional[OAuthSessionModel]:
|
||||
async def get_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> Optional[OAuthSessionModel]:
|
||||
"""Get OAuth session by ID"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
session = db.query(OAuthSession).filter_by(id=session_id).first()
|
||||
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:
|
||||
db.expunge(session)
|
||||
session.token = self._decrypt_token(session.token)
|
||||
return OAuthSessionModel.model_validate(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
|
||||
|
||||
def get_session_by_id_and_user_id(
|
||||
self, session_id: str, user_id: str, db: Optional[Session] = 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:
|
||||
with get_db_context(db) as db:
|
||||
session = db.query(OAuthSession).filter_by(id=session_id, user_id=user_id).first()
|
||||
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:
|
||||
db.expunge(session)
|
||||
session.token = self._decrypt_token(session.token)
|
||||
return OAuthSessionModel.model_validate(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
|
||||
|
||||
def get_session_by_provider_and_user_id(
|
||||
self, provider: str, user_id: str, db: Optional[Session] = 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:
|
||||
with get_db_context(db) as db:
|
||||
session = (
|
||||
db.query(OAuthSession)
|
||||
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())
|
||||
.first()
|
||||
)
|
||||
session = result.scalars().first()
|
||||
if session:
|
||||
db.expunge(session)
|
||||
session.token = self._decrypt_token(session.token)
|
||||
return OAuthSessionModel.model_validate(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
|
||||
|
||||
def get_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> List[OAuthSessionModel]:
|
||||
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:
|
||||
with get_db_context(db) as db:
|
||||
sessions = db.query(OAuthSession).filter_by(user_id=user_id).all()
|
||||
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:
|
||||
db.expunge(session)
|
||||
session.token = self._decrypt_token(session.token)
|
||||
results.append(OAuthSessionModel.model_validate(session))
|
||||
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}'
|
||||
)
|
||||
db.query(OAuthSession).filter_by(id=session.id).delete()
|
||||
db.commit()
|
||||
await db.execute(delete(OAuthSession).filter_by(id=session.id))
|
||||
await db.commit()
|
||||
|
||||
return results
|
||||
|
||||
@@ -221,62 +257,69 @@ class OAuthSessionTable:
|
||||
log.error(f'Error getting OAuth sessions by user ID: {e}')
|
||||
return []
|
||||
|
||||
def update_session_by_id(
|
||||
self, session_id: str, token: dict, db: Optional[Session] = None
|
||||
async def update_session_by_id(
|
||||
self, session_id: str, token: dict, db: Optional[AsyncSession] = None
|
||||
) -> Optional[OAuthSessionModel]:
|
||||
"""Update OAuth session tokens"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
async with get_async_db_context(db) as db:
|
||||
current_time = int(time.time())
|
||||
|
||||
db.query(OAuthSession).filter_by(id=session_id).update(
|
||||
{
|
||||
'token': self._encrypt_token(token),
|
||||
'expires_at': token.get('expires_at'),
|
||||
'updated_at': current_time,
|
||||
}
|
||||
await db.execute(
|
||||
update(OAuthSession).filter_by(id=session_id).values(
|
||||
token=self._encrypt_token(token),
|
||||
expires_at=token.get('expires_at'),
|
||||
updated_at=current_time,
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
session = db.query(OAuthSession).filter_by(id=session_id).first()
|
||||
await db.commit()
|
||||
result = await db.execute(select(OAuthSession).filter_by(id=session_id))
|
||||
session = result.scalars().first()
|
||||
|
||||
if session:
|
||||
db.expunge(session)
|
||||
session.token = self._decrypt_token(session.token)
|
||||
return OAuthSessionModel.model_validate(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
|
||||
|
||||
def delete_session_by_id(self, session_id: str, db: Optional[Session] = None) -> bool:
|
||||
async def delete_session_by_id(self, session_id: str, db: Optional[AsyncSession] = None) -> bool:
|
||||
"""Delete an OAuth session"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
result = db.query(OAuthSession).filter_by(id=session_id).delete()
|
||||
db.commit()
|
||||
return result > 0
|
||||
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
|
||||
|
||||
def delete_sessions_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
|
||||
async def delete_sessions_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> bool:
|
||||
"""Delete all OAuth sessions for a user"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
result = db.query(OAuthSession).filter_by(user_id=user_id).delete()
|
||||
db.commit()
|
||||
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
|
||||
|
||||
def delete_sessions_by_provider(self, provider: str, db: Optional[Session] = None) -> bool:
|
||||
async def delete_sessions_by_provider(self, provider: str, db: Optional[AsyncSession] = None) -> bool:
|
||||
"""Delete all OAuth sessions for a provider"""
|
||||
try:
|
||||
with get_db_context(db) as db:
|
||||
db.query(OAuthSession).filter_by(provider=provider).delete()
|
||||
db.commit()
|
||||
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}')
|
||||
|
||||
Reference in New Issue
Block a user