feat(database): per-operation statement_timeout_ms

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "noverde-serpens"
-version = "2.8.1"
+version = "2.9.0"
 description = "A set of Python utilities, recipes and snippets"
 readme = {file = "README.md", content-type = "text/markdown"}
 authors = [

serpens/database/__init__.py

@@ -9,7 +9,7 @@
 from datetime import datetime
 from typing import AsyncIterator, Iterator, Optional
 
-from sqlalchemy import DateTime, Engine, MetaData, create_engine, event
+from sqlalchemy import DateTime, Engine, MetaData, create_engine, event, text
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
     AsyncSession,
@@ -85,6 +85,15 @@ def _engine_args(url, pool_use_lifo=None):
     }
 
 
+def _statement_timeout_sql(dialect_name: str, statement_timeout_ms: Optional[int]) -> Optional[str]:
+    if statement_timeout_ms is None or dialect_name != "postgresql":
+        return None
+    timeout = int(statement_timeout_ms)
+    if timeout <= 0:
+        return None
+    return f"SET LOCAL statement_timeout = {timeout}"
+
+
 def _normalize_sync_url(url):
     if url and url.startswith("postgres://"):
         return url.replace("postgres://", "postgresql+psycopg2://", 1)
@@ -161,11 +170,14 @@ async def async_dispose() -> None:
 
 
 @contextmanager
-def db_session() -> Iterator[Session]:
+def db_session(statement_timeout_ms: Optional[int] = None) -> Iterator[Session]:
     if SessionLocal is None:
         bind()
     sess = SessionLocal()
     try:
+        timeout_sql = _statement_timeout_sql(_engine.dialect.name, statement_timeout_ms)
+        if timeout_sql:
+            sess.execute(text(timeout_sql))
         yield sess
         sess.commit()
     except BaseException:
@@ -176,11 +188,16 @@ def db_session() -> Iterator[Session]:
 
 
 @asynccontextmanager
-async def async_db_session() -> AsyncIterator[AsyncSession]:
+async def async_db_session(
+    statement_timeout_ms: Optional[int] = None,
+) -> AsyncIterator[AsyncSession]:
     if AsyncSessionLocal is None:
         async_bind()
     sess = AsyncSessionLocal()
     try:
+        timeout_sql = _statement_timeout_sql(_async_engine.dialect.name, statement_timeout_ms)
+        if timeout_sql:
+            await sess.execute(text(timeout_sql))
         yield sess
         await sess.commit()
     except BaseException:

tests/test_database.py

@@ -1,6 +1,6 @@
 import asyncio
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from sqlalchemy import Column, Integer, String, select
 
@@ -107,6 +107,92 @@ def test_timestamp_mixin_on_update(self):
             self.assertGreaterEqual(row.updated_at, original_updated)
 
 
+class TestStatementTimeout(unittest.TestCase):
+    def test_sql_built_for_postgres(self):
+        self.assertEqual(
+            database._statement_timeout_sql("postgresql", 3000),
+            "SET LOCAL statement_timeout = 3000",
+        )
+
+    def test_sql_none_for_non_postgres(self):
+        self.assertIsNone(database._statement_timeout_sql("sqlite", 3000))
+
+    def test_sql_none_when_timeout_unset(self):
+        self.assertIsNone(database._statement_timeout_sql("postgresql", None))
+
+    def test_sql_coerces_to_int(self):
+        self.assertEqual(
+            database._statement_timeout_sql("postgresql", "1500"),
+            "SET LOCAL statement_timeout = 1500",
+        )
+
+    def test_sql_none_for_non_positive(self):
+        self.assertIsNone(database._statement_timeout_sql("postgresql", 0))
+        self.assertIsNone(database._statement_timeout_sql("postgresql", -100))
+
+    def test_db_session_accepts_timeout_on_non_postgres(self):
+        database.dispose()
+        engine = database.bind("sqlite:///:memory:")
+        Base.metadata.create_all(engine)
+        try:
+            with db_session(statement_timeout_ms=3000) as sess:
+                sess.add(_Item(name="x"))
+            with db_session() as sess:
+                self.assertEqual(len(sess.execute(select(_Item)).scalars().all()), 1)
+        finally:
+            database.dispose()
+
+    def test_db_session_executes_set_local_for_postgres(self):
+        mock_sess = MagicMock()
+        engine = MagicMock()
+        engine.dialect.name = "postgresql"
+        with patch.object(database, "SessionLocal", MagicMock(return_value=mock_sess)):
+            with patch.object(database, "_engine", engine):
+                with db_session(statement_timeout_ms=2500):
+                    pass
+
+        executed_sql = str(mock_sess.execute.call_args[0][0])
+        self.assertEqual(executed_sql, "SET LOCAL statement_timeout = 2500")
+
+    def test_async_db_session_executes_set_local_for_postgres(self):
+        mock_sess = MagicMock()
+        mock_sess.execute = AsyncMock()
+        mock_sess.commit = AsyncMock()
+        mock_sess.close = AsyncMock()
+        engine = MagicMock()
+        engine.dialect.name = "postgresql"
+
+        async def run():
+            with patch.object(database, "AsyncSessionLocal", MagicMock(return_value=mock_sess)):
+                with patch.object(database, "_async_engine", engine):
+                    async with async_db_session(statement_timeout_ms=2500):
+                        pass
+
+        asyncio.run(run())
+
+        mock_sess.execute.assert_awaited_once()
+        executed_sql = str(mock_sess.execute.await_args[0][0])
+        self.assertEqual(executed_sql, "SET LOCAL statement_timeout = 2500")
+
+    def test_async_db_session_no_set_local_for_non_postgres(self):
+        mock_sess = MagicMock()
+        mock_sess.execute = AsyncMock()
+        mock_sess.commit = AsyncMock()
+        mock_sess.close = AsyncMock()
+        engine = MagicMock()
+        engine.dialect.name = "sqlite"
+
+        async def run():
+            with patch.object(database, "AsyncSessionLocal", MagicMock(return_value=mock_sess)):
+                with patch.object(database, "_async_engine", engine):
+                    async with async_db_session(statement_timeout_ms=2500):
+                        pass
+
+        asyncio.run(run())
+
+        mock_sess.execute.assert_not_awaited()
+
+
 class TestDeclarativeBaseFactory(unittest.TestCase):
     def test_no_schema(self):
         b = database.declarative_base()