Spaces:
Runtime error
Runtime error
File size: 2,722 Bytes
18ade12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """
Database connection test utility.
This module provides a simple connection test function that
verifies database connectivity by executing a SELECT 1 query.
"""
import asyncio
import logging
from typing import Tuple
from sqlalchemy import text
from sqlmodel.ext.asyncio.session import AsyncSession
from src.models.base import get_engine
logger = logging.getLogger(__name__)
async def test_connection() -> Tuple[bool, str]:
"""
Test database connection by executing a simple query.
Returns:
Tuple[bool, str]: (success, message)
- success: True if connection succeeded
- message: Descriptive message about the result
Example:
>>> import asyncio
>>> success, message = asyncio.run(test_connection())
>>> print(f"Connection {'succeeded' if success else 'failed'}: {message}")
"""
try:
engine = get_engine()
async with AsyncSession(engine) as session:
result = await session.exec(text("SELECT 1"))
value = result.scalar()
if value == 1:
logger.info("Database connection test successful")
return True, "Database connection successful (SELECT 1 = 1)"
else:
logger.warning(f"Unexpected result from SELECT 1: {value}")
return False, f"Unexpected result from SELECT 1: {value}"
except Exception as e:
error_type = type(e).__name__
# Don't expose connection details in error message
logger.error(f"Database connection test failed: {error_type}")
return False, f"Database connection failed: {error_type}"
async def test_connection_with_timeout(timeout_seconds: float = 5.0) -> Tuple[bool, str]:
"""
Test database connection with a timeout.
Args:
timeout_seconds: Maximum time to wait for connection (default 5 seconds)
Returns:
Tuple[bool, str]: (success, message)
"""
try:
return await asyncio.wait_for(
test_connection(),
timeout=timeout_seconds
)
except asyncio.TimeoutError:
logger.error(f"Database connection test timed out after {timeout_seconds}s")
return False, f"Connection timed out after {timeout_seconds} seconds"
def run_connection_test() -> None:
"""
CLI entry point to test database connection.
Usage:
python -m backend.src.db.connection_test
"""
print("Testing database connection...")
success, message = asyncio.run(test_connection_with_timeout())
if success:
print(f"SUCCESS: {message}")
else:
print(f"FAILED: {message}")
exit(1)
if __name__ == "__main__":
run_connection_test()
|