streamlit_healthcheck.server

  1# -*- coding: utf-8 -*-
  2"""
  3This module implements a FastAPI server for monitoring the health of a Streamlit application.
  4
  5Features:
  6
  7- Provides REST API endpoints to report health status of system resources, dependencies, and Streamlit pages.
  8- Uses a configurable health check service to gather health data.
  9- Supports startup and shutdown events for proper initialization and cleanup of the health check service.
 10- Includes endpoints:
 11    - `/health`: Returns complete health status.
 12    - `/health/system`: Returns system resource health (CPU, Memory, Disk).
 13    - `/health/dependencies`: Returns dependencies health status.
 14    - `/health/pages`: Returns Streamlit pages health status.
 15- Configurable via command-line arguments for host, port, and health check config file.
 16- Uses logging for operational visibility.
 17
 18Global Variables:
 19
 20- `health_service`: Instance of HealthCheckService, manages health checks.
 21- `config_file`: Path to health check configuration file.
 22
 23Usage:
 24
 25Run as a standalone server or import as a module.
 26"""
 27
 28from fastapi import FastAPI, HTTPException
 29from fastapi.responses import JSONResponse
 30from typing import Dict, Any, Optional
 31import uvicorn
 32from .healthcheck import HealthCheckService
 33import logging
 34import argparse
 35from contextlib import asynccontextmanager
 36
 37
 38
 39# Set up logging
 40logging.basicConfig(
 41    level=logging.INFO,
 42    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
 43    handlers=[
 44        logging.StreamHandler()  # Prints to console
 45    ]
 46)
 47logger = logging.getLogger(__name__)
 48
 49# Initialize health check service as global variable
 50health_service: Optional[HealthCheckService] = None
 51config_file: Optional[str] = None
 52
 53@asynccontextmanager
 54async def lifespan(app: FastAPI):
 55    """Lifespan context manager for FastAPI application"""
 56    # Startup
 57    global health_service
 58    logger.info("Initializing health check service")
 59    config_path = config_file if config_file is not None else "health_check_config.json"
 60    health_service = HealthCheckService(config_path=config_path)
 61    health_service.start()
 62    
 63    yield
 64    
 65    # Shutdown
 66    if health_service:
 67        logger.info("Stopping health check service")
 68        health_service.stop()
 69
 70
 71app = FastAPI(
 72    title="Streamlit Health Check API",
 73    description="API endpoints for monitoring Streamlit application health",
 74    version="1.0.0",
 75    lifespan=lifespan
 76)
 77
 78# Root endpoint providing service metadata and available endpoints
 79@app.get("/", response_model=Dict[str, Any])
 80async def root():
 81    """Asynchronous handler for the application root ("/") endpoint.
 82    Returns a JSONResponse containing basic service metadata intended for health
 83    checks and simple API discovery. The returned JSON includes the service name,
 84    version, a short description, and a mapping of available health-related
 85    endpoints.
 86    
 87    Returns:
 88    
 89        JSONResponse: HTTP 200 response with a JSON body.
 90        
 91    Notes:
 92    
 93        - This function is asynchronous and suitable for use with Starlette or FastAPI.
 94        - No input parameters are required.
 95    """
 96    
 97    return JSONResponse(
 98        content={
 99            "service": "streamlit-healthcheck",
100            "version": "1.0.0",
101            "description": "API for monitoring Streamlit application health",
102            "endpoints": {
103                "health": "/health",
104                "system": "/health/system",
105                "dependencies": "/health/dependencies",
106                "pages": "/health/pages"
107            }
108        }
109    )
110
111# Initialize health check service and config file path as global variables
112# health_service and config_file are already defined above
113
114
115@app.get("/health", response_model=Dict[str, Any])
116async def get_health_status():
117    """
118    Asynchronous endpoint that retrieves and returns the application's health status as a JSON response.
119    This coroutine uses the module-level `health_service` to obtain health data and returns it wrapped in a
120    fastapi.responses.JSONResponse. It logs unexpected errors and maps them to appropriate HTTP error responses.
121    
122    Behavior:
123    
124    - If the global `health_service` is not initialized or falsy, raises HTTPException(status_code=503).
125    - If `health_service.get_health_data()` succeeds, returns a JSONResponse with the health data (typically a dict).
126    - If any exception occurs while obtaining health data, the exception is logged and an HTTPException(status_code=500) is raised with the original error message.
127    
128    Returns:
129    
130            fastapi.responses.JSONResponse: A JSONResponse containing the health data returned by `health_service.get_health_data()`.
131    
132    Raises:
133    
134            fastapi.HTTPException: With status_code=503 when the health service is not initialized.
135            fastapi.HTTPException: With status_code=500 when an unexpected error occurs while retrieving health data.
136    """
137    
138    global health_service
139    if not health_service:
140        raise HTTPException(status_code=503, detail="Health check service not initialized")
141    
142    try:
143        health_data = health_service.get_health_data()
144        return JSONResponse(content=health_data)
145    except Exception as e:
146        logger.error(f"Error getting health data: {str(e)}")
147        raise HTTPException(status_code=500, detail=str(e))
148
149@app.get("/health/system", response_model=Dict[str, Any])
150async def get_system_health():
151    """
152    Asynchronously retrieve the system health payload from the global health service and return it
153    as a JSONResponse suitable for use in a FastAPI/Starlette endpoint.
154    
155    Behavior:
156    
157    - Verifies that the module-level 'health_service' is initialized; if not, raises HTTPException(503).
158    - Calls health_service.get_health_data() to obtain health information.
159    - Extracts the "system" sub-dictionary from the returned health data (defaults to an empty dict).
160    - Returns a fastapi.responses.JSONResponse with content {"system": <system_data>}.
161    - If any unexpected error occurs while obtaining or processing health data, logs the error and
162        raises HTTPException(500) with the error message.
163        
164    Returns:
165    
166            fastapi.responses.JSONResponse: JSON response containing the "system" health data.
167            
168    Raises:
169    
170            HTTPException: with status_code=503 if the health service is not initialized.
171            HTTPException: with status_code=500 if an unexpected error occurs while retrieving health data.
172            
173    Notes:
174    
175    - Relies on the module-level 'health_service' and 'logger' variables.
176    - Designed to be used as an async request handler in a web application.
177    """
178    global health_service
179    if not health_service:
180        raise HTTPException(status_code=503, detail="Health check service not initialized")
181    
182    try:
183        health_data = health_service.get_health_data()
184        return JSONResponse(content={"system": health_data.get("system", {})})
185    except Exception as e:
186        logger.error(f"Error getting system health data: {str(e)}")
187        raise HTTPException(status_code=500, detail=str(e))
188
189@app.get("/health/dependencies", response_model=Dict[str, Any])
190async def get_dependencies_health():
191    """
192    Asynchronously retrieve and return the health status of external dependencies.
193    This function relies on a module-level `health_service` object. If the service is
194    not initialized, it raises an HTTPException with status 503. Otherwise it calls
195    `health_service.get_health_data()` and returns a JSONResponse containing the
196    "dependencies" entry from the returned health data (an empty dict is returned
197    if that key is missing). Any unexpected error during retrieval is logged and
198    propagated as an HTTPException with status 500.
199    
200    Returns:
201    
202        JSONResponse: A JSON response with the shape {"dependencies": {...}}.
203        
204    Raises:
205    
206        HTTPException: 503 if the global health_service is not initialized.
207        HTTPException: 500 for any unexpected error while fetching health data.
208        
209    Notes:
210    
211        - The function is asynchronous but calls a synchronous `get_health_data()` on
212          the health service; the service implementation should be safe to call from
213          an async context.
214        - Errors are logged using the module-level `logger`.
215    """
216    
217    global health_service
218    if not health_service:
219        raise HTTPException(status_code=503, detail="Health check service not initialized")
220    
221    try:
222        health_data = health_service.get_health_data()
223        return JSONResponse(content={"dependencies": health_data.get("dependencies", {})})
224    except Exception as e:
225        logger.error(f"Error getting dependencies health data: {str(e)}")
226        raise HTTPException(status_code=500, detail=str(e))
227
228@app.get("/health/pages", response_model=Dict[str, Any])
229async def get_pages_health():
230    """
231    Asynchronously run health checks and return the health status for Streamlit pages.
232    This coroutine depends on a module-level `health_service` which must be initialized
233    before calling. It triggers a fresh health check by calling `health_service.run_all_checks()`
234    and then retrieves the aggregated health data via `health_service.get_health_data()`.
235    Only the "streamlit_pages" portion of the health data is returned in a FastAPI
236    JSONResponse (defaults to an empty dict if the key is absent).
237    
238    Returns:
239    
240        fastapi.responses.JSONResponse: A JSON response with the shape
241            {"streamlit_pages": {...}}.
242            
243    Raises:
244    
245        fastapi.HTTPException: If `health_service` is not initialized (status_code=503).
246        fastapi.HTTPException: If any unexpected error occurs while running checks or
247            retrieving data (status_code=500). The exception detail contains the original
248            error message.
249            
250    Side effects:
251    
252        - Calls `health_service.run_all_checks()` which may perform I/O or long-running checks.
253        - Logs errors to the module logger when exceptions occur.
254        
255    Notes:
256    
257        - This function is intended to be used in an ASGI/FastAPI context and must be awaited.
258        - Only the `"streamlit_pages"` key from the health payload is exposed to the caller.
259    """
260    
261    global health_service
262    if not health_service:
263        raise HTTPException(status_code=503, detail="Health check service not initialized")
264    
265    try:
266        health_service.run_all_checks()
267        health_data = health_service.get_health_data()
268        return JSONResponse(content={"streamlit_pages": health_data.get("streamlit_pages", {})})
269    except Exception as e:
270        logger.error(f"Error getting pages health data: {str(e)}")
271        raise HTTPException(status_code=500, detail=str(e))
272
273def start_api_server(host: str = "0.0.0.0", port: int = 8000, config: str = "health_check_config.json"):
274    """
275    Start the API server using uvicorn and set the global configuration file.
276    Sets the module-level variable `config_file` to the provided `config` path
277    and then starts the ASGI server by calling `uvicorn.run(app, host=host, port=port)`.
278    The call is blocking and will run until the server is stopped.
279    
280    Parameters
281    ----------
282    
283    host : str, optional
284        Host interface to bind the server to. Defaults to "0.0.0.0" (all interfaces).
285    port : int, optional
286        TCP port to listen on. Defaults to 8000.
287    config : str, optional
288        Path to the health check configuration file. Defaults to "health_check_config.json".
289        This value is assigned to the module-level `config_file` variable before the server starts.
290        
291    Returns
292    -------
293    
294    None
295        This function does not return; it blocks while the server is running.
296        
297    Raises
298    ------
299    
300    OSError
301        If the server cannot bind to the given host/port (for example, if the port is already in use).
302    
303    Exception
304        Any exceptions raised by uvicorn or the ASGI application are propagated.
305    
306    Example
307    -------
308    >>> start_api_server(host="127.0.0.1", port=8080, config="my_config.json")
309    """
310    
311    global config_file
312    config_file = config
313    uvicorn.run(app, host=host, port=port)
314
315def parse_args():
316    """
317    Parse command-line arguments for the Streamlit Health Check API Server.
318    
319    Defines and parses the following command-line options:
320        --host   (str)  Host address to bind. Default: "0.0.0.0".
321        --port   (int)  Port to run the server on. Default: 8000.
322        --config (str)  Path to the health check configuration file. Default: "health_check_config.json".
323        
324    Returns:
325    
326            argparse.Namespace: The parsed arguments with attributes `host`, `port`, and `config`.
327    """
328    
329    parser = argparse.ArgumentParser(description="Streamlit Health Check API Server")
330    parser.add_argument("--host", default="0.0.0.0", help="Host address to bind")
331    parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
332    parser.add_argument(
333        "--config", 
334        default="health_check_config.json",
335        help="Path to health check configuration file"
336    )
337    return parser.parse_args()
338
339if __name__ == "__main__":
340    args = parse_args()
341    logger.info(f"Starting server with config file: {args.config}")
342    start_api_server(host=args.host, port=args.port, config=args.config)