Merge remote-tracking branch 'origin/enterprise' into fix/TG-1120-misc-fixes

Author: testgen-ci-bot

testgen/api/runs.py

@@ -35,7 +35,7 @@
 )
 def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_tests)):  # noqa: B008
     """Get a test run by the job execution ID that created it."""
-    test_run = TestRun.get_by_id_or_job(job.id)
+    test_run = TestRun.get(job.id)
 
     result = None
     if test_run:
@@ -76,7 +76,7 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key ==
 )
 def get_profiling_run(job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_profile)):  # noqa: B008
     """Get a profiling run by the job execution ID that created it."""
-    profiling_run = ProfilingRun.get_by_id_or_job(job.id)
+    profiling_run = ProfilingRun.get(job.id)
 
     result = None
     if profiling_run:

testgen/commands/job_registry.py

@@ -80,7 +80,7 @@ def run_final_callbacks(job_exec: JobExecution) -> None:
 def _notify_profiling_run(job_exec: JobExecution) -> None:
     with database_session() as session:
         profiling_run = session.scalars(
-            select(ProfilingRun).where(ProfilingRun.job_execution_id == job_exec.id)
+            select(ProfilingRun).where(ProfilingRun.id == job_exec.id)
         ).first()
         if not profiling_run:
             LOG.warning("No profiling_run found for job %s; skipping notification", job_exec.id)
@@ -90,7 +90,7 @@ def _notify_profiling_run(job_exec: JobExecution) -> None:
 
 def _notify_test_run(job_exec: JobExecution) -> None:
     with database_session() as session:
-        test_run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_exec.id)).first()
+        test_run = session.scalars(select(TestRun).where(TestRun.id == job_exec.id)).first()
         if not test_run:
             LOG.warning("No test_run found for job %s; skipping notification", job_exec.id)
             return
@@ -99,7 +99,7 @@ def _notify_test_run(job_exec: JobExecution) -> None:
 
 def _notify_monitor_run(job_exec: JobExecution) -> None:
     with database_session() as session:
-        test_run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_exec.id)).first()
+        test_run = session.scalars(select(TestRun).where(TestRun.id == job_exec.id)).first()
         if not test_run:
             LOG.warning("No test_run found for job %s; skipping monitor notification", job_exec.id)
             return

testgen/commands/job_runner.py

@@ -71,12 +71,12 @@ def _print_run_summary(job_id: UUID, job_key: str) -> None:
         session = get_current_session()
         match job_key:
             case "run-profile":
-                run = session.scalars(select(ProfilingRun).where(ProfilingRun.job_execution_id == job_id)).first()
+                run = session.scalars(select(ProfilingRun).where(ProfilingRun.id == job_id)).first()
                 if run:
                     status_msg = "Profiling encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Profiling completed."
                     click.echo(f"\n        {status_msg}\n        Run ID: {run.id}\n    ")
             case "run-tests" | "run-monitors":
-                run = session.scalars(select(TestRun).where(TestRun.job_execution_id == job_id)).first()
+                run = session.scalars(select(TestRun).where(TestRun.id == job_id)).first()
                 if run:
                     status_msg = "Test execution encountered an error. Check log for details." if run.job_execution.status == JobStatus.ERROR else "Test execution completed."
                     click.echo(f"\n        {status_msg}\n        Run ID: {run.id}\n    ")

testgen/commands/run_data_cleanup.py

@@ -42,13 +42,9 @@ def run_data_cleanup(project_code: str, retention_days: int) -> None:
     with database_session():
         protected_profiling_ids = ProfilingRun.find_latest_per_table_group(project_code)
         protected_test_run_ids = TestRun.find_latest_per_test_suite(project_code)
-        # Translate protected run ids → their job_execution_ids so the JE sweep
-        # can carve them out. Nulls (older runs without a JE) are filtered here.
-        je_map = {
-            **ProfilingRun.get_job_execution_ids(list(protected_profiling_ids)),
-            **TestRun.get_job_execution_ids(list(protected_test_run_ids)),
-        }
-        protected_job_execution_ids = {je for je in je_map.values() if je is not None}
+        # The run id is the job execution id, so the protected run ids are
+        # already the job executions the sweep must carve out.
+        protected_job_execution_ids = protected_profiling_ids | protected_test_run_ids
 
     LOG.info(
         "Protected latest runs: profiling=%d test=%d job_executions=%d",

testgen/commands/run_profiling.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from datetime import UTC, datetime, timedelta
 from uuid import UUID
 
@@ -52,14 +51,12 @@ def run_profiling(
 
     LOG.info("Creating profiling run record")
     profiling_run = ProfilingRun(
+        id=job_context.get().job_id,
         project_code=table_group.project_code,
         connection_id=connection.connection_id,
         table_groups_id=table_group.id,
         profiling_starttime=datetime.now(UTC) + time_delta,
-        process_id=os.getpid(),
     )
-    if job_id := job_context.get().job_id:
-        profiling_run.job_execution_id = job_id
 
     # This runs in a subprocess — commit after every save so progress is visible
     # to the UI (separate session) and to execute_db_queries (independent connection).
@@ -96,19 +93,12 @@ def run_profiling(
             #     run_pairwise_contingency_check(profiling_run.id, table_group.profile_pair_rule_pct)
         else:
             LOG.info("No columns were selected to profile.")
-    except Exception as e:
+    except Exception:
         LOG.exception("Profiling encountered an error.")
-        LOG.info("Setting profiling run status to Error")
-        profiling_run.log_message = get_exception_message(e)
-        profiling_run.profiling_endtime = datetime.now(UTC) + time_delta
-        profiling_run.status = "Error"
-        profiling_run.save()
-        session.commit()
+        end_time = datetime.now(UTC) + time_delta
         raise
     else:
-        LOG.info("Setting profiling run status to Completed")
-        profiling_run.profiling_endtime = datetime.now(UTC) + time_delta
-        profiling_run.status = "Complete"
+        end_time = datetime.now(UTC) + time_delta
         profiling_run.save()
         session.commit()
 
@@ -122,7 +112,7 @@ def run_profiling(
             sampling=table_group.profile_use_sampling,
             table_count=profiling_run.table_ct or 0,
             column_count=profiling_run.column_ct or 0,
-            run_duration=(profiling_run.profiling_endtime - profiling_run.profiling_starttime).total_seconds(),
+            run_duration=(end_time - profiling_run.profiling_starttime.replace(tzinfo=UTC)).total_seconds(),
         )
 
     return profiling_run.id

testgen/commands/run_score_update.py

@@ -36,7 +36,7 @@ def run_score_update(parent_job_id: str, parent_job_key: str) -> None:
 def _rollup_profiling(parent_je_id: UUID) -> None:
     with database_session() as session:
         profiling_run = session.scalars(
-            select(ProfilingRun).where(ProfilingRun.job_execution_id == parent_je_id)
+            select(ProfilingRun).where(ProfilingRun.id == parent_je_id)
         ).first()
         if not profiling_run:
             LOG.error("No profiling_run found for job execution %s; skipping score rollup", parent_je_id)
@@ -60,7 +60,7 @@ def _rollup_test(parent_je_id: UUID) -> None:
         row = session.execute(
             select(TestRun.id, TestRun.test_starttime, TestSuite.table_groups_id, TestSuite.project_code)
             .join(TestSuite, TestRun.test_suite_id == TestSuite.id)
-            .where(TestRun.job_execution_id == parent_je_id)
+            .where(TestRun.id == parent_je_id)
         ).first()
         if not row:
             LOG.error("No test_run found for job execution %s; skipping score rollup", parent_je_id)

testgen/commands/run_test_execution.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from collections import defaultdict
 from datetime import UTC, datetime, timedelta
 from functools import partial
@@ -24,7 +23,6 @@
 from testgen.common.models.table_group import TableGroup
 from testgen.common.models.test_run import TestRun
 from testgen.common.models.test_suite import TestSuite
-from testgen.utils import get_exception_message
 
 from .run_refresh_data_chars import run_data_chars_refresh
 from .run_test_validation import run_test_validation
@@ -52,12 +50,10 @@ def run_test_execution(
 
     LOG.info("Creating test run record")
     test_run = TestRun(
+        id=job_context.get().job_id,
         test_suite_id=test_suite_id,
         test_starttime=datetime.now(UTC) + time_delta,
-        process_id=os.getpid(),
     )
-    if job_id := job_context.get().job_id:
-        test_run.job_execution_id = job_id
 
     # This runs in a subprocess — commit after every save so progress is visible
     # to the UI (separate session) and to execute_db_queries (independent connection).
@@ -134,22 +130,12 @@ def run_test_execution(
         # Refresh needed because previous query updates the test run too
         test_run.refresh()
 
-    except Exception as e:
+    except Exception:
         LOG.exception("Test execution encountered an error.")
-        LOG.info("Setting test run status to Error")
-        test_run.log_message = get_exception_message(e)
-        test_run.test_endtime = datetime.now(UTC) + time_delta
-        test_run.status = "Error"
-        test_run.save()
-        session.commit()
+        end_time = datetime.now(UTC) + time_delta
         raise
     else:
-        LOG.info("Setting test run status to Completed")
-        test_run.test_endtime = datetime.now(UTC) + time_delta
-        test_run.status = "Complete"
-        test_run.save()
-        session.commit()
-
+        end_time = datetime.now(UTC) + time_delta
         LOG.info("Updating latest run for test suite")
         test_suite.last_complete_test_run_id = test_run.id
         test_suite.save()
@@ -167,7 +153,7 @@ def run_test_execution(
             username=username,
             sql_flavor=connection.sql_flavor_code,
             test_count=test_run.test_ct,
-            run_duration=(test_run.test_endtime - test_run.test_starttime.replace(tzinfo=UTC)).total_seconds(),
+            run_duration=(end_time - test_run.test_starttime.replace(tzinfo=UTC)).total_seconds(),
             prediction_duration=(datetime.now(UTC) + time_delta - prediction_start).total_seconds(),
         )
 

testgen/common/models/data_column.py

@@ -22,6 +22,7 @@
 from testgen.common.models import get_current_session
 from testgen.common.models.entity import Entity, EntityMinimal
 from testgen.common.models.hygiene_issue import HygieneIssue
+from testgen.common.models.job_execution import JobExecution
 from testgen.common.models.profile_result import ProfileResult
 from testgen.common.models.profiling_run import ProfilingRun
 
@@ -559,11 +560,11 @@ def get_column_detail(
                 cls.dq_score_testing,
                 func.coalesce(hygiene_subq.c.hygiene_issue_count, 0).label("hygiene_issue_count"),
                 ProfilingRun.id.label("profile_run_id"),
-                ProfilingRun.job_execution_id.label("profile_run_je_id"),
-                ProfilingRun.status.label("profile_run_status"),
+                ProfilingRun.id.label("profile_run_je_id"),
+                JobExecution.status.label("profile_run_status"),
                 ProfilingRun.profiling_starttime.label("profile_run_started_at"),
-                ProfilingRun.profiling_endtime.label("profile_run_ended_at"),
-                ProfilingRun.log_message.label("profile_run_log_message"),
+                JobExecution.completed_at.label("profile_run_ended_at"),
+                JobExecution.error_message.label("profile_run_log_message"),
             )
             .outerjoin(DataTable, DataTable.id == cls.table_id)
             .outerjoin(
@@ -585,6 +586,7 @@ def get_column_detail(
                 ),
             )
             .outerjoin(ProfilingRun, ProfilingRun.id == ProfileResult.profile_run_id)
+            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id)
             .where(
                 cls.table_groups_id == table_groups_id,
                 cls.table_name == table_name,

testgen/common/models/data_table.py

@@ -134,7 +134,7 @@ def get_profiling_overview(
                 JobExecution.id.label("latest_profile_job_execution_id"),
             )
             .outerjoin(ProfilingRun, ProfilingRun.id == cls.last_complete_profile_run_id)
-            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id)
+            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id)
             .where(
                 cls.table_groups_id == table_groups_id,
                 cls.table_name == table_name,

testgen/common/models/hygiene_issue.py

@@ -275,7 +275,7 @@ def list_for_run(
                     ProfileResult.column_name == cls.column_name,
                 ),
             )
-            .where(ProfilingRun.job_execution_id == job_execution_id, *clauses)
+            .where(ProfilingRun.id == job_execution_id, *clauses)
             .order_by(cls._priority_order(), cls.table_name, cls.column_name, cls.id)
         )
         return cls._paginate(query, page=page, limit=limit, data_class=HygieneIssueListRow)
@@ -299,7 +299,7 @@ def search(
                 cls.project_code.label("project_code"),
                 HygieneIssueType.name.label("issue_type_name"),
                 TableGroup.table_groups_name.label("table_groups_name"),
-                ProfilingRun.job_execution_id.label("job_execution_id"),
+                ProfilingRun.id.label("job_execution_id"),
                 JobExecution.started_at.label("started_at"),
                 cls.schema_name.label("schema_name"),
                 cls.table_name.label("table_name"),
@@ -314,7 +314,7 @@ def search(
             )
             .join(HygieneIssueType, HygieneIssueType.id == cls.type_id)
             .join(ProfilingRun, ProfilingRun.id == cls.profile_run_id)
-            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id)
+            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id)
             .join(TableGroup, TableGroup.id == cls.table_groups_id)
             .outerjoin(
                 ProfileResult,
@@ -358,7 +358,7 @@ def get_with_context(cls, issue_id: UUID, *clauses) -> HygieneIssueDetail | None
                 cls.detail.label("detail"),
                 HygieneIssueType.detail_redactable.label("detail_redactable"),
                 ProfileResult.pii_flag.label("pii_flag"),
-                ProfilingRun.job_execution_id.label("job_execution_id"),
+                ProfilingRun.id.label("job_execution_id"),
                 JobExecution.started_at.label("started_at"),
                 ProfileResult.general_type.label("column_general_type"),
                 ProfileResult.db_data_type.label("column_db_data_type"),
@@ -368,7 +368,7 @@ def get_with_context(cls, issue_id: UUID, *clauses) -> HygieneIssueDetail | None
             )
             .join(HygieneIssueType, HygieneIssueType.id == cls.type_id)
             .join(ProfilingRun, ProfilingRun.id == cls.profile_run_id)
-            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.job_execution_id)
+            .outerjoin(JobExecution, JobExecution.id == ProfilingRun.id)
             .outerjoin(
                 ProfileResult,
                 and_(