Skip to content

Commit 4a7785d

Browse files
committed
feat(database): add database URL retrieval and health check functionality
- Implemented a method to retrieve the current database URL from configuration. - Added a health check method to verify database connectivity and return status information. - Enhanced connection configuration for SQLite and PostgreSQL with specific optimizations. - Updated connection pooling settings to improve performance and reliability.
1 parent 54ba641 commit 4a7785d

File tree

1 file changed

+57
-1
lines changed

1 file changed

+57
-1
lines changed

src/tux/database/service.py

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from collections.abc import AsyncGenerator, Callable
1616
from contextlib import asynccontextmanager
17+
from datetime import UTC, datetime
1718
from typing import Any
1819

1920
import sentry_sdk
@@ -54,6 +55,10 @@ def __init__(self, *, echo: bool = False):
5455
self._session_factory: async_sessionmaker[AsyncSession] | None = None
5556
self._echo = echo
5657

58+
def get_database_url(self) -> str:
59+
"""Get the current database URL from configuration."""
60+
return get_database_url()
61+
5762
# =====================================================================
5863
# Connection & Session Management
5964
# =====================================================================
@@ -84,12 +89,35 @@ async def connect(self, database_url: str | None = None, *, echo: bool | None =
8489
echo_setting = echo if echo is not None else self._echo
8590

8691
logger.debug(f"Creating async SQLAlchemy engine (echo={echo_setting})")
92+
93+
# Enhanced connection configuration based on SQLModel best practices
94+
connect_args = {}
95+
if "sqlite" in database_url:
96+
# SQLite-specific optimizations
97+
connect_args = {
98+
"check_same_thread": False,
99+
"timeout": 30,
100+
}
101+
elif "postgresql" in database_url:
102+
# PostgreSQL-specific optimizations
103+
connect_args = {
104+
"server_settings": {
105+
"timezone": "UTC",
106+
"application_name": "TuxBot",
107+
},
108+
}
109+
87110
self._engine = create_async_engine(
88111
database_url,
89112
echo=echo_setting,
90-
pool_pre_ping=True,
113+
future=True, # Enable SQLAlchemy 2.0 style
114+
# Connection pooling configuration
115+
pool_pre_ping=True, # Verify connections before use
91116
pool_size=10,
92117
max_overflow=20,
118+
pool_timeout=30, # Connection timeout
119+
pool_recycle=1800, # Recycle connections after 30 minutes
120+
connect_args=connect_args,
93121
)
94122
self._session_factory = async_sessionmaker(
95123
self._engine,
@@ -121,6 +149,34 @@ async def disconnect(self) -> None:
121149
self._session_factory = None
122150
logger.info("Disconnected from database")
123151

152+
async def health_check(self) -> dict[str, Any]:
153+
"""Perform a database health check."""
154+
if not self.is_connected():
155+
return {"status": "disconnected", "error": "Database engine not connected"}
156+
157+
try:
158+
async with self.session() as session:
159+
# Simple query to test connectivity
160+
from sqlalchemy import text # noqa: PLC0415
161+
162+
result = await session.execute(text("SELECT 1"))
163+
value = result.scalar()
164+
165+
if value == 1:
166+
return {
167+
"status": "healthy",
168+
"pool_size": getattr(self._engine.pool, "size", "unknown") if self._engine else "unknown",
169+
"checked_connections": getattr(self._engine.pool, "checkedin", "unknown")
170+
if self._engine
171+
else "unknown",
172+
"timestamp": datetime.now(UTC).isoformat(),
173+
}
174+
return {"status": "unhealthy", "error": "Unexpected query result"}
175+
176+
except Exception as exc:
177+
logger.error(f"Database health check failed: {exc}")
178+
return {"status": "unhealthy", "error": str(exc)}
179+
124180
@asynccontextmanager
125181
async def session(self) -> AsyncGenerator[AsyncSession]:
126182
"""Return an async SQLAlchemy session context-manager."""

0 commit comments

Comments
 (0)