diff --git a/src/app_readonly.py b/src/app_readonly.py index 10fb031..fb737b5 100644 --- a/src/app_readonly.py +++ b/src/app_readonly.py @@ -30,6 +30,7 @@ async def lifespan(app: FastAPI): # Shutdown logger.info("Shutting down Hell Divers 2 API") scraper.close() + db.close_pool() # Initialize FastAPI app @@ -394,9 +395,17 @@ async def get_config(): } +@app.get("/api/livez", tags=["Health"]) +async def liveness_probe(): + """Lightweight liveness probe - no DB, for Kubernetes liveness checks. + Returns 200 if the process is alive. Use /api/health for readiness (includes DB check). + """ + return {"status": "ok"} + + @app.get("/api/health", tags=["Health"]) async def health_check(): - """Health check endpoint""" + """Health check endpoint - includes DB and upstream status (for readiness).""" # Returns local service status and upstream API status upstream_status = db.get_upstream_status() return { diff --git a/src/database.py b/src/database.py index 5a76a2c..edb0a57 100644 --- a/src/database.py +++ b/src/database.py @@ -1,19 +1,24 @@ import json import logging import psycopg2 +from psycopg2 import pool from datetime import datetime, timezone from typing import Dict, List, Optional logger = logging.getLogger(__name__) +# Connection pool defaults +DEFAULT_POOL_MIN_CONN = 2 +DEFAULT_POOL_MAX_CONN = 10 + class Database: - """PostgreSQL database manager for Hell Divers 2 API data""" + """PostgreSQL database manager for Hell Divers 2 API data with connection pooling""" def __init__(self, database_url: Optional[str] = None): """ - Initialize database connection - + Initialize database connection pool + Args: database_url: PostgreSQL connection string (postgresql://user:pass@host:port/db) If None, will try to get from DATABASE_URL env var @@ -21,17 +26,41 @@ def __init__(self, database_url: Optional[str] = None): if database_url is None: import os database_url = os.getenv("DATABASE_URL", "") - + if not database_url: raise ValueError("DATABASE_URL must be provided") - + self.database_url = database_url self._initialized = False - # Lazy initialization - only connect when first needed + self._pool: Optional[pool.ThreadedConnectionPool] = None + + def _get_pool(self) -> pool.ThreadedConnectionPool: + """Get or create the connection pool (lazy init)""" + if self._pool is None: + self._pool = pool.ThreadedConnectionPool( + minconn=DEFAULT_POOL_MIN_CONN, + maxconn=DEFAULT_POOL_MAX_CONN, + dsn=self.database_url, + ) + logger.info("Database connection pool initialized") + return self._pool def _get_connection(self): - """Get a database connection""" - return psycopg2.connect(self.database_url) + """Get a database connection from the pool (returns to pool on close).""" + conn = self._get_pool().getconn() + # Wrap so conn.close() returns to pool instead of closing the underlying connection + def _close(): + self._get_pool().putconn(conn) + + conn.close = _close + return conn + + def close_pool(self): + """Close the connection pool. Call on application shutdown.""" + if self._pool is not None: + self._pool.closeall() + self._pool = None + logger.info("Database connection pool closed") @staticmethod def _parse_expiration_time(expiration_time: str) -> Optional[datetime]: diff --git a/src/main.py b/src/main.py index 35ad8fd..1113df3 100644 --- a/src/main.py +++ b/src/main.py @@ -46,39 +46,39 @@ def run_collector(): from src.database import Database from src.collector import DataCollector from src.config import Config - + global _collector - + logger.info("Starting High Command API Collector") - + # Initialize database database_url = Config.DATABASE_URL if not database_url: raise ValueError("DATABASE_URL environment variable must be set") - + db = Database(database_url) logger.info("Database initialized with connection string") - + # Initialize collector interval = Config.SCRAPE_INTERVAL _collector = DataCollector(db, interval=interval) - + def signal_handler(sig, frame): """Handle shutdown signals""" logger.info("Received shutdown signal, stopping collector...") if _collector: _collector.stop() sys.exit(0) - + # Register signal handlers signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) - + # Start collector try: _collector.start() logger.info(f"Collector started with {interval}s interval") - + # Keep running while _collector.is_running: time.sleep(1) @@ -89,6 +89,7 @@ def signal_handler(sig, frame): finally: if _collector: _collector.stop() + db.close_pool() logger.info("Collector stopped") diff --git a/tests/test_app_readonly.py b/tests/test_app_readonly.py new file mode 100644 index 0000000..d3f97fa --- /dev/null +++ b/tests/test_app_readonly.py @@ -0,0 +1,47 @@ +"""Unit tests for app_readonly (read-only API used in production).""" + +import pytest +from fastapi.testclient import TestClient +from unittest.mock import patch +from src.app_readonly import app + + +@pytest.fixture +def client(): + """Create test client for readonly app""" + return TestClient(app) + + +class TestLivezEndpoint: + """Test lightweight liveness probe - no DB dependency.""" + + def test_livez_returns_200(self, client): + """Liveness probe returns 200 without hitting DB.""" + response = client.get("/api/livez") + assert response.status_code == 200 + data = response.json() + assert data == {"status": "ok"} + + +class TestHealthEndpoint: + """Test health endpoint (includes DB check).""" + + @patch("src.app_readonly.db.get_upstream_status") + def test_health_upstream_online(self, mock_status, client): + """Health check when upstream is online.""" + mock_status.return_value = True + response = client.get("/api/health") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert data["mode"] == "read-only" + assert data["upstream_api"] == "online" + + @patch("src.app_readonly.db.get_upstream_status") + def test_health_upstream_offline(self, mock_status, client): + """Health check when upstream is offline.""" + mock_status.return_value = False + response = client.get("/api/health") + assert response.status_code == 200 + data = response.json() + assert data["upstream_api"] == "offline"