Merge remote-tracking branch 'origin/enterprise' into feat/TG-1121-normalize-run-summary-count-shapes

Author: rboni-dk

deploy/testgen-base.dockerfile

@@ -27,9 +27,7 @@ RUN apk update && apk upgrade && apk add --no-cache \
     openblas=0.3.30-r2 \
     openblas-dev=0.3.30-r2 \
     unixodbc=2.3.14-r0 \
-    unixodbc-dev=2.3.14-r0 \
-    libarrow=21.0.0-r4 \
-    apache-arrow-dev=21.0.0-r4
+    unixodbc-dev=2.3.14-r0
 
 COPY --chmod=775 ./deploy/install_linuxodbc.sh /tmp/dk/install_linuxodbc.sh
 RUN /tmp/dk/install_linuxodbc.sh
@@ -39,7 +37,7 @@ COPY ./pyproject.toml /tmp/dk/pyproject.toml
 RUN mkdir /dk
 
 # Upgrading pip for security
-RUN python3 -m pip install --no-cache-dir --upgrade pip==26.0
+RUN python3 -m pip install --no-cache-dir --upgrade pip==26.1.2
 
 # hdbcli only ships manylinux wheels (no musl). pip 26+ correctly rejects these on Alpine.
 # We download the wheel for the correct arch, then extract it directly into site-packages
@@ -73,8 +71,7 @@ RUN apk del \
     openssl \
     linux-headers \
     openblas-dev \
-    unixodbc-dev \
-    apache-arrow-dev
+    unixodbc-dev
 
 # Remove interactive ODBC tools — not needed at runtime, and iusql triggers
 # false-positive secret detection in security scanners (SECRET-3010)

pyproject.toml

@@ -61,7 +61,7 @@ dependencies = [
     "xlsxwriter==3.2.0",
     "psutil==5.9.8",
     "concurrent_log_handler==0.9.25",
-    "cryptography==46.0.6",
+    "cryptography==46.0.7",
     "validators==0.33.0",
     "reportlab==4.2.2",
     "cron-converter==1.2.1",
@@ -72,7 +72,7 @@ dependencies = [
     "holidays~=0.89",
 
     # Pinned to match the manually compiled libs or for security
-    "pyarrow==21.0.0",
+    "pyarrow==23.0.1",
     "matplotlib==3.9.2",
     "scipy==1.14.1",
     "jinja2==3.1.6",
@@ -82,7 +82,7 @@ dependencies = [
     # MCP server
     "mcp[cli]==1.26.0",
     "uvicorn==0.41.0",
-    "PyJWT==2.12.0",
+    "PyJWT==2.13.0",
     "bcrypt==5.0.0",
 
     # API & OAuth server

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:
@@ -69,7 +69,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,