Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions sqlspec/base.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
# ruff: noqa: PLR6301
import atexit
import contextlib
import re
from abc import ABC, abstractmethod
from collections.abc import Awaitable
Expand All @@ -19,6 +21,7 @@
from sqlspec.exceptions import NotFoundError
from sqlspec.statement import SQLStatement
from sqlspec.typing import ModelDTOT, StatementParameterType
from sqlspec.utils.sync_tools import maybe_async_

if TYPE_CHECKING:
from contextlib import AbstractAsyncContextManager, AbstractContextManager
Expand Down Expand Up @@ -202,6 +205,15 @@ class SQLSpec:

def __init__(self) -> None:
self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
# Register the cleanup handler to run at program exit
atexit.register(self._cleanup_pools)

def _cleanup_pools(self) -> None:
"""Clean up all open database pools at program exit."""
for config in self._configs.values():
if config.support_connection_pooling and config.pool_instance is not None:
with contextlib.suppress(Exception):
maybe_async_(config.close_pool)()

@overload
def add_config(self, config: "SyncConfigT") -> "type[SyncConfigT]": ...
Expand Down