Merge branch 'monitor-test-suite-1' into 'enterprise'

feat(tests): generate a monitor suite for new table groups

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

Author: ci bot

testgen/commands/run_execute_tests.py

@@ -125,9 +125,12 @@ def run_execution_steps_in_background(project_code, test_suite):
     if settings.IS_DEBUG:
         LOG.info(msg + ". Running in debug mode (new thread instead of new process).")
         empty_cache()
+        username = None
+        if session.auth:
+            username = session.auth.user_display
         background_thread = threading.Thread(
             target=run_execution_steps,
-            args=(project_code, test_suite, session.auth.user_display),
+            args=(project_code, test_suite, username),
         )
         background_thread.start()
     else:

testgen/commands/run_profiling_bridge.py

@@ -10,6 +10,8 @@
 import testgen.common.process_service as process_service
 from testgen import settings
 from testgen.commands.queries.profiling_query import CProfilingSQL
+from testgen.commands.run_execute_tests import run_execution_steps_in_background
+from testgen.commands.run_generate_tests import run_test_gen_queries
 from testgen.commands.run_refresh_score_cards_results import run_refresh_score_cards_results
 from testgen.common import (
     date_service,
@@ -25,6 +27,7 @@
 from testgen.common.mixpanel_service import MixpanelService
 from testgen.common.models import with_database_session
 from testgen.common.models.connection import Connection
+from testgen.common.models.test_suite import TestSuite
 from testgen.ui.session import session
 
 LOG = logging.getLogger("testgen")
@@ -238,6 +241,9 @@ def run_profiling_queries(table_group_id: str, username: str | None = None, spin
     profiling_run_id = str(uuid.uuid4())
 
     params = get_profiling_params(table_group_id)
+    needs_monitor_tests_generated = (
+        bool(params["monitor_test_suite_id"]) and not params["last_complete_profile_run_id"]
+    )
 
     LOG.info("CurrentStep: Initializing Query Generator")
     clsProfiling = CProfilingSQL(params["project_code"], connection.sql_flavor, minutes_offset=minutes_offset)
@@ -471,7 +477,24 @@ def run_profiling_queries(table_group_id: str, username: str | None = None, spin
             scoring_duration=(datetime.now(UTC) - end_time).total_seconds(),
         )
 
+        if needs_monitor_tests_generated:
+            _generate_monitor_tests(params["project_code"], table_group_id, params["monitor_test_suite_id"])
+
     return f"""
         Profiling completed {"with errors. Check log for details." if has_errors else "successfully."}
         Run ID: {profiling_run_id}
     """
+
+
+@with_database_session
+def _generate_monitor_tests(project_code: str, table_group_id: str, test_suite_id: str) -> None:
+    try:
+        monitor_test_suite = TestSuite.get(test_suite_id)
+        if not monitor_test_suite:
+            LOG.info("Skipping test generation on missing monitor test suite")
+        else:
+            LOG.info("Generating monitor tests")
+            run_test_gen_queries(table_group_id, monitor_test_suite.test_suite, "Monitor")
+            run_execution_steps_in_background(project_code, monitor_test_suite.test_suite)
+    except Exception:
+        LOG.exception("Error generating monitor tests")

testgen/common/get_pipeline_parms.py

@@ -21,6 +21,8 @@ class ProfilingParams(BaseParams):
     profile_sample_min_count: int
     profile_do_pair_rules: str
     profile_pair_rule_pct: int
+    monitor_test_suite_id: str | None
+    last_complete_profile_run_id: str | None
 
 
 class TestGenerationParams(BaseParams):

testgen/common/models/scheduler.py

@@ -4,11 +4,16 @@
 from uuid import UUID, uuid4
 
 from cron_converter import Cron
-from sqlalchemy import Column, String, select
+from sqlalchemy import Column, String, func, select
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.orm import InstrumentedAttribute
 
 from testgen.common.models import Base, get_current_session
+from testgen.common.models.test_definition import TestDefinition
+from testgen.common.models.test_suite import TestSuite
+
+RUN_TESTS_JOB_KEY = "run-tests"
+RUN_PROFILE_JOB_KEY = "run-profile"
 
 
 class JobSchedule(Base):
@@ -25,7 +30,23 @@ class JobSchedule(Base):
 
     @classmethod
     def select_where(cls, *clauses, order_by: str | InstrumentedAttribute | None = None) -> Iterable[Self]:
-        query = select(cls).where(*clauses).order_by(order_by)
+        test_definitions_count = (
+            select(cls.id)
+            .join(TestSuite, TestSuite.test_suite == cls.kwargs["test_suite_key"].astext)
+            .join(TestDefinition, TestDefinition.test_suite_id == TestSuite.id)
+            .where(cls.key == RUN_TESTS_JOB_KEY)
+            .group_by(cls.id, TestSuite.test_suite)
+            .having(func.count(TestDefinition.id) > 0)
+            .subquery()
+        )
+        test_runs_query = (
+            select(cls)
+            .join(test_definitions_count, test_definitions_count.c.id == cls.id)
+            .where(*clauses)
+        )
+        non_test_runs_query = select(cls).where(cls.key != RUN_TESTS_JOB_KEY, *clauses)
+        query = test_runs_query.union_all(non_test_runs_query).order_by(order_by)
+
         return get_current_session().execute(query)
 
     @classmethod

testgen/common/models/table_group.py

@@ -11,6 +11,7 @@
 from testgen.common.models import get_current_session
 from testgen.common.models.custom_types import NullIfEmptyString, YNString
 from testgen.common.models.entity import ENTITY_HASH_FUNCS, Entity, EntityMinimal
+from testgen.common.models.scheduler import RUN_TESTS_JOB_KEY, JobSchedule
 from testgen.common.models.scores import ScoreDefinition
 from testgen.common.models.test_suite import TestSuite
 
@@ -52,6 +53,11 @@ class TableGroup(Entity):
     id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, default=uuid4)
     project_code: str = Column(String, ForeignKey("projects.project_code"))
     connection_id: int = Column(BigInteger, ForeignKey("connections.connection_id"))
+    monitor_test_suite_id: UUID | None = Column(
+        postgresql.UUID(as_uuid=True),
+        ForeignKey("test_suites.id"),
+        default=None,
+    )
     table_groups_name: str = Column(String)
     table_group_schema: str = Column(String)
     profiling_table_set: str = Column(NullIfEmptyString)
@@ -260,7 +266,12 @@ def clear_cache(cls) -> bool:
         cls.select_minimal_where.clear()
         cls.select_summary.clear()
 
-    def save(self, add_scorecard_definition: bool = False) -> None:
+    def save(
+        self,
+        add_scorecard_definition: bool = False,
+        add_monitor_test_suite: bool = False,
+        monitor_schedule_timezone: str = "UTC",
+    ) -> None:
         if self.id:
             values = {
                 column.key: getattr(self, column.key, None)
@@ -273,7 +284,38 @@ def save(self, add_scorecard_definition: bool = False) -> None:
             db_session.commit()
         else:
             super().save()
+            db_session = get_current_session()
+
             if add_scorecard_definition:
                 ScoreDefinition.from_table_group(self).save()
 
+            if add_monitor_test_suite:
+                test_suite = TestSuite(
+                    project_code=self.project_code,
+                    test_suite=f"{self.table_groups_name} Monitor",
+                    connection_id=self.connection_id,
+                    table_groups_id=self.id,
+                    export_to_observability=False,
+                    dq_score_exclude=True,
+                    view_mode="Monitor",
+                )
+                test_suite.save()
+
+                schedule_job = JobSchedule(
+                    project_code=self.project_code,
+                    key=RUN_TESTS_JOB_KEY,
+                    cron_expr="0 * * * *",
+                    cron_tz=monitor_schedule_timezone,
+                    args=[],
+                    kwargs={"project_key": self.project_code, "test_suite_key": test_suite.test_suite},
+                )
+                db_session.add(schedule_job)
+
+                self.monitor_test_suite_id = test_suite.id
+                db_session.execute(
+                    update(TableGroup)
+                    .where(TableGroup.id == self.id).values(monitor_test_suite_id=test_suite.id)
+                )
+                db_session.commit()
+
         TableGroup.clear_cache()

testgen/common/models/test_suite.py

@@ -65,6 +65,7 @@ class TestSuite(Entity):
     component_name: str = Column(NullIfEmptyString)
     last_complete_test_run_id: UUID = Column(postgresql.UUID(as_uuid=True))
     dq_score_exclude: bool = Column(Boolean, default=False)
+    view_mode: str | None = Column(NullIfEmptyString, default=None)
 
     _default_order_by = (asc(func.lower(test_suite)),)
     _minimal_columns = TestSuiteMinimal.__annotations__.keys()

testgen/scheduler/cli_scheduler.py

@@ -46,7 +46,7 @@ def get_jobs(self) -> Iterable[CliJob]:
         self.reload_timer.start()
 
         jobs = {}
-        for (job_model,) in JobSchedule.select_where():
+        for job_model in JobSchedule.select_where():
             if job_model.key not in JOB_REGISTRY:
                 LOG.error("Job '%s' scheduled but not registered", job_model.key)
                 continue

testgen/template/dbsetup/030_initialize_new_schema_structure.sql

@@ -84,14 +84,15 @@ CREATE TABLE connections (
 CREATE TABLE table_groups
 (
     id                       UUID DEFAULT gen_random_uuid()
-      CONSTRAINT pk_tg_id
+         CONSTRAINT pk_tg_id
          PRIMARY KEY,
     project_code             VARCHAR(30)
-      CONSTRAINT table_groups_projects_project_code_fk
-           REFERENCES projects,
+         CONSTRAINT table_groups_projects_project_code_fk
+         REFERENCES projects,
     connection_id            BIGINT
-          CONSTRAINT table_groups_connections_connection_id_fk
-           REFERENCES connections,
+         CONSTRAINT table_groups_connections_connection_id_fk
+         REFERENCES connections,
+    monitor_test_suite_id    UUID DEFAULT NULL,
     table_groups_name        VARCHAR(100),
     table_group_schema       VARCHAR(100),
     profiling_table_set      VARCHAR(2000),
@@ -161,10 +162,14 @@ CREATE TABLE test_suites (
    component_name          VARCHAR(100),
    last_complete_test_run_id UUID,
    dq_score_exclude        BOOLEAN default FALSE,
+   view_mode               VARCHAR(20) DEFAULT NULL,
    CONSTRAINT test_suites_id_pk
       PRIMARY KEY (id)
 );
 
+ALTER TABLE table_groups ADD CONSTRAINT table_groups_test_suites_monitor_test_suite_id_fk
+    FOREIGN KEY (monitor_test_suite_id) REFERENCES test_suites ON DELETE SET NULL;
+
 CREATE TABLE test_definitions (
    id                     UUID DEFAULT gen_random_uuid(),
    cat_test_id            BIGINT GENERATED BY DEFAULT AS IDENTITY

testgen/template/dbupgrade/0150_incremental_upgrade.sql

@@ -0,0 +1,10 @@
+SET SEARCH_PATH TO {SCHEMA_NAME};
+
+ALTER TABLE test_suites
+  ADD COLUMN view_mode VARCHAR(20) DEFAULT NULL;
+
+ALTER TABLE table_groups
+  ADD COLUMN monitor_test_suite_id UUID DEFAULT NULL;
+
+ALTER TABLE table_groups ADD CONSTRAINT table_groups_test_suites_monitor_test_suite_id_fk
+    FOREIGN KEY (monitor_test_suite_id) REFERENCES test_suites ON DELETE SET NULL;

testgen/template/parms/parms_profiling.sql

@@ -1,6 +1,7 @@
 SELECT tg.project_code,
        tg.id::VARCHAR(50) as table_groups_id,
        tg.table_group_schema,
+       tg.table_group_schema,
        CASE
          WHEN tg.profiling_table_set ILIKE '''%''' THEN tg.profiling_table_set
          ELSE fn_format_csv_quotes(tg.profiling_table_set)
@@ -14,6 +15,14 @@ SELECT tg.project_code,
        tg.profile_sample_percent,
        tg.profile_sample_min_count,
        tg.profile_do_pair_rules,
-       tg.profile_pair_rule_pct
+       tg.profile_pair_rule_pct,
+       CASE
+        WHEN tg.monitor_test_suite_id IS NULL THEN NULL
+        ELSE tg.monitor_test_suite_id::VARCHAR(50)
+       END as monitor_test_suite_id,
+       CASE
+        WHEN tg.last_complete_profile_run_id is NULL THEN NULL
+        ELSE tg.last_complete_profile_run_id::VARCHAR(50)
+       END as last_complete_profile_run_id
   FROM table_groups tg
  WHERE tg.id = :TABLE_GROUP_ID;