From 54cfca3dd6c990d0d1eebc01689366f82b9ebb87 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 1 Sep 2023 17:20:55 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(api=5Fkey/crud.py):=20fix=20?= =?UTF-8?q?session=20concurrency=20issue=20when=20updating=20total=20uses?= =?UTF-8?q?=20and=20last=20used=20at=20fields?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../services/database/models/api_key/crud.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/backend/langflow/services/database/models/api_key/crud.py b/src/backend/langflow/services/database/models/api_key/crud.py index abc84f108..d78da5a05 100644 --- a/src/backend/langflow/services/database/models/api_key/crud.py +++ b/src/backend/langflow/services/database/models/api_key/crud.py @@ -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