From 7c5f61d6a61e1a3cb6d20d0918610abafa2cd326 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Thu, 18 May 2017 15:25:09 -0700
Subject: [PATCH] Adding some STATSD logging (#2715)

---
 superset/config.py                            |  5 ++
 .../a6c18f869a4e_query_start_running_time.py  |  1 -
 superset/models/core.py                       |  2 +
 superset/stats_logger.py                      | 57 +++++++++++++++++++
 4 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 superset/stats_logger.py

diff --git a/superset/config.py b/superset/config.py
index 2e981921df..dcf39d18cf 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -17,6 +17,11 @@ from collections import OrderedDict
 from dateutil import tz
 from flask_appbuilder.security.manager import AUTH_DB
 
+from superset.stats_logger import DummyStatsLogger
+
+# Realtime stats logger, a StatsD implementation exists
+STATS_LOGGER = DummyStatsLogger()
+
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
 DATA_DIR = os.path.join(os.path.expanduser('~'), '.superset')
 if not os.path.exists(DATA_DIR):
diff --git a/superset/migrations/versions/a6c18f869a4e_query_start_running_time.py b/superset/migrations/versions/a6c18f869a4e_query_start_running_time.py
index 3be8ae06b2..0f89b3f516 100644
--- a/superset/migrations/versions/a6c18f869a4e_query_start_running_time.py
+++ b/superset/migrations/versions/a6c18f869a4e_query_start_running_time.py
@@ -13,7 +13,6 @@ revision = 'a6c18f869a4e'
 down_revision = '979c03af3341'
 
 
-
 def upgrade():
     op.add_column(
         'query',
diff --git a/superset/models/core.py b/superset/models/core.py
index d2a8816ca3..49f95d7669 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -42,6 +42,7 @@ install_aliases()
 from urllib import parse  # noqa
 
 config = app.config
+stats_logger = config.get('STATS_LOGGER')
 metadata = Model.metadata  # pylint: disable=no-member
 
 
@@ -751,6 +752,7 @@ class Log(Model):
                 params = json.dumps(d)
             except:
                 pass
+            stats_logger.incr(f.__name__)
             value = f(*args, **kwargs)
 
             sesh = db.session()
diff --git a/superset/stats_logger.py b/superset/stats_logger.py
new file mode 100644
index 0000000000..0203adc330
--- /dev/null
+++ b/superset/stats_logger.py
@@ -0,0 +1,57 @@
+
+class BaseStatsLogger(object):
+    """Base class for logging realtime events"""
+
+    def __init__(self, prefix='superset'):
+        self.prefix = prefix
+
+    def key(self, key):
+        if self.prefix:
+            return self.prefix + key
+        return key
+
+    def incr(self, key):
+        """Increment a counter"""
+        raise NotImplementedError()
+
+    def decr(self, key):
+        """Decrement a counter"""
+        raise NotImplementedError()
+
+    def gauge(self, key):
+        """Setup a gauge"""
+        raise NotImplementedError()
+
+
+class DummyStatsLogger(BaseStatsLogger):
+
+    def incr(self, key):
+        pass
+
+    def decr(self, key):
+        pass
+
+    def gauge(self, key):
+        pass
+
+
+try:
+    from statsd import StatsClient
+
+    class StatsdStatsLogger(BaseStatsLogger):
+        def __init__(self, host, port, prefix='superset'):
+            self.client = StatsClient(
+                  host=host,
+                  port=port,
+                  prefix=prefix)
+
+        def incr(self, key):
+            self.client.incr(key)
+
+        def decr(self, key):
+            self.client.decr(key)
+
+        def gauge(self, key):
+            self.client.gauge(key)
+except Exception as e:
+    pass