feat: update some stats metrics names and add tags (#1208)

querybook/server/app/datasource.py

@@ -19,11 +19,7 @@
 )
 from flask_login import current_user
 from lib.event_logger import event_logger
-from lib.stats_logger import (
-    stats_logger,
-    API_REQUEST_COUNTER,
-    API_REQUEST_LATENCY_TIMER,
-)
+from lib.stats_logger import API_REQUESTS, stats_logger
 from lib.logger import get_logger
 from logic.impression import create_impression
 from werkzeug.exceptions import Forbidden, NotFound
@@ -61,8 +57,6 @@ def wrapper(fn):
         @flask_app.route(r"%s%s" % (DS_PATH, url), methods=methods)
         @functools.wraps(fn)
         def handler(**kwargs):
-            # increment the number of api request counter
-            stats_logger.incr(API_REQUEST_COUNTER.format(fn.__name__))
             # start the timer for api request duration
             start_time = time.time()
 
@@ -92,7 +86,9 @@ def handler(**kwargs):
                 # stop the timer and record the duration
                 duration_ms = (time.time() - start_time) * 1000.0
                 stats_logger.timing(
-                    API_REQUEST_LATENCY_TIMER.format(fn.__name__), duration_ms
+                    API_REQUESTS,
+                    duration_ms,
+                    tags={"endpoint": fn.__name__, "method": flask.request.method},
                 )
 
                 if not custom_response:

querybook/server/app/db.py

@@ -7,7 +7,8 @@
 from sqlalchemy.exc import SQLAlchemyError, DisconnectionError
 from sqlalchemy import create_engine, event
 from sqlalchemy.orm import sessionmaker, scoped_session
-from lib.stats_logger import stats_logger, SQL_SESSION_FAILURE_COUNTER
+
+from lib.stats_logger import SQL_SESSION_FAILURES, stats_logger
 
 try:
     from greenlet import getcurrent as _get_ident
@@ -125,7 +126,7 @@ def func(*args, **kwargs):
                 LOG.error(traceback.format_exc())
 
                 # increment sql session failure counter
-                stats_logger.incr(SQL_SESSION_FAILURE_COUNTER)
+                stats_logger.incr(SQL_SESSION_FAILURES)
 
                 raise e
             finally:
@@ -158,7 +159,7 @@ def DBSession():
         LOG.error(traceback.format_exc())
 
         # increment sql session failure counter
-        stats_logger.incr(SQL_SESSION_FAILURE_COUNTER)
+        stats_logger.incr(SQL_SESSION_FAILURES)
 
         raise e
     finally:

querybook/server/clients/redis_client.py

@@ -3,7 +3,7 @@
 
 import redis
 from env import QuerybookSettings
-from lib.stats_logger import REDIS_LATENCY_TIMER, stats_logger
+from lib.stats_logger import REDIS_OPERATIONS, stats_logger
 
 __redis = None
 
@@ -34,7 +34,9 @@ def func(*args, **kwargs):
 
         # stop the timer and record the duration
         duration_ms = (time.time() - start_time) * 1000.0
-        stats_logger.timing(REDIS_LATENCY_TIMER.format(fn.__name__), duration_ms)
+        stats_logger.timing(
+            REDIS_OPERATIONS, duration_ms, tags={"operation": fn.__name__}
+        )
 
         return result
 

querybook/server/datasources_socketio/connect.py

@@ -1,17 +1,22 @@
 from flask_login import current_user
 from flask_socketio import ConnectionRefusedError
 
-from app.flask_app import socketio
 from const.data_doc import DATA_DOC_NAMESPACE
 from const.query_execution import QUERY_EXECUTION_NAMESPACE
-from lib.stats_logger import stats_logger, WS_CONNECTIONS_COUNTER
 
+from .helper import register_socket
 
-def connect():
-    stats_logger.incr(WS_CONNECTIONS_COUNTER)
+
+def on_connect():
     if not current_user.is_authenticated:
         raise ConnectionRefusedError("User is not logged in, please refresh the page.")
 
 
-socketio.on_event("connect", connect, namespace=DATA_DOC_NAMESPACE)
-socketio.on_event("connect", connect, namespace=QUERY_EXECUTION_NAMESPACE)
+@register_socket("connect", namespace=QUERY_EXECUTION_NAMESPACE)
+def connect_query_execution(auth):
+    on_connect()
+
+
+@register_socket("connect", namespace=DATA_DOC_NAMESPACE)
+def connect_datadoc(auth):
+    on_connect()

querybook/server/datasources_socketio/helper.py

@@ -6,7 +6,7 @@
 from app.flask_app import socketio
 from lib.event_logger import event_logger
 from lib.logger import get_logger
-from lib.stats_logger import stats_logger, WS_CONNECTIONS_COUNTER
+from lib.stats_logger import WS_CONNECTIONS, stats_logger
 
 LOG = get_logger(__file__)
 
@@ -20,7 +20,7 @@ def handler(*args, **kwargs):
                 LOG.error("Unauthorized websocket access")
                 disconnect()
                 # decrement ws connections counter on disconnect
-                stats_logger.decr(WS_CONNECTIONS_COUNTER)
+                stats_logger.decr(WS_CONNECTIONS)
             else:
                 try:
                     if websocket_logging:
@@ -39,9 +39,12 @@ def handler(*args, **kwargs):
                         room=flask.request.sid,
                     )
 
+                # increment ws connections counter on connect
+                if url == "connect":
+                    stats_logger.incr(WS_CONNECTIONS, tags={"namespace": namespace})
                 # decrement ws connections counter on disconnect
-                if url == "disconnect":
-                    stats_logger.decr(WS_CONNECTIONS_COUNTER)
+                elif url == "disconnect":
+                    stats_logger.decr(WS_CONNECTIONS, tags={"namespace": namespace})
 
         handler.__raw__ = fn
         return handler

querybook/server/lib/stats_logger/__init__.py

@@ -1,17 +1,15 @@
 from env import QuerybookSettings
 from lib.stats_logger.all_stats_loggers import get_stats_logger_class
-from .base_stats_logger import BaseStatsLogger
 
 
-# metrics name templates
-API_REQUEST_COUNTER = "api.{}"
-API_REQUEST_LATENCY_TIMER = "api.duration.ms.{}"
-WS_CONNECTIONS_COUNTER = "ws.connections"
-SQL_SESSION_FAILURE_COUNTER = "sql.session.failure"
-SYSTEM_TASK_FAILURE_COUNTER = "task.failure.system"
-DATADOC_TASK_FAILURE_COUNTER = "task.failure.datadoc"
-REDIS_LATENCY_TIMER = "redis.duration.ms.{}"
-QUERY_EXECUTION_COUNTER = "query_execution.{}"
+# metrics name
+API_REQUESTS = "api.requests"
+WS_CONNECTIONS = "ws.connections"
+SQL_SESSION_FAILURES = "sql_session.failures"
+TASK_FAILURES = "task.failures"
+REDIS_OPERATIONS = "redis.operations"
+QUERY_EXECUTIONS = "query.executions"
+
 
 logger_name = QuerybookSettings.STATS_LOGGER_NAME
-stats_logger: BaseStatsLogger = get_stats_logger_class(logger_name)
+stats_logger = get_stats_logger_class(logger_name)

querybook/server/lib/stats_logger/all_stats_loggers.py

@@ -15,4 +15,4 @@ def get_stats_logger_class(name: str):
     for logger in ALL_STATS_LOGGERS:
         if logger.logger_name == name:
             return logger
-    raise ValueError(f"Unknown event logger name {name}")
+    raise ValueError(f"Unknown stats logger name {name}")

querybook/server/lib/stats_logger/base_stats_logger.py

@@ -18,20 +18,20 @@ def prefix(self) -> str:
         return "querybook."
 
     @abstractmethod
-    def incr(self, key: str) -> None:
+    def incr(self, key: str, tags: dict[str, str] = None) -> None:
         """Increment a counter"""
         raise NotImplementedError()
 
     @abstractmethod
-    def decr(self, key: str) -> None:
+    def decr(self, key: str, tags: dict[str, str] = None) -> None:
         """Decrement a counter"""
         raise NotImplementedError()
 
     @abstractmethod
-    def timing(self, key: str, value: float) -> None:
+    def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         raise NotImplementedError()
 
     @abstractmethod
-    def gauge(self, key: str, value: float) -> None:
+    def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         """Setup a gauge"""
         raise NotImplementedError()

querybook/server/lib/stats_logger/loggers/console_stats_logger.py

@@ -14,18 +14,22 @@ class ConsoleStatsLogger(BaseStatsLogger):
     def logger_name(self) -> str:
         return "console"
 
-    def incr(self, key: str) -> None:
-        LOG.debug(COLOR_YELLOW + "[stats_logger] (incr) " + key + COLOR_RESET)
+    def incr(self, key: str, tags: dict[str, str] = None) -> None:
+        LOG.debug(COLOR_YELLOW + f"[stats_logger] (incr) {key} | {tags}" + COLOR_RESET)
 
-    def decr(self, key: str) -> None:
-        LOG.debug(COLOR_YELLOW + "[stats_logger] (decr) " + key + COLOR_RESET)
+    def decr(self, key: str, tags: dict[str, str] = None) -> None:
+        LOG.debug(COLOR_YELLOW + f"[stats_logger] (decr) {key} | {tags}" + COLOR_RESET)
 
-    def timing(self, key: str, value: float) -> None:
+    def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         LOG.debug(
-            COLOR_YELLOW + f"[stats_logger] (timing) {key} | {value} " + COLOR_RESET
+            COLOR_YELLOW
+            + f"[stats_logger] (timing) {key} | {value} | {tags}"
+            + COLOR_RESET
         )
 
-    def gauge(self, key: str, value: float) -> None:
+    def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         LOG.debug(
-            COLOR_YELLOW + f"[stats_logger] (gauge) {key} | {value} " + COLOR_RESET
+            COLOR_YELLOW
+            + f"[stats_logger] (gauge) {key} | {value} | {tags}"
+            + COLOR_RESET
         )

querybook/server/lib/stats_logger/loggers/null_stats_logger.py

@@ -8,14 +8,14 @@ class NullStatsLogger(BaseStatsLogger):
     def logger_name(self) -> str:
         return "null"
 
-    def incr(self, key: str) -> None:
+    def incr(self, key: str, tags: dict[str, str] = None) -> None:
         pass
 
-    def decr(self, key: str) -> None:
+    def decr(self, key: str, tags: dict[str, str] = None) -> None:
         pass
 
-    def timing(self, key: str, value: float) -> None:
+    def timing(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         pass
 
-    def gauge(self, key: str, value: float) -> None:
+    def gauge(self, key: str, value: float, tags: dict[str, str] = None) -> None:
         pass

querybook/server/tasks/all_tasks.py

@@ -5,11 +5,8 @@
 from app.flask_app import celery
 from env import QuerybookSettings
 from lib.logger import get_logger
-from lib.stats_logger import (
-    stats_logger,
-    DATADOC_TASK_FAILURE_COUNTER,
-    SYSTEM_TASK_FAILURE_COUNTER,
-)
+from lib.stats_logger import TASK_FAILURES, stats_logger
+from logic.schedule import get_schedule_task_type
 
 from .export_query_execution import export_query_execution_task
 from .run_query import run_query_task
@@ -63,8 +60,5 @@ def configure_workers(sender=None, conf=None, **kwargs):
 
 @task_failure.connect
 def handle_task_failure(sender, signal, *args, **kwargs):
-    task_name = sender.name
-    if task_name.startswith("tasks.run_datadoc"):
-        stats_logger.incr(DATADOC_TASK_FAILURE_COUNTER)
-    else:
-        stats_logger.incr(SYSTEM_TASK_FAILURE_COUNTER)
+    task_type = get_schedule_task_type(sender.name)
+    stats_logger.incr(TASK_FAILURES, tags={"task_type": task_type})

querybook/server/tasks/run_query.py

@@ -11,7 +11,7 @@
 from lib.query_executor.executor_factory import create_executor_from_execution
 from lib.query_executor.exc import QueryExecutorException
 from lib.query_executor.utils import format_error_message
-from lib.stats_logger import stats_logger, QUERY_EXECUTION_COUNTER
+from lib.stats_logger import QUERY_EXECUTIONS, stats_logger
 
 from logic import query_execution as qe_logic
 from logic.elasticsearch import update_query_execution_by_id
@@ -32,7 +32,7 @@
 def run_query_task(
     self, query_execution_id, execution_type=QueryExecutionType.ADHOC.value
 ):
-    stats_logger.incr(QUERY_EXECUTION_COUNTER.format(execution_type))
+    stats_logger.incr(QUERY_EXECUTIONS, tags={"execution_type": execution_type})
 
     executor = None
     error_message = None
@@ -57,7 +57,7 @@ def run_query_task(
             7406, "{}\n{}".format(e, traceback.format_exc())
         )
     finally:
-        stats_logger.decr(QUERY_EXECUTION_COUNTER.format(execution_type))
+        stats_logger.decr(QUERY_EXECUTIONS, tags={"execution_type": execution_type})
 
         # When the finally block is reached, it is expected
         # that the executor should be in one of the end state