feat: add websocket logging

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.15.3",
+    "version": "3.15.4",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

querybook/server/const/event_log.py

@@ -4,6 +4,8 @@
 class EventType(Enum):
     # an api request
     API = "api"
+    # websocket event
+    WEBSOCKET = "websocket"
     # a UI element gets viewed
     VIEW = "view"
     # a UI element gets clicked

querybook/server/datasources_socketio/helper.py

@@ -4,12 +4,13 @@
 from flask_socketio import disconnect
 
 from app.flask_app import socketio
+from lib.event_logger import event_logger
 from lib.logger import get_logger
 
 LOG = get_logger(__file__)
 
 
-def register_socket(url, namespace=None):
+def register_socket(url, namespace=None, websocket_logging=True):
     def wrapper(fn):
         @socketio.on(url, namespace=namespace)
         @functools.wraps(fn)
@@ -19,6 +20,12 @@ def handler(*args, **kwargs):
                 disconnect()
             else:
                 try:
+                    if websocket_logging:
+                        event_logger.log_websocket_event(
+                            route=namespace + "/" + url,
+                            args=args,
+                            kwargs=kwargs,
+                        )
                     fn(*args, **kwargs)
                 except Exception as e:
                     LOG.error(e, exc_info=True)

querybook/server/lib/event_logger/__init__.py

@@ -36,5 +36,13 @@ def log_api_request(self, route: str, method: str, params: dict):
         except Exception as e:
             LOG.error(e, exc_info=True)
 
+    def log_websocket_event(self, route: str, args: list, kwargs: dict):
+        try:
+            self.logger.log_websocket_event(
+                uid=current_user.id, route=route, args=args, kwargs=kwargs
+            )
+        except Exception as e:
+            LOG.error(e, exc_info=True)
+
 
 event_logger = EventLogger()

querybook/server/lib/event_logger/base_event_logger.py

@@ -89,6 +89,7 @@ def log_api_request(self, uid: int, route: str, method: str, params: dict) -> No
         """Log an API request.
 
         Args:
+            uid (int): id of the user who performed the action
             method (str): request method, e.g. GET, POST
             route (str): route of the api endpoint which serves the request
             params (dict): params of the request, includes path params,
@@ -101,6 +102,20 @@ def log_api_request(self, uid: int, route: str, method: str, params: dict) -> No
         event_data = {"method": method, "route": route, "params": params}
         self.log(uid=uid, event_type=EventType.API, event_data=event_data)
 
+    def log_websocket_event(
+        self, uid: int, route: str, args: list, kwargs: dict
+    ) -> None:
+        """Log websocket event.
+
+        Args:
+            uid (int): id of the user who performed the action
+            route (str): route is namespace joined with event name.
+            args (list):  non-keyworded arguments passed over from client.
+            params (dict): keyworded arguments passed over from client.
+        """
+        event_data = {"route": route, "args": args, "kwargs": kwargs}
+        self.log(uid=uid, event_type=EventType.WEBSOCKET, event_data=event_data)
+
     def __match_filter_rule(self, rule: ApiFilterRule, route: str, method: str) -> bool:
         route_matched = (
             route.startswith(rule["route"])