To fix the rate limiting for /healthz, reorganize the app/__init__ a bunch to support decorators (#280)

* To fix the rate limiting for /healthz, reorganize the app/__init__ a bunch to support decorators

Signed-off-by: Irving Popovetsky <[email protected]>

* Making the linter happy broke the app, so quiet the linter instead

Signed-off-by: Irving Popovetsky <[email protected]>

* nobody found my requirements.txt joke funny

Signed-off-by: Irving Popovetsky <[email protected]>

* Fix the /metrics endpoint

Signed-off-by: Irving Popovetsky <[email protected]>

* test the healthz endpoint

Signed-off-by: Irving Popovetsky <[email protected]>

* Revert "Fix the /metrics endpoint"

This reverts commit ddfdcc5.

* Comment out the /environment endpoint

Co-authored-by: Aaron Suarez <[email protected]>

Author: irvingpop

app/__init__.py

@@ -1,11 +1,16 @@
 from algoliasearch.search_client import SearchClient
 from configs import Config
+
 from flask import Flask
 from flask_migrate import Migrate
 from flask_sqlalchemy import SQLAlchemy
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
-from app.healthcheck import add_health_check
+
+from healthcheck import HealthCheck
+# from healthcheck import EnvironmentDump
+
+from app.versioning import versioned
 
 db = SQLAlchemy()
 migrate = Migrate()
@@ -14,30 +19,50 @@
 search_client = SearchClient.create(Config.ALGOLIA_APP_ID, Config.ALGOLIA_API_KEY)
 index = search_client.init_index(Config.INDEX_NAME)
 
+app = Flask(__name__, static_folder='app/static')
+app.config.from_object(Config)
+app.url_map.strict_slashes = False
 
-def create_app(config_class=Config):
-    app = Flask(__name__, static_folder=None)
-    app.config.from_object(config_class)
-    app.url_map.strict_slashes = False
+limiter = Limiter(
+    app,
+    key_func=get_remote_address,
+    default_limits=["200 per day", "50 per hour"]
+)
 
-    Limiter(
-        app,
-        key_func=get_remote_address,
-        default_limits=["200 per day", "50 per hour"]
-    )
+db.init_app(app)
+migrate.init_app(app, db)
 
-    db.init_app(app)
-    migrate.init_app(app, db)
+from app.api import bp as api_bp # noqa
+app.register_blueprint(api_bp, url_prefix='/api/v1')
 
-    from app.api import bp as api_bp
-    app.register_blueprint(api_bp, url_prefix='/api/v1')
+from app.views import bp as view_bp # noqa
+app.register_blueprint(view_bp)
 
-    from app.views import bp as view_bp
-    app.register_blueprint(view_bp)
+from app.errors import bp as error_bp # noqa
+app.register_blueprint(error_bp)
 
-    from app.errors import bp as error_bp
-    app.register_blueprint(error_bp)
 
-    add_health_check(app)
+@app.route("/healthz")
+@limiter.exempt
+def healthz():
+    health = HealthCheck()
+    health.add_section("application", application_data)
+    return health.run()
 
-    return app
+
+# @app.route("/environment")
+# @limiter.limit("1 per hour")
+# def environment():
+#     envdump = EnvironmentDump()
+#     envdump.add_section("application", application_data)
+#     return envdump.run()
+
+
+@versioned
+def application_data(version):
+    return dict(
+        apiVersion=version,
+        status="ok",
+        status_code=200,
+        data=None
+    )

app/healthcheck.py

@@ -1,9 +1,12 @@
-from app import cli, create_app
+from app import app, cli
 from app.models import Category, Language, Resource, db
-from werkzeug.wsgi import DispatcherMiddleware
+from werkzeug.middleware.dispatcher import DispatcherMiddleware
 from prometheus_client import make_wsgi_app
 
-app = create_app()
+
+if __name__ == "__main__":
+    app.run()
+
 cli.register(app, db)
 
 # Add prometheus wsgi middleware to route /metrics requests

run.py

@@ -1,10 +1,9 @@
 import pytest
 from algoliasearch.exceptions import (AlgoliaException,
                                       AlgoliaUnreachableHostException)
-from app import create_app
+from app import app
 from app import db as _db
 from app.utils import standardize_response
-from configs import Config
 
 TEST_DATABASE_URI = 'sqlite:///:memory:'
 
@@ -16,19 +15,19 @@
 # away the need for a persistent DB and be able to use this method again.
 @pytest.fixture(scope='session')
 def session_app(request):
-    Config.SQLALCHEMY_DATABASE_URI = TEST_DATABASE_URI
-    Config.TESTING = True
-    app = create_app(Config)
+    flask_app = app
+    flask_app.config['TESTING'] = True
+    flask_app.config['SQLALCHEMY_DATABASE_URI'] = TEST_DATABASE_URI
 
     # Establish an application context before running the tests.
-    ctx = app.app_context()
+    ctx = flask_app.app_context()
     ctx.push()
 
     def teardown():
         ctx.pop()
 
     request.addfinalizer(teardown)
-    return app
+    return flask_app
 
 
 # TODO: session_db is not currently used because we need a
@@ -72,9 +71,9 @@ def teardown():
 
 @pytest.fixture(scope='module')
 def module_client():
-    Config.SQLALCHEMY_DATABASE_URI = TEST_DATABASE_URI
-    Config.TESTING = True
-    flask_app = create_app(Config)
+    flask_app = app
+    flask_app.config['TESTING'] = True
+    flask_app.config['SQLALCHEMY_DATABASE_URI'] = TEST_DATABASE_URI
 
     # Flask provides a way to test your application by exposing the Werkzeug test Client
     # and handling the context locals for you.

tests/conftest.py

@@ -113,3 +113,14 @@ def test_rate_limit(module_client, module_db):
     # Response should be a failure on request 51
     response = client.get('api/v1/resources')
     assert_correct_response(response, 429)
+
+
+# Ensure the healthz endpoint is never rate limited
+def test_rate_limit_healthz(module_client, module_db):
+    client = module_client
+
+    for _ in range(50):
+        client.get('/healthz')
+
+    response = client.get('/healthz')
+    assert (response.status_code == 200)