From f5a7533e7fdaefa7ee5ddad51205a55a0a27320b Mon Sep 17 00:00:00 2001
From: Cristhian Garcia <crisgarta8@gmail.com>
Date: Mon, 30 Sep 2024 10:06:22 -0500
Subject: [PATCH] feat(metrics): add organization filter option

---
 tutoraspects/commands_v1.py                           | 11 +++++++++--
 .../apps/superset/pythonpath/performance_metrics.py   |  8 +++++++-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/tutoraspects/commands_v1.py b/tutoraspects/commands_v1.py
index 2723b6f58..9dbf55c6d 100644
--- a/tutoraspects/commands_v1.py
+++ b/tutoraspects/commands_v1.py
@@ -148,6 +148,11 @@ def init_clickhouse() -> list[tuple[str, str]]:
 
 # Ex: "tutor local do performance-metrics "
 @click.command(context_settings={"ignore_unknown_options": True})
+@click.option(
+    "--org",
+    default="",
+    help="An organization to apply as a filter.",
+)
 @click.option(
     "--course_name",
     default="",
@@ -169,12 +174,14 @@ def init_clickhouse() -> list[tuple[str, str]]:
     "--fail_on_error", is_flag=True, default=False, help="Allow errors to fail the run."
 )
 def performance_metrics(
-    course_name, dashboard_slug, slice_name, print_sql, fail_on_error
+    org, course_name, dashboard_slug, slice_name, print_sql, fail_on_error
 ) -> (list)[tuple[str, str]]:
     """
     Job to measure performance metrics of charts and its queries in Superset and ClickHouse.
     """
-    options = f"--course_name '{course_name}'" if course_name else ""
+    options = ""
+    options += f"--org '{org}' " if org else ""
+    options += f"--course_name '{course_name}' " if course_name else ""
     options += f" --dashboard_slug {dashboard_slug}" if dashboard_slug else ""
     options += f' --slice_name "{slice_name}"' if slice_name else ""
     options += " --print_sql" if print_sql else ""
diff --git a/tutoraspects/templates/aspects/apps/superset/pythonpath/performance_metrics.py b/tutoraspects/templates/aspects/apps/superset/pythonpath/performance_metrics.py
index ba681a3a7..3d709fce7 100644
--- a/tutoraspects/templates/aspects/apps/superset/pythonpath/performance_metrics.py
+++ b/tutoraspects/templates/aspects/apps/superset/pythonpath/performance_metrics.py
@@ -49,6 +49,10 @@
 )
 
 @click.command()
+@click.option(
+    "--org",
+    default="",
+    help="An organization to apply as a filter.")
 @click.option(
     "--course_name",
     default="",
@@ -71,7 +75,7 @@
 @click.option(
     "--fail_on_error", is_flag=True, default=False, help="Allow errors to fail the run."
 )
-def performance_metrics(course_name, dashboard_slug, slice_name, print_sql,
+def performance_metrics(org, course_name, dashboard_slug, slice_name, print_sql,
                         fail_on_error):
     """
     Measure the performance of the dashboard.
@@ -81,6 +85,8 @@ def performance_metrics(course_name, dashboard_slug, slice_name, print_sql,
     extra_filters = []
     if course_name:
         extra_filters += [{"col": "course_name", "op": "IN", "val": course_name}]
+    if org:
+        extra_filters += [{"col": "org", "op": "IN", "val": org}]
 
     with patch("clickhouse_connect.common.build_client_name") as mock_build_client_name:
         mock_build_client_name.return_value = RUN_ID