Merge pull request alphagov#143 from alphagov/add-performance-platform-collect

timmow · timmow · commit d60d137079e0 · 2015-07-22T12:02:05.000+01:00
Add collector execution task to fabric
diff --git a/performanceplatform.py b/performanceplatform.py
@@ -1,4 +1,10 @@
-from fabric.api import (task, hosts, run)
+import json
+
+from fabric.api import (
+    task, hosts, run, settings,
+    sudo, get, shell_env, cd
+)
+from StringIO import StringIO
 
 
 @task
@@ -23,5 +29,51 @@ def unpublish_dashboard(slug):
     run_stagecraft_postgres_command(sql_command)
 
 
+@task
+@hosts('performance-backend-1.backend')
+def collect(data_group, data_type):
+    """Run a collector config"""
+    base_path = '/data/apps/performanceplatform-collector/current'
+
+    with cd(base_path):
+        query_path = './config/queries/{}/{}.json'.format(
+            data_group, data_type)
+
+        query = get_file_contents(query_path)
+
+        command = get_command(query_path, query)
+        print('Executing: {}'.format(command))
+        with settings(sudo_user='deploy'):
+            with shell_env(LOGLEVEL='debug'):
+                sudo(command)
+
+
+def get_command(query_path, query):
+    """
+    >>> get_command('/tmp/foo/bar', {'entrypoint': 'foo.bar.monkey', 'token': 'foo'})
+    './venv/bin/pp-collector -q /tmp/foo/bar -c ./config/credentials/monkey.json -t ./config/tokens/foo.json -b ./config/performanceplatform.json --console-logging'
+    """
+    credentials = query['entrypoint'].split('.')[-1]
+
+    # realtime uses GA credentials, gcloud needs none but pp-collector fails with none
+    if credentials in ['realtime', 'gcloud', 'trending']:
+        credentials = 'ga'
+
+    pp_collector = './venv/bin/pp-collector'
+    credentials_path = './config/credentials/{}.json'.format(credentials)
+    token_path = './config/tokens/{}.json'.format(
+    query['token'])
+    platform_path = './config/performanceplatform.json'
+
+    return '{} -q {} -c {} -t {} -b {} --console-logging'.format(
+        pp_collector, query_path, credentials_path, token_path, platform_path)
+
+
+def get_file_contents(path):
+    fd = StringIO()
+    get(path, fd)
+    return json.loads(fd.getvalue())
+
+
 def run_stagecraft_postgres_command(sql_command):
     run('sudo -iu postgres psql stagecraft -c "{0}"'.format(sql_command))
