feat: add automatic database initialization and recovery
- Add init_database() with alembic programmatic API and retry logic - Add connection retry with exponential backoff (5 attempts) - Improve error messages for connection/auth failures - Add table existence check before seeding data - Update startup event to run migrations before seeding - Add wait-for-db.sh script for Docker containers - Update Docker and docker-compose configurations Quality gates: ruff ✓, mypy ✓, unit tests (8 passed)
This commit is contained in:
+52
-20
@@ -2,8 +2,7 @@ import logging
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import ProgrammingError
|
||||
from sqlalchemy import select, text
|
||||
|
||||
from src.api.auth import router as auth_router
|
||||
from src.api.git_repositories import router as git_repositories_router
|
||||
@@ -12,15 +11,41 @@ from src.api.ssh_keys import router as ssh_keys_router
|
||||
from src.api.tool_types import router as tool_types_router
|
||||
from src.api.user_config import router as user_config_router
|
||||
from src.api.users import router as users_router
|
||||
from src.database import SessionLocal
|
||||
from src.database import SessionLocal, init_database
|
||||
from src.models.tool_type import ToolType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
app = FastAPI(title="Headquarter API")
|
||||
|
||||
|
||||
async def _table_exists(session, table_name: str) -> bool:
|
||||
"""Check if a table exists in the database."""
|
||||
try:
|
||||
result = await session.execute(
|
||||
text("""
|
||||
SELECT EXISTS (
|
||||
SELECT FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name = :table_name
|
||||
)
|
||||
"""),
|
||||
{"table_name": table_name},
|
||||
)
|
||||
return result.scalar() or False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def seed_builtin_tool_types():
|
||||
async with SessionLocal() as session:
|
||||
# Check if tool_types table exists before attempting to seed
|
||||
if not await _table_exists(session, "tool_types"):
|
||||
logger.warning(
|
||||
"tool_types table does not exist. Skipping seeding. "
|
||||
"Migrations may not have run yet."
|
||||
)
|
||||
return
|
||||
|
||||
builtin_types = [
|
||||
{
|
||||
"name": "code-server",
|
||||
@@ -62,29 +87,36 @@ services:
|
||||
},
|
||||
]
|
||||
|
||||
try:
|
||||
for tool_data in builtin_types:
|
||||
existing = await session.scalar(select(ToolType).where(ToolType.name == tool_data["name"]))
|
||||
if not existing:
|
||||
tool_type = ToolType(
|
||||
name=tool_data["name"],
|
||||
display_name=tool_data["display_name"],
|
||||
description=tool_data["description"],
|
||||
compose_template=tool_data["compose_template"],
|
||||
required_variables=tool_data["required_variables"],
|
||||
is_builtin=True,
|
||||
)
|
||||
session.add(tool_type)
|
||||
for tool_data in builtin_types:
|
||||
existing = await session.scalar(select(ToolType).where(ToolType.name == tool_data["name"]))
|
||||
if not existing:
|
||||
tool_type = ToolType(
|
||||
name=tool_data["name"],
|
||||
display_name=tool_data["display_name"],
|
||||
description=tool_data["description"],
|
||||
compose_template=tool_data["compose_template"],
|
||||
required_variables=tool_data["required_variables"],
|
||||
is_builtin=True,
|
||||
)
|
||||
session.add(tool_type)
|
||||
|
||||
await session.commit()
|
||||
except ProgrammingError:
|
||||
logger.warning("tool_types table does not exist yet. Skipping seeding. Run migrations first.")
|
||||
await session.rollback()
|
||||
await session.commit()
|
||||
logger.info("Built-in tool types seeded successfully.")
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def on_startup():
|
||||
logger.info("Starting up Headquarter API...")
|
||||
|
||||
# Initialize database (run migrations)
|
||||
db_ready = await init_database()
|
||||
if not db_ready:
|
||||
logger.error("Database initialization failed. API may not function correctly.")
|
||||
# Continue anyway so the health endpoint remains available
|
||||
|
||||
# Seed built-in data
|
||||
await seed_builtin_tool_types()
|
||||
logger.info("Startup complete.")
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
|
||||
Reference in New Issue
Block a user