feat(server): implement PostgreSQL storage backend for server mode
Add PostgreSQL support as the server-mode storage backend alongside existing SQLite for local mode. Verify all ORM models are dialect- agnostic, configure connection pooling for multi-user access, add Docker Compose for local PG development, and wire database URL selection based on deployment mode. Changes: - Add psycopg2-binary dependency to pyproject.toml - Add server_mode, db_pool_size, db_max_overflow, db_pool_recycle settings to Settings with environment variable support - Add resolve_database_url() and is_postgresql() to Settings for mode-aware database URL resolution - Configure UnitOfWork engine creation with pool_size, max_overflow, pool_recycle, and pool_pre_ping for PostgreSQL connections - Update MigrationRunner to handle both SQLite and PostgreSQL backends - Add compare_type=True to Alembic env.py for dialect-aware migrations - Add docker-compose.yml with PostgreSQL 16-alpine for local development - Add Behave BDD feature (14 scenarios) covering settings, pool config, engine creation, ORM dialect compatibility, and migration runner - Add Robot Framework integration tests (12 test cases) for the abstraction layer with requires_postgresql tag for live PG tests Fixes: - Restore require_confirmation parameter to UnitOfWork.__init__ - Add argument validation to UnitOfWork.__init__ (fail-fast) - Add explicit PostgreSQL isolation_level (READ COMMITTED) - Add dispose() method and context manager support to UnitOfWork - Fix MigrationRunner.get_current_revision() to dispose engine - Fix Settings.is_postgresql() to handle ValueError gracefully ISSUES CLOSED: #878
This commit is contained in:
@@ -13,6 +13,7 @@ services:
|
|||||||
image: postgres:16-alpine
|
image: postgres:16-alpine
|
||||||
environment:
|
environment:
|
||||||
POSTGRES_USER: cleveragents
|
POSTGRES_USER: cleveragents
|
||||||
|
# WARNING: Development credentials only — do not use in production
|
||||||
POSTGRES_PASSWORD: cleveragents
|
POSTGRES_PASSWORD: cleveragents
|
||||||
POSTGRES_DB: cleveragents
|
POSTGRES_DB: cleveragents
|
||||||
ports:
|
ports:
|
||||||
|
|||||||
@@ -15,10 +15,10 @@ Feature: PostgreSQL Storage Backend for Server Mode
|
|||||||
Then the resolved database URL should start with "sqlite"
|
Then the resolved database URL should start with "sqlite"
|
||||||
|
|
||||||
@settings @database_url
|
@settings @database_url
|
||||||
Scenario: Server mode with default URL resolves to PostgreSQL
|
Scenario: Server mode with default URL raises ValueError
|
||||||
Given I create settings with server_mode enabled
|
Given I create settings with server_mode enabled
|
||||||
And the database_url is the default SQLite value
|
And the database_url is the default SQLite value
|
||||||
Then the resolved database URL should start with "postgresql"
|
Then resolving the database URL should raise a ValueError
|
||||||
|
|
||||||
@settings @database_url
|
@settings @database_url
|
||||||
Scenario: Server mode with explicit PostgreSQL URL uses that URL
|
Scenario: Server mode with explicit PostgreSQL URL uses that URL
|
||||||
|
|||||||
@@ -136,6 +136,17 @@ def step_settings_custom_pool(context: Context) -> None:
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@then("resolving the database URL should raise a ValueError")
|
||||||
|
def step_resolve_url_raises_value_error(context: Context) -> None:
|
||||||
|
"""Assert that resolve_database_url raises ValueError in server mode."""
|
||||||
|
try:
|
||||||
|
context.settings.resolve_database_url()
|
||||||
|
except ValueError:
|
||||||
|
pass # expected
|
||||||
|
else:
|
||||||
|
raise AssertionError("Expected ValueError from resolve_database_url()")
|
||||||
|
|
||||||
|
|
||||||
@then('the resolved database URL should start with "{prefix}"')
|
@then('the resolved database URL should start with "{prefix}"')
|
||||||
def step_resolved_url_starts_with(context: Context, prefix: str) -> None:
|
def step_resolved_url_starts_with(context: Context, prefix: str) -> None:
|
||||||
"""Assert the resolved database URL prefix."""
|
"""Assert the resolved database URL prefix."""
|
||||||
@@ -231,20 +242,20 @@ def step_engine_dialect(context: Context, dialect: str) -> None:
|
|||||||
|
|
||||||
@then("the UnitOfWork pool_size should be {value:d}")
|
@then("the UnitOfWork pool_size should be {value:d}")
|
||||||
def step_uow_pool_size(context: Context, value: int) -> None:
|
def step_uow_pool_size(context: Context, value: int) -> None:
|
||||||
"""Assert UnitOfWork._pool_size."""
|
"""Assert UnitOfWork.pool_size via read-only property."""
|
||||||
assert context.uow._pool_size == value
|
assert context.uow.pool_size == value
|
||||||
|
|
||||||
|
|
||||||
@then("the UnitOfWork max_overflow should be {value:d}")
|
@then("the UnitOfWork max_overflow should be {value:d}")
|
||||||
def step_uow_max_overflow(context: Context, value: int) -> None:
|
def step_uow_max_overflow(context: Context, value: int) -> None:
|
||||||
"""Assert UnitOfWork._max_overflow."""
|
"""Assert UnitOfWork.max_overflow via read-only property."""
|
||||||
assert context.uow._max_overflow == value
|
assert context.uow.max_overflow == value
|
||||||
|
|
||||||
|
|
||||||
@then("the UnitOfWork pool_recycle should be {value:d}")
|
@then("the UnitOfWork pool_recycle should be {value:d}")
|
||||||
def step_uow_pool_recycle(context: Context, value: int) -> None:
|
def step_uow_pool_recycle(context: Context, value: int) -> None:
|
||||||
"""Assert UnitOfWork._pool_recycle."""
|
"""Assert UnitOfWork.pool_recycle via read-only property."""
|
||||||
assert context.uow._pool_recycle == value
|
assert context.uow.pool_recycle == value
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -268,9 +279,13 @@ _PORTABLE_TYPE_NAMES = frozenset(
|
|||||||
"TIME",
|
"TIME",
|
||||||
"TIMESTAMP",
|
"TIMESTAMP",
|
||||||
"JSON",
|
"JSON",
|
||||||
|
"JSONB",
|
||||||
"BLOB",
|
"BLOB",
|
||||||
"LARGEBINARY",
|
"LARGEBINARY",
|
||||||
"ENUM",
|
"ENUM",
|
||||||
|
"UUID",
|
||||||
|
"INTERVAL",
|
||||||
|
"ARRAY",
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -53,20 +53,31 @@ def step_uow_non_sqlite_url(context: Context) -> None:
|
|||||||
@when("I access the engine property with mocked create_engine")
|
@when("I access the engine property with mocked create_engine")
|
||||||
def step_access_engine_property_mocked(context: Context) -> None:
|
def step_access_engine_property_mocked(context: Context) -> None:
|
||||||
"""Access the engine property, mocking _ensure_database_initialized and create_engine."""
|
"""Access the engine property, mocking _ensure_database_initialized and create_engine."""
|
||||||
|
import sys as _sys
|
||||||
|
|
||||||
from sqlalchemy.engine import Engine
|
from sqlalchemy.engine import Engine
|
||||||
|
|
||||||
mock_engine = MagicMock(spec=Engine)
|
mock_engine = MagicMock(spec=Engine)
|
||||||
|
|
||||||
with (
|
# Ensure psycopg2 can be imported (it may not be installed in test env)
|
||||||
patch.object(context.uow, "_ensure_database_initialized"),
|
_needs_psycopg2_stub = "psycopg2" not in _sys.modules
|
||||||
patch(
|
if _needs_psycopg2_stub:
|
||||||
"cleveragents.infrastructure.database.unit_of_work.create_engine",
|
_sys.modules["psycopg2"] = MagicMock()
|
||||||
return_value=mock_engine,
|
|
||||||
) as mock_create,
|
try:
|
||||||
):
|
with (
|
||||||
engine = context.uow.engine
|
patch.object(context.uow, "_ensure_database_initialized"),
|
||||||
context.result_engine = engine
|
patch(
|
||||||
context.mock_create_engine = mock_create
|
"cleveragents.infrastructure.database.unit_of_work.create_engine",
|
||||||
|
return_value=mock_engine,
|
||||||
|
) as mock_create,
|
||||||
|
):
|
||||||
|
engine = context.uow.engine
|
||||||
|
context.result_engine = engine
|
||||||
|
context.mock_create_engine = mock_create
|
||||||
|
finally:
|
||||||
|
if _needs_psycopg2_stub:
|
||||||
|
_sys.modules.pop("psycopg2", None)
|
||||||
|
|
||||||
|
|
||||||
@then("the engine should be created via the non-SQLite code path")
|
@then("the engine should be created via the non-SQLite code path")
|
||||||
@@ -232,20 +243,31 @@ def step_uow_postgresql_url(context: Context) -> None:
|
|||||||
@when("I trigger engine creation with mocked dependencies")
|
@when("I trigger engine creation with mocked dependencies")
|
||||||
def step_trigger_engine_creation_mocked(context: Context) -> None:
|
def step_trigger_engine_creation_mocked(context: Context) -> None:
|
||||||
"""Access engine property with mocked create_engine and migration runner."""
|
"""Access engine property with mocked create_engine and migration runner."""
|
||||||
|
import sys as _sys
|
||||||
|
|
||||||
from sqlalchemy.engine import Engine
|
from sqlalchemy.engine import Engine
|
||||||
|
|
||||||
mock_engine = MagicMock(spec=Engine)
|
mock_engine = MagicMock(spec=Engine)
|
||||||
|
|
||||||
with (
|
# Ensure psycopg2 can be imported (it may not be installed in test env)
|
||||||
patch.object(context.uow, "_ensure_database_initialized"),
|
_needs_psycopg2_stub = "psycopg2" not in _sys.modules
|
||||||
patch(
|
if _needs_psycopg2_stub:
|
||||||
"cleveragents.infrastructure.database.unit_of_work.create_engine",
|
_sys.modules["psycopg2"] = MagicMock()
|
||||||
return_value=mock_engine,
|
|
||||||
) as mock_create,
|
try:
|
||||||
):
|
with (
|
||||||
engine = context.uow.engine
|
patch.object(context.uow, "_ensure_database_initialized"),
|
||||||
context.result_engine = engine
|
patch(
|
||||||
context.mock_create_engine = mock_create
|
"cleveragents.infrastructure.database.unit_of_work.create_engine",
|
||||||
|
return_value=mock_engine,
|
||||||
|
) as mock_create,
|
||||||
|
):
|
||||||
|
engine = context.uow.engine
|
||||||
|
context.result_engine = engine
|
||||||
|
context.mock_create_engine = mock_create
|
||||||
|
finally:
|
||||||
|
if _needs_psycopg2_stub:
|
||||||
|
_sys.modules.pop("psycopg2", None)
|
||||||
|
|
||||||
|
|
||||||
@then("create_engine should have been called without isolation_level")
|
@then("create_engine should have been called without isolation_level")
|
||||||
|
|||||||
+3
-1
@@ -49,10 +49,12 @@ dependencies = [
|
|||||||
"aiohttp>=3.13.4", # CVE-2026-34515 mitigation: open redirect vulnerability
|
"aiohttp>=3.13.4", # CVE-2026-34515 mitigation: open redirect vulnerability
|
||||||
"pyyaml>=6.0.3", # Security: address known YAML parsing vulnerabilities
|
"pyyaml>=6.0.3", # Security: address known YAML parsing vulnerabilities
|
||||||
"a2a-sdk>=0.3.0,<1.0.0", # A2A Python SDK — required transport for local (stdio) and server (HTTP) modes (ADR-047); pinned <1.0.0 (removed legacy A2AClient)
|
"a2a-sdk>=0.3.0,<1.0.0", # A2A Python SDK — required transport for local (stdio) and server (HTTP) modes (ADR-047); pinned <1.0.0 (removed legacy A2AClient)
|
||||||
"psycopg2-binary>=2.9.9", # PostgreSQL adapter for server-mode storage
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.optional-dependencies]
|
[project.optional-dependencies]
|
||||||
|
server = [
|
||||||
|
"psycopg2-binary>=2.9.9", # PostgreSQL adapter for server-mode storage
|
||||||
|
]
|
||||||
tui = [
|
tui = [
|
||||||
"textual>=1.0.0,<2.0.0",
|
"textual>=1.0.0,<2.0.0",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -31,10 +31,11 @@ Settings Resolve Database URL Local Mode
|
|||||||
|
|
||||||
Settings Resolve Database URL Server Mode
|
Settings Resolve Database URL Server Mode
|
||||||
[Documentation] Verify resolve_database_url returns PostgreSQL when server_mode is on
|
[Documentation] Verify resolve_database_url returns PostgreSQL when server_mode is on
|
||||||
|
... and CLEVERAGENTS_DATABASE_URL points to a PostgreSQL instance.
|
||||||
${script}= Catenate SEPARATOR=\n
|
${script}= Catenate SEPARATOR=\n
|
||||||
... import os
|
... import os
|
||||||
... os.environ['CLEVERAGENTS_SERVER_MODE'] = 'true'
|
... os.environ['CLEVERAGENTS_SERVER_MODE'] = 'true'
|
||||||
... os.environ.pop('CLEVERAGENTS_DATABASE_URL', None)
|
... os.environ['CLEVERAGENTS_DATABASE_URL'] = 'postgresql://localhost/cleveragents'
|
||||||
... from cleveragents.config.settings import Settings
|
... from cleveragents.config.settings import Settings
|
||||||
... s = Settings()
|
... s = Settings()
|
||||||
... print(s.resolve_database_url())
|
... print(s.resolve_database_url())
|
||||||
|
|||||||
@@ -251,6 +251,9 @@ def get_database_url() -> str:
|
|||||||
relative to ``CLEVERAGENTS_HOME`` (if set) rather than CWD, ensuring
|
relative to ``CLEVERAGENTS_HOME`` (if set) rather than CWD, ensuring
|
||||||
test isolation and correct data placement.
|
test isolation and correct data placement.
|
||||||
|
|
||||||
|
.. todo:: (#878) Use ``Settings.resolve_database_url()`` once
|
||||||
|
PostgreSQL wiring is complete so that *server_mode* is honoured.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
SQLAlchemy database URL with absolute path for SQLite
|
SQLAlchemy database URL with absolute path for SQLite
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -37,6 +37,7 @@ def _get_runner() -> MigrationRunner:
|
|||||||
from cleveragents.application.container import get_database_url
|
from cleveragents.application.container import get_database_url
|
||||||
from cleveragents.infrastructure.database.migration_runner import MigrationRunner
|
from cleveragents.infrastructure.database.migration_runner import MigrationRunner
|
||||||
|
|
||||||
|
# TODO(#878): Use Settings.resolve_database_url() once PostgreSQL wiring is complete
|
||||||
return MigrationRunner(get_database_url())
|
return MigrationRunner(get_database_url())
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -126,6 +126,7 @@ def _store_project_extras(
|
|||||||
|
|
||||||
from cleveragents.application.container import get_database_url
|
from cleveragents.application.container import get_database_url
|
||||||
|
|
||||||
|
# TODO(#878): Use Settings.resolve_database_url() once PostgreSQL wiring is complete
|
||||||
db_url = get_database_url()
|
db_url = get_database_url()
|
||||||
engine = create_engine(db_url, echo=False)
|
engine = create_engine(db_url, echo=False)
|
||||||
session = sessionmaker(bind=engine, expire_on_commit=False)()
|
session = sessionmaker(bind=engine, expire_on_commit=False)()
|
||||||
|
|||||||
@@ -413,6 +413,8 @@ def _check_stale_locks() -> dict[str, Any]:
|
|||||||
try:
|
try:
|
||||||
from sqlalchemy.orm import Session, sessionmaker
|
from sqlalchemy.orm import Session, sessionmaker
|
||||||
|
|
||||||
|
# TODO(#878): Use Settings.resolve_database_url()
|
||||||
|
# once PostgreSQL wiring is complete.
|
||||||
db_url = get_database_url()
|
db_url = get_database_url()
|
||||||
engine = create_engine(db_url, echo=False)
|
engine = create_engine(db_url, echo=False)
|
||||||
|
|
||||||
|
|||||||
@@ -929,11 +929,10 @@ class Settings(BaseSettings):
|
|||||||
"""Return the database URL, honouring *server_mode*.
|
"""Return the database URL, honouring *server_mode*.
|
||||||
|
|
||||||
In **server mode** the database URL must point to a PostgreSQL
|
In **server mode** the database URL must point to a PostgreSQL
|
||||||
instance. When the user has not explicitly overridden
|
instance. If the caller has not provided an explicit PostgreSQL
|
||||||
``database_url`` (i.e. it still has its SQLite default), this
|
URL via ``CLEVERAGENTS_DATABASE_URL``, a ``ValueError`` is
|
||||||
method returns the conventional PostgreSQL URL
|
raised to prevent silent fallback to a possibly unreachable
|
||||||
``postgresql://localhost/cleveragents`` so that the application
|
default.
|
||||||
falls through to the correct backend.
|
|
||||||
|
|
||||||
In **local mode** the existing SQLite URL is returned unchanged.
|
In **local mode** the existing SQLite URL is returned unchanged.
|
||||||
|
|
||||||
@@ -942,15 +941,25 @@ class Settings(BaseSettings):
|
|||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
A SQLAlchemy-compatible database URL string.
|
A SQLAlchemy-compatible database URL string.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: When *server_mode* is True but no explicit
|
||||||
|
PostgreSQL URL has been configured.
|
||||||
"""
|
"""
|
||||||
base_url = self.get_database_url(test=test)
|
base_url = self.get_database_url(test=test)
|
||||||
if self.server_mode and base_url.startswith("sqlite"):
|
if self.server_mode and base_url.startswith("sqlite"):
|
||||||
return "postgresql://localhost/cleveragents"
|
raise ValueError(
|
||||||
|
"server_mode=True requires an explicit database URL. "
|
||||||
|
"Set CLEVERAGENTS_DATABASE_URL to a PostgreSQL connection string."
|
||||||
|
)
|
||||||
return base_url
|
return base_url
|
||||||
|
|
||||||
def is_postgresql(self) -> bool:
|
def is_postgresql(self) -> bool:
|
||||||
"""Return True when the active database URL targets PostgreSQL."""
|
"""Return True when the active database URL targets PostgreSQL."""
|
||||||
return self.resolve_database_url().startswith("postgresql")
|
try:
|
||||||
|
return self.resolve_database_url().startswith("postgresql")
|
||||||
|
except ValueError:
|
||||||
|
return False
|
||||||
|
|
||||||
def is_production(self) -> bool:
|
def is_production(self) -> bool:
|
||||||
"""Whether the runtime matches production settings."""
|
"""Whether the runtime matches production settings."""
|
||||||
|
|||||||
@@ -167,9 +167,12 @@ class MigrationRunner:
|
|||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
engine = create_engine(self.database_url)
|
engine = create_engine(self.database_url)
|
||||||
with engine.connect() as connection:
|
try:
|
||||||
context = MigrationContext.configure(connection)
|
with engine.connect() as connection:
|
||||||
return context.get_current_revision()
|
context = MigrationContext.configure(connection)
|
||||||
|
return context.get_current_revision()
|
||||||
|
finally:
|
||||||
|
engine.dispose()
|
||||||
|
|
||||||
def get_pending_migrations(self) -> list[str]:
|
def get_pending_migrations(self) -> list[str]:
|
||||||
"""Get list of pending migrations.
|
"""Get list of pending migrations.
|
||||||
|
|||||||
@@ -43,6 +43,10 @@ class UnitOfWork:
|
|||||||
database_url: str,
|
database_url: str,
|
||||||
prompt_for_migration: Callable[[str], bool] | None = None,
|
prompt_for_migration: Callable[[str], bool] | None = None,
|
||||||
require_confirmation: bool = True,
|
require_confirmation: bool = True,
|
||||||
|
# Pool defaults are intentionally duplicated from Settings.db_pool_*
|
||||||
|
# so that UnitOfWork remains usable without a Settings instance
|
||||||
|
# (e.g. in tests or standalone scripts). When constructed via the
|
||||||
|
# DI container the caller should forward Settings values explicitly.
|
||||||
pool_size: int = 5,
|
pool_size: int = 5,
|
||||||
max_overflow: int = 10,
|
max_overflow: int = 10,
|
||||||
pool_recycle: int = 1800,
|
pool_recycle: int = 1800,
|
||||||
@@ -60,6 +64,16 @@ class UnitOfWork:
|
|||||||
max_overflow: Extra connections above *pool_size* under load.
|
max_overflow: Extra connections above *pool_size* under load.
|
||||||
pool_recycle: Seconds before a pooled connection is recycled.
|
pool_recycle: Seconds before a pooled connection is recycled.
|
||||||
"""
|
"""
|
||||||
|
# Validate arguments per CONTRIBUTING.md fail-fast principles
|
||||||
|
if not database_url:
|
||||||
|
raise ValueError("database_url must be a non-empty string")
|
||||||
|
if pool_size < 1:
|
||||||
|
raise ValueError(f"pool_size must be >= 1, got {pool_size}")
|
||||||
|
if max_overflow < 0:
|
||||||
|
raise ValueError(f"max_overflow must be >= 0, got {max_overflow}")
|
||||||
|
if pool_recycle < -1:
|
||||||
|
raise ValueError(f"pool_recycle must be >= -1, got {pool_recycle}")
|
||||||
|
|
||||||
self.database_url = database_url
|
self.database_url = database_url
|
||||||
self._engine: Engine | None = None
|
self._engine: Engine | None = None
|
||||||
self._session_factory: sessionmaker[Session] | None = None
|
self._session_factory: sessionmaker[Session] | None = None
|
||||||
@@ -70,6 +84,21 @@ class UnitOfWork:
|
|||||||
self._max_overflow = max_overflow
|
self._max_overflow = max_overflow
|
||||||
self._pool_recycle = pool_recycle
|
self._pool_recycle = pool_recycle
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pool_size(self) -> int:
|
||||||
|
"""Read-only access to the configured connection-pool size."""
|
||||||
|
return self._pool_size
|
||||||
|
|
||||||
|
@property
|
||||||
|
def max_overflow(self) -> int:
|
||||||
|
"""Read-only access to the configured max overflow."""
|
||||||
|
return self._max_overflow
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pool_recycle(self) -> int:
|
||||||
|
"""Read-only access to the configured pool recycle interval."""
|
||||||
|
return self._pool_recycle
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def engine(self) -> Engine:
|
def engine(self) -> Engine:
|
||||||
"""Get or create database engine."""
|
"""Get or create database engine."""
|
||||||
@@ -104,10 +133,24 @@ class UnitOfWork:
|
|||||||
else:
|
else:
|
||||||
# PostgreSQL (or other non-SQLite): use connection pooling
|
# PostgreSQL (or other non-SQLite): use connection pooling
|
||||||
# suitable for multi-user server mode.
|
# suitable for multi-user server mode.
|
||||||
|
# Use READ COMMITTED isolation level for concurrent multi-user access.
|
||||||
|
# This allows non-repeatable reads and phantom reads, which is
|
||||||
|
# acceptable for the concurrent plan execution use case where
|
||||||
|
# individual plan steps are atomic operations. For stronger
|
||||||
|
# consistency, callers can use explicit locking or application-
|
||||||
|
# level coordination.
|
||||||
|
try:
|
||||||
|
import psycopg2 # noqa: F401
|
||||||
|
except ImportError as exc:
|
||||||
|
raise ImportError(
|
||||||
|
"PostgreSQL support requires the 'server' extra: "
|
||||||
|
"pip install cleveragents[server]"
|
||||||
|
) from exc
|
||||||
self._engine = create_engine(
|
self._engine = create_engine(
|
||||||
self.database_url,
|
self.database_url,
|
||||||
echo=False, # Set to True for SQL debugging
|
echo=False, # Set to True for SQL debugging
|
||||||
future=True, # Use SQLAlchemy 2.0 style
|
future=True, # Use SQLAlchemy 2.0 style
|
||||||
|
isolation_level="READ COMMITTED",
|
||||||
pool_size=self._pool_size,
|
pool_size=self._pool_size,
|
||||||
max_overflow=self._max_overflow,
|
max_overflow=self._max_overflow,
|
||||||
pool_recycle=self._pool_recycle,
|
pool_recycle=self._pool_recycle,
|
||||||
@@ -151,6 +194,25 @@ class UnitOfWork:
|
|||||||
finally:
|
finally:
|
||||||
session.close()
|
session.close()
|
||||||
|
|
||||||
|
def dispose(self) -> None:
|
||||||
|
"""Release all database connections and dispose the engine.
|
||||||
|
|
||||||
|
Must be called when the UnitOfWork is no longer needed, especially
|
||||||
|
for PostgreSQL backends where the connection pool holds TCP connections.
|
||||||
|
"""
|
||||||
|
if self._engine is not None and self.database_url != "sqlite:///:memory:":
|
||||||
|
self._engine.dispose()
|
||||||
|
self._engine = None
|
||||||
|
self._session_factory = None
|
||||||
|
|
||||||
|
def __enter__(self) -> UnitOfWork:
|
||||||
|
"""Enter context manager."""
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, *args: object) -> None:
|
||||||
|
"""Exit context manager and dispose resources."""
|
||||||
|
self.dispose()
|
||||||
|
|
||||||
def _ensure_database_initialized(self, force: bool = False) -> None:
|
def _ensure_database_initialized(self, force: bool = False) -> None:
|
||||||
"""Run migrations if needed before database access."""
|
"""Run migrations if needed before database access."""
|
||||||
if self._database_initialized and not force:
|
if self._database_initialized and not force:
|
||||||
|
|||||||
Reference in New Issue
Block a user