Skip to content

Commit c2a9ad2

Browse files
febus982claude
andauthored
Remove obsolete async helpers and update engine/session cleanup logic (#84)
* Remove obsolete async helpers and update engine/session cleanup logic * Add test for atexit cleanup handler Ensures _cleanup_all_managers() is covered by tests, restoring 100% test coverage. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 2eaad26 commit c2a9ad2

File tree

5 files changed

+67
-162
lines changed

5 files changed

+67
-162
lines changed

sqlalchemy_bind_manager/_async_helpers.py

Lines changed: 0 additions & 42 deletions
This file was deleted.

sqlalchemy_bind_manager/_bind_manager.py

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@
1818
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
1919
# DEALINGS IN THE SOFTWARE.
2020

21-
from typing import Mapping, MutableMapping, Union
21+
import atexit
22+
import weakref
23+
from typing import ClassVar, Mapping, MutableMapping, Union
2224

2325
from pydantic import BaseModel, ConfigDict, StrictBool
2426
from sqlalchemy import MetaData, create_engine
@@ -32,7 +34,6 @@
3234
from sqlalchemy.orm import Session, sessionmaker
3335
from sqlalchemy.orm.decl_api import DeclarativeMeta, registry
3436

35-
from sqlalchemy_bind_manager._async_helpers import run_async_from_sync
3637
from sqlalchemy_bind_manager.exceptions import (
3738
InvalidConfigError,
3839
NotInitializedBindError,
@@ -73,6 +74,7 @@ class SQLAlchemyAsyncBind(BaseModel):
7374

7475
class SQLAlchemyBindManager:
7576
__binds: MutableMapping[str, Union[SQLAlchemyBind, SQLAlchemyAsyncBind]]
77+
_instances: ClassVar[weakref.WeakSet["SQLAlchemyBindManager"]] = weakref.WeakSet()
7678

7779
def __init__(
7880
self,
@@ -87,14 +89,24 @@ def __init__(
8789
self.__init_bind(name, conf)
8890
else:
8991
self.__init_bind(DEFAULT_BIND_NAME, config)
92+
SQLAlchemyBindManager._instances.add(self)
9093

91-
def __del__(self):
94+
def _dispose_sync(self) -> None:
95+
"""Dispose all engines synchronously.
96+
97+
This method is safe to call from any context, including __del__
98+
and atexit handlers. For async engines, it uses the underlying
99+
sync_engine to perform synchronous disposal.
100+
"""
92101
for bind in self.__binds.values():
93102
if isinstance(bind, SQLAlchemyAsyncBind):
94-
run_async_from_sync(bind.engine.dispose())
103+
bind.engine.sync_engine.dispose()
95104
else:
96105
bind.engine.dispose()
97106

107+
def __del__(self) -> None:
108+
self._dispose_sync()
109+
98110
def __init_bind(self, name: str, config: SQLAlchemyConfig):
99111
if not isinstance(config, SQLAlchemyConfig):
100112
raise InvalidConfigError(
@@ -210,3 +222,15 @@ def get_session(
210222
:return: The SQLAlchemy Session object
211223
"""
212224
return self.get_bind(bind_name).session_class()
225+
226+
227+
@atexit.register
228+
def _cleanup_all_managers() -> None:
229+
"""Cleanup handler that runs during interpreter shutdown.
230+
231+
This ensures all SQLAlchemyBindManager instances have their engines
232+
disposed before the interpreter exits, even if __del__ hasn't been
233+
called yet due to reference cycles or other GC timing issues.
234+
"""
235+
for manager in list(SQLAlchemyBindManager._instances):
236+
manager._dispose_sync()

sqlalchemy_bind_manager/_session_handler.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
)
2929
from sqlalchemy.orm import Session, scoped_session
3030

31-
from sqlalchemy_bind_manager._async_helpers import run_async_from_sync
3231
from sqlalchemy_bind_manager._bind_manager import (
3332
SQLAlchemyAsyncBind,
3433
SQLAlchemyBind,
@@ -69,12 +68,20 @@ def commit(self, session: Session) -> None:
6968
"""
7069
try:
7170
session.commit()
72-
except:
71+
except Exception:
7372
session.rollback()
7473
raise
7574

7675

7776
class AsyncSessionHandler:
77+
"""Async session handler for managing async scoped sessions.
78+
79+
Note: Unlike SessionHandler, this class does not implement __del__ cleanup
80+
because async_scoped_session.remove() is an async operation that cannot be
81+
safely executed during garbage collection. Sessions should be properly
82+
closed via the get_session() context manager.
83+
"""
84+
7885
scoped_session: async_scoped_session
7986

8087
def __init__(self, bind: SQLAlchemyAsyncBind):
@@ -85,11 +92,6 @@ def __init__(self, bind: SQLAlchemyAsyncBind):
8592
bind.session_class, asyncio.current_task
8693
)
8794

88-
def __del__(self):
89-
if not getattr(self, "scoped_session", None):
90-
return
91-
run_async_from_sync(self.scoped_session.remove())
92-
9395
@asynccontextmanager
9496
async def get_session(self, read_only: bool = False) -> AsyncIterator[AsyncSession]:
9597
session = self.scoped_session()
@@ -110,6 +112,6 @@ async def commit(self, session: AsyncSession) -> None:
110112
"""
111113
try:
112114
await session.commit()
113-
except:
115+
except Exception:
114116
await session.rollback()
115117
raise

tests/session_handler/test_session_lifecycle.py

Lines changed: 8 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,14 @@
1111
from sqlalchemy_bind_manager._session_handler import AsyncSessionHandler, SessionHandler
1212

1313

14-
async def test_session_is_removed_on_cleanup(session_handler_class, sa_bind):
15-
sh = session_handler_class(sa_bind)
14+
def test_sync_session_is_removed_on_cleanup(sa_manager):
15+
"""Test that sync SessionHandler removes session on garbage collection.
16+
17+
Note: AsyncSessionHandler does not implement __del__ cleanup because
18+
async_scoped_session.remove() is an async operation that cannot be
19+
safely executed during garbage collection.
20+
"""
21+
sh = SessionHandler(sa_manager.get_bind("sync"))
1622
original_session_remove = sh.scoped_session.remove
1723

1824
with patch.object(
@@ -26,53 +32,6 @@ async def test_session_is_removed_on_cleanup(session_handler_class, sa_bind):
2632
mocked_remove.assert_called_once()
2733

2834

29-
def test_session_is_removed_on_cleanup_even_if_loop_is_not_running(sa_manager):
30-
# Running the test without a loop will trigger the loop creation
31-
sh = AsyncSessionHandler(sa_manager.get_bind("async"))
32-
original_session_remove = sh.scoped_session.remove
33-
original_get_event_loop = asyncio.get_event_loop
34-
35-
with (
36-
patch.object(
37-
sh.scoped_session,
38-
"remove",
39-
wraps=original_session_remove,
40-
) as mocked_close,
41-
patch(
42-
"asyncio.get_event_loop",
43-
wraps=original_get_event_loop,
44-
) as mocked_get_event_loop,
45-
):
46-
# This should trigger the garbage collector and close the session
47-
sh = None
48-
49-
mocked_get_event_loop.assert_called_once()
50-
mocked_close.assert_called_once()
51-
52-
53-
def test_session_is_removed_on_cleanup_even_if_loop_search_errors_out(sa_manager):
54-
# Running the test without a loop will trigger the loop creation
55-
sh = AsyncSessionHandler(sa_manager.get_bind("async"))
56-
original_session_remove = sh.scoped_session.remove
57-
58-
with (
59-
patch.object(
60-
sh.scoped_session,
61-
"remove",
62-
wraps=original_session_remove,
63-
) as mocked_close,
64-
patch(
65-
"asyncio.get_event_loop",
66-
side_effect=RuntimeError(),
67-
) as mocked_get_event_loop,
68-
):
69-
# This should trigger the garbage collector and close the session
70-
sh = None
71-
72-
mocked_get_event_loop.assert_called_once()
73-
mocked_close.assert_called_once()
74-
75-
7635
@pytest.mark.parametrize("read_only_flag", [True, False])
7736
async def test_commit_is_called_only_if_not_read_only(
7837
read_only_flag,

tests/test_sqlalchemy_bind_manager.py

Lines changed: 21 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -75,39 +75,19 @@ def test_multiple_binds(multiple_config):
7575
assert isinstance(sa_manager.get_session("async"), AsyncSession)
7676

7777

78-
async def test_engine_is_disposed_on_cleanup(multiple_config):
79-
sa_manager = SQLAlchemyBindManager(multiple_config)
80-
sync_engine = sa_manager.get_bind("default").engine
81-
async_engine = sa_manager.get_bind("async").engine
82-
83-
original_sync_dispose = sync_engine.dispose
84-
original_async_dispose = async_engine.dispose
85-
86-
with (
87-
patch.object(
88-
sync_engine,
89-
"dispose",
90-
wraps=original_sync_dispose,
91-
) as mocked_dispose,
92-
patch.object(
93-
type(async_engine),
94-
"dispose",
95-
wraps=original_async_dispose,
96-
) as mocked_async_dispose,
97-
):
98-
sa_manager = None
99-
100-
mocked_dispose.assert_called_once()
101-
mocked_async_dispose.assert_called()
78+
def test_engine_is_disposed_on_cleanup(multiple_config):
79+
"""Test that engines are disposed synchronously during garbage collection.
10280
103-
104-
def test_engine_is_disposed_on_cleanup_even_if_no_loop(multiple_config):
81+
This test verifies that both sync and async engines are properly disposed
82+
using synchronous disposal (sync_engine.dispose() for async engines).
83+
"""
10584
sa_manager = SQLAlchemyBindManager(multiple_config)
10685
sync_engine = sa_manager.get_bind("default").engine
10786
async_engine = sa_manager.get_bind("async").engine
10887

10988
original_sync_dispose = sync_engine.dispose
110-
original_async_dispose = async_engine.dispose
89+
# For async engines, we now use sync_engine.dispose() for safe cleanup
90+
original_async_sync_dispose = async_engine.sync_engine.dispose
11191

11292
with (
11393
patch.object(
@@ -116,45 +96,27 @@ def test_engine_is_disposed_on_cleanup_even_if_no_loop(multiple_config):
11696
wraps=original_sync_dispose,
11797
) as mocked_dispose,
11898
patch.object(
119-
type(async_engine),
99+
async_engine.sync_engine,
120100
"dispose",
121-
wraps=original_async_dispose,
122-
) as mocked_async_dispose,
101+
wraps=original_async_sync_dispose,
102+
) as mocked_async_sync_dispose,
123103
):
124104
sa_manager = None
125105

126106
mocked_dispose.assert_called_once()
127-
mocked_async_dispose.assert_called()
107+
mocked_async_sync_dispose.assert_called_once()
108+
128109

110+
def test_atexit_cleanup_disposes_all_managers(multiple_config):
111+
"""Test that the atexit handler disposes all tracked manager instances."""
112+
from sqlalchemy_bind_manager._bind_manager import _cleanup_all_managers
129113

130-
def test_engine_is_disposed_on_cleanup_even_if_loop_search_errors_out(
131-
multiple_config,
132-
):
133114
sa_manager = SQLAlchemyBindManager(multiple_config)
134-
sync_engine = sa_manager.get_bind("default").engine
135-
async_engine = sa_manager.get_bind("async").engine
136115

137-
original_sync_dispose = sync_engine.dispose
138-
original_async_dispose = async_engine.dispose
116+
with patch.object(
117+
sa_manager,
118+
"_dispose_sync",
119+
) as mocked_dispose_sync:
120+
_cleanup_all_managers()
139121

140-
with (
141-
patch.object(
142-
sync_engine,
143-
"dispose",
144-
wraps=original_sync_dispose,
145-
) as mocked_dispose,
146-
patch.object(
147-
type(async_engine),
148-
"dispose",
149-
wraps=original_async_dispose,
150-
) as mocked_async_dispose,
151-
patch(
152-
"asyncio.get_event_loop",
153-
side_effect=RuntimeError(),
154-
) as mocked_get_event_loop,
155-
):
156-
sa_manager = None
157-
158-
mocked_get_event_loop.assert_called_once()
159-
mocked_dispose.assert_called_once()
160-
mocked_async_dispose.assert_called()
122+
mocked_dispose_sync.assert_called_once()

0 commit comments

Comments
 (0)