diff --git a/api/main.py b/api/main.py index 010e6b4..9d902b3 100644 --- a/api/main.py +++ b/api/main.py @@ -3,8 +3,9 @@ from contextlib import asynccontextmanager from time import perf_counter -from fastapi import FastAPI, HTTPException, Request +from fastapi import FastAPI, HTTPException, Request, Depends, status from fastapi.responses import JSONResponse, PlainTextResponse +from fastapi.security import APIKeyHeader from prometheus_client import CONTENT_TYPE_LATEST, Counter, Histogram, generate_latest from mlops_nlp.config import load_config @@ -20,9 +21,21 @@ REQUEST_LATENCY = Histogram("request_latency_seconds", "HTTP request latency", ["path", "method"]) PROMETHEUS_ENABLED = APP_CONFIG.monitoring.enable_prometheus +API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False) + pipeline: InferencePipeline | None = None +def get_api_key(api_key_header: str = Depends(API_KEY_HEADER)): + if APP_CONFIG.api.api_key: + if api_key_header != APP_CONFIG.api.api_key: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Could not validate credentials", + ) + return api_key_header + + @asynccontextmanager async def lifespan(app: FastAPI): global pipeline @@ -81,7 +94,10 @@ def health() -> JSONResponse: @app.post("/predict", response_model=PredictionResponse) -def predict(payload: PredictionRequest) -> PredictionResponse: +def predict( + payload: PredictionRequest, + _api_key: str = Depends(get_api_key) +) -> PredictionResponse: if pipeline is None: raise HTTPException(status_code=503, detail="Model is not loaded. Train model first.") prediction, confidence = pipeline.predict(payload.text) diff --git a/src/mlops_nlp/config.py b/src/mlops_nlp/config.py index 17487ce..0fb5444 100644 --- a/src/mlops_nlp/config.py +++ b/src/mlops_nlp/config.py @@ -49,6 +49,7 @@ class ApiConfig(BaseModel): title: str version: str description: str + api_key: str | None = None class MonitoringConfig(BaseModel): @@ -88,4 +89,8 @@ def load_config(config_path: str | Path = "configs/config.yaml") -> AppConfig: if enable_prometheus is not None: config.monitoring.enable_prometheus = _parse_bool(enable_prometheus) + api_key = os.getenv("MLOPS_API_KEY") + if api_key: + config.api.api_key = api_key + return config diff --git a/tests/test_api_auth.py b/tests/test_api_auth.py new file mode 100644 index 0000000..cc75cd2 --- /dev/null +++ b/tests/test_api_auth.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import pytest +from fastapi.testclient import TestClient +from api.main import app +import os + +client = TestClient(app) + +def test_predict_no_auth(): + # If MLOPS_API_KEY is not set, it should allow access (default behavior for local dev if not configured) + # But for this test, we want to ensure that IF it is set, it blocks. + # However, the current implementation in main.py checks APP_CONFIG.api.api_key which is loaded at start. + pass + +@pytest.fixture +def mock_api_key(monkeypatch): + monkeypatch.setenv("MLOPS_API_KEY", "test-secret-key") + # We need to reload the config or the app to pick up the env var if it's already loaded. + # In api/main.py, APP_CONFIG = load_config() happens at module level. + # For a robust test, we might need to override the dependency. + return "test-secret-key" + +def test_predict_with_auth_dependency_override(): + from api.main import get_api_key, app + + # Mocking the actual config value for the test + import api.main + original_key = api.main.APP_CONFIG.api.api_key + api.main.APP_CONFIG.api.api_key = "test-secret-key" + + try: + # 1. No key provided + response = client.post("/predict", json={"text": "hello"}) + assert response.status_code == 403 + + # 2. Wrong key provided + response = client.post("/predict", json={"text": "hello"}, headers={"X-API-Key": "wrong-key"}) + assert response.status_code == 403 + + # 3. Correct key provided + # We need a model loaded for a full 200, otherwise it might be 503 if no model exists. + response = client.post("/predict", json={"text": "hello"}, headers={"X-API-Key": "test-secret-key"}) + # We expect 503 (Model not loaded) or 200 (Success), but NOT 403. + assert response.status_code != 403 + + finally: + api.main.APP_CONFIG.api.api_key = original_key