refactor: replace generic Exception handling with SQLAlchemyError in CRUD and session management
This commit is contained in:
@@ -5,7 +5,7 @@ from typing import Any, TypeVar, cast
|
||||
|
||||
from sqlalchemy import delete as sql_delete
|
||||
from sqlalchemy import update as sql_update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.exc import IntegrityError, SQLAlchemyError
|
||||
from sqlmodel import SQLModel, select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
from sqlmodel.sql.expression import SelectOfScalar
|
||||
@@ -24,7 +24,7 @@ class MultipleObjectsReturned(LookupError):
|
||||
async def _flush_or_rollback(session: AsyncSession) -> None:
|
||||
try:
|
||||
await session.flush()
|
||||
except Exception:
|
||||
except SQLAlchemyError:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
@@ -32,7 +32,7 @@ async def _flush_or_rollback(session: AsyncSession) -> None:
|
||||
async def _commit_or_rollback(session: AsyncSession) -> None:
|
||||
try:
|
||||
await session.commit()
|
||||
except Exception:
|
||||
except SQLAlchemyError:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
@@ -268,7 +268,7 @@ async def get_or_create(
|
||||
if existing is not None:
|
||||
return existing, False
|
||||
raise
|
||||
except Exception:
|
||||
except SQLAlchemyError:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from pathlib import Path
|
||||
import anyio
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
|
||||
from sqlmodel import SQLModel
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
@@ -69,9 +70,15 @@ async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
async with async_session_maker() as session:
|
||||
try:
|
||||
yield session
|
||||
except Exception:
|
||||
finally:
|
||||
try:
|
||||
in_txn = bool(session.in_transaction())
|
||||
except SQLAlchemyError:
|
||||
logger.exception("Failed to inspect session transaction state.")
|
||||
return
|
||||
if not in_txn:
|
||||
return
|
||||
try:
|
||||
await session.rollback()
|
||||
except Exception:
|
||||
except SQLAlchemyError:
|
||||
logger.exception("Failed to rollback session after request error.")
|
||||
raise
|
||||
|
||||
@@ -5,6 +5,7 @@ from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlmodel import Field, SQLModel
|
||||
|
||||
from app.db import crud
|
||||
@@ -43,12 +44,17 @@ class _Maker:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_session_rolls_back_on_dependency_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
fake_session = SimpleNamespace(rollbacks=0)
|
||||
@dataclass
|
||||
class _FakeDependencySession:
|
||||
rollbacks: int = 0
|
||||
|
||||
async def _rollback() -> None:
|
||||
fake_session.rollbacks += 1
|
||||
def in_transaction(self) -> bool:
|
||||
return True
|
||||
|
||||
fake_session.rollback = _rollback
|
||||
async def rollback(self) -> None:
|
||||
self.rollbacks += 1
|
||||
|
||||
fake_session = _FakeDependencySession()
|
||||
ctx = _SessionCtx(fake_session)
|
||||
monkeypatch.setattr(db_session, "async_session_maker", _Maker(ctx))
|
||||
|
||||
@@ -66,6 +72,9 @@ async def test_get_session_rolls_back_on_dependency_error(monkeypatch: pytest.Mo
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_rolls_back_when_commit_fails() -> None:
|
||||
class _CommitError(SQLAlchemyError):
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class _FailCommitSession:
|
||||
rollback_calls: int = 0
|
||||
@@ -82,7 +91,7 @@ async def test_create_rolls_back_when_commit_fails() -> None:
|
||||
return None
|
||||
|
||||
async def commit(self) -> None:
|
||||
raise RuntimeError("commit failed")
|
||||
raise _CommitError("commit failed")
|
||||
|
||||
async def rollback(self) -> None:
|
||||
self.rollback_calls += 1
|
||||
@@ -92,7 +101,7 @@ async def test_create_rolls_back_when_commit_fails() -> None:
|
||||
|
||||
session = _FailCommitSession()
|
||||
|
||||
with pytest.raises(RuntimeError, match="commit failed"):
|
||||
with pytest.raises(SQLAlchemyError, match="commit failed"):
|
||||
await crud.create(session, _Thing, name="demo")
|
||||
|
||||
assert session.rollback_calls == 1
|
||||
@@ -101,6 +110,9 @@ async def test_create_rolls_back_when_commit_fails() -> None:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_where_rolls_back_when_commit_fails() -> None:
|
||||
class _CommitError(SQLAlchemyError):
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class _FailCommitDmlSession:
|
||||
rollback_calls: int = 0
|
||||
@@ -111,14 +123,14 @@ async def test_delete_where_rolls_back_when_commit_fails() -> None:
|
||||
return SimpleNamespace(rowcount=3)
|
||||
|
||||
async def commit(self) -> None:
|
||||
raise RuntimeError("commit failed")
|
||||
raise _CommitError("commit failed")
|
||||
|
||||
async def rollback(self) -> None:
|
||||
self.rollback_calls += 1
|
||||
|
||||
session = _FailCommitDmlSession()
|
||||
|
||||
with pytest.raises(RuntimeError, match="commit failed"):
|
||||
with pytest.raises(SQLAlchemyError, match="commit failed"):
|
||||
await crud.delete_where(session, _Thing, commit=True)
|
||||
|
||||
assert session.exec_calls == 1
|
||||
|
||||
Reference in New Issue
Block a user