From d0964bdc476d4a34f813f71d9d9f0a904c4ae9cf Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Tue, 9 Jan 2024 10:33:56 +0500 Subject: [PATCH] Fix long-running db transactions --- src/dstack/_internal/server/db.py | 5 ++--- src/dstack/_internal/server/services/projects.py | 1 + src/dstack/_internal/server/services/repos.py | 2 ++ src/dstack/_internal/server/services/runs.py | 1 + src/dstack/_internal/server/services/users.py | 3 +++ src/tests/_internal/server/conftest.py | 5 ++--- 6 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/dstack/_internal/server/db.py b/src/dstack/_internal/server/db.py index 285aeb1f67..dec0a7bd7a 100644 --- a/src/dstack/_internal/server/db.py +++ b/src/dstack/_internal/server/db.py @@ -46,9 +46,8 @@ async def migrate(): async def get_session(): async with db.get_session() as session: - async with session as s: - yield s - await s.commit() + yield session + await session.commit() get_session_ctx = asynccontextmanager(get_session) diff --git a/src/dstack/_internal/server/services/projects.py b/src/dstack/_internal/server/services/projects.py index 2f8b12ff5b..386a84449f 100644 --- a/src/dstack/_internal/server/services/projects.py +++ b/src/dstack/_internal/server/services/projects.py @@ -109,6 +109,7 @@ async def delete_projects( name=new_project_name, ) ) + await session.commit() async def add_project_member( diff --git a/src/dstack/_internal/server/services/repos.py b/src/dstack/_internal/server/services/repos.py index 4a996666c0..cc9163d6a7 100644 --- a/src/dstack/_internal/server/services/repos.py +++ b/src/dstack/_internal/server/services/repos.py @@ -115,6 +115,7 @@ async def update_repo( creds=repo.creds, ) ) + await session.commit() return repo @@ -126,6 +127,7 @@ async def delete_repos( await session.execute( delete(RepoModel).where(RepoModel.project_id == project.id, RepoModel.name.in_(repos_ids)) ) + await session.commit() async def upload_code( diff --git a/src/dstack/_internal/server/services/runs.py b/src/dstack/_internal/server/services/runs.py index 0c8674da55..8f79ca9f08 100644 --- a/src/dstack/_internal/server/services/runs.py +++ b/src/dstack/_internal/server/services/runs.py @@ -262,6 +262,7 @@ async def delete_runs( ) .values(deleted=True) ) + await session.commit() def run_model_to_run(run_model: RunModel, include_job_submissions: bool = True) -> Run: diff --git a/src/dstack/_internal/server/services/users.py b/src/dstack/_internal/server/services/users.py index b3e8058434..4c3d86acd6 100644 --- a/src/dstack/_internal/server/services/users.py +++ b/src/dstack/_internal/server/services/users.py @@ -85,6 +85,7 @@ async def update_user( .where(UserModel.name == username) .values(global_role=global_role, email=email) ) + await session.commit() return await get_user_model_by_name_or_error(session=session, username=username) @@ -92,6 +93,7 @@ async def refresh_user_token(session: AsyncSession, username: str) -> Optional[U await session.execute( update(UserModel).where(UserModel.name == username).values(token=uuid.uuid4()) ) + await session.commit() return await get_user_model_by_name(session=session, username=username) @@ -101,6 +103,7 @@ async def delete_users( usernames: List[str], ): await session.execute(delete(UserModel).where(UserModel.name.in_(usernames))) + await session.commit() async def get_user_model_by_name(session: AsyncSession, username: str) -> Optional[UserModel]: diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index c700a4830d..5a46803c5a 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -18,6 +18,5 @@ async def test_db(): @pytest_asyncio.fixture async def session(): async with db.get_session() as session: - async with session as s: - yield s - await s.commit() + yield session + await session.commit()