Add option to manually set maintenance mode

Use `python -m pyramid_heroku.maintenance` to manually set Heroku
maintenance mode.

Also, some cleanup and refactor.

CHANGES.rst

@@ -9,6 +9,14 @@ Changes
    This helps when we want to enable/disable the maintenance mode manually or externally.
    [sayanarijit]
 
+ * Add options to manually enable or disable Heroku maintenance mode.
+   Use `pyramid_heroku.maintenance` script to manage the maintenance state.
+   [sayanarijit]
+
+ * Heroku API client related code has been moved from `pyramid_heroku.migrate` to
+   `pyramid_heroku.heroku`, while the `shell` function is now in `pyramid_heroku`.
+   [sayanarijit]
+
 
 0.7.0
 -----

README.rst

@@ -14,6 +14,7 @@ It provides the following:
   ``<app>.herokuapp.com`` domain for any non-whitelisted IPs.
 * ``migrate.py`` script for automatically running alembic migrations on
   deploy.
+* ``maintenance.py`` script for controlling Heroku maintenance mode.
 
 
 Installation
@@ -46,9 +47,10 @@ Usage example for tweens::
         config.include('pyramid_heroku.herokuapp_access')
         return config.make_wsgi_app()
 
-The `pyramid_heroku.herokuapp_access` tween depends on
-`pyramid_heroku.client_addr` tween and it requires you to list whitelisted IPs
-in the `pyramid_heroku.herokuapp_whitelist` setting.
+The ``pyramid_heroku.herokuapp_access`` tween depends on
+``pyramid_heroku.client_addr`` tween and it requires you to list whitelisted IPs
+in the ``pyramid_heroku.herokuapp_whitelist`` setting.
+
 
 Usage example for automatic alembic migration script::
 
@@ -65,13 +67,27 @@ Usage example for automatic alembic migration script::
 For migration script to work, you need to set the ``MIGRATE_API_SECRET_HEROKU``
 env var in Heroku. This allows the migration script to use the Heroku API.
 
-See tests for more examples.
+
+Before running DB migration, the script will enable `Heroku maintenance mode <https://devcenter.heroku.com/articles/maintenance-mode>`_
+if the app is not already in maintenance mode. After the migration, maintenance mode will
+be disabled only if it was enabled by the migration script.
+
+Maintenance mode can also be enabled/disabled using the ``pyramid_heroku.maintenance`` script.
+
+Usage example for enabling the Heroku maintenance mode::
+
+    python -m pyramid_heroku.maintenance on my_app etc/production.ini
+
 
 If you use structlog, add the following configuration setting to your INI file to enable structlog-like logging::
 
     pyramid_heroku.structlog = true
 
 
+See tests for more examples.
+
+
+
 Releasing
 ---------
 

pyramid_heroku/__init__.py

@@ -3,9 +3,27 @@
 from ast import literal_eval
 from expandvars import expandvars
 
+import shlex
+import subprocess
+import sys
 import typing as t
 
 
+def shell(cmd: str) -> str:
+    """
+    Run shell command.
+
+    :param cmd: shell command to run
+    :return: stdout of command
+    """
+
+    p = subprocess.run(shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    print(p.stdout.decode())
+    print(p.stderr.decode(), file=sys.stderr)
+    p.check_returncode()
+    return p.stdout.decode()
+
+
 def safe_eval(text: str) -> t.Optional[str]:
     """Safely evaluate `text` argument.
 

pyramid_heroku/heroku.py

@@ -0,0 +1,99 @@
+"""Heroku API client."""
+
+from requests import Response
+from requests import Session
+from typing import Optional
+
+import os
+
+
+class Heroku(object):
+
+    api_endpoint = "https://api.heroku.com"
+
+    def __init__(self, app_name: str, ini_file: str) -> None:
+        """
+        :param app_name: Name of Heroku app or id.
+        :param ini_file: development.ini or production.ini filename.
+        """
+        self._formation = None
+        self.app_name = app_name
+        self.ini_file = ini_file
+
+        headers = {
+            "Authorization": f"Bearer {self.auth_key}",
+            "Accept": "application/vnd.heroku+json; version=3",
+            "Content-Type": "application/json",
+        }
+        self.session = Session()
+        self.session.headers.update(headers)
+
+    @property
+    def auth_key(self) -> Optional[str]:
+        """Heroku API secret.
+
+        https://devcenter.heroku.com/articles/platform-api-quickstart#authentication.
+        """
+        return os.environ.get("MIGRATE_API_SECRET_HEROKU")
+
+    def scale_down(self):
+        """Scale all app workers to 0."""
+        updates = [dict(type=t, quantity=0) for t in self.formation.keys()]
+        res = self.session.patch(
+            f"{self.api_endpoint}/apps/{self.app_name}/formation",
+            json=dict(updates=updates),
+        )
+        self.parse_response(res)
+        print("Scaled down to:")
+        for x in res.json():
+            print(f'{x["type"]}={x["quantity"]}')
+
+    def scale_up(self):
+        """Scale app back to initial state."""
+        updates = [dict(type=t, quantity=s) for t, s in self.formation.items()]
+        res = self.session.patch(
+            f"{self.api_endpoint}/apps/{self.app_name}/formation",
+            json=dict(updates=updates),
+        )
+        self.parse_response(res)
+        print("Scaled up to:")
+        for x in res.json():
+            print(f'{x["type"]}={x["quantity"]}')
+
+    @property
+    def formation(self):
+        """Get current app status and configuration.
+
+        :return: Heroku app status as dict.
+        """
+        if not self._formation:
+            res = self.session.get(
+                f"{self.api_endpoint}/apps/{self.app_name}/formation"
+            )
+            self.parse_response(res)
+            self._formation = {x["type"]: x["quantity"] for x in res.json()}
+        return self._formation
+
+    def get_maintenance(self) -> bool:
+        res = self.session.get(f"{self.api_endpoint}/apps/{self.app_name}")
+        self.parse_response(res)
+        return res.json()["maintenance"]
+
+    def set_maintenance(self, state: bool) -> None:
+        res = self.session.patch(
+            f"{self.api_endpoint}/apps/{self.app_name}", json=dict(maintenance=state)
+        )
+        if self.parse_response(res):
+            print("Maintenance {}".format("enabled" if state else "disabled"))
+
+    def parse_response(self, res: Response) -> Optional[bool]:
+        """
+        Parses Heroku API response.
+
+        :param res: requests object
+        :return: true if request succeeded
+        """
+        if res.status_code != 200:
+            print(res.json())
+        res.raise_for_status()
+        return True

pyramid_heroku/maintenance.py

@@ -0,0 +1,44 @@
+"""Enable or disable Heroku app maintenance mode."""
+
+from enum import Enum
+from pyramid_heroku.heroku import Heroku
+
+import argparse
+
+
+class Mode(Enum):
+    """Heroku maintenance modes."""
+
+    on = "ON"
+    off = "OFF"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        usage=(
+            "usage: maintenance.py [-h] [app_name] [ini_file] [on|off]"
+            "\nexample: python -m pyramid_heroku.maintenance on my_app etc/production.ini"
+        ),
+    )
+    parser.add_argument("mode", choices=[x.name for x in Mode], help="Maintenance mode")
+    parser.add_argument("app_name", help="Heroku app name")
+    parser.add_argument(
+        "ini_file",
+        nargs="?",
+        default="etc/production.ini",
+        help="Path to Pyramid configuration file ",
+    )
+
+    options = parser.parse_args()
+
+    h = Heroku(options.app_name, options.ini_file)
+
+    if Mode[options.mode] == Mode.on:
+        h.set_maintenance(True)
+    else:
+        h.set_maintenance(False)
+
+
+if __name__ == "__main__":
+    main()