Merge branch 'telemetry' into 'enterprise'

misc(analytics): add events for nav, execution, and dialogs

See merge request dkinternal/testgen/dataops-testgen!209

Author: ci bot

testgen/__main__.py

@@ -113,12 +113,19 @@ def cli(ctx: Context, verbose: bool):
     help="The identifier for the table group used during a profile run. Use a table_group_id shown in list-table-groups.",
     default=None,
 )
-def run_profile(configuration: Configuration, table_group_id: str):
+@click.option(
+    "-s",
+    "--source",
+    required=False,
+    type=click.STRING,
+    default="cli",
+)
+def run_profile(configuration: Configuration, table_group_id: str, source: str):
     click.echo(f"run-profile with table_group_id: {table_group_id}")
     spinner = None
     if not configuration.verbose:
         spinner = MoonSpinner("Processing ... ")
-    message = run_profiling_queries(table_group_id, spinner=spinner)
+    message = run_profiling_queries(table_group_id, spinner=spinner, source=source)
     click.echo("\n" + message)
 
 
@@ -145,10 +152,17 @@ def run_profile(configuration: Configuration, table_group_id: str):
     required=False,
     default=None,
 )
+@click.option(
+    "-s",
+    "--source",
+    required=False,
+    type=click.STRING,
+    default="cli",
+)
 @pass_configuration
-def run_test_generation(configuration: Configuration, table_group_id: str, test_suite_key: str, generation_set: str):
+def run_test_generation(configuration: Configuration, table_group_id: str, test_suite_key: str, generation_set: str, source: str):
     LOG.info("CurrentStep: Generate Tests - Main Procedure")
-    message = run_test_gen_queries(table_group_id, test_suite_key, generation_set)
+    message = run_test_gen_queries(table_group_id, test_suite_key, generation_set, source)
     LOG.info("Current Step: Generate Tests - Main Procedure Complete")
     click.echo("\n" + message)
 
@@ -170,13 +184,20 @@ def run_test_generation(configuration: Configuration, table_group_id: str, test_
     required=False,
     default=settings.DEFAULT_TEST_SUITE_KEY,
 )
+@click.option(
+    "-s",
+    "--source",
+    required=False,
+    type=click.STRING,
+    default="cli",
+)
 @pass_configuration
-def run_tests(configuration: Configuration, project_key: str, test_suite_key: str):
+def run_tests(configuration: Configuration, project_key: str, test_suite_key: str, source: str):
     click.echo(f"run-tests for suite: {test_suite_key}")
     spinner = None
     if not configuration.verbose:
         spinner = MoonSpinner("Processing ... ")
-    message = run_execution_steps(project_key, test_suite_key, spinner=spinner)
+    message = run_execution_steps(project_key, test_suite_key, spinner=spinner, source=source)
     click.echo("\n" + message)
 
 

testgen/commands/run_execute_cat_tests.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import UTC, datetime
 
 from testgen.commands.queries.execute_cat_tests_query import CCATExecutionSQL
 from testgen.commands.run_refresh_score_cards_results import run_refresh_score_cards_results
@@ -9,6 +10,7 @@
     WriteListToDB,
     date_service,
 )
+from testgen.common.mixpanel_service import MixpanelService
 
 LOG = logging.getLogger("testgen")
 
@@ -60,14 +62,26 @@ def ParseCATResults(clsCATExecute):
     RunActionQueryList("DKTG", [strQuery])
 
 
-def FinalizeTestRun(clsCATExecute: CCATExecutionSQL):
-    lstQueries = [clsCATExecute.FinalizeTestResultsSQL(),
-                  clsCATExecute.PushTestRunStatusUpdateSQL(),
-                  clsCATExecute.FinalizeTestSuiteUpdateSQL(),
-                  clsCATExecute.CalcPrevalenceTestResultsSQL(),
-                  clsCATExecute.TestScoringRollupRunSQL(),
-                  clsCATExecute.TestScoringRollupTableGroupSQL()]
-    RunActionQueryList(("DKTG"), lstQueries)
+def FinalizeTestRun(clsCATExecute: CCATExecutionSQL, source: str):
+    _, row_counts = RunActionQueryList(("DKTG"), [
+        clsCATExecute.FinalizeTestResultsSQL(),
+        clsCATExecute.PushTestRunStatusUpdateSQL(),
+        clsCATExecute.FinalizeTestSuiteUpdateSQL(),
+    ])
+
+    MixpanelService().send_event(
+        "run-tests",
+        source=source,
+        sql_flavor=clsCATExecute.flavor,
+        test_count=row_counts[0],
+        duration=(datetime.now(UTC) - date_service.parse_now(clsCATExecute.run_date)).total_seconds(),
+    )
+    
+    RunActionQueryList(("DKTG"), [
+        clsCATExecute.CalcPrevalenceTestResultsSQL(),
+        clsCATExecute.TestScoringRollupRunSQL(),
+        clsCATExecute.TestScoringRollupTableGroupSQL(),
+    ])
     run_refresh_score_cards_results(
         project_code=clsCATExecute.project_code,
         add_history_entry=True,
@@ -76,7 +90,7 @@ def FinalizeTestRun(clsCATExecute: CCATExecutionSQL):
 
 
 def run_cat_test_queries(
-    dctParms, strTestRunID, strTestTime, strProjectCode, strTestSuite, error_msg, minutes_offset=0, spinner=None
+    dctParms, strTestRunID, strTestTime, strProjectCode, strTestSuite, error_msg, minutes_offset=0, spinner=None, source=None
 ):
     booErrors = False
 
@@ -146,4 +160,4 @@ def run_cat_test_queries(
 
     finally:
         LOG.info("Finalizing test run")
-        FinalizeTestRun(clsCATExecute)
+        FinalizeTestRun(clsCATExecute, source)

testgen/commands/run_execute_tests.py

@@ -103,10 +103,8 @@ def run_execution_steps_in_background(project_code, test_suite):
         empty_cache()
         background_thread = threading.Thread(
             target=run_execution_steps,
-            args=(
-                project_code,
-                test_suite
-            ),
+            args=(project_code, test_suite),
+            kwargs={"source": "ui"},
         )
         background_thread.start()
     else:
@@ -115,7 +113,13 @@ def run_execution_steps_in_background(project_code, test_suite):
         subprocess.Popen(script)  # NOQA S603
 
 
-def run_execution_steps(project_code: str, test_suite: str, minutes_offset: int=0, spinner: Spinner=None) -> str:
+def run_execution_steps(
+    project_code: str,
+    test_suite: str,
+    minutes_offset: int=0,
+    spinner: Spinner=None,
+    source: str | None=None,
+) -> str:
     # Initialize required parms for all steps
     has_errors = False
     error_msg = ""
@@ -165,7 +169,7 @@ def run_execution_steps(project_code: str, test_suite: str, minutes_offset: int=
 
     LOG.info("CurrentStep: Execute Step - CAT Test Execution")
     if run_cat_test_queries(
-        test_exec_params, test_run_id, test_time, project_code, test_suite, error_msg, minutes_offset, spinner
+        test_exec_params, test_run_id, test_time, project_code, test_suite, error_msg, minutes_offset, spinner, source
     ):
         has_errors = True
 

testgen/commands/run_generate_tests.py

@@ -2,11 +2,12 @@
 
 from testgen.commands.queries.generate_tests_query import CDeriveTestsSQL
 from testgen.common import AssignConnectParms, RetrieveDBResultsToDictList, RetrieveTestGenParms, RunActionQueryList
+from testgen.common.mixpanel_service import MixpanelService
 
 LOG = logging.getLogger("testgen")
 
 
-def run_test_gen_queries(strTableGroupsID, strTestSuite, strGenerationSet=None):
+def run_test_gen_queries(strTableGroupsID, strTestSuite, strGenerationSet=None, source=None):
     if strTableGroupsID is None:
         raise ValueError("Table Group ID was not specified")
 
@@ -56,7 +57,8 @@ def run_test_gen_queries(strTableGroupsID, strTestSuite, strGenerationSet=None):
         LOG.info("CurrentStep: Creating new Test Suite")
         strQuery = clsTests.GetInsertTestSuiteSQL(booClean)
         if strQuery:
-            clsTests.test_suite_id, = RunActionQueryList("DKTG", [strQuery])
+            insert_ids, _ = RunActionQueryList("DKTG", [strQuery])
+            clsTests.test_suite_id = insert_ids[0]
         else:
             raise ValueError("Test Suite not found and could not be created")
 
@@ -104,6 +106,15 @@ def run_test_gen_queries(strTableGroupsID, strTestSuite, strGenerationSet=None):
     if lstQueries:
         LOG.info("Running Test Generation Template Queries")
         RunActionQueryList("DKTG", lstQueries)
-        return "Test generation completed successfully."
+        message = "Test generation completed successfully."
     else:
-        return "No TestGen Queries were compiled."
+        message = "No TestGen Queries were compiled."
+    
+    MixpanelService().send_event(
+        "generate-tests",
+        source=source,
+        sql_flavor=clsTests.sql_flavor,
+        generation_set=clsTests.generation_set,
+    )
+
+    return message

testgen/commands/run_profiling_bridge.py

@@ -2,6 +2,7 @@
 import subprocess
 import threading
 import uuid
+from datetime import UTC, datetime
 
 import pandas as pd
 
@@ -20,6 +21,7 @@
     date_service,
 )
 from testgen.common.database.database_service import empty_cache
+from testgen.common.mixpanel_service import MixpanelService
 
 booClean = True
 LOG = logging.getLogger("testgen")
@@ -234,15 +236,19 @@ def run_profiling_in_background(table_group_id):
     if settings.IS_DEBUG:
         LOG.info(msg + ". Running in debug mode (new thread instead of new process).")
         empty_cache()
-        background_thread = threading.Thread(target=run_profiling_queries, args=(table_group_id,))
+        background_thread = threading.Thread(
+            target=run_profiling_queries,
+            args=(table_group_id),
+            kwargs={"source": "ui"},
+        )
         background_thread.start()
     else:
         LOG.info(msg)
         script = ["testgen", "run-profile", "-tg", table_group_id]
         subprocess.Popen(script)  # NOQA S603
 
 
-def run_profiling_queries(strTableGroupsID, spinner=None):
+def run_profiling_queries(strTableGroupsID, spinner=None, source=None):
     if strTableGroupsID is None:
         raise ValueError("Table Group ID was not specified")
 
@@ -308,28 +314,29 @@ def run_profiling_queries(strTableGroupsID, spinner=None):
     if spinner:
         spinner.next()
 
+    table_count = 0
+    column_count = 0
     try:
         # Retrieve Column Metadata
         LOG.info("CurrentStep: Getting DDF from project")
 
         strQuery = clsProfiling.GetDDFQuery()
         lstResult = RetrieveDBResultsToDictList("PROJECT", strQuery)
-
-        if len(lstResult) == 0:
-            LOG.warning("SQL retrieved 0 records")
+        column_count = len(lstResult)
 
         if lstResult:
-            if clsProfiling.profile_use_sampling == "Y":
-                # Get distinct tables
-                distinct_tables = set()
-                for item in lstResult:
-                    schema_name = item["table_schema"]
-                    table_name = item["table_name"]
-                    distinct_tables.add(f"{schema_name}.{table_name}")
+            # Get distinct tables
+            distinct_tables = set()
+            for item in lstResult:
+                schema_name = item["table_schema"]
+                table_name = item["table_name"]
+                distinct_tables.add(f"{schema_name}.{table_name}")
 
-                # Convert the set to a list
-                distinct_tables_list = list(distinct_tables)
+            # Convert the set to a list
+            distinct_tables_list = list(distinct_tables)
+            table_count = len(distinct_tables_list)
 
+            if clsProfiling.profile_use_sampling == "Y":
                 # Sampling tables
                 lstQueries = []
                 for parm_sampling_table in distinct_tables_list:
@@ -494,12 +501,24 @@ def run_profiling_queries(strTableGroupsID, spinner=None):
         raise
     finally:
         LOG.info("Updating the profiling run record")
-        lstProfileRunQuery = [
+        RunActionQueryList("DKTG", [
             clsProfiling.GetProfileRunInfoRecordUpdateQuery(),
+        ])
+
+        MixpanelService().send_event(
+            "run-profiling",
+            source=source,
+            sql_flavor=clsProfiling.flavor,
+            sampling=clsProfiling.profile_use_sampling == "Y",
+            table_count=table_count,
+            column_count=column_count,
+            duration=(datetime.now(UTC) - date_service.parse_now(clsProfiling.run_date)).total_seconds(),
+        )
+
+        RunActionQueryList("DKTG", [
             clsProfiling.GetAnomalyScoringRollupRunQuery(),
             clsProfiling.GetAnomalyScoringRollupTableGroupQuery(),
-        ]
-        RunActionQueryList("DKTG", lstProfileRunQuery)
+        ])
         run_refresh_score_cards_results(
             project_code=dctParms["project_code"],
             add_history_entry=True,

testgen/common/database/database_service.py

@@ -393,7 +393,8 @@ def RunActionQueryList(strCredentialSet, lstQueries, strAdminNDS="N", user_overr
     ) as con:
         i = 0
         n = len(lstQueries)
-        lstInsertedIds = []
+        insert_ids = []
+        row_counts = []
         if n == 0:
             LOG.info("No queries to process")
         for q in lstQueries:
@@ -402,20 +403,21 @@ def RunActionQueryList(strCredentialSet, lstQueries, strAdminNDS="N", user_overr
             LOG.info(f"(Processing {i} of {n})")
             tx = con.begin()
             exQ = con.execute(text(q))
+            row_counts.append(exQ.rowcount)
             if exQ.rowcount == -1:
                 strMsg = "Action query processed no records."
             else:
                 strMsg = str(exQ.rowcount) + " records processed."
 
                 try:
-                    lstInsertedIds.append(exQ.fetchone()[0])
+                    insert_ids.append(exQ.fetchone()[0])
                 except Exception:
-                    lstInsertedIds.append(None)
+                    insert_ids.append(None)
 
             tx.commit()
             LOG.info(strMsg)
 
-    return lstInsertedIds
+    return insert_ids, row_counts
 
 
 

testgen/common/mixpanel_service.py

@@ -35,7 +35,7 @@ def instance_id(self):
 
     @cached_property
     def distinct_id(self):
-        return self._hash_value(session.username)
+        return self._hash_value(session.username or self.instance_id)
 
     def _hash_value(self, value: bytes | str, digest_size: int = 8) -> str:
         if isinstance(value, str):

testgen/ui/navigation/router.py

@@ -6,6 +6,7 @@
 import streamlit as st
 
 import testgen.ui.navigation.page
+from testgen.common.mixpanel_service import MixpanelService
 from testgen.ui.session import session
 from testgen.utils.singleton import Singleton
 
@@ -83,6 +84,8 @@ def navigate(self, /, to: str, with_args: dict = {}) -> None:  # noqa: B006
                 len((st.query_params or {}).keys()) != len(final_args.keys())
                 or any(st.query_params.get(name) != value for name, value in final_args.items())
             )
+            if is_different_page:
+                MixpanelService().send_event(f"nav-{to}")
             if is_different_page or query_params_changed:
                 route = self._routes[to]
                 session.page_args_pending_router = {

testgen/ui/views/dialogs/generate_tests_dialog.py

@@ -71,7 +71,7 @@ def generate_tests_dialog(test_suite: pd.Series) -> None:
         status_container.info("Starting test generation ...")
 
         try:
-            run_test_gen_queries(table_group_id, test_suite_name, selected_set)
+            run_test_gen_queries(table_group_id, test_suite_name, selected_set, "ui")
         except Exception as e:
             status_container.error(f"Test generation encountered errors: {e!s}.")