Make webserver resilient to database not starting up (#649)

* Update changelog * Add try except around database config * Add database health checks to health endpoint * Attempt to add tests * Create engine without system exit * Fix test mock by using the starlette client * Fix linting errors * add the PR link to the changelog * fix the visualize test * fix linting, fix the visualize test Co-authored-by: ThomasLaPiana <[email protected]>
ethyca · May 19, 2022 · 6b3f4c7 · 6b3f4c7
1 parent e3cc072
commit 6b3f4c7
Show file tree

Hide file tree

Showing 5 changed files with 220 additions and 138 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -45,9 +45,9 @@ The types of changes are:
 * Resolved a failure with populating applicable data subject rights to a data map
 * Updated `fideslog` to v1.1.5, resolving an issue where some exceptions thrown by the SDK were not handled as expected
 * Host static files via fidesapi [#621](https://github.com/ethyca/fides/pull/621)
+* Updated the webserver so that it won't fail if the database is inaccessible [#649](https://github.com/ethyca/fides/pull/649)
 * Handle complex characters in external tests  [#661](https://github.com/ethyca/fides/pull/661)
 
-
 ## [1.6.0](https://github.com/ethyca/fides/compare/1.5.3...1.6.0) - 2022-05-02
 
 ### Added

diff --git a/src/fidesapi/database/database.py b/src/fidesapi/database/database.py
@@ -3,15 +3,21 @@
 """
 from os import path
 
-from alembic import command
+from alembic import command, script
 from alembic.config import Config
 from alembic.migration import MigrationContext
+from alembic.runtime import migration
 from fideslang import DEFAULT_TAXONOMY
 from loguru import logger as log
+from sqlalchemy import create_engine
 from sqlalchemy_utils.functions import create_database, database_exists
 
 from fidesapi.sql_models import SqlAlchemyBase, sql_model_map
-from fidesapi.utils.errors import AlreadyExistsError, QueryError
+from fidesapi.utils.errors import (
+    AlreadyExistsError,
+    QueryError,
+    get_full_exception_name,
+)
 from fidesctl.core.utils import get_db_engine
 
 from .crud import create_resource, upsert_resources
@@ -108,3 +114,23 @@ def reset_db(database_url: str) -> None:
     version = migration_context._version  # pylint: disable=protected-access
     if version.exists(connection):
         version.drop(connection)
+
+
+def get_db_health(database_url: str) -> str:
+    """Checks if the db is reachable and up to date in alembic migrations"""
+    try:
+        engine = create_engine(database_url)
+        alembic_config = get_alembic_config(database_url)
+        alembic_script_directory = script.ScriptDirectory.from_config(alembic_config)
+        with engine.begin() as conn:
+            context = migration.MigrationContext.configure(conn)
+            if (
+                context.get_current_revision()
+                != alembic_script_directory.get_current_head()
+            ):
+                return "needs migration"
+        return "healthy"
+    except Exception as error:  # pylint: disable=broad-except
+        error_type = get_full_exception_name(error)
+        log.error(f"Unable to reach the database: {error_type}: {error}")
+        return "unhealthy"
diff --git a/src/fidesapi/main.py b/src/fidesapi/main.py
@@ -2,6 +2,7 @@
 Contains the code that sets up the API.
 """
 
+
 from datetime import datetime
 from enum import Enum
 from logging import WARNING
@@ -17,8 +18,10 @@
 import fidesctl
 from fidesapi import view
 from fidesapi.database import database
+from fidesapi.database.database import get_db_health
 from fidesapi.routes import crud, visualize
 from fidesapi.routes.util import API_PREFIX, WEBAPP_DIRECTORY, WEBAPP_INDEX
+from fidesapi.utils.errors import get_full_exception_name
 from fidesapi.utils.logger import setup as setup_logging
 from fidesctl.core.config import FidesctlConfig, get_config
 
@@ -42,8 +45,12 @@ def configure_routes() -> None:
 
 async def configure_db(database_url: str) -> None:
     "Set up the db to be used by the app."
-    database.create_db_if_not_exists(database_url)
-    await database.init_db(database_url)
+    try:
+        database.create_db_if_not_exists(database_url)
+        await database.init_db(database_url)
+    except Exception as error:  # pylint: disable=broad-except
+        error_type = get_full_exception_name(error)
+        log.error(f"Unable to configure database: {error_type}: {error}")
 
 
 @app.on_event("startup")
@@ -96,20 +103,44 @@ async def log_request(request: Request, call_next: Callable) -> Response:
                     "example": {
                         "status": "healthy",
                         "version": "1.0.0",
+                        "database": "healthy",
+                    }
+                }
+            }
+        },
+        status.HTTP_503_SERVICE_UNAVAILABLE: {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "detail": {
+                            "status": "healthy",
+                            "version": "1.0.0",
+                            "database": "unhealthy",
+                        }
                     }
                 }
             }
-        }
+        },
     },
     tags=["Health"],
 )
 async def health() -> Dict:
     "Confirm that the API is running and healthy."
-    return {
+    database_health = get_db_health(CONFIG.api.sync_database_url)
+    response = {
         "status": "healthy",
         "version": str(fidesctl.__version__),
+        "database": database_health,
     }
 
+    for key in response:
+        if response[key] == "unhealthy":
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=response
+            )
+
+    return response
+
 
 class DBActions(str, Enum):
     "The available path parameters for the `/admin/db/{action}` endpoint."
@@ -152,7 +183,9 @@ def read_other_paths(request: Request) -> FileResponse:
 
     # raise 404 for anything that should be backend endpoint but we can't find it
     if path.startswith(API_PREFIX[1:]):
-        raise HTTPException(status_code=404, detail="Item not found")
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Item not found"
+        )
 
     # otherwise return the index
     return FileResponse(WEBAPP_INDEX)

diff --git a/src/fidesapi/utils/errors.py b/src/fidesapi/utils/errors.py
@@ -54,3 +54,13 @@ def __init__(self) -> None:
             status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={"error": "a database query failed"},
         )
+
+
+def get_full_exception_name(exception: Exception) -> str:
+    """Get the full exception name
+    i.e. get sqlalchemy.exc.IntegrityError instead of just IntegrityError
+    """
+    module = exception.__class__.__module__
+    if module is None or module == str.__class__.__module__:
+        return exception.__class__.__name__
+    return module + "." + exception.__class__.__name__