Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
118 commits
Select commit Hold shift + click to select a range
9c0b076
chore: e2e tests update
luis-dk Apr 10, 2026
571bf41
refactor(ui): add data-value to help e2e tests
luis-dk Apr 17, 2026
63786f8
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 5, 2026
c16f2ad
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 5, 2026
bd07969
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 5, 2026
12e8f8c
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 6, 2026
b259f0e
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 6, 2026
6ad1b8c
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 6, 2026
4257369
feat(server): harden API + MCP server for production deployments (TG-…
rboni-dk May 7, 2026
5c8e479
Merge branch 'feat/TG-1065-harden-api-mcp-server-for-production' into…
May 7, 2026
ac8baa2
feat(mcp): add run status & history tools (TG-1050)
rboni-dk May 8, 2026
3fa16dd
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 7, 2026
6250ebb
Merge branch 'feat/TG-1050-mcp-run-status-and-history' into 'enterprise'
May 11, 2026
cb23b50
Merge remote-tracking branch 'origin/enterprise' into ui-page-tests-e2e
May 11, 2026
6b2c390
feat(mcp): add test definition CRUD tools (TG-1054)
rboni-dk May 11, 2026
3b47f0b
Merge branch 'ui-page-tests-e2e' into 'enterprise'
luis-dk May 11, 2026
89fa359
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1054-mc…
May 11, 2026
55cca1a
refactor(mcp): add get_column_profile_detail tool
luis-dk May 8, 2026
26127eb
Merge branch 'tg-1052-column-profile' into 'enterprise'
May 11, 2026
e02afc3
refactor(mcp): apply TG-1054 review feedback
rboni-dk May 12, 2026
a0ca7a3
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1054-mc…
rboni-dk May 12, 2026
079331d
refactor: consolidate row-limiting clauses into FlavorService
rboni-dk May 12, 2026
600df67
Merge branch 'feat/TG-1054-mcp-test-definition-crud' into 'enterprise'
May 12, 2026
11bba62
feat(mcp): profiling L3 — cross-column search, frequent values, patte…
rboni-dk May 11, 2026
c44ec72
fix(monitors): freshness-gate Volume_Trend/Metric_Trend prediction
aarthy-dk May 11, 2026
e5f6ac0
refactor: centralize /api/v1 prefix in api package router
rboni-dk May 13, 2026
61509f3
refactor: extract _check_access helper for API resolvers
rboni-dk May 13, 2026
158331d
refactor: drop vestigial args column from job_executions and job_sche…
rboni-dk May 13, 2026
453203b
refactor: consolidate cross-cutting enums into common.enums
rboni-dk May 13, 2026
207f8e4
refactor: gate public job exposure by job_key allowlist, not source
rboni-dk May 13, 2026
470fc1e
fix(scorecards): filter categories by CDE
luis-dk May 13, 2026
5f8c106
Merge branch 'fix/cde-quality-scores' into 'enterprise'
May 13, 2026
4a1669e
Merge remote-tracking branch 'origin/enterprise' into refactor/light-…
May 13, 2026
3f2297c
Merge remote-tracking branch 'origin/enterprise' into feat/freshness-…
May 13, 2026
32a2085
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1067-mc…
May 13, 2026
39d4798
fix: drop args column from quick-start seed insert
rboni-dk May 13, 2026
e957c70
Merge branch 'refactor/light-cleanup' into 'enterprise'
May 13, 2026
c967ce7
Merge remote-tracking branch 'origin/enterprise' into feat/freshness-…
May 13, 2026
6e16024
Merge branch 'feat/freshness-gated-monitors' into 'enterprise'
May 13, 2026
c7bd8de
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1067-mc…
rboni-dk May 13, 2026
f3a1582
fix(standalone): resolve embedded host/port at connection-build time
aarthy-dk May 15, 2026
1a6150d
fix(standalone): revert Windows signal forwarding to TerminateProcess
aarthy-dk May 15, 2026
8d9b600
feat: add server-side pagination for test definitions (TG-1041)
FernandezAstor Apr 28, 2026
75eafe5
refactor(TG-1041): address reviewer feedback on pagination implementa…
FernandezAstor May 15, 2026
4e848c3
Merge branch 'fix/TG-1083-standalone-windows-port' into 'enterprise'
May 15, 2026
442cad5
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 15, 2026
3c3f5f3
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1067-mc…
rboni-dk May 15, 2026
e33ef2f
fix(scoring): accept leading-dot decimals in fn_eval
rboni-dk May 15, 2026
c006257
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1079-te…
rboni-dk May 15, 2026
ce8cca3
refactor(mcp): apply TG-1067 review feedback
rboni-dk May 15, 2026
967ce26
Merge branch 'feat/TG-1067-mcp-profiling-l3-cross-column-search' into…
May 16, 2026
5f73a74
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1079-te…
May 16, 2026
863d2c6
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 16, 2026
7eb4c2c
feat(mcp): profiling L4 — cross-run comparison, trends, schema histor…
rboni-dk May 14, 2026
49c3d2c
refactor(mcp): apply TG-1068 review feedback
rboni-dk May 17, 2026
11332c8
Merge branch 'feat/TG-1068-mcp-profiling-l4-cross-run-comparison' int…
May 18, 2026
628ad33
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 18, 2026
af35cf5
Merge branch 'feat/TG-1079-test-type-validation-cat-table-level-tests…
rboni-dk May 18, 2026
354aa95
feat(salesforce): add Salesforce Data 360 flavor
aarthy-dk Apr 2, 2026
918088c
feat(mcp): schedule CRUD tools (TG-1071)
rboni-dk May 19, 2026
55d6a79
fix(salesforce): apply MR review feedback
aarthy-dk May 20, 2026
a4037d8
refactor(mcp): apply TG-1071 review feedback
rboni-dk May 20, 2026
34f3dca
Merge branch 'feat/TG-994-salesforce' into 'enterprise'
May 20, 2026
f927231
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1071-mc…
May 20, 2026
8cd8edb
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 20, 2026
476a7d0
ci: bump base image to v16
May 20, 2026
78b36df
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1071-mc…
May 20, 2026
207410c
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 20, 2026
5f4b3b8
feat(TG-1001): exclude monitor suites from all queries
FernandezAstor May 20, 2026
843f914
Merge branch 'feat/TG-1001-exclude-monitor-suites' into 'enterprise'
May 20, 2026
a227dd3
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1071-mc…
May 20, 2026
0208934
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 20, 2026
183805c
fix(TG-1080): cross-flavor template fixes for QUERY-style tests
rboni-dk May 21, 2026
8a869fa
Merge branch 'feat/TG-1071-mcp-schedules-list-get-crud' into 'enterpr…
May 21, 2026
a760a70
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1080-te…
May 21, 2026
9310d33
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 21, 2026
6ae73dd
Merge branch 'feat/TG-1080-test-type-validation-query-tests' into 'en…
rboni-dk May 21, 2026
3e54466
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 21, 2026
167a7b1
refactor(mcp): update inventory tool to display scorecards
luis-dk May 12, 2026
cec8098
feat(mcp): add new tool get_quality_scores
luis-dk May 13, 2026
3818ff0
feat(mcp): add CRUD tools for quality scores
luis-dk May 15, 2026
8a34d01
Merge branch 'mcp-quality-scores' into 'enterprise'
May 22, 2026
3b2dd49
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1041-te…
May 22, 2026
b104cb1
refactor(TG-1041): address second round of reviewer feedback
FernandezAstor May 22, 2026
4cec0bb
Merge branch 'feat/TG-1041-test-definitions-db-pagination' into 'ente…
FernandezAstor May 22, 2026
9f0a452
refactor(models): decouple Streamlit cache from common layer
luis-dk May 26, 2026
3e800d9
fix(common-models): get_previous returns self in TestRun and Profilin…
rboni-dk May 27, 2026
fe37c41
feat(mcp): single-arg compare_test_runs (TG-1056)
rboni-dk May 27, 2026
ca685ee
Merge branch 'fix/st-cache-on-mcp' into 'enterprise'
May 27, 2026
161cc1c
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1056-mc…
May 27, 2026
c0670fc
Merge branch 'feat/TG-1056-mcp-compare-runs-single-arg-diff' into 'en…
May 28, 2026
077c70d
feat(retention): add per-project data retention cleanup (TG-1063)
diogodk May 11, 2026
b95ccea
feat(mcp): test definition note CRUD tools (TG-1086)
rboni-dk May 28, 2026
f29dd01
Merge branch 'feat/TG-1063-data-retention' into 'enterprise'
aarthy-dk May 28, 2026
0b651a2
misc: remove noisy streamlit logs in debug mode
aarthy-dk May 28, 2026
0e1a20e
Merge branch 'misc-fix' into 'enterprise'
May 28, 2026
48c04f3
Merge remote-tracking branch 'origin/enterprise' into feat/TG-1086-mc…
May 28, 2026
a36ee7d
refactor(mcp): apply TG-1086 review feedback
rboni-dk May 29, 2026
80f13f2
Merge branch 'feat/TG-1086-mcp-test-definition-notes-crud' into 'ente…
May 29, 2026
e9b3c0e
feat: add feedback popup and help item
FernandezAstor May 29, 2026
d92b7a5
Merge branch 'astor/TG-1006_feedback_widget' into 'enterprise'
May 29, 2026
8e5d3ae
feat(mcp): add CRUD mcp tools for notifications
luis-dk May 19, 2026
e4fef2c
refactor(mcp): remove redundant session flush in schedule tools
luis-dk May 29, 2026
b2a0a0e
Merge branch 'mcp-notifications-crud' into 'enterprise'
Jun 1, 2026
b26e147
fix(ui): handle out-of-range dates when serializing results to JSON
aarthy-dk Jun 1, 2026
1049e1d
fix(generation): correct Freshness_Trend tran_date_cols filter preced…
aarthy-dk Jun 1, 2026
a2ece91
fix(profiling): guard empty SPLIT_PART casts in pattern anomaly criteria
aarthy-dk Jun 1, 2026
a3173bf
feat(ui): log UI render errors and show a custom error page
aarthy-dk Jun 1, 2026
f9c9da2
Merge branch 'main' into fix/TG-1101-sqlserver-profiling-and-freshness
aarthy-dk Jun 1, 2026
2704c7a
fix(scorecard): improve category layout
aarthy-dk Jun 1, 2026
3a09de1
fix(reports): correct Column Tags and link layout in test issue repor…
aarthy-dk Jun 2, 2026
753fc8e
fix(source-data): preserve datetimes for source-data queries and reports
aarthy-dk Jun 2, 2026
563dc0c
docs(mcp): change doc group for test definitions
aarthy-dk Jun 2, 2026
73b66ef
fix(source-data): handle fractional-second timestamps in parse_fuzzy_…
aarthy-dk Jun 2, 2026
6c95948
fix: address review feedback
aarthy-dk Jun 2, 2026
b2e9ca6
Merge branch 'fix/TG-1101-sqlserver-profiling-and-freshness' into 'en…
Jun 2, 2026
ed8ffa7
release: 5.33.3 -> 5.48.0
aarthy-dk Jun 2, 2026
e474fe7
fix(project-settings): disable Save button correctly
aarthy-dk Jun 3, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
  •  
  •  
  •  
2 changes: 2 additions & 0 deletions deploy/build_mcp_docs.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@
DocGroup.INVESTIGATE,
DocGroup.BROWSE_PROFILING,
DocGroup.TRIGGER,
DocGroup.SCORING,
DocGroup.MANAGE,
]
_FALLBACK_GROUP = "Other tools"

Expand Down
2 changes: 1 addition & 1 deletion deploy/testgen.dockerfile
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
ARG TESTGEN_BASE_LABEL=v15
ARG TESTGEN_BASE_LABEL=v16

FROM datakitchen/dataops-testgen-base:${TESTGEN_BASE_LABEL} AS release-image

Expand Down
10 changes: 9 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"

[project]
name = "dataops-testgen"
version = "5.33.3"
version = "5.48.0"
description = "DataKitchen's Data Quality DataOps TestGen"
authors = [
{ "name" = "DataKitchen, Inc.", "email" = "info@datakitchen.io" },
Expand Down Expand Up @@ -41,6 +41,7 @@ dependencies = [
"oracledb==3.4.0",
"hdbcli==2.25.31",
"sqlalchemy-hana==4.4.0",
"salesforce-cdp-connector>=1.0.19",
"pyodbc==5.2.0",
"psycopg2-binary==2.9.11",
"pycryptodome==3.21",
Expand Down Expand Up @@ -117,6 +118,9 @@ release = [
testgen = "testgen.__main__:cli"
tg-patch-streamlit = "testgen.ui.scripts.patch_streamlit:patch"

[project.entry-points."sqlalchemy.dialects"]
salesforce_data360 = "testgen.common.database.salesforce_data360_dialect:SalesforceData360Dialect"

[project.urls]
"Source Code" = "https://github.com/DataKitchen/dataops-testgen"
"Bug Tracker" = "https://github.com/DataKitchen/dataops-testgen/issues"
Expand Down Expand Up @@ -397,3 +401,7 @@ asset_dir = "ui/components/frontend/js"
[[tool.streamlit.component.components]]
name = "sidebar"
asset_dir = "ui/components/frontend/js"

[[tool.streamlit.component.components]]
name = "feedback_widget"
asset_dir = "ui/components/frontend/js"
2 changes: 1 addition & 1 deletion testgen/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -982,7 +982,7 @@ def init_ui():
"run",
app_file,
"--browser.gatherUsageStats=false",
f"--logger.level={'debug' if settings.IS_DEBUG else 'error'}",
"--logger.level=error",
"--client.showErrorDetails=none",
"--client.toolbarMode=minimal",
"--server.enableStaticServing=true",
Expand Down
12 changes: 12 additions & 0 deletions testgen/api/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
from fastapi import APIRouter

from testgen.api.app import router as _app_router
from testgen.api.jobs import router as _jobs_router
from testgen.api.runs import router as _runs_router
from testgen.api.test_definitions import router as _test_definitions_router

router = APIRouter(prefix="/api/v1")
router.include_router(_app_router)
router.include_router(_jobs_router)
router.include_router(_runs_router)
router.include_router(_test_definitions_router)
2 changes: 1 addition & 1 deletion testgen/api/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from testgen.api.deps import db_session
from testgen.common import version_service

router = APIRouter(prefix="/api/v1", tags=["API"], dependencies=[Depends(db_session)])
router = APIRouter(tags=["API"], dependencies=[Depends(db_session)])


@router.get("/health")
Expand Down
51 changes: 26 additions & 25 deletions testgen/api/deps.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,14 @@

from fastapi import Depends, HTTPException, Security, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from sqlalchemy import select

from testgen.common.auth import authorize_token, decode_jwt_token
from testgen.common.models import Session, _current_session_wrapper, get_current_session
from testgen.common.models.job_execution import PUBLIC_JOB_KEYS, JobExecution
from testgen.common.models.project_membership import ProjectMembership
from testgen.common.models.table_group import TableGroup
from testgen.common.models.test_suite import TestSuite
from testgen.common.models.user import User
from testgen.utils.plugins import PluginHook

Expand Down Expand Up @@ -73,14 +77,25 @@ def has_project_permission(user: User, project_code: str, permission: str) -> bo

# --- Resolver dependency factories ---
# Each factory takes a permission string and returns Depends(). The entity ID
# comes from a URL path parameter (FastAPI resolves it natively).
# Entity not found and insufficient permission both raise the same 404
# with a stable code/message — no variation that could leak the cause.
# comes from a URL path parameter (FastAPI resolves it natively, including
# UUID validation that yields a 422 for malformed inputs).

_require_user = Depends(get_authorized_user)
_not_found = api_error(404, "not_found", "Not found")


def _check_access(entity, user: User, permission: str):
"""Return ``entity`` if the user has ``permission`` on its project, else raise 404.

Entity-not-found and insufficient-permission both surface as the same 404
with a stable code/message — no variation that could leak the cause to an
unauthorized caller.
"""
if entity and has_project_permission(user, entity.project_code, permission):
return entity
raise _not_found


def resolve_project_code(permission: str):
"""Verify the user has ``permission`` on the project identified by ``project_code`` path param."""
def dependency(project_code: str, user: User = _require_user) -> str:
Expand All @@ -92,45 +107,31 @@ def dependency(project_code: str, user: User = _require_user) -> str:

def resolve_table_group(permission: str):
"""Resolve a TableGroup by ``table_group_id`` path param and verify project permission."""
from testgen.common.models.table_group import TableGroup

def dependency(table_group_id: UUID, user: User = _require_user) -> TableGroup:
if (table_group := TableGroup.get(table_group_id)) and has_project_permission(user, table_group.project_code, permission):
return table_group
raise _not_found
return _check_access(TableGroup.get(table_group_id), user, permission)
return Depends(dependency)


def resolve_test_suite(permission: str):
"""Resolve a non-monitor TestSuite by ``test_suite_id`` path param and verify project permission."""
from testgen.common.models.test_suite import TestSuite

def dependency(test_suite_id: UUID, user: User = _require_user) -> TestSuite:
if (test_suite := TestSuite.get_regular(test_suite_id)) and has_project_permission(user, test_suite.project_code, permission):
return test_suite
raise _not_found
return _check_access(TestSuite.get_regular(test_suite_id), user, permission)
return Depends(dependency)


def resolve_job(permission: str, *extra_filters):
"""Resolve a JobExecution by ``job_id`` path param and verify project permission.

Internally-submitted jobs (source='system') are never exposed via the API.
Extra ORM clauses are appended to the WHERE clause, e.g. to restrict by job_key.
Mismatches surface as the same 404 — no information leakage.
Only jobs whose ``job_key`` is in ``PUBLIC_JOB_KEYS`` are exposed via the API.
Internal kinds (score rollups, recalculations, monitor runs) are filtered out
by construction. Extra ORM clauses are appended to the WHERE clause to further
restrict by job_key when a caller wants a single kind.
"""
from sqlalchemy import select

from testgen.common.models.job_execution import JobExecution

def dependency(job_id: UUID, user: User = _require_user) -> JobExecution:
query = select(JobExecution).where(
JobExecution.id == job_id,
JobExecution.source != "system",
JobExecution.job_key.in_(PUBLIC_JOB_KEYS),
*extra_filters,
)
job = get_current_session().scalars(query).first()
if job and has_project_permission(user, job.project_code, permission):
return job
raise _not_found
return _check_access(get_current_session().scalars(query).first(), user, permission)
return Depends(dependency)
9 changes: 5 additions & 4 deletions testgen/api/jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,17 @@
resolve_table_group,
resolve_test_suite,
)
from testgen.api.schemas import ErrorResponse, JobKey, JobListResponse, JobResponse, JobSource, JobSubmittedResponse
from testgen.common.models.job_execution import JobExecution, JobStatus
from testgen.api.schemas import ErrorResponse, JobListResponse, JobResponse, JobSubmittedResponse
from testgen.common.enums import JobKey, JobSource, JobStatus
from testgen.common.models.job_execution import PUBLIC_JOB_KEYS, JobExecution
from testgen.common.models.table_group import TableGroup
from testgen.common.models.test_suite import TestSuite

_error_responses = {
404: {"model": ErrorResponse, "description": "Not found"},
}

router = APIRouter(prefix="/api/v1", tags=["Jobs"], dependencies=[Depends(db_session)], responses=_error_responses)
router = APIRouter(tags=["Jobs"], dependencies=[Depends(db_session)], responses=_error_responses)


@router.post(
Expand Down Expand Up @@ -105,7 +106,7 @@ def list_jobs(
"""List job executions for a project, with optional filters and pagination."""
items, total = JobExecution.list_for_project(
project_code,
JobExecution.source != "system",
JobExecution.job_key.in_(PUBLIC_JOB_KEYS),
job_key=job_key,
status=status,
page=page,
Expand Down
2 changes: 1 addition & 1 deletion testgen/api/runs.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
404: {"model": ErrorResponse, "description": "Not found"},
}

router = APIRouter(prefix="/api/v1", tags=["runs"], dependencies=[Depends(db_session)], responses=_error_responses)
router = APIRouter(tags=["runs"], dependencies=[Depends(db_session)], responses=_error_responses)


@router.get(
Expand Down
18 changes: 1 addition & 17 deletions testgen/api/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,27 +6,11 @@

from pydantic import BaseModel, field_validator

from testgen.common.models.job_execution import JobStatus
from testgen.common.enums import JobKey, JobSource, JobStatus

# --- Jobs ---


class JobKey(StrEnum):
run_profile = "run-profile"
run_tests = "run-tests"
run_monitors = "run-monitors"
run_test_generation = "run-test-generation"


class JobSource(StrEnum):
api = "api"
ui = "ui"
scheduler = "scheduler"
mcp = "mcp"
cli = "cli"
backfill = "backfill"


class JobSubmittedResponse(BaseModel):
"""Returned on 202 Accepted after successful job submission."""

Expand Down
1 change: 0 additions & 1 deletion testgen/api/test_definitions.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
}

router = APIRouter(
prefix="/api/v1",
tags=["Test Definitions"],
dependencies=[Depends(db_session)],
responses=_error_responses,
Expand Down
9 changes: 5 additions & 4 deletions testgen/commands/exec_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,10 @@
from uuid import UUID

from testgen.commands.job_registry import JOB_DISPATCH, run_final_callbacks
from testgen.common.enums import JobStatus
from testgen.common.job_context import JobContext, job_context
from testgen.common.models import database_session
from testgen.common.models.job_execution import JobExecution, JobStatus
from testgen.common.models.job_execution import JobExecution
from testgen.utils import get_exception_message

LOG = logging.getLogger("testgen")
Expand All @@ -35,8 +36,8 @@ def exec_job(job_execution_id: UUID) -> None:
LOG.error("Job execution %s not found", job_execution_id)
sys.exit(1)

handler = JOB_DISPATCH.get(job_exec.job_key)
if not handler:
job_config = JOB_DISPATCH.get(job_exec.job_key)
if not job_config:
job_exec.mark_interrupted(f"Unknown job key: {job_exec.job_key}")
return

Expand All @@ -48,7 +49,7 @@ def exec_job(job_execution_id: UUID) -> None:
with database_session():
job_exec = JobExecution.get(job_execution_id)
job_context.set(JobContext(job_id=job_execution_id, source=job_exec.source))
handler(**job_exec.kwargs)
job_config.handler(**job_exec.kwargs)

with database_session():
job_exec = JobExecution.get(job_execution_id)
Expand Down
52 changes: 37 additions & 15 deletions testgen/commands/job_registry.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
"""Wiring between the JobExecution engine and the concrete job handlers.

Two registries keyed by `job_key`:
- `JOB_DISPATCH`: maps a job to its handler (`exec_job` resolves this).
- `JOB_DISPATCH`: maps a job to its `JobConfig` (handler + per-job metadata).
`exec_job` and the scheduler resolve this.
- `JOB_FINAL_CALLBACKS`: maps a job to post-terminal-transition callbacks
(notifications, follow-up job submissions). `run_final_callbacks` iterates.

Expand All @@ -12,16 +13,19 @@

import logging
from collections.abc import Callable
from dataclasses import dataclass

from sqlalchemy import select

from testgen.commands.run_data_cleanup import run_data_cleanup
from testgen.commands.run_profiling import run_profiling
from testgen.commands.run_recalculate_project_scores import run_recalculate_project_scores
from testgen.commands.run_score_update import run_score_update
from testgen.commands.run_test_execution import run_test_execution
from testgen.commands.test_generation import run_test_generation
from testgen.common.enums import JobKey, JobSource, JobStatus
from testgen.common.models import database_session
from testgen.common.models.job_execution import JobExecution, JobStatus
from testgen.common.models.job_execution import JobExecution
from testgen.common.models.profiling_run import ProfilingRun
from testgen.common.models.test_run import TestRun
from testgen.common.notifications.monitor_run import send_monitor_notifications
Expand All @@ -32,13 +36,31 @@

FinalCallback = Callable[[JobExecution], None]

JOB_DISPATCH: dict[str, Callable] = {
"run-profile": run_profiling,
"run-tests": run_test_execution,
"run-monitors": run_test_execution,
"run-test-generation": run_test_generation,
"run-score-update": run_score_update,
"recalculate-project-scores": run_recalculate_project_scores,

@dataclass(frozen=True)
class JobConfig:
"""Per-job-key registration metadata.

`scheduler_source` is the value the scheduler tags `JobExecution.source`
with when it spawns this job key — ``"scheduler"`` for user-facing jobs,
``"system"`` for system-internal jobs (e.g., retention cleanup). Read
only by the scheduler; direct `JobExecution.submit(source=...)` callers
(UI, follow-up enqueues, CLI) set their own source independently and do
not consult this field.
"""

handler: Callable
scheduler_source: JobSource = JobSource.scheduler


JOB_DISPATCH: dict[JobKey, JobConfig] = {
JobKey.run_profile: JobConfig(handler=run_profiling),
JobKey.run_tests: JobConfig(handler=run_test_execution),
JobKey.run_monitors: JobConfig(handler=run_test_execution),
JobKey.run_test_generation: JobConfig(handler=run_test_generation),
JobKey.run_score_update: JobConfig(handler=run_score_update, scheduler_source=JobSource.system),
JobKey.recalculate_project_scores: JobConfig(handler=run_recalculate_project_scores, scheduler_source=JobSource.system),
JobKey.run_data_cleanup: JobConfig(handler=run_data_cleanup, scheduler_source=JobSource.system),
}


Expand Down Expand Up @@ -91,18 +113,18 @@ def _enqueue_score_update(job_exec: JobExecution) -> None:

with database_session():
JobExecution.submit(
job_key="run-score-update",
job_key=JobKey.run_score_update,
kwargs={
"parent_job_id": str(job_exec.id),
"parent_job_key": job_exec.job_key,
},
source="system",
source=JobSource.system,
project_code=job_exec.project_code,
)


JOB_FINAL_CALLBACKS: dict[str, list[FinalCallback]] = {
"run-profile": [_notify_profiling_run, _enqueue_score_update],
"run-tests": [_notify_test_run, _enqueue_score_update],
"run-monitors": [_notify_monitor_run],
JOB_FINAL_CALLBACKS: dict[JobKey, list[FinalCallback]] = {
JobKey.run_profile: [_notify_profiling_run, _enqueue_score_update],
JobKey.run_tests: [_notify_test_run, _enqueue_score_update],
JobKey.run_monitors: [_notify_monitor_run],
}
Loading
Loading