🐛 fix(api_key/crud.py): fix session concurrency issue when updating total uses and last used at fields

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-09-01 17:20:55 -03:00
commit 54cfca3dd6

View file

@ -3,6 +3,8 @@ import secrets
import threading
from uuid import UUID
from typing import List, Optional
from langflow.services.database.utils import session_getter
from langflow.services.utils import get_db_manager
from sqlmodel import Session, select
from langflow.services.database.models.api_key import (
ApiKey,
@ -63,9 +65,13 @@ def check_key(session: Session, api_key: str) -> Optional[ApiKey]:
def update_total_uses(session, api_key: ApiKey):
"""Update the total uses and last used at."""
api_key.total_uses += 1
api_key.last_used_at = datetime.datetime.now(datetime.timezone.utc)
session.add(api_key)
session.commit()
session.refresh(api_key)
# This is running in a separate thread to avoid slowing down the request
# but session is not thread safe so we need to create a new session
db_manager = get_db_manager()
with session_getter(db_manager) as new_session:
api_key = new_session.get(ApiKey, api_key.id)
api_key.total_uses += 1
api_key.last_used_at = datetime.datetime.now(datetime.timezone.utc)
new_session.add(api_key)
new_session.commit()
return api_key