From 721afd20d1eb0211a9bd7b0b9e4bebe4cc255405 Mon Sep 17 00:00:00 2001
From: Jonathan Le <jonathan.le@gemini.com>
Date: Mon, 25 Nov 2024 10:59:38 -0800
Subject: [PATCH] Add ability to dynamically extend flask cli management
commands (#201)
---
Dockerfile | 2 +-
api/app.py | 14 ++++-
api/config.py | 3 +
.../plugins/health_check_plugin/README.md | 50 ++++++++++++++++
.../plugins/health_check_plugin/__init__.py | 7 +++
examples/plugins/health_check_plugin/cli.py | 59 +++++++++++++++++++
examples/plugins/health_check_plugin/setup.py | 16 +++++
7 files changed, 149 insertions(+), 2 deletions(-)
create mode 100644 examples/plugins/health_check_plugin/README.md
create mode 100644 examples/plugins/health_check_plugin/__init__.py
create mode 100644 examples/plugins/health_check_plugin/cli.py
create mode 100644 examples/plugins/health_check_plugin/setup.py
diff --git a/Dockerfile b/Dockerfile
index 92754ea..5173c28 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,7 +3,7 @@
ARG PUSH_SENTRY_RELEASE="false"
# Build step #1: build the React front end
-FROM node:23-alpine AS build-step
+FROM node:22-alpine AS build-step
ARG SENTRY_RELEASE=""
WORKDIR /app
ENV PATH /app/node_modules/.bin:$PATH
diff --git a/api/app.py b/api/app.py
index 38557d0..8c041c7 100644
--- a/api/app.py
+++ b/api/app.py
@@ -5,6 +5,7 @@
import logging
import sys
import warnings
+from importlib.metadata import entry_points
from os import environ
from typing import Optional
@@ -183,6 +184,7 @@ def add_headers(response: Response) -> ResponseReturnValue:
##########################################
# Configure flask cli commands
##########################################
+ # Register static commands
app.cli.add_command(manage.init)
app.cli.add_command(manage.import_from_okta)
app.cli.add_command(manage.init_builtin_apps)
@@ -191,6 +193,16 @@ def add_headers(response: Response) -> ResponseReturnValue:
app.cli.add_command(manage.fix_role_memberships)
app.cli.add_command(manage.notify)
+ # Register dynamically loaded commands
+ flask_commands = entry_points(group="flask.commands")
+
+ for entry_point in flask_commands:
+ try:
+ command = entry_point.load()
+ app.cli.add_command(command)
+ except Exception as e:
+ logger.warning(f"Failed to load command '{entry_point.name}': {e}")
+
###########################################
# Configure APISpec for swagger support
###########################################
@@ -204,7 +216,7 @@ def add_headers(response: Response) -> ResponseReturnValue:
# https://github.com/marshmallow-code/apispec/issues/444
warnings.filterwarnings("ignore", message="Multiple schemas resolved to the name ")
# Ignore the following warning because nested schemas may declare less fields via only tuples
- # than the actual schema has specfieid in the fields tuple
+ # than the actual schema has specified in the fields tuple
warnings.filterwarnings("ignore", message="Only explicitly-declared fields will be included in the Schema Object")
app.register_blueprint(exception_views.bp)
diff --git a/api/config.py b/api/config.py
index 2615682..558f1d0 100644
--- a/api/config.py
+++ b/api/config.py
@@ -83,3 +83,6 @@ def default_user_search() -> list[str]:
FLASK_SENTRY_DSN = os.getenv("FLASK_SENTRY_DSN")
REACT_SENTRY_DSN = os.getenv("REACT_SENTRY_DSN")
+
+# Add APP_VERSION, defaulting to 'Not Defined' if not set
+APP_VERSION = os.getenv("APP_VERSION", "Not Defined")
diff --git a/examples/plugins/health_check_plugin/README.md b/examples/plugins/health_check_plugin/README.md
new file mode 100644
index 0000000..dfecaf6
--- /dev/null
+++ b/examples/plugins/health_check_plugin/README.md
@@ -0,0 +1,50 @@
+# Health Check Plugin
+
+This is an example plugin that demonstrates how to extend Flask CLI commands using plugins. The `health_check_plugin` adds a custom `health` command to the Flask CLI, which performs a health check of the application, including verifying database connectivity.
+
+## Overview
+
+The plugin consists of the following files:
+
+- **`__init__.py`**: Initializes the plugin by defining an `init_app` function that registers the CLI commands.
+- **`cli.py`**: Contains the implementation of the `health` command.
+- **`setup.py`**: Defines the plugin's setup configuration and registers the entry point for the CLI command.
+
+## Installation
+
+To install the plugin the App container Dockerfile
+
+```
+WORKDIR /app/plugins
+ADD ./examples/plugins/health_check_plugin ./health_check_plugin
+RUN pip install ./health_check_plugin
+
+# Reset working directory
+WORKDIR /app
+```
+
+## Usage
+
+After installing the plugin, the `health` command becomes available in the Flask CLI:
+
+```bash
+flask health
+```
+
+This command outputs the application's health status in JSON format, indicating the database connection status and the application version.
+
+## Purpose
+
+This plugin serves as an example of how to extend Flask CLI commands using plugins and entry points. It demonstrates:
+
+- How to create a custom CLI command in a plugin.
+- How to register the command using entry points in `setup.py`.
+
+By following this example, you can create your own plugins to extend the functionality of your Flask application's CLI in a modular and scalable way.
+
+## Files
+
+- **[`__init__.py`](./__init__.py)**: Plugin initialization code.
+- **[`cli.py`](./cli.py)**: Implementation of the `health` CLI command.
+- **[`setup.py`](./setup.py)**: Setup script defining the plugin metadata and entry points.
+
diff --git a/examples/plugins/health_check_plugin/__init__.py b/examples/plugins/health_check_plugin/__init__.py
new file mode 100644
index 0000000..bbaee7b
--- /dev/null
+++ b/examples/plugins/health_check_plugin/__init__.py
@@ -0,0 +1,7 @@
+from flask import Flask
+
+
+def init_app(app: Flask) -> None:
+ from .cli import health_command
+
+ app.cli.add_command(health_command)
diff --git a/examples/plugins/health_check_plugin/cli.py b/examples/plugins/health_check_plugin/cli.py
new file mode 100644
index 0000000..f570c83
--- /dev/null
+++ b/examples/plugins/health_check_plugin/cli.py
@@ -0,0 +1,59 @@
+import logging
+
+import click
+from flask.cli import with_appcontext
+from sqlalchemy import text
+
+
+@click.command("health")
+@with_appcontext
+def health_command() -> None:
+ """Displays application database health and metrics in JSON format."""
+ from flask import current_app, json
+
+ from api.extensions import db
+
+ logger = logging.getLogger(__name__)
+
+ try:
+ # Perform a simple database health check using SQLAlchemy
+ db.session.execute(text("SELECT 1"))
+ db_status = "connected"
+ error = None
+ logger.info("Database connection successful.")
+
+ # Retrieve all table names and their row counts
+ tables_query = text("""
+ SELECT table_name
+ FROM information_schema.tables
+ WHERE table_schema = 'public';
+ """)
+ tables = db.session.execute(tables_query).fetchall()
+
+ table_sizes = {}
+ for table in tables:
+ table_name = table[0]
+ row_count_query = text(f"SELECT COUNT(*) FROM {table_name}")
+ row_count = db.session.execute(row_count_query).scalar()
+ table_sizes[table_name] = row_count
+
+ except Exception as e:
+ db_status = "disconnected"
+ error = str(e)
+ table_sizes = {}
+ logger.error(f"Database connection error: {error}")
+
+ # Prepare the health status response
+ status = {
+ "status": "ok" if db_status == "connected" else "error",
+ "database": db_status,
+ "tables": table_sizes,
+ "version": current_app.config.get("APP_VERSION", "Not Defined"),
+ **({"error": error} if error else {}),
+ }
+
+ # Log the health status
+ logger.info(f"Health status: {status}")
+
+ # Output the health status as a JSON string
+ click.echo(json.dumps(status))
diff --git a/examples/plugins/health_check_plugin/setup.py b/examples/plugins/health_check_plugin/setup.py
new file mode 100644
index 0000000..dea404e
--- /dev/null
+++ b/examples/plugins/health_check_plugin/setup.py
@@ -0,0 +1,16 @@
+from setuptools import setup
+
+setup(
+ name="health_check_plugin",
+ version="0.1.0",
+ packages=["health_check_plugin"],
+ package_dir={"health_check_plugin": "."}, # Map package to current directory
+ install_requires=[
+ "Flask",
+ ],
+ entry_points={
+ "flask.commands": [
+ "health=health_check_plugin.cli:health_command",
+ ],
+ },
+)