Integrate websockets into the Flask webapp

11 files changed, 214 insertions, 93 deletions

diff --git a/pysite/base_route.py b/pysite/base_route.py
index 8e6648ee..36b7fef5 100644
--- a/pysite/base_route.py
+++ b/pysite/base_route.py
@@ -3,15 +3,10 @@ import os
 import random
 import string
 
-from _weakref import ref
-
 from flask import Blueprint, jsonify, render_template
 from flask.views import MethodView
 
-from rethinkdb.ast import Table
-
 from pysite.constants import ErrorCodes
-from pysite.database import RethinkDB
 
 
 class BaseView(MethodView):
@@ -124,53 +119,6 @@ class APIView(RouteView):
         return response
 
 
-class DBViewMixin:
-    """
-    Mixin for views that make use of RethinkDB. It can automatically create a table with the specified primary
-    key using the attributes set at class-level.
-
-    This class is intended to be mixed in alongside one of the other view classes. For example:
-
-    >>> class MyView(APIView, DBViewMixin):
-    ...     name = "my_view"  # Flask internal name for this route
-    ...     path = "/my_view"  # Actual URL path to reach this route
-    ...     table_name = "my_table"  # Name of the table to create
-    ...     table_primary_key = "username"  # Primary key to set for this table
-
-    You may omit `table_primary_key` and it will be defaulted to RethinkDB's default column - "id".
-    """
-
-    table_name = ""  # type: str
-    table_primary_key = "id"  # type: str
-
-    @classmethod
-    def setup(cls: "DBViewMixin", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
-        """
-        Set up the view by creating the table specified by the class attributes - this will also deal with multiple
-        inheritance by calling `super().setup()` as appropriate.
-
-        :param manager: Instance of the current RouteManager (used to get a handle for the database object)
-        :param blueprint: Current Flask blueprint
-        """
-
-        if hasattr(super(), "setup"):
-            super().setup(manager, blueprint)
-
-        if not cls.table_name:
-            raise RuntimeError("Routes using DBViewMixin must define `table_name`")
-
-        cls._db = ref(manager.db)
-        manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key)
-
-    @property
-    def table(self) -> Table:
-        return self.db.query(self.table_name)
-
-    @property
-    def db(self) -> RethinkDB:
-        return self._db()
-
-
 class ErrorView(BaseView):
     """
     Error view, shown for a specific HTTP status code, as defined in the class attributes.
diff --git a/pysite/mixins.py b/pysite/mixins.py
new file mode 100644
index 00000000..891aa185
--- /dev/null
+++ b/pysite/mixins.py
@@ -0,0 +1,63 @@
+# coding=utf-8
+from _weakref import ref
+
+from flask import Blueprint
+
+from rethinkdb.ast import Table
+
+from pysite.database import RethinkDB
+
+
+class DBMixin:
+    """
+    Mixin for classes that make use of RethinkDB. It can automatically create a table with the specified primary
+    key using the attributes set at class-level.
+
+    This class is intended to be mixed in alongside one of the other view classes. For example:
+
+    >>> class MyView(APIView, DBMixin):
+    ...     name = "my_view"  # Flask internal name for this route
+    ...     path = "/my_view"  # Actual URL path to reach this route
+    ...     table_name = "my_table"  # Name of the table to create
+    ...     table_primary_key = "username"  # Primary key to set for this table
+
+    This class will also work with Websockets:
+
+    >>> class MyWeboscket(Websocket, DBMixin):
+    ...     name = "my_websocket"
+    ...     path = "/my_websocket"
+    ...     table_name = "my_table"
+    ...     table_primary_key = "username"
+
+    You may omit `table_primary_key` and it will be defaulted to RethinkDB's default column - "id".
+    """
+
+    table_name = ""  # type: str
+    table_primary_key = "id"  # type: str
+
+    @classmethod
+    def setup(cls: "DBMixin", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
+        """
+        Set up the view by creating the table specified by the class attributes - this will also deal with multiple
+        inheritance by calling `super().setup()` as appropriate.
+
+        :param manager: Instance of the current RouteManager (used to get a handle for the database object)
+        :param blueprint: Current Flask blueprint
+        """
+
+        if hasattr(super(), "setup"):
+            super().setup(manager, blueprint)
+
+        if not cls.table_name:
+            raise RuntimeError("Routes using DBViewMixin must define `table_name`")
+
+        cls._db = ref(manager.db)
+        manager.db.create_table(cls.table_name, primary_key=cls.table_primary_key)
+
+    @property
+    def table(self) -> Table:
+        return self.db.query(self.table_name)
+
+    @property
+    def db(self) -> RethinkDB:
+        return self._db()
diff --git a/pysite/route_manager.py b/pysite/route_manager.py
index 3a23619e..b3f71643 100644
--- a/pysite/route_manager.py
+++ b/pysite/route_manager.py
@@ -5,8 +5,11 @@ import os
 
 from flask import Blueprint, Flask
 
+from flask_sockets import Sockets
+
 from pysite.base_route import APIView, BaseView, ErrorView, RouteView
 from pysite.database import RethinkDB
+from pysite.websockets import Websocket
 
 TEMPLATES_PATH = "../templates"
 STATIC_PATH = "../static"
@@ -19,6 +22,8 @@ class RouteManager:
         self.app = Flask(
             __name__, template_folder=TEMPLATES_PATH, static_folder=STATIC_PATH, static_url_path="/static",
         )
+        self.sockets = Sockets(self.app)
+
         self.db = RethinkDB()
         self.app.secret_key = os.environ.get("WEBPAGE_SECRET_KEY", "super_secret")
         self.app.config["SERVER_NAME"] = os.environ.get("SERVER_NAME", "pythondiscord.com:8080")
@@ -36,17 +41,29 @@ class RouteManager:
         self.subdomains = ['api', 'staff']
 
         for sub in self.subdomains:
-            self.sub_blueprint = Blueprint(sub, __name__, subdomain=sub)
+            sub_blueprint = Blueprint(sub, __name__, subdomain=sub)
 
-            print(f"Loading Blueprint: {self.sub_blueprint.name}")
-            self.load_views(self.sub_blueprint, f"pysite/views/{sub}")
-            self.app.register_blueprint(self.sub_blueprint)
+            print(f"Loading Blueprint: {sub_blueprint.name}")
+            self.load_views(sub_blueprint, f"pysite/views/{sub}")
+            self.app.register_blueprint(sub_blueprint)
             print("")
 
+        # Load the websockets
+        self.ws_blueprint = Blueprint("ws", __name__)
+
+        print("Loading websocket routes...")
+        self.load_views(self.ws_blueprint, "pysite/views/ws")
+        self.sockets.register_blueprint(self.ws_blueprint, url_prefix="/ws")
+
     def run(self):
-        self.app.run(
-            port=int(os.environ.get("WEBPAGE_PORT", 8080)), debug="FLASK_DEBUG" in os.environ
+        from gevent.pywsgi import WSGIServer
+        from geventwebsocket.handler import WebSocketHandler
+
+        server = WSGIServer(
+            ("", int(os.environ.get("WEBPAGE_PORT", 8080))),
+            self.app, handler_class=WebSocketHandler
         )
+        server.serve_forever()
 
     def load_views(self, blueprint, location="pysite/views"):
         for filename in os.listdir(location):
@@ -65,7 +82,11 @@ class RouteManager:
                             cls is not ErrorView and
                             cls is not RouteView and
                             cls is not APIView and
-                            BaseView in cls.__mro__
+                            cls is not Websocket and
+                            (
+                                BaseView in cls.__mro__ or
+                                Websocket in cls.__mro__
+                            )
                     ):
                         cls.setup(self, blueprint)
                         print(f">> View loaded: {cls.name: <15} ({module.__name__}.{cls_name})")
diff --git a/pysite/views/api/bot/tag.py b/pysite/views/api/bot/tag.py
index 363f98fe..2117d948 100644
--- a/pysite/views/api/bot/tag.py
+++ b/pysite/views/api/bot/tag.py
@@ -2,12 +2,13 @@
 
 from flask import jsonify, request
 
-from pysite.base_route import APIView, DBViewMixin
+from pysite.base_route import APIView
 from pysite.constants import ErrorCodes
 from pysite.decorators import api_key
+from pysite.mixins import DBMixin
 
 
-class TagView(APIView, DBViewMixin):
+class TagView(APIView, DBMixin):
     path = "/tag"
     name = "tag"
     table_name = "tag"
diff --git a/pysite/views/api/bot/user.py b/pysite/views/api/bot/user.py
index 8c2d8149..174407b8 100644
--- a/pysite/views/api/bot/user.py
+++ b/pysite/views/api/bot/user.py
@@ -4,10 +4,10 @@ from flask import jsonify
 
 from schema import Schema
 
-from pysite.base_route import APIView, DBViewMixin
+from pysite.base_route import APIView
 from pysite.constants import ValidationTypes
 from pysite.decorators import api_key, api_params
-
+from pysite.mixins import DBMixin
 
 SCHEMA = Schema([
     {
@@ -22,7 +22,7 @@ REQUIRED_KEYS = [
 ]
 
 
-class UserView(APIView, DBViewMixin):
+class UserView(APIView, DBMixin):
     path = "/user"
     name = "user"
     table_name = "users"
diff --git a/pysite/views/ws/__init__.py b/pysite/views/ws/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/pysite/views/ws/__init__.py
diff --git a/pysite/views/ws/echo.py b/pysite/views/ws/echo.py
new file mode 100644
index 00000000..135adfcf
--- /dev/null
+++ b/pysite/views/ws/echo.py
@@ -0,0 +1,18 @@
+# coding=utf-8
+from pysite.websockets import Websocket
+
+
+class EchoWebsocket(Websocket):
+    path = "/echo"
+    name = "ws_echo"
+
+    def on_open(self):
+        print("Echo | WS opened.")
+        self.send("Hey, welcome!")
+
+    def on_message(self, message):
+        print(f"Echo | Message: {message}")
+        self.send(message)
+
+    def on_close(self):
+        print("Echo | WS closed.")
diff --git a/pysite/websockets.py b/pysite/websockets.py
new file mode 100644
index 00000000..f82217ef
--- /dev/null
+++ b/pysite/websockets.py
@@ -0,0 +1,97 @@
+# coding=utf-8
+from flask import Blueprint
+
+from geventwebsocket.websocket import WebSocket
+
+
+class Websocket:
+    """
+    Base class for representing a Websocket.
+
+    At minimum, you must implement the `on_message(self, message)` function. Without it, you won't be able to handle
+    any messages, and an error will be thrown!
+
+    If you need access to the database, you can mix-in DBMixin, just like any view class:
+
+    >>> class DBWebsocket(Websocket, DBMixin):
+    ...     name = "db_websocket"
+    ...     path = "/db_websocket"  # This will be prefixed with "/ws" by the blueprint
+    ...     table = "ws"
+    ...
+    ...     def on_message(self, message):
+    ...         self.send(
+    ...             json.loads(self.db.get(self.table_name, message))
+    ...         )
+
+    Please note that an instance of this class is created for every websocket connected to the path. This does, however,
+    mean that you can store any state required by your websocket.
+    """
+
+    path = ""  # type: str
+    name = ""  # type: str
+
+    def __init__(self, socket: WebSocket):
+        self.socket = socket
+
+    def on_open(self):
+        """
+        Called once when the websocket is opened. Optional.
+        """
+
+    def on_message(self, message: str):
+        """
+        Called when a message is received by the websocket.
+        """
+
+        raise NotImplementedError()
+
+    def on_close(self):
+        """
+        Called once when the websocket is closed. Optional.
+        """
+
+    def send(self, message, binary=None):
+        """
+        Send a message to the currently-connected websocket, if it's open.
+
+        Nothing will happen if the websocket is closed.
+        """
+
+        if not self.socket.closed:
+            self.socket.send(message, binary=binary)
+
+    @classmethod
+    def setup(cls: "type(Websocket)", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
+        """
+        Set up the websocket object, calling `setup()` on any superclasses as necessary (for example, on the DB
+        mixin).
+
+        This function will set up a websocket handler so that it behaves in a class-oriented way. It's up to you to
+        deal with message handling yourself, however.
+        """
+
+        if hasattr(super(), "setup"):
+            super().setup(manager, blueprint)
+
+        if not cls.path or not cls.name:
+            raise RuntimeError("Websockets must have both `path` and `name` defined")
+
+        def handle(socket: WebSocket):
+            """
+            Wrap the current Websocket class, dispatching events to it as necessary. We're using gevent, so there's
+            no need to worry about blocking here.
+            """
+
+            ws = cls(socket)  # Instantiate the current class, passing it the WS object
+
+            ws.on_open()  # Call the "on_open" handler
+
+            while not socket.closed:  # As long as the socket is open...
+                message = socket.receive()  # Wait for a message
+
+                if not socket.closed:  # If the socket didn't just close (there's always a None message on closing)
+                    ws.on_message(message)  # Call the "on_message" handler
+
+            ws.on_close()  # The socket just closed, call the "on_close" handler
+
+        blueprint.route(cls.path)(handle)  # Register the handling function to the WS blueprint
diff --git a/requirements.txt b/requirements.txt
index 205154bb..cc2dc0e9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,3 +6,4 @@ gevent-websocket
 wsaccel
 ujson
 schema
+flask_sockets
diff --git a/templates/ws_test.html b/templates/ws_test.html
index 921226c0..36be61d6 100644
--- a/templates/ws_test.html
+++ b/templates/ws_test.html
@@ -5,7 +5,7 @@
         <h1>Open your JS console to test</h1>
 
         <script type="application/javascript">
-            let ws = new WebSocket("wss://api.{{ server_name }}/ws/echo");
+            let ws = new WebSocket("ws://{{ server_name }}/ws/echo");
 
             ws.onopen = function(event) {
                 console.log("WS opened! Use send() to send a message.");
diff --git a/ws_app.py b/ws_app.py
deleted file mode 100644
index cab43ac0..00000000
--- a/ws_app.py
+++ /dev/null
@@ -1,28 +0,0 @@
-# coding=utf-8
-import os
-
-from geventwebsocket import Resource, WebSocketApplication, WebSocketServer
-
-
-class EchoApplication(WebSocketApplication):
-    def on_open(self):
-        print("Connection opened")
-
-    def on_message(self, message, **kwargs):
-        print(f"<- {message}")
-        self.ws.send(message)
-        print(f"-> {message}")
-
-    def on_close(self, reason):
-        print(reason)
-
-
-app = WebSocketServer(
-    ('', os.environ.get("WS_PORT", 8080)),
-    Resource({
-        "/ws/echo": EchoApplication  # Dicts are ordered in Python 3.6
-    })
-)
-
-if __name__ == "__main__":
-    app.serve_forever()