From 66c0942c0e7dfeffa0ed0b9671e7b8c2580373af Mon Sep 17 00:00:00 2001
From: Ilya Shakhat <ishakhat@mirantis.com>
Date: Fri, 19 Feb 2016 16:47:44 +0300
Subject: [PATCH] Allow to tag data

---
 performa/engine/config.py |  4 ++++
 performa/engine/main.py   |  9 +++++++--
 performa/engine/player.py |  8 +++++++-
 performa/engine/report.py | 11 ++++++++---
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/performa/engine/config.py b/performa/engine/config.py
index bd52878..1e49db5 100644
--- a/performa/engine/config.py
+++ b/performa/engine/config.py
@@ -35,6 +35,7 @@ class Endpoint(types.String):
         return "Endpoint host[:port]"
 
 
+
 MAIN_OPTS = [
     cfg.StrOpt('scenario',
                default=utils.env('PERFORMA_SCENARIO'),
@@ -60,6 +61,9 @@ MAIN_OPTS = [
                default=utils.env('PERFORMA_BOOK'),
                help='Generate report in ReST format and store it into the '
                     'specified folder, defaults to env[PERFORMA_BOOK]. '),
+    cfg.StrOpt('tag',
+               default=utils.env('PERFORMA_TAG'),
+               help='Tag the execution, defaults to env[PERFORMA_TAG].'),
 ]
 
 
diff --git a/performa/engine/main.py b/performa/engine/main.py
index 2f2118d..7eb48eb 100644
--- a/performa/engine/main.py
+++ b/performa/engine/main.py
@@ -37,12 +37,17 @@ def main():
     scenario = utils.read_yaml_file(scenario_file_path)
     base_dir = os.path.dirname(scenario_file_path)
 
-    records = player.play_scenario(scenario)
+    tag = cfg.CONF.tag
+    if not tag:
+        tag = utils.random_string()
+        LOG.info('Using auto-generated tag "%s"', tag)
+
+    records = player.play_scenario(scenario, tag)
 
     storage.store_data(records, cfg.CONF.mongo_url, cfg.CONF.mongo_db)
 
     report.generate_report(scenario, base_dir, cfg.CONF.mongo_url,
-                           cfg.CONF.mongo_db, cfg.CONF.book)
+                           cfg.CONF.mongo_db, cfg.CONF.book, tag)
 
 
 if __name__ == "__main__":
diff --git a/performa/engine/player.py b/performa/engine/player.py
index 77b6a0f..e06625c 100644
--- a/performa/engine/player.py
+++ b/performa/engine/player.py
@@ -80,7 +80,12 @@ def play_execution(execution):
     return records
 
 
-def play_scenario(scenario):
+def tag_records(records, tag):
+    for r in records:
+        r['tag'] = tag
+
+
+def play_scenario(scenario, tag):
     records = {}
 
     if 'preparation' in scenario:
@@ -90,5 +95,6 @@ def play_scenario(scenario):
         execution = scenario['execution']
 
         records = play_execution(execution)
+        tag_records(records, tag)
 
     return records
diff --git a/performa/engine/report.py b/performa/engine/report.py
index 21a1685..623ed6f 100644
--- a/performa/engine/report.py
+++ b/performa/engine/report.py
@@ -32,12 +32,15 @@ from performa.engine import utils
 LOG = logging.getLogger(__name__)
 
 
-def generate_chart(chart_str, records_collection, doc_folder):
+def generate_chart(chart_str, records_collection, doc_folder, tag):
     chart = yaml.safe_load(chart_str)
     pipeline = chart.get('pipeline')
     title = chart.get('title')
     axes = chart.get('axes') or dict(x='x', y='y')
 
+    if tag:
+        pipeline.insert(0, {'$match': {'tag': tag}})
+
     chart_data = records_collection.aggregate(pipeline)
 
     line = []
@@ -81,7 +84,8 @@ def _make_dir(name):
             raise
 
 
-def generate_report(scenario, base_dir, mongo_url, db_name, doc_folder):
+def generate_report(scenario, base_dir, mongo_url, db_name, doc_folder,
+                    tag=None):
     LOG.info('Generate report')
 
     doc_folder = doc_folder or tempfile.mkdtemp(prefix='performa')
@@ -101,7 +105,8 @@ def generate_report(scenario, base_dir, mongo_url, db_name, doc_folder):
     jinja_env.filters['chart'] = functools.partial(
         generate_chart,
         records_collection=records_collection,
-        doc_folder=doc_folder)
+        doc_folder=doc_folder,
+        tag=tag)
 
     template = utils.read_file(report_template, base_dir=base_dir)
     compiled_template = jinja_env.from_string(template)