From ad1f9e7d89f628e562d3976d65c8996aca489a23 Mon Sep 17 00:00:00 2001 From: gogoncalves Date: Sun, 28 Jun 2026 23:09:42 -0300 Subject: [PATCH 1/2] feat(database): per-operation statement_timeout_ms for db_session/async_db_session --- pyproject.toml | 2 +- serpens/database/__init__.py | 20 +++++++-- tests/test_database.py | 84 +++++++++++++++++++++++++++++++++++- 3 files changed, 101 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ace1022..af0d041 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/serpens/database/__init__.py b/serpens/database/__init__.py index 0f4f70b..f0b3870 100644 --- a/serpens/database/__init__.py +++ b/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,12 @@ 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 + return f"SET LOCAL statement_timeout = {int(statement_timeout_ms)}" + + def _normalize_sync_url(url): if url and url.startswith("postgres://"): return url.replace("postgres://", "postgresql+psycopg2://", 1) @@ -161,11 +167,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 +185,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: diff --git a/tests/test_database.py b/tests/test_database.py index 24e740c..a317b3b 100644 --- a/tests/test_database.py +++ b/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,88 @@ 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_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() From 5816f9851af810044dcfceac0ab608737b62bfeb Mon Sep 17 00:00:00 2001 From: gogoncalves Date: Mon, 29 Jun 2026 00:52:10 -0300 Subject: [PATCH 2/2] fix(database): ignore non-positive statement_timeout_ms --- serpens/database/__init__.py | 5 ++++- tests/test_database.py | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/serpens/database/__init__.py b/serpens/database/__init__.py index f0b3870..2be2503 100644 --- a/serpens/database/__init__.py +++ b/serpens/database/__init__.py @@ -88,7 +88,10 @@ 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 - return f"SET LOCAL statement_timeout = {int(statement_timeout_ms)}" + timeout = int(statement_timeout_ms) + if timeout <= 0: + return None + return f"SET LOCAL statement_timeout = {timeout}" def _normalize_sync_url(url): diff --git a/tests/test_database.py b/tests/test_database.py index a317b3b..3e6d7c7 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -126,6 +126,10 @@ def test_sql_coerces_to_int(self): "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:")