task/WG-236: Websocket notifications (#185)

devops/docker-compose.local.yml

@@ -102,6 +102,6 @@ services:
     tty: true
     container_name: geoapi
     hostname: geoapi
-    command: "gunicorn -w 4 -b 0.0.0.0:8000 geoapi.app:app -k gevent --reload --timeout 300"
+    command: "gunicorn -w 1 -b 0.0.0.0:8000 geoapi.app:app -k gevent --reload --timeout 300"
     extra_hosts:
       - "host.docker.internal:host-gateway"

devops/local_conf/nginx.conf

@@ -36,6 +36,19 @@ http {
             proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
         }
 
+        location /socket.io {
+            proxy_pass http://geoapi:8000;
+            proxy_http_version 1.1;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection 'upgrade';
+            proxy_set_header Host $host;
+            proxy_cache_bypass $http_upgrade;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            proxy_redirect off;
+        }
+
         location /assets {
             max_ranges 0;
             expires 30d;

devops/pyproject.toml

@@ -23,6 +23,7 @@ Fiona = "1.8.21"
 gunicorn = "^20.1.0"
 gevent = "^21.12.0"
 Flask = "^2.1.2"
+flask-socketio = "^5.3.6"
 Werkzeug = "2.1.2"
 flask-restx = "^0.5.1"
 psycopg2 = "^2.9.3"

geoapi/app.py

@@ -1,5 +1,6 @@
 
 from flask import Flask
+from flask_socketio import SocketIO
 from geoapi.routes import api
 from geoapi.settings import settings as app_settings
 from geoapi.db import db_session
@@ -8,15 +9,58 @@
                                StreetviewLimitException)
 
 import logging
+from geoapi.utils.decorators import jwt_socket_decoder
 
 logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
 
 
 app = Flask(__name__)
-api.init_app(app)
 app.config.from_object(app_settings)
 
 
+# Initialize SocketIO
+socketio = SocketIO(app, cors_allowed_origins='http://localhost:4200')
+api.init_app(app)
+
+
+@socketio.on('connect')
+@jwt_socket_decoder
+def handle_connect():
+    logging.info('Client connected')
+
+
+@socketio.on('trigger_notification')
+def handle_notification(data):
+    logging.info('Received trigger notification event: %s', data)
+    # Emit a new notification with the message sent from the client
+    # socketio.emit('new_notification', {'message': data.get('message')})
+    # Otherwise, emit a default message
+    socketio.emit('new_notification', {'message': 'This is a toast message!'})
+
+
+@socketio.on('trigger_asset_success')
+def handle_asset_success(data):
+    logging.info('Received trigger notification event: %s', data)
+    socketio.emit('asset_success', {'message': 'Asset was successfully added!'})
+
+
+@socketio.on('trigger_asset_failure')
+def handle_asset_failure(data):
+    logging.info('Received trigger notification event: %s', data)
+    socketio.emit('asset_failure', {'message': 'Asset failed to be added!'})
+
+
+@socketio.on('client_message')
+def handle_client_message(message):
+    logging.info('Received message from client: %s', message)
+    socketio.emit('server_message', {'message': 'This is a server message!'})
+
+
+@socketio.on('disconnect')
+def handle_disconnect():
+    logging.info('Client disconnected')
+
+
 @api.errorhandler(InvalidGeoJSON)
 def handle_geojson_exception(error: Exception):
     '''Return a custom message and 400 status code'''
@@ -67,3 +111,7 @@ def handle_streetview_limit_exception(error: Exception):
 @app.teardown_appcontext
 def shutdown_session(exception=None):
     db_session.remove()
+
+
+if __name__ == '__main__':
+    socketio.run(app, port=8000, debug=True)

geoapi/celery_app.py

@@ -9,14 +9,8 @@
     vhost=settings.RABBITMQ_VHOST
 )
 
-app = Celery('hello',
-             backend='rpc',
-             broker=CELERY_CONNECTION_STRING,
-             include=['geoapi.tasks'])
+app = Celery('geoapi', backend='rpc', broker=CELERY_CONNECTION_STRING, include=['geoapi.tasks'])
 
-app.conf.beat_schedule = {
-    'refresh_observable_projects': {
+app.conf.beat_schedule = {'refresh_observable_projects': {
         'task': 'geoapi.tasks.external_data.refresh_observable_projects',
-        'schedule': crontab(hour='*', minute='0')
-    }
-}
+        'schedule': crontab(hour='*', minute='0')}}

geoapi/services/notifications.py

@@ -40,6 +40,8 @@ def create(database_session, user: User, status: AnyStr, message: AnyStr) -> Not
         try:
             database_session.add(note)
             database_session.commit()
+            from geoapi.app import socketio
+            socketio.emit('new_notification', {'message': message})
             return note
         except Exception:
             database_session.rollback()
@@ -114,3 +116,9 @@ def deleteAllDoneProgress(database_session):
         for pn in note:
             database_session.delete(pn)
         database_session.commit()
+
+    @staticmethod
+    def emit_socketio_event(event: str, message: str, socketio):
+        logger.info('Emitting socketio event: %s', event)
+        socketio.emit(event, {'message': message})
+        logger.info('Socketio event emitted: %s', event)

geoapi/utils/decorators.py

@@ -15,6 +15,7 @@
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 import base64
+from flask_socketio import disconnect
 
 
 def get_pub_key():
@@ -63,6 +64,18 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def jwt_socket_decoder(fn):
+    @wraps(fn)
+    def wrapper(auth=None):
+        token = auth.get('token') if auth else None
+
+        if not token:
+            logger.error('No token provided.')
+            disconnect()
+            return
+    return wrapper
+
+
 def check_access_and_get_project(current_user, allow_public_use=False, project_id=None, uuid=None):
     """
     Check if user (authenticated or anonymous) can access a project id and *aborts* if there is no access.