diff --git a/.gitignore b/.gitignore index 62b713e7..f1b1b4ab 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,8 @@ __pycache__/ work_area/ .envrc environment.list +.claude/settings.local.json +.claude/repo.yml # C extensions *.so diff --git a/deploy/testgen-base.dockerfile b/deploy/testgen-base.dockerfile index 2fafe213..9b32f85a 100644 --- a/deploy/testgen-base.dockerfile +++ b/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) diff --git a/deploy/testgen.dockerfile b/deploy/testgen.dockerfile index 800cfa4f..d60654aa 100644 --- a/deploy/testgen.dockerfile +++ b/deploy/testgen.dockerfile @@ -1,4 +1,4 @@ -ARG TESTGEN_BASE_LABEL=v16 +ARG TESTGEN_BASE_LABEL=v17 FROM datakitchen/dataops-testgen-base:${TESTGEN_BASE_LABEL} AS release-image diff --git a/pyproject.toml b/pyproject.toml index 7af55506..b64fa8fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" [project] name = "dataops-testgen" -version = "5.48.0" +version = "5.70.2" description = "DataKitchen's Data Quality DataOps TestGen" authors = [ { "name" = "DataKitchen, Inc.", "email" = "info@datakitchen.io" }, @@ -61,7 +61,7 @@ dependencies = [ "xlsxwriter==3.2.0", "psutil==5.9.8", "concurrent_log_handler==0.9.25", - "cryptography==46.0.6", + "cryptography==48.0.1", "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 @@ -145,7 +145,7 @@ include-package-data = true include = [ "testgen*", ] -exclude = [ "*.tests", "tests*", "deploy*", "invocations*", "testlib*"] +exclude = [ "*.tests", "tests*", "deploy*", "invocations*", "testgen.testing*"] [tool.pytest.ini_options] minversion = "7.0" diff --git a/testgen/api/deps.py b/testgen/api/deps.py index 1807d68d..e59d783d 100644 --- a/testgen/api/deps.py +++ b/testgen/api/deps.py @@ -6,9 +6,10 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from sqlalchemy import select -from testgen.common.auth import authorize_token, decode_jwt_token +from testgen.common.auth import AuthError, authorize_token, decode_jwt_token +from testgen.common.enums import PublicJobKey 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.job_execution import 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 @@ -47,7 +48,7 @@ def get_authorized_user(credentials: HTTPAuthorizationCredentials = _bearer_secu try: payload = decode_jwt_token(credentials.credentials) - except ValueError: + except AuthError: raise _invalid from None username = payload.get("username") @@ -57,7 +58,7 @@ def get_authorized_user(credentials: HTTPAuthorizationCredentials = _bearer_secu session = get_current_session() try: return authorize_token(credentials.credentials, username, session) - except ValueError: + except AuthError: raise _invalid from None @@ -122,7 +123,7 @@ def dependency(test_suite_id: UUID, user: User = _require_user) -> TestSuite: def resolve_job(permission: str, *extra_filters): """Resolve a JobExecution by ``job_id`` path param and verify project permission. - Only jobs whose ``job_key`` is in ``PUBLIC_JOB_KEYS`` are exposed via the API. + Only externally-triggerable jobs (``PublicJobKey``) 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. @@ -130,7 +131,7 @@ def resolve_job(permission: str, *extra_filters): def dependency(job_id: UUID, user: User = _require_user) -> JobExecution: query = select(JobExecution).where( JobExecution.id == job_id, - JobExecution.job_key.in_(PUBLIC_JOB_KEYS), + JobExecution.job_key.in_(list(PublicJobKey)), *extra_filters, ) return _check_access(get_current_session().scalars(query).first(), user, permission) diff --git a/testgen/api/enums.py b/testgen/api/enums.py new file mode 100644 index 00000000..7fc4eef0 --- /dev/null +++ b/testgen/api/enums.py @@ -0,0 +1,52 @@ +"""Lowercase presentation enums for API v1 request filters and responses. + +Shared home for StrEnums that map a DB-stored value to the lowercase snake_case +form exposed through the API. Several DB columns store title-case values +(``test_results.result_status``, ``*.disposition``); the mapping dicts here are the +single seam that normalizes them to the API surface — the DB is never changed. +""" + +from enum import StrEnum + +from testgen.common.enums import Disposition as DbDisposition +from testgen.common.models.test_result import TestResultStatus + + +class ResultStatus(StrEnum): + """Outcome of a single test result.""" + + passed = "passed" + failed = "failed" + warning = "warning" + error = "error" + log = "log" + + +class Disposition(StrEnum): + """Triage state of a test result. ``no_decision`` is the state of a result no one + has triaged yet. Omitting the filter returns active results (``confirmed`` and + ``no_decision``); pass an explicit value to filter to a single state.""" + + confirmed = "confirmed" + dismissed = "dismissed" + muted = "muted" + no_decision = "no_decision" + + +RESULT_STATUS_TO_DB: dict[ResultStatus, TestResultStatus] = { + ResultStatus.passed: TestResultStatus.Passed, + ResultStatus.failed: TestResultStatus.Failed, + ResultStatus.warning: TestResultStatus.Warning, + ResultStatus.error: TestResultStatus.Error, + ResultStatus.log: TestResultStatus.Log, +} +RESULT_STATUS_FROM_DB: dict[TestResultStatus, ResultStatus] = {v: k for k, v in RESULT_STATUS_TO_DB.items()} + +# ``no_decision`` has no stored DB value — it corresponds to a NULL ``disposition`` +# column, handled explicitly at the API boundary (not present in these dicts). +DISPOSITION_TO_DB: dict[Disposition, DbDisposition] = { + Disposition.confirmed: DbDisposition.CONFIRMED, + Disposition.dismissed: DbDisposition.DISMISSED, + Disposition.muted: DbDisposition.INACTIVE, +} +DISPOSITION_FROM_DB: dict[DbDisposition, Disposition] = {v: k for k, v in DISPOSITION_TO_DB.items()} diff --git a/testgen/api/jobs.py b/testgen/api/jobs.py index 3ab291cf..c4f28f2f 100644 --- a/testgen/api/jobs.py +++ b/testgen/api/jobs.py @@ -11,8 +11,8 @@ resolve_test_suite, ) 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.enums import JobKey, JobSource, JobStatus, PublicJobKey +from testgen.common.models.job_execution import JobExecution from testgen.common.models.table_group import TableGroup from testgen.common.models.test_suite import TestSuite @@ -98,7 +98,7 @@ def cancel_job(job: JobExecution = resolve_job("edit")): # noqa: B008 ) def list_jobs( project_code: str = resolve_project_code("view"), - job_key: JobKey | None = Query(default=None), # noqa: B008 + job_key: PublicJobKey | None = Query(default=None), # noqa: B008 status: JobStatus | None = Query(default=None), # noqa: B008 page: int = Query(default=1, ge=1), limit: int = Query(default=20, ge=1, le=100), @@ -106,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.job_key.in_(PUBLIC_JOB_KEYS), + JobExecution.job_key.in_(list(PublicJobKey)), job_key=job_key, status=status, page=page, diff --git a/testgen/api/oauth/metadata.py b/testgen/api/oauth/metadata.py index 55d6fb28..d54d3dfe 100644 --- a/testgen/api/oauth/metadata.py +++ b/testgen/api/oauth/metadata.py @@ -12,7 +12,7 @@ def authorization_server_metadata(): """Return OAuth 2.1 Authorization Server Metadata per RFC 8414. - MCP clients use this for server discovery. + OAuth clients use this for server discovery. """ base_url = settings.BASE_URL.rstrip("/") diff --git a/testgen/api/oauth/models.py b/testgen/api/oauth/models.py deleted file mode 100644 index 49b2c961..00000000 --- a/testgen/api/oauth/models.py +++ /dev/null @@ -1,45 +0,0 @@ -import time - -from authlib.integrations.sqla_oauth2 import ( - OAuth2AuthorizationCodeMixin, - OAuth2ClientMixin, - OAuth2TokenMixin, -) -from sqlalchemy import Column, ForeignKey, String -from sqlalchemy.dialects import postgresql - -from testgen import settings -from testgen.common.models import Base - - -class OAuth2Client(Base, OAuth2ClientMixin): - __tablename__ = "oauth2_clients" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="SET NULL"), nullable=True) - - # Override to widen — JWTs can exceed 255 chars - # (the mixin defines client_id as VARCHAR(48) which is fine) - - -class OAuth2AuthorizationCode(Base, OAuth2AuthorizationCodeMixin): - __tablename__ = "oauth2_authorization_codes" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=False) - - -class OAuth2Token(Base, OAuth2TokenMixin): - __tablename__ = "oauth2_tokens" - - id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") - user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=True) - - # Override to allow longer JWTs as access tokens - access_token = Column(String(2048), unique=True, nullable=False) - - def is_refresh_token_active(self) -> bool: - if self.refresh_token_revoked_at: - return False - expires_at = self.issued_at + settings.REFRESH_TOKEN_EXPIRES_IN - return expires_at >= time.time() diff --git a/testgen/api/oauth/routes.py b/testgen/api/oauth/routes.py index d2ea828b..da228383 100644 --- a/testgen/api/oauth/routes.py +++ b/testgen/api/oauth/routes.py @@ -21,10 +21,10 @@ from testgen import settings from testgen.api.deps import db_session from testgen.api.oauth.login import render_login_page -from testgen.api.oauth.models import OAuth2Client from testgen.api.oauth.server import TestGenAuthorizationServer from testgen.common.auth import create_jwt_token, decode_jwt_token, verify_password from testgen.common.models import get_current_session +from testgen.common.models.oauth import OAuth2Client from testgen.common.models.user import User LOG = logging.getLogger("testgen") diff --git a/testgen/api/oauth/server.py b/testgen/api/oauth/server.py index 32ca120a..0d7859cb 100644 --- a/testgen/api/oauth/server.py +++ b/testgen/api/oauth/server.py @@ -2,7 +2,6 @@ Grant types: - Authorization Code + PKCE (for MCP clients) -- Client Credentials (for automation scripts) - Refresh Token (for token renewal) All DB operations use get_current_session() for thread-local session access. @@ -13,15 +12,14 @@ from typing import ClassVar from authlib.oauth2.rfc6749 import AuthorizationServer, JsonRequest, OAuth2Request, grants -from authlib.oauth2.rfc6749.errors import InvalidGrantError from authlib.oauth2.rfc7009 import RevocationEndpoint from authlib.oauth2.rfc7636 import CodeChallenge from sqlalchemy import select from testgen import settings -from testgen.api.oauth.models import OAuth2AuthorizationCode, OAuth2Client, OAuth2Token from testgen.common.auth import create_jwt_token from testgen.common.models import get_current_session +from testgen.common.models.oauth import OAuth2AuthorizationCode, OAuth2Client, OAuth2Token from testgen.common.models.user import User @@ -85,24 +83,6 @@ def revoke_old_credential(self, credential): credential.access_token_revoked_at = int(time.time()) -class ClientCredentialsGrant(grants.ClientCredentialsGrant): - """Client credentials grant that resolves the client's owner as the token user. - - Ensures every token has a real User identity — no "ghost" usernames. - """ - - def validate_token_request(self): - super().validate_token_request() - client = self.request.client - if not client.user_id: - raise InvalidGrantError(description="Client has no registered owner.") - session = get_current_session() - owner = session.scalars(select(User).where(User.id == client.user_id)).first() - if owner is None: - raise InvalidGrantError(description="Client owner no longer exists.") - self.request.user = owner - - class TestGenRevocationEndpoint(RevocationEndpoint): def query_token(self, token_string, token_type_hint): session = get_current_session() @@ -183,7 +163,6 @@ def create_authorization_server() -> TestGenAuthorizationServer: """Create and configure the authorization server with all grant types.""" server = TestGenAuthorizationServer() server.register_grant(AuthorizationCodeGrant, [CodeChallenge(required=True)]) - server.register_grant(ClientCredentialsGrant) server.register_grant(RefreshTokenGrant) server.register_endpoint(TestGenRevocationEndpoint) server.register_token_generator("default", _generate_bearer_token) diff --git a/testgen/api/runs.py b/testgen/api/runs.py index fcbe0445..277e982b 100644 --- a/testgen/api/runs.py +++ b/testgen/api/runs.py @@ -1,24 +1,35 @@ """API v1 — test run and profiling run retrieval.""" -from fastapi import APIRouter, Depends -from sqlalchemy import select +from fastapi import APIRouter, Depends, Query +from sqlalchemy import or_, select from testgen.api.deps import db_session, resolve_job +from testgen.api.enums import ( + DISPOSITION_FROM_DB, + DISPOSITION_TO_DB, + RESULT_STATUS_FROM_DB, + RESULT_STATUS_TO_DB, + Disposition, + ResultStatus, +) from testgen.api.schemas import ( ErrorResponse, - IssueBreakdown, - JobKey, + IssueCounts, ProfilingRunResponse, ProfilingRunResult, - TestBreakdown, + ResultCounts, + TestResultItem, + TestResultListResponse, TestRunResponse, TestRunResult, ) +from testgen.common.enums import Disposition as DbDisposition +from testgen.common.enums import JobKey from testgen.common.models import get_current_session from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.job_execution import JobExecution from testgen.common.models.profiling_run import ProfilingRun -from testgen.common.models.test_result import TestResult +from testgen.common.models.test_result import TestResult, TestRunResultRow from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite @@ -35,21 +46,14 @@ ) 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: counts = TestResult.count_by_status(test_run.id) result = TestRunResult( score=test_run.dq_score_test_run, - tests=TestBreakdown( - passed=counts.passed, - failed=counts.failed, - warning=counts.warning, - error=counts.error, - log=counts.log, - dismissed=counts.dismissed, - ), + result_counts=ResultCounts.model_validate(counts, from_attributes=True), ) test_suite_id = test_run.test_suite_id if test_run else None @@ -70,28 +74,99 @@ def get_test_run(job: JobExecution = resolve_job("view", JobExecution.job_key == ) +def _disposition_from_db(value: str | None) -> Disposition: + """Map a stored ``disposition`` to the API enum, degrading unknown values to ``no_decision``. + + A NULL or unmapped value resolves to ``no_decision`` rather than raising, so a single odd + row never fails serialization of the whole results page. + """ + if not value: + return Disposition.no_decision + try: + return DISPOSITION_FROM_DB[DbDisposition(value)] + except (ValueError, KeyError): + return Disposition.no_decision + + +def _to_item(row: TestRunResultRow) -> TestResultItem: + """Map a DB-valued result row to the API item, normalizing enum casing.""" + return TestResultItem( + test_definition_id=row.test_definition_id, + test_type=row.test_type, + schema_name=row.schema_name, + table_name=row.table_name, + column_names=row.column_names, + result_status=RESULT_STATUS_FROM_DB.get(row.status), + result_measure=row.result_measure, + threshold_value=row.threshold_value, + result_message=row.message, + test_time=row.test_time, + disposition=_disposition_from_db(row.disposition), + ) + + +@router.get( + "/test-runs/{job_id}/results", + response_model=TestResultListResponse, +) +def list_test_run_results( + job: JobExecution = resolve_job("view", JobExecution.job_key == JobKey.run_tests), # noqa: B008 + status: ResultStatus | None = Query(default=None), # noqa: B008 + table_name: str | None = Query(default=None), + column_name: str | None = Query(default=None), + test_type: str | None = Query(default=None), + disposition: Disposition | None = Query(default=None), # noqa: B008 + page: int = Query(default=1, ge=1), + limit: int = Query(default=20, ge=1, le=100), +): + """List individual results for a test run. + + Omitting ``disposition`` returns active results — confirmed and no_decision + (excludes dismissed and muted); pass an explicit value to filter to one state. + """ + clauses = [] + if status: + clauses.append(TestResult.status == RESULT_STATUS_TO_DB[status]) + if table_name: + clauses.append(TestResult.table_name == table_name) + if column_name: + clauses.append(TestResult.column_names == column_name) + if test_type: + clauses.append(TestResult.test_type == test_type) + if disposition is None: + # Active: confirmed plus no_decision (NULL). Dismissed/muted excluded. + clauses.append(or_(TestResult.disposition.is_(None), TestResult.disposition == DbDisposition.CONFIRMED.value)) + elif disposition == Disposition.no_decision: + clauses.append(TestResult.disposition.is_(None)) + else: + clauses.append(TestResult.disposition == DISPOSITION_TO_DB[disposition].value) + + items, total = TestResult.list_for_run(job.id, *clauses, page=page, limit=limit) + return TestResultListResponse( + items=[_to_item(row) for row in items], + page=page, + limit=limit, + total=total, + ) + + @router.get( "/profiling-runs/{job_id}", response_model=ProfilingRunResponse, ) 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: - counts = HygieneIssue.count_by_likelihood(profiling_run.id) + counts = HygieneIssue.count_for_run(profiling_run.id) result = ProfilingRunResult( score=profiling_run.dq_score_profiling, table_ct=profiling_run.table_ct, column_ct=profiling_run.column_ct, record_ct=profiling_run.record_ct, - issues=IssueBreakdown( - definite=counts.definite, - likely=counts.likely, - possible=counts.possible, - dismissed=counts.dismissed, - ), + issue_counts=IssueCounts.model_validate(counts, from_attributes=True), ) return ProfilingRunResponse( diff --git a/testgen/api/schemas.py b/testgen/api/schemas.py index 2543f227..4bfd0f40 100644 --- a/testgen/api/schemas.py +++ b/testgen/api/schemas.py @@ -1,12 +1,13 @@ """Pydantic request/response models for API v1 endpoints.""" from datetime import datetime -from enum import StrEnum from uuid import UUID -from pydantic import BaseModel, field_validator +from pydantic import BaseModel -from testgen.common.enums import JobKey, JobSource, JobStatus +from testgen.api.enums import Disposition, ResultStatus +from testgen.common.enums import JobSource, JobStatus, PublicJobKey +from testgen.common.test_definition_export_import_service import ImportConfig, ImportPayload, ImportResponse # --- Jobs --- @@ -24,7 +25,7 @@ class JobResponse(BaseModel): """Full job execution record returned by status and cancel endpoints.""" id: UUID - job_key: JobKey + job_key: PublicJobKey status: JobStatus source: JobSource created_at: datetime @@ -48,8 +49,8 @@ class JobListResponse(BaseModel): # --- Test Runs --- -class TestBreakdown(BaseModel): - """Counts of test results by outcome status.""" +class ResultCounts(BaseModel): + """Counts of test results by outcome status, with dismissed results separated.""" passed: int = 0 failed: int = 0 @@ -63,7 +64,7 @@ class TestRunResult(BaseModel): """Run-specific data populated when execution completes.""" score: float | None = None - tests: TestBreakdown + result_counts: ResultCounts class TestRunResponse(BaseModel): @@ -78,15 +79,54 @@ class TestRunResponse(BaseModel): result: TestRunResult | None = None +class TestResultItem(BaseModel): + """One individual test result within a test run.""" + + test_definition_id: UUID + test_type: str + schema_name: str + table_name: str | None = None + column_names: str | None = None + result_status: ResultStatus | None = None + result_measure: str | None = None + threshold_value: str | None = None + result_message: str | None = None + test_time: datetime | None = None + disposition: Disposition + + +class TestResultListResponse(BaseModel): + """Paginated list of individual test results.""" + + items: list[TestResultItem] + page: int + limit: int + total: int + + # --- Profiling Runs --- -class IssueBreakdown(BaseModel): - """Counts of hygiene issues by likelihood category.""" +class HygieneIssueCounts(BaseModel): + """Counts of active data-quality hygiene issues by likelihood category.""" definite: int = 0 likely: int = 0 possible: int = 0 + + +class PotentialPiiCounts(BaseModel): + """Counts of active Potential PII findings by risk level.""" + + high: int = 0 + moderate: int = 0 + + +class IssueCounts(BaseModel): + """Profiling-finding breakdown: active counts by kind, plus a single dismissed total.""" + + hygiene_issues: HygieneIssueCounts + potential_pii: PotentialPiiCounts dismissed: int = 0 @@ -97,7 +137,7 @@ class ProfilingRunResult(BaseModel): table_ct: int | None = None column_ct: int | None = None record_ct: int | None = None - issues: IssueBreakdown + issue_counts: IssueCounts class ProfilingRunResponse(BaseModel): @@ -127,153 +167,7 @@ class ErrorResponse(BaseModel): errors: list[ErrorDetail] -# --- Test Definition Export/Import --- - - -class Origin(StrEnum): - manual = "manual" - auto = "auto" - both = "both" - - -class ImportMode(StrEnum): - preview = "preview" - apply = "apply" - apply_strict = "apply_strict" - - -class OnMatch(StrEnum): - overwrite_all = "overwrite_all" - overwrite_unlocked = "overwrite_unlocked" - skip = "skip" - - -class OnNew(StrEnum): - skip = "skip" - create = "create" - create_and_lock = "create_and_lock" - - -class OnAbsence(StrEnum): - do_nothing = "do_nothing" - delete_all = "delete_all" - delete_unlocked = "delete_unlocked" - - -class ImportAction(StrEnum): - create = "create" - update = "update" - skip = "skip" - delete = "delete" - - -class ImportReason(StrEnum): - matched = "matched" - no_match = "no_match" - policy = "policy" - locked = "locked" - invalid_test_type = "invalid_test_type" - invalid_table = "invalid_table" - missing_external_id = "missing_external_id" - absent = "absent" - - -# Non-None defaults must match the ORM column defaults in TestDefinition: -# test_active=True (YNString default="Y"), lock_refresh=False (YNString default="N"), -# skip_errors=0 (ZeroIfEmptyInteger), window_days=0 (ZeroIfEmptyInteger), -# history_lookback=0 (Column default=0). -# On export, the model_serializer omits fields matching these defaults to keep the file compact. -# On import, model_fields_set distinguishes explicit from defaulted. -class TestDefinitionExport(BaseModel): - """Test definition fields included in the export/import file.""" - - model_config = {"from_attributes": True} - - # Matching / identity - test_type: str - external_id: UUID | None = None - last_auto_gen_date: datetime | None = None - - # Definition fields - table_name: str | None = None - column_name: str | None = None - test_description: str | None = None - test_active: bool = True - severity: str | None = None - lock_refresh: bool = False - export_to_observability: bool | None = None - skip_errors: int = 0 - - # Calibration fields - baseline_ct: str | None = None - baseline_unique_ct: str | None = None - baseline_value: str | None = None - baseline_value_ct: str | None = None - threshold_value: str | None = None - baseline_sum: str | None = None - baseline_avg: str | None = None - baseline_sd: str | None = None - lower_tolerance: str | None = None - upper_tolerance: str | None = None - - # Subset / grouping - subset_condition: str | None = None - groupby_names: str | None = None - having_condition: str | None = None - window_date_column: str | None = None - window_days: int = 0 - - # Referential - match_schema_name: str | None = None - match_table_name: str | None = None - match_column_names: str | None = None - match_subset_condition: str | None = None - match_groupby_names: str | None = None - match_having_condition: str | None = None - - # Query / history - custom_query: str | None = None - history_calculation: str | None = None - history_calculation_upper: str | None = None - history_lookback: int = 0 - - @field_validator("skip_errors", "window_days", "history_lookback", mode="before") - @classmethod - def _coerce_none_to_zero(cls, v: int | None) -> int: - return v if v is not None else 0 - - -class ExportSource(BaseModel): - project_code: str - test_suite: str - table_group: str - table_group_schema: str - exported_at: datetime - testgen_version: str | None = None - - -class ExportDocument(BaseModel): - version: int = 1 - source: ExportSource - definitions: list[TestDefinitionExport] - - -# --- Import --- - - -class ImportConfig(BaseModel): - mode: ImportMode - on_match: OnMatch - on_new: OnNew - on_absence: OnAbsence - - -class ImportPayload(BaseModel): - """Import payload — same structure as an export document, but definitions are typed.""" - - version: int = 1 - source: ExportSource | None = None - definitions: list[TestDefinitionExport] +# --- Test Definition Export/Import (wire types) --- class ImportRequest(BaseModel): @@ -281,29 +175,6 @@ class ImportRequest(BaseModel): payload: ImportPayload -class ImportItemTD(BaseModel): - idx: int | None = None - target_id: UUID | None = None - - -class ImportItem(BaseModel): - action: ImportAction - reason: ImportReason - tds: list[ImportItemTD] - - -class ImportSummary(BaseModel): - created: int = 0 - updated: int = 0 - skipped: int = 0 - deleted: int = 0 - - -class ImportResponse(BaseModel): - summary: ImportSummary - items: list[ImportItem] - - class ImportStrictError(ErrorResponse): """400 response for apply_strict when entries would be skipped.""" diff --git a/testgen/api/test_definitions.py b/testgen/api/test_definitions.py index 6d2d0e41..eadb0b92 100644 --- a/testgen/api/test_definitions.py +++ b/testgen/api/test_definitions.py @@ -2,19 +2,23 @@ from fastapi import APIRouter, Depends, HTTPException, Query -from testgen.api import test_definition_service -from testgen.api.deps import db_session, resolve_test_suite +from testgen.api.deps import api_error, db_session, resolve_test_suite from testgen.api.schemas import ( ErrorDetail, ErrorResponse, - ExportDocument, - ImportMode, ImportRequest, - ImportResponse, ImportStrictError, - Origin, ) from testgen.common.models.test_suite import TestSuite +from testgen.common.test_definition_export_import_service import ( + ExportDocument, + ImportResponse, + ImportStrictViolation, + InvalidImportPayload, + Origin, + export_definitions, + import_definitions, +) _error_responses = { 404: {"model": ErrorResponse, "description": "Not found"}, @@ -39,14 +43,17 @@ def export_test_definitions( test_type: str | None = Query(default=None), ) -> ExportDocument: """Export test definitions from a test suite as a portable JSON document.""" - return test_definition_service.export_definitions(test_suite, origin, table_name, test_type) + return export_definitions(test_suite, origin, table_name, test_type) @router.post( "/test-suites/{test_suite_id}/test-definition-import", response_model=ImportResponse, responses={ - 400: {"model": ImportStrictError, "description": "Invalid request or strict validation failed"}, + 400: { + "model": ImportStrictError | ErrorResponse, + "description": "Strict validation failed (includes the projected import result) or invalid payload", + }, }, ) def import_test_definitions( @@ -54,18 +61,15 @@ def import_test_definitions( test_suite: TestSuite = resolve_test_suite("edit"), # noqa: B008 ) -> ImportResponse: """Import test definitions into a test suite from a portable JSON document.""" - result = test_definition_service.import_definitions(test_suite, body.config, body.payload) - - if body.config.mode == ImportMode.apply_strict and result.summary.skipped > 0: + try: + return import_definitions(test_suite, body.config, body.payload) + except InvalidImportPayload as err: + raise api_error(400, err.code, str(err)) from err + except ImportStrictViolation as err: raise HTTPException( status_code=400, detail=ImportStrictError( - errors=[ErrorDetail( - code="strict_validation_failed", - detail=f"{result.summary.skipped} test definition(s) would be skipped", - )], - import_result=result, + errors=[ErrorDetail(code="strict_validation_failed", detail=str(err))], + import_result=err.result, ).model_dump(mode="json"), - ) - - return result + ) from err diff --git a/testgen/commands/job_registry.py b/testgen/commands/job_registry.py index 3fa3dceb..f41ee4e1 100644 --- a/testgen/commands/job_registry.py +++ b/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 diff --git a/testgen/commands/job_runner.py b/testgen/commands/job_runner.py index f95584e6..5643999e 100644 --- a/testgen/commands/job_runner.py +++ b/testgen/commands/job_runner.py @@ -71,14 +71,14 @@ 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.status == "Error" else "Profiling completed." + 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.status == "Error" else "Test execution completed." + 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 ") case "run-test-generation": click.echo("Test generation completed.") diff --git a/testgen/commands/queries/execute_tests_query.py b/testgen/commands/queries/execute_tests_query.py index 03eab489..1316f6b6 100644 --- a/testgen/commands/queries/execute_tests_query.py +++ b/testgen/commands/queries/execute_tests_query.py @@ -7,7 +7,7 @@ import pandas as pd from testgen.common import read_template_sql_file -from testgen.common.clean_sql import concat_columns +from testgen.common.clean_sql import concat_columns, null_if_empty from testgen.common.database.database_service import ( fetch_dict_from_db, get_flavor_service, @@ -21,7 +21,7 @@ is_excluded_day, resolve_holiday_dates, ) -from testgen.common.models.connection import Connection +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup from testgen.common.models.test_definition import TestRunType, TestScope @@ -348,16 +348,27 @@ def _get_params(self, test_def: TestExecutionDef | None = None) -> dict: "CONCAT_COLUMNS": concat_columns(test_def.column_name, self.null_value) if test_def.column_name else "", "SKIP_ERRORS": test_def.skip_errors or 0, "CUSTOM_QUERY": test_def.custom_query, - "BASELINE_CT": test_def.baseline_ct, - "BASELINE_UNIQUE_CT": test_def.baseline_unique_ct, + # Numeric baseline params feed division/arithmetic in CAT measures. + # Render empty values as NULL so the measure evaluates to NULL instead of producing + # invalid SQL like CAST( AS FLOAT). Cannot default to 0 — several are denominators. + "BASELINE_CT": null_if_empty(test_def.baseline_ct), + "BASELINE_UNIQUE_CT": null_if_empty(test_def.baseline_unique_ct), + # BASELINE_VALUE is a quoted literal, an unquoted number, or an IN-list depending on + # test type — no single empty-substitution is valid SQL, so it is templated as-is. "BASELINE_VALUE": test_def.baseline_value, - "BASELINE_VALUE_CT": test_def.baseline_value_ct, + "BASELINE_VALUE_CT": null_if_empty(test_def.baseline_value_ct), "THRESHOLD_VALUE": test_def.threshold_value or 0, - "BASELINE_SUM": test_def.baseline_sum, - "BASELINE_AVG": test_def.baseline_avg, - "BASELINE_SD": test_def.baseline_sd, - "LOWER_TOLERANCE": "NULL" if test_def.lower_tolerance in (None, "") else test_def.lower_tolerance, - "UPPER_TOLERANCE": "NULL" if test_def.upper_tolerance in (None, "") else test_def.upper_tolerance, + # Freshness_Trend uses BASELINE_SUM as a quoted timestamp (the template's + # NULLIF('{...}', '') handles empty); Other tests use it as a numeric sum. + "BASELINE_SUM": ( + test_def.baseline_sum + if test_def.test_type == "Freshness_Trend" + else null_if_empty(test_def.baseline_sum) + ), + "BASELINE_AVG": null_if_empty(test_def.baseline_avg), + "BASELINE_SD": null_if_empty(test_def.baseline_sd), + "LOWER_TOLERANCE": null_if_empty(test_def.lower_tolerance), + "UPPER_TOLERANCE": null_if_empty(test_def.upper_tolerance), # SUBSET_CONDITION should be replaced after CUSTOM_QUERY # since the latter may contain the former "SUBSET_CONDITION": test_def.subset_condition or "1=1", @@ -534,7 +545,7 @@ def aggregate_cat_tests( null_value=self.null_value, ) - max_query_chars = self.connection.max_query_chars - 400 + max_query_chars = (self.connection.max_query_chars or DEFAULT_MAX_QUERY_CHARS) - 400 groups = group_cat_tests(test_defs, max_query_chars, concat_operator, single) aggregate_queries: list[tuple[str, None]] = [] diff --git a/testgen/commands/queries/refresh_data_chars_query.py b/testgen/commands/queries/refresh_data_chars_query.py index 762e7261..68eb388f 100644 --- a/testgen/commands/queries/refresh_data_chars_query.py +++ b/testgen/commands/queries/refresh_data_chars_query.py @@ -5,7 +5,7 @@ from testgen.common import read_template_sql_file from testgen.common.database.column_chars import ColumnChars from testgen.common.database.database_service import get_flavor_service, replace_params -from testgen.common.models.connection import Connection +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection from testgen.common.models.table_group import TableGroup from testgen.utils import chunk_queries, to_sql_timestamp @@ -30,6 +30,7 @@ class RefreshDataCharsSQL: "general_type", "column_type", "db_data_type", + "object_type", "approx_record_ct", "record_ct", ) @@ -133,7 +134,8 @@ def get_row_counts(self, table_names: Iterable[str]) -> list[tuple[str, None]]: f"SELECT '{table}' AS table_name, COUNT(*) AS row_count FROM {self.flavor_service.get_table_ref(schema, table)}" for table in table_names ] - chunked_queries = chunk_queries(count_queries, " UNION ALL ", self.connection.max_query_chars) + max_query_chars = self.connection.max_query_chars or DEFAULT_MAX_QUERY_CHARS + chunked_queries = chunk_queries(count_queries, " UNION ALL ", max_query_chars) return [ (query, None) for query in chunked_queries ] def verify_access(self, table_name: str) -> tuple[str, None]: @@ -156,6 +158,7 @@ def get_staging_data_chars(self, data_chars: list[ColumnChars], run_date: dateti column.general_type, column.column_type, column.db_data_type, + column.object_type, column.approx_record_ct, column.record_ct, ] diff --git a/testgen/commands/run_data_cleanup.py b/testgen/commands/run_data_cleanup.py index 08e6341c..0dbb0799 100644 --- a/testgen/commands/run_data_cleanup.py +++ b/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", diff --git a/testgen/commands/run_profiling.py b/testgen/commands/run_profiling.py index ffd1a58f..0748f57b 100644 --- a/testgen/commands/run_profiling.py +++ b/testgen/commands/run_profiling.py @@ -1,5 +1,4 @@ import logging -import os from datetime import UTC, datetime, timedelta from uuid import UUID @@ -19,7 +18,7 @@ write_to_app_db, ) from testgen.common.database.column_chars import ColumnChars -from testgen.common.database.database_service import ThreadedProgress +from testgen.common.database.database_service import ThreadedProgress, get_flavor_service from testgen.common.job_context import job_context from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import get_current_session, with_database_session @@ -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). @@ -86,8 +83,9 @@ def run_profiling( if data_chars: sql_generator = ProfilingSQL(connection, table_group, profiling_run) - _run_column_profiling(sql_generator, data_chars) - _run_frequency_analysis(sql_generator) + sampling_params = _compute_sampling_params(sql_generator, data_chars) + _run_column_profiling(sql_generator, data_chars, sampling_params) + _run_frequency_analysis(sql_generator, sampling_params) _run_hygiene_issue_detection(sql_generator) # if table_group.profile_do_pair_rules == "Y": @@ -95,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() @@ -121,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 @@ -143,26 +134,40 @@ def _exclude_xde_columns(data_chars: list[ColumnChars], table_group_id: UUID) -> return filtered -def _run_column_profiling(sql_generator: ProfilingSQL, data_chars: list[ColumnChars]) -> None: +def _compute_sampling_params( + sql_generator: ProfilingSQL, data_chars: list[ColumnChars] +) -> dict[str, TableSampling]: + table_group = sql_generator.table_group + sampling_params: dict[str, TableSampling] = {} + if not table_group.profile_use_sampling: + return sampling_params + + sampleable_types = get_flavor_service(sql_generator.flavor).sampleable_object_types + for column in data_chars: + if sampling_params.get(column.table_name): + continue + if sampleable_types is not None and column.object_type not in sampleable_types: + continue + result = calculate_sampling_params( + table_name=column.table_name, + record_count=column.record_ct, + sample_percent_raw=table_group.profile_sample_percent, + min_sample=table_group.profile_sample_min_count, + ) + if result: + sampling_params[column.table_name] = result + return sampling_params + + +def _run_column_profiling( + sql_generator: ProfilingSQL, data_chars: list[ColumnChars], sampling_params: dict[str, TableSampling] +) -> None: profiling_run = sql_generator.profiling_run profiling_run.set_progress("col_profiling", "Running") profiling_run.save() get_current_session().commit() LOG.info(f"Running column profiling queries: {len(data_chars)}") - table_group = sql_generator.table_group - sampling_params: dict[str, TableSampling] = {} - if table_group.profile_use_sampling: - for column in data_chars: - if not sampling_params.get(column.table_name): - result = calculate_sampling_params( - table_name=column.table_name, - record_count=column.record_ct, - sample_percent_raw=table_group.profile_sample_percent, - min_sample=table_group.profile_sample_min_count, - ) - if result: - sampling_params[column.table_name] = result def update_column_progress(progress: ThreadedProgress) -> None: profiling_run.set_progress( @@ -218,7 +223,7 @@ def update_column_progress(progress: ThreadedProgress) -> None: ) -def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None: +def _run_frequency_analysis(sql_generator: ProfilingSQL, sampling_params: dict[str, TableSampling]) -> None: profiling_run = sql_generator.profiling_run profiling_run.set_progress("freq_analysis", "Running") profiling_run.save() @@ -227,7 +232,7 @@ def _run_frequency_analysis(sql_generator: ProfilingSQL) -> None: error_data = None try: LOG.info("Selecting columns for frequency analysis") - frequency_columns = fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns()) + frequency_columns = [ColumnChars(**column) for column in fetch_dict_from_db(*sql_generator.get_frequency_analysis_columns())] if frequency_columns: LOG.info(f"Running frequency analysis queries: {len(frequency_columns)}") @@ -240,7 +245,10 @@ def update_frequency_progress(progress: ThreadedProgress) -> None: get_current_session().commit() frequency_results, result_columns, error_data = fetch_from_db_threaded( - [sql_generator.run_frequency_analysis(ColumnChars(**column)) for column in frequency_columns], + [ + sql_generator.run_frequency_analysis(column, sampling_params.get(column.table_name)) + for column in frequency_columns + ], use_target_db=True, max_threads=sql_generator.connection.max_threads, progress_callback=update_frequency_progress, diff --git a/testgen/commands/run_score_update.py b/testgen/commands/run_score_update.py index 230a4a9d..caa9125e 100644 --- a/testgen/commands/run_score_update.py +++ b/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) diff --git a/testgen/commands/run_test_execution.py b/testgen/commands/run_test_execution.py index 65d759df..381dc2f0 100644 --- a/testgen/commands/run_test_execution.py +++ b/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(), ) diff --git a/testgen/commands/test_thresholds_prediction.py b/testgen/commands/test_thresholds_prediction.py index 7c501e98..9089551e 100644 --- a/testgen/commands/test_thresholds_prediction.py +++ b/testgen/commands/test_thresholds_prediction.py @@ -89,7 +89,7 @@ def run(self) -> None: # Freshness update events are fetched as secondary data only when the suite # is a monitor — Volume_Trend / Metric_Trend in monitor suites couple to the - # Freshness_Trend signal to avoid stairstep false positives. + # Freshness_Trend signal to avoid stairstep false positives. freshness_updates_by_table: dict[tuple[str, str], list[str]] = ( self._fetch_freshness_updates_by_table() if self.test_suite.is_monitor else {} ) @@ -268,10 +268,16 @@ def compute_freshness_threshold( window_end=schedule.window_end if has_window else None, ) lower, upper = result.lower, result.upper - staleness = result.staleness except NotEnoughData: pass # Keep first-pass thresholds + # An active schedule certifies the cadence is regular, so the tight staleness + # threshold (median * factor) applies. Full-week schedules skip the recompute + # above and use the first-pass result — with no excluded days, those gaps are + # already in business minutes. Non-active stages keep staleness = None so the + # prediction SQL falls back to upper_tolerance (avoids false weekend anomalies). + staleness = result.staleness + # Override upper threshold with schedule-based deadline (daily/weekly only) if schedule.frequency != "sub_daily": holiday_dates = resolve_holiday_dates(holiday_codes, history.index) if holiday_codes else None @@ -293,9 +299,12 @@ def compute_sarimax_threshold( exclude_weekends: bool = False, holiday_codes: list[str] | None = None, schedule_tz: str | None = None, + event_space: bool = False, ) -> tuple[float | None, float | None, str | None]: """Compute SARIMAX-based thresholds for the next forecast point. + `event_space=True` fits in event-space (one step per observation, no interpolation). + Returns (lower, upper, forecast_json) or (None, None, None) if insufficient data. """ try: @@ -305,6 +314,7 @@ def compute_sarimax_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, tz=schedule_tz, + event_space=event_space, ) num_points = len(history) @@ -344,7 +354,7 @@ def compute_volume_or_metric_threshold( This avoids the "stairstep" false-positive shape where inter-change plateaus collapse the SE estimate. The returned prediction JSON is augmented with `freshness_gated` and `baseline_value` so that test execution can apply dual-branch evaluation. - + If the filtered fit fails for any reason, falls back to fit SARIMAX on the raw value series and emits a prediction JSON without the freshness-gating markers. @@ -353,6 +363,9 @@ def compute_volume_or_metric_threshold( Freshness_Trend detected a fingerprint change. """ filtered_history = history.loc[history["test_run_id"].astype(str).isin(freshness_updates)] + # Event-space fit: the filtered series has one point per refresh and is irregularly spaced. + # Calendar resampling + interpolation would smooth the refresh jumps into small uniform + # increments and collapse the jump variance, biasing the forecast low and the band too tight. lower, upper, prediction = compute_sarimax_threshold( filtered_history, sensitivity=sensitivity, @@ -360,6 +373,7 @@ def compute_volume_or_metric_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, schedule_tz=schedule_tz, + event_space=True, ) if prediction is not None: # Pull the baseline value from the most-recent filtered row. @@ -381,5 +395,5 @@ def compute_volume_or_metric_threshold( exclude_weekends=exclude_weekends, holiday_codes=holiday_codes, schedule_tz=schedule_tz, - ) + ) return lower, upper, None, prediction diff --git a/testgen/common/auth.py b/testgen/common/auth.py index 1dcdc479..1df03684 100644 --- a/testgen/common/auth.py +++ b/testgen/common/auth.py @@ -4,12 +4,19 @@ import bcrypt import jwt +from sqlalchemy import func, select from testgen import settings +from testgen.common.models.oauth import OAuth2Token +from testgen.common.models.user import User LOG = logging.getLogger("testgen") +class AuthError(Exception): + """Token authentication failed — invalid/expired token, unknown user, or revoked token.""" + + def get_jwt_signing_key() -> bytes: """Decode the base64-encoded JWT signing key from settings.""" return base64.b64decode(settings.JWT_HASHING_KEY_B64.encode("ascii")) @@ -27,13 +34,13 @@ def create_jwt_token(username: str, expiry_seconds: int = 86400) -> str: def decode_jwt_token(token_str: str) -> dict: """Decode and validate a JWT token. Returns the payload dict. - Raises ValueError if the token is invalid or expired. + Raises ``AuthError`` if the token is invalid or expired. PyJWT auto-validates the standard ``exp`` claim during decode. """ try: return jwt.decode(token_str, get_jwt_signing_key(), algorithms=["HS256"]) except jwt.InvalidTokenError as e: - raise ValueError(f"Invalid token: {e}") from e + raise AuthError(f"Invalid token: {e}") from e def authorize_token(token_str: str, username: str, session): @@ -41,18 +48,13 @@ def authorize_token(token_str: str, username: str, session): Shared implementation for API and MCP authorization. """ - from sqlalchemy import func, select - - from testgen.api.oauth.models import OAuth2Token - from testgen.common.models.user import User - user = session.scalars(select(User).where(func.lower(User.username) == func.lower(username))).first() if user is None: - raise ValueError("User not found") + raise AuthError("User not found") token_record = session.scalars(select(OAuth2Token).where(OAuth2Token.access_token == token_str)).first() if token_record and token_record.access_token_revoked_at: - raise ValueError("Token has been revoked") + raise AuthError("Token has been revoked") return user diff --git a/testgen/common/clean_sql.py b/testgen/common/clean_sql.py index 8d1856e3..397fc9fc 100644 --- a/testgen/common/clean_sql.py +++ b/testgen/common/clean_sql.py @@ -50,6 +50,19 @@ def quote_identifiers(identifiers: str, flavor: str) -> str: return ", ".join(quoted_values) +def null_if_empty(value: object) -> object: + """Return the literal ``"NULL"`` when ``value`` is empty (``None`` or ``""``), else ``value``. + + For numeric test parameters substituted into SQL templates (baseline counts, + averages, standard deviations, tolerances). An empty substitution produces invalid SQL + such as ``CAST( AS FLOAT)``; ``"NULL"`` makes the surrounding expression evaluate to NULL + instead. A real ``0`` is preserved (``0 in (None, "")`` is ``False``). Not suitable for + params used as quoted literals or IN-lists (e.g. ``BASELINE_VALUE``, ``BASELINE_SUM`` for + Freshness_Trend), where ``"NULL"`` would not be valid SQL. + """ + return "NULL" if value in (None, "") else value + + def concat_columns(columns: str, null_value: str): # Prepares SQL expression to concatenate comma-separated column list expression = "" diff --git a/testgen/common/cron_service.py b/testgen/common/cron_service.py index afe75485..fd9a2984 100644 --- a/testgen/common/cron_service.py +++ b/testgen/common/cron_service.py @@ -5,6 +5,8 @@ import cron_converter import cron_descriptor +from testgen.common import date_service + class CronSample(TypedDict, total=False): id: str | None @@ -26,7 +28,7 @@ def get_cron_sample( cron_schedule = cron_obj.schedule(reference_time or datetime.now(zoneinfo.ZoneInfo(cron_tz))) readable_cron_schedule = cron_descriptor.get_description(cron_expr) if formatted: - samples = [cron_schedule.next().strftime("%a %b %-d, %-I:%M %p") for _ in range(sample_count)] + samples = [date_service.format_friendly_datetime(cron_schedule.next(), "%a %b %-d, %-I:%M %p") for _ in range(sample_count)] else: samples = [int(cron_schedule.next().timestamp()) for _ in range(sample_count)] except zoneinfo.ZoneInfoNotFoundError: diff --git a/testgen/common/data_catalog_service.py b/testgen/common/data_catalog_service.py new file mode 100644 index 00000000..ba385411 --- /dev/null +++ b/testgen/common/data_catalog_service.py @@ -0,0 +1,191 @@ +"""Shared Data Catalog service. + +Generates CREATE TABLE scripts from profiled column metadata, fetches sample rows +from source tables, and reads/writes table/column catalog metadata. Used by both +the Streamlit UI and the MCP tools so they share one code path. +""" +import logging +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Literal +from uuid import UUID + +import pandas as pd + +from testgen.common.database.database_service import get_flavor_service +from testgen.common.database.flavor.flavor_service import FlavorService +from testgen.common.models.connection import Connection +from testgen.common.models.data_column import CreateScriptColumn, DataColumnChars +from testgen.common.models.data_table import DataTable +from testgen.common.models.table_group import TableGroup +from testgen.common.pii_masking import get_pii_columns, mask_source_data_pii +from testgen.ui.services.database_service import fetch_from_target_db +from testgen.utils import to_dataframe + +LOG = logging.getLogger("testgen") + +DESCRIPTION_MAX_LENGTH = 1000 +TAG_MAX_LENGTH = 40 + +TAG_FIELDS = [ + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", + "data_classification", +] + +# Metadata fields settable per target type, keyed by their data_*_chars column names. +_TABLE_FIELDS = ("description", *TAG_FIELDS, "critical_data_element") +_COLUMN_FIELDS = ("description", *TAG_FIELDS, "critical_data_element", "excluded_data_element", "pii_flag") + + +@dataclass +class TableSampleResult: + status: Literal["OK", "ND", "ERR"] + message: str | None = None + df: pd.DataFrame | None = None + pii_redacted: bool = False + + +def render_create_table_script( + schema_name: str, + table_name: str, + columns: list[CreateScriptColumn], + flavor_service: FlavorService, + *, + annotate_changes: bool = False, +) -> str: + """Render CREATE TABLE DDL from profiled columns, quoting identifiers for the flavor. + + Column types use the profiling-derived suggestion, falling back to the original + database type. ``annotate_changes`` appends ``-- WAS `` comments where the + suggestion differs from the original type. + """ + quote = flavor_service.quote_character + table_ref = flavor_service.get_table_ref(schema_name, table_name) + quoted_names = [f"{quote}{col.column_name}{quote}" for col in columns] + + name_width = max(len(name) for name in quoted_names) + type_width = max(len(col.datatype_suggestion or col.db_data_type or "") for col in columns) + + col_defs = [] + for index, (col, name) in enumerate(zip(columns, quoted_names, strict=True)): + col_type = col.datatype_suggestion or col.db_data_type or "" + separator = "" if index == len(columns) - 1 else "," + line = f"{name:<{name_width}} {col_type:<{type_width}}{separator}" + if ( + annotate_changes + and col.db_data_type + and col.datatype_suggestion + and col.db_data_type.lower() != col.datatype_suggestion.lower() + ): + line = f"{line} -- WAS {col.db_data_type}" + col_defs.append(line.rstrip()) + + body = "\n ".join(col_defs) + return f"CREATE TABLE {table_ref} (\n {body}\n);" + + +def build_create_table_script( + table_group_id: UUID, table_name: str, *, annotate_changes: bool = False, +) -> str | None: + """Build a CREATE TABLE script for a profiled table, or ``None`` if it is not in the catalog.""" + schema_name, columns = DataColumnChars.list_for_create_script(table_group_id, table_name) + if not columns or schema_name is None: + return None + connection = Connection.get_by_table_group(table_group_id) + if connection is None: + return None + flavor_service = get_flavor_service(connection.sql_flavor) + return render_create_table_script( + schema_name, table_name, columns, flavor_service, annotate_changes=annotate_changes, + ) + + +def fetch_table_sample( + connection: Connection, + table_group_id: UUID, + schema_name: str, + table_name: str, + *, + limit: int, + mask_pii: bool, + column_name: str | None = None, +) -> TableSampleResult: + """Fetch distinct sample rows from a source table, masking PII columns when requested.""" + flavor_service = get_flavor_service(connection.sql_flavor) + prefix, suffix = flavor_service.row_limit_clauses(limit) + quote = flavor_service.quote_character + table_ref = flavor_service.get_table_ref(schema_name, table_name) + columns_expr = f"{quote}{column_name}{quote}" if column_name else "*" + query = f"SELECT DISTINCT {prefix} {columns_expr} FROM {table_ref} {suffix}".strip() + + try: + results = fetch_from_target_db(connection, query) + except Exception: + LOG.exception("Table sample fetch encountered an error.") + return TableSampleResult("ERR", message="The sample data could not be loaded.") + + if not results: + return TableSampleResult("ND") + + df = to_dataframe(results) + pii_redacted = False + if mask_pii: + pii_columns = get_pii_columns(str(table_group_id), schema_name, table_name) + pii_redacted = mask_source_data_pii(df, pii_columns) + return TableSampleResult("OK", df=df, pii_redacted=pii_redacted) + + +def validate_metadata_fields(fields: Mapping[str, Any]) -> list[str]: + """Validate free-text field lengths, returning one message per violation. + + Only checks length: ``description`` and the tag fields. ``None`` values clear the + field and are always valid; absent keys are skipped. + """ + errors: list[str] = [] + + description = fields.get("description") + if description is not None and len(description) > DESCRIPTION_MAX_LENGTH: + errors.append(f"description must be {DESCRIPTION_MAX_LENGTH} characters or fewer.") + + for tag in TAG_FIELDS: + value = fields.get(tag) + if value is not None and len(value) > TAG_MAX_LENGTH: + errors.append(f"{tag} must be {TAG_MAX_LENGTH} characters or fewer.") + + return errors + + +def apply_table_metadata(table: DataTable, fields: Mapping[str, Any]) -> None: + """Set table metadata attributes for every present key (None clears, value sets).""" + for field in _TABLE_FIELDS: + if field in fields: + setattr(table, field, fields[field]) + + +def apply_column_metadata(column: DataColumnChars, fields: Mapping[str, Any]) -> None: + """Set column metadata attributes for every present key (None clears, value sets).""" + for field in _COLUMN_FIELDS: + if field in fields: + setattr(column, field, fields[field]) + + +def disable_autoflags(table_group: TableGroup, *, wrote_cde: bool, wrote_pii: bool) -> list[str]: + """Turn off auto-detect flags so the next profiling pass preserves manual marks. + + Returns the names of the flags that were turned off (skips flags already off). + """ + disabled: list[str] = [] + if wrote_cde and table_group.profile_flag_cdes: + table_group.profile_flag_cdes = False + disabled.append("profile_flag_cdes") + if wrote_pii and table_group.profile_flag_pii: + table_group.profile_flag_pii = False + disabled.append("profile_flag_pii") + return disabled diff --git a/testgen/common/database/column_chars.py b/testgen/common/database/column_chars.py index 6faa08f7..b9c79d08 100644 --- a/testgen/common/database/column_chars.py +++ b/testgen/common/database/column_chars.py @@ -1,4 +1,13 @@ import dataclasses +from enum import StrEnum + + +class ObjectType(StrEnum): + TABLE = "TABLE" + VIEW = "VIEW" + MATERIALIZED_VIEW = "MATERIALIZED_VIEW" + EXTERNAL = "EXTERNAL" + OTHER = "OTHER" @dataclasses.dataclass @@ -11,6 +20,7 @@ class ColumnChars: column_type: str | None = None db_data_type: str | None = None is_decimal: bool = False + object_type: ObjectType | None = None approx_record_ct: int | None = None # This should not default to 0 since we don't always retrieve actual row counts # UI relies on the null value to know that the approx_record_ct should be displayed instead diff --git a/testgen/common/database/connection_service.py b/testgen/common/database/connection_service.py new file mode 100644 index 00000000..879f6208 --- /dev/null +++ b/testgen/common/database/connection_service.py @@ -0,0 +1,145 @@ +"""Shared connection domain helpers — moved out of the Streamlit connections page +so MCP and any future caller can reuse identical test/normalize logic without +forking SQL templates or per-flavor rules. + +Two pieces live here, both flavor-aware and both operating on a ``Connection`` by +column / flag (no user-facing labels — those are presentation): + +* ``ConnectionStatus`` + ``test_connection_status`` — open an external connection + and report success / failure with driver-text details. +* ``normalize_auth_fields`` — pre-save scrub of mutually-exclusive credential + columns so flipping auth modes doesn't leave orphan values behind. +* ``apply_connection_defaults`` — pre-save fill of flavor-dependent defaults the + caller didn't supply. + +The MCP-facing connection-parameter contract (field labels, modes, requirement +rules, validation) is presentation/input vocabulary and lives in +``mcp/tools/common.py``; flavor display labels live in ``common/flavors.py``. +""" + +from __future__ import annotations + +import logging +import random +from dataclasses import dataclass, field + +from sqlalchemy.exc import DatabaseError, DBAPIError + +from testgen.common.database.database_service import empty_cache, get_flavor_service +from testgen.common.models.connection import DEFAULT_MAX_QUERY_CHARS, Connection +from testgen.ui.services.database_service import fetch_from_target_db + +try: + from pyodbc import Error as PyODBCError +except ImportError: + PyODBCError = None + +LOG = logging.getLogger("testgen") + + +# --------------------------------------------------------------------------- +# ConnectionStatus + test runner +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class ConnectionStatus: + message: str + successful: bool + details: str | None = field(default=None) + # Streamlit's reactive system suppresses re-renders when a frozen dataclass + # compares equal to its previous value. A random per-instance field forces + # ``__eq__`` to differ so repeated "Test Connection" clicks producing the + # same error message still re-render the alert. This field has no consumer + # outside the UI re-render path; do NOT delete it as cleanup. + _: float = field(default_factory=random.random) + + +def is_open_ssl_error(error: Exception) -> bool: + return ( + bool(error.args) + and len(error.args) > 1 + and isinstance(error.args[1], list) + and len(error.args[1]) > 0 + and type(error.args[1][0]).__name__ == "OpenSSLError" + ) + + +def test_connection_status(connection: Connection) -> ConnectionStatus: + """Open the connection, run the flavor's smoke query, classify the outcome.""" + # Drop pooled engines so the next checkout reflects any credential changes. + empty_cache() + try: + flavor_service = get_flavor_service(connection.sql_flavor) + results = fetch_from_target_db(connection, flavor_service.test_query) + connection_successful = len(results) == 1 and next(iter(results[0].values())) == 1 + + if not connection_successful: + return ConnectionStatus(message="Error completing a query to the database server.", successful=False) + return ConnectionStatus(message="The connection was successful.", successful=True) + except KeyError: + return ConnectionStatus( + message="Error attempting the connection. ", + details="Complete all the required fields.", + successful=False, + ) + except DatabaseError as error: + LOG.exception("Error testing database connection") + return ConnectionStatus(message="Error attempting the connection.", details=str(error.orig), successful=False) + except DBAPIError as error: + LOG.exception("Error testing database connection") + details = str(error.orig) + if PyODBCError and isinstance(error.orig, PyODBCError) and error.orig.args: + details = error.orig.args[1] + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + except (TypeError, ValueError) as error: + LOG.exception("Error testing database connection") + details = str(error) + if is_open_ssl_error(error): + details = error.args[0] + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + except Exception: + details = "Something went wrong while testing the connection." + if connection.connect_by_key and not connection.private_key: + details = "The private key is missing." + LOG.exception("Error testing database connection") + return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + + +# --------------------------------------------------------------------------- +# Auth-path normalization +# --------------------------------------------------------------------------- + + +def normalize_auth_fields(connection: Connection) -> None: + """Clear credential columns the active auth mode doesn't use. + + Mirrors the UI's pre-save scrub at ``ui/views/connections.py`` so flipping + auth modes doesn't leave stale opposite-mode values behind. Databricks OAuth + stores the client secret in ``project_pw_encrypted`` despite + ``connect_by_key=True``, so it sticks with the password path. + """ + code = connection.sql_flavor_code + + uses_private_key = bool(connection.connect_by_key) and code != "databricks" + if uses_private_key: + connection.project_pw_encrypted = None + else: + connection.private_key = None + connection.private_key_passphrase = None + + if connection.connect_with_identity: + connection.project_user = None + connection.project_pw_encrypted = None + + +# --------------------------------------------------------------------------- +# Flavor-dependent defaults +# --------------------------------------------------------------------------- + + +def apply_connection_defaults(connection: Connection) -> None: + """Fill flavor-dependent defaults for fields the caller didn't supply.""" + if connection.max_query_chars is None: + # Salesforce Data 360's Hyper engine has a lower expression-depth limit + connection.max_query_chars = 15000 if connection.sql_flavor_code == "salesforce_data360" else DEFAULT_MAX_QUERY_CHARS diff --git a/testgen/common/database/flavor/flavor_service.py b/testgen/common/database/flavor/flavor_service.py index becb8b38..7e09406b 100644 --- a/testgen/common/database/flavor/flavor_service.py +++ b/testgen/common/database/flavor/flavor_service.py @@ -6,7 +6,7 @@ from sqlalchemy import create_engine as sqlalchemy_create_engine from sqlalchemy.engine.base import Engine -from testgen.common.database.column_chars import ColumnChars +from testgen.common.database.column_chars import ColumnChars, ObjectType from testgen.common.encrypt import DecryptText SQLFlavor = Literal["redshift", "redshift_spectrum", "snowflake", "mssql", "postgresql", "databricks", "bigquery", "oracle", "sap_hana", "salesforce_data360"] @@ -116,6 +116,7 @@ def row_limit_clauses(self, n: int) -> tuple[str, str]: qualifies_table_refs_with_schema = True metadata_via_api = False + sampleable_object_types: frozenset[ObjectType] | None = None def get_schema_columns(self, _params: ResolvedConnectionParams, _schema: str) -> list[ColumnChars] | None: """Return column metadata without querying information_schema. diff --git a/testgen/common/database/flavor/mssql_flavor_service.py b/testgen/common/database/flavor/mssql_flavor_service.py index 7f248e43..356cba09 100644 --- a/testgen/common/database/flavor/mssql_flavor_service.py +++ b/testgen/common/database/flavor/mssql_flavor_service.py @@ -1,6 +1,7 @@ from sqlalchemy.engine import URL from testgen import settings +from testgen.common.database.column_chars import ObjectType from testgen.common.database.flavor.flavor_service import FlavorService, ResolvedConnectionParams @@ -10,6 +11,8 @@ class MssqlFlavorService(FlavorService): escaped_underscore = "[_]" row_limiting_clause = "top" url_scheme = "mssql+pyodbc" + # TABLESAMPLE is rejected on views; SQL Server has no materialized views, so only base tables. + sampleable_object_types = frozenset({ObjectType.TABLE}) def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str: connection_url = URL.create( diff --git a/testgen/common/database/flavor/postgresql_flavor_service.py b/testgen/common/database/flavor/postgresql_flavor_service.py index 011ab05a..5f187ee0 100644 --- a/testgen/common/database/flavor/postgresql_flavor_service.py +++ b/testgen/common/database/flavor/postgresql_flavor_service.py @@ -1,5 +1,6 @@ from urllib.parse import quote_plus +from testgen.common.database.column_chars import ObjectType from testgen.common.database.flavor.flavor_service import ResolvedConnectionParams from testgen.common.database.flavor.redshift_flavor_service import RedshiftFlavorService @@ -8,6 +9,8 @@ class PostgresqlFlavorService(RedshiftFlavorService): escaped_underscore = "\\_" url_scheme = "postgresql" + # TABLESAMPLE applies only to tables and materialized views + sampleable_object_types = frozenset({ObjectType.TABLE, ObjectType.MATERIALIZED_VIEW}) def get_connection_string_from_fields(self, params: ResolvedConnectionParams) -> str: if params.host.startswith("/"): diff --git a/testgen/common/database/table_group_service.py b/testgen/common/database/table_group_service.py new file mode 100644 index 00000000..f4c09882 --- /dev/null +++ b/testgen/common/database/table_group_service.py @@ -0,0 +1,260 @@ +"""Shared table-group helpers — common-layer logic shared by MCP, the UI, and +any future caller. Two pieces live here, both flavor-aware: + +* ``validate_table_group_fields`` — required-field + format checks. Callers + surface the returned bullets verbatim. +* ``preview_table_group`` — schema-introspection preview. Returns + ``(preview, data_chars, sql_generator)``; the ``save_data_chars`` callback + is built by callers via ``make_save_data_chars``. + +None of these belong on the ``TableGroup`` model: preview opens an external +DB connection (I/O against an arbitrary target system), and the field-required +rules are form-shaped per-flavor logic, not model invariants. +""" + +from __future__ import annotations + +from collections.abc import Callable +from datetime import UTC, datetime +from typing import TypedDict +from uuid import UUID + +from testgen.commands.queries.refresh_data_chars_query import RefreshDataCharsSQL +from testgen.commands.run_refresh_data_chars import write_data_chars +from testgen.common.database.column_chars import ColumnChars +from testgen.common.database.flavor.flavor_service import resolve_connection_params +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup +from testgen.ui.services.database_service import fetch_from_target_db + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +_NAME_MIN = 3 +_NAME_MAX = 40 +_SAMPLE_PCT_MIN = 1 +_SAMPLE_PCT_MAX = 100 + + +def _missing(value: object) -> bool: + return value is None or (isinstance(value, str) and not value.strip()) + + +def _coerce_int(value: object) -> int | None: + """Return ``int(value)`` for ints / numeric strings, ``None`` otherwise. + + The model stores ``profiling_delay_days`` and ``profile_sample_percent`` as + ``String`` columns; callers may pass either ints or numeric strings. Return + ``None`` to signal "not parseable" so the validator can emit the right error. + """ + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, str): + try: + return int(value.strip()) + except (ValueError, AttributeError): + return None + return None + + +def validate_table_group_fields(table_group: TableGroup) -> list[str]: + """Return every validation error (empty list = valid). + + Mirrors the per-field form validators on the UI ``Add Table Group`` wizard. + The MCP tools call this and raise an ``MCPUserError`` containing the bullets + — they never duplicate a rule themselves. + """ + errors: list[str] = [] + + name = table_group.table_groups_name + if _missing(name): + errors.append("`table_group_name` is required.") + elif not (_NAME_MIN <= len(name.strip()) <= _NAME_MAX): + errors.append(f"`table_group_name` must be between {_NAME_MIN} and {_NAME_MAX} characters.") + + if _missing(table_group.table_group_schema): + errors.append("`schema` is required.") + + delay = _coerce_int(table_group.profiling_delay_days) + if delay is None or delay < 0: + errors.append("`profiling_delay_days` must be a non-negative integer.") + + pct = _coerce_int(table_group.profile_sample_percent) + if pct is None or not (_SAMPLE_PCT_MIN <= pct <= _SAMPLE_PCT_MAX): + errors.append(f"`profile_sample_percent` must be between {_SAMPLE_PCT_MIN} and {_SAMPLE_PCT_MAX}.") + + min_count = table_group.profile_sample_min_count + if not isinstance(min_count, int) or isinstance(min_count, bool) or min_count < 0: + errors.append("`profile_sample_min_count` must be a non-negative integer.") + + return errors + + +# --------------------------------------------------------------------------- +# Preview +# --------------------------------------------------------------------------- + + +class StatsPreview(TypedDict, total=False): + id: UUID | None + table_groups_name: str + table_group_schema: str + table_ct: int | None + column_ct: int | None + approx_record_ct: int | None + approx_data_point_ct: int | None + + +class TablePreview(TypedDict): + column_ct: int + approx_record_ct: int | None + approx_data_point_ct: int | None + can_access: bool | None + + +class TableGroupPreview(TypedDict): + stats: StatsPreview + tables: dict[str, TablePreview] + success: bool + message: str | None + + +_NO_CONNECTION_MESSAGE = "No connection selected. Please select a connection to preview the Table Group." +_NO_TABLES_MESSAGE = ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." +) +_INACCESSIBLE_MESSAGE = "Some tables were not accessible. Please the check the database permissions." + + +def preview_table_group( + table_group: TableGroup, + *, + connection: Connection | None = None, + verify_access: bool = True, +) -> tuple[TableGroupPreview, list[ColumnChars] | None, RefreshDataCharsSQL | None]: + """Probe the connected target DB for tables matching ``table_group``'s filters. + + Returns ``(preview, data_chars, sql_generator)`` — three picklable values. + On the error path (no connection, DDF failure, or empty schema) the second + and third tuple elements are ``None``. + + Use ``make_save_data_chars(data_chars, sql_generator)`` in callers that + need to record the introspected metadata in ``data_chars``. + """ + preview: TableGroupPreview = { + "stats": { + "id": table_group.id, + "table_groups_name": table_group.table_groups_name, + "table_group_schema": table_group.table_group_schema, + }, + "tables": {}, + "success": True, + "message": None, + } + + if not (connection or table_group.connection_id): + preview["success"] = False + preview["message"] = _NO_CONNECTION_MESSAGE + return preview, None, None + + data_chars: list[ColumnChars] | None = None + sql_generator: RefreshDataCharsSQL | None = None + try: + if connection is None: + connection = Connection.get(table_group.connection_id) + preview, data_chars, sql_generator = _build_preview(table_group, connection) + + if verify_access and preview["success"]: + for table_name in list(preview["tables"]): + try: + results = fetch_from_target_db(connection, *sql_generator.verify_access(table_name)) + except Exception: + preview["tables"][table_name]["can_access"] = False + else: + preview["tables"][table_name]["can_access"] = bool(results) and len(results) > 0 + if not all(t["can_access"] for t in preview["tables"].values()): + preview["message"] = _INACCESSIBLE_MESSAGE + except Exception as error: + preview["success"] = False + preview["message"] = error.args[0] if error.args else str(error) + data_chars = None + sql_generator = None + + return preview, data_chars, sql_generator + + +def make_save_data_chars( + data_chars: list[ColumnChars], + sql_generator: RefreshDataCharsSQL, +) -> Callable[[UUID], None]: + """Build the ``save_data_chars(table_group_id)`` callback for the caller. + + Kept out of ``preview_table_group`` so the service's return value stays + picklable; local closures don't pickle. + """ + def save(table_group_id: UUID) -> None: + # Unsaved table groups won't have an ID yet; sync it before writing. + sql_generator.table_group.id = table_group_id + write_data_chars(data_chars, sql_generator, datetime.now(UTC)) + return save + + +def _build_preview( + table_group: TableGroup, + connection: Connection, +) -> tuple[TableGroupPreview, list[ColumnChars], RefreshDataCharsSQL]: + sql_generator = RefreshDataCharsSQL(connection, table_group) + if sql_generator.flavor_service.metadata_via_api: + params = resolve_connection_params(connection.__dict__) + api_columns = sql_generator.flavor_service.get_schema_columns(params, table_group.table_group_schema) or [] + data_chars = sql_generator.filter_schema_columns(api_columns) + else: + rows = fetch_from_target_db(connection, *sql_generator.get_schema_ddf()) + data_chars = [ColumnChars(**column) for column in rows] + + preview: TableGroupPreview = { + "stats": { + "id": table_group.id, + "table_groups_name": table_group.table_groups_name, + "table_group_schema": table_group.table_group_schema, + "table_ct": 0, + "column_ct": 0, + "approx_record_ct": None, + "approx_data_point_ct": None, + }, + "tables": {}, + "success": True, + "message": None, + } + stats = preview["stats"] + tables = preview["tables"] + + for column in data_chars: + if not tables.get(column.table_name): + tables[column.table_name] = { + "column_ct": 0, + "approx_record_ct": column.approx_record_ct, + "approx_data_point_ct": None, + "can_access": None, + } + stats["table_ct"] += 1 + if column.approx_record_ct is not None: + stats["approx_record_ct"] = (stats["approx_record_ct"] or 0) + column.approx_record_ct + + stats["column_ct"] += 1 + tables[column.table_name]["column_ct"] += 1 + if column.approx_record_ct is not None: + stats["approx_data_point_ct"] = (stats["approx_data_point_ct"] or 0) + column.approx_record_ct + tables[column.table_name]["approx_data_point_ct"] = ( + tables[column.table_name]["approx_data_point_ct"] or 0 + ) + column.approx_record_ct + + if len(data_chars) <= 0: + preview["success"] = False + preview["message"] = _NO_TABLES_MESSAGE + + return preview, data_chars, sql_generator diff --git a/testgen/common/date_service.py b/testgen/common/date_service.py index eefaf131..54c04faf 100644 --- a/testgen/common/date_service.py +++ b/testgen/common/date_service.py @@ -97,12 +97,19 @@ def accommodate_dataframe_to_timezone(df, streamlit_session, time_columns=None): df[time_column] = df[time_column].dt.strftime("%Y-%m-%d %H:%M:%S") +def format_friendly_datetime(dt: datetime, fmt: str) -> str: + """Cross-platform strftime: replaces Linux-only %-d and %-I before calling strftime.""" + return dt.strftime(fmt.replace("%-d", str(dt.day)).replace("%-I", str(dt.hour % 12 or 12))) + + def get_timezoned_timestamp(streamlit_session, value, dateformat="%b %-d, %-I:%M %p"): ret = None if value and "browser_timezone" in streamlit_session: data = {"value": [value]} df = pd.DataFrame(data) timezone = streamlit_session["browser_timezone"] - df["value"] = df["value"].dt.tz_localize("UTC").dt.tz_convert(timezone).dt.strftime(dateformat) + df["value"] = df["value"].dt.tz_localize("UTC").dt.tz_convert(timezone).apply( + lambda dt: format_friendly_datetime(dt, dateformat) if pd.notna(dt) else "" + ) ret = df.iloc[0, 0] return ret diff --git a/testgen/common/enums.py b/testgen/common/enums.py index 3679f709..84fffb65 100644 --- a/testgen/common/enums.py +++ b/testgen/common/enums.py @@ -40,6 +40,14 @@ class JobKey(StrEnum): run_data_cleanup = "run-data-cleanup" +class PublicJobKey(StrEnum): + """``job_key`` values exposed through the public API — the externally-triggerable + subset of ``JobKey``. Internal maintenance kinds are intentionally absent.""" + run_profile = JobKey.run_profile.value + run_tests = JobKey.run_tests.value + run_test_generation = JobKey.run_test_generation.value + + class JobSource(StrEnum): """``source`` column values for ``job_executions``.""" api = "api" @@ -47,7 +55,6 @@ class JobSource(StrEnum): scheduler = "scheduler" mcp = "mcp" cli = "cli" - backfill = "backfill" system = "system" @@ -63,6 +70,18 @@ class JobStatus(StrEnum): CANCELED = "canceled" +# User-facing display labels for JobStatus values. +JOB_STATUS_LABEL: dict[str, str] = { + JobStatus.COMPLETED: "Completed", + JobStatus.CANCELED: "Canceled", + JobStatus.CANCEL_REQUESTED: "Canceling", + JobStatus.PENDING: "Pending", + JobStatus.CLAIMED: "Starting", + JobStatus.RUNNING: "Running", + JobStatus.ERROR: "Error", +} + + class Disposition(StrEnum): """Stored disposition values for ``profile_anomaly_results.disposition`` and ``test_results.disposition``. The user-facing label for ``INACTIVE`` is "Muted".""" @@ -83,3 +102,12 @@ class PiiRisk(StrEnum): """Risk level extracted from PII issue ``detail`` strings via ``priority`` hybrid.""" HIGH = "High" MODERATE = "Moderate" + + +class MonitorType(StrEnum): + """Stored ``test_type`` values for the four monitor test types. Surfaced to users + as the lowercase short labels (freshness / volume / schema / metric).""" + FRESHNESS = "Freshness_Trend" + VOLUME = "Volume_Trend" + SCHEMA = "Schema_Drift" + METRIC = "Metric_Trend" diff --git a/testgen/common/flavors.py b/testgen/common/flavors.py new file mode 100644 index 00000000..277bc7da --- /dev/null +++ b/testgen/common/flavors.py @@ -0,0 +1,71 @@ +"""Flavor identity — the single source of truth for database-flavor display labels +and the code→family mapping. + +Flavor *type* labels (e.g. "PostgreSQL", "Snowflake") are shared presentation: the +Streamlit connections page feeds them to its JS component via ``FLAVOR_OPTIONS`` and +the MCP tools accept/return them as the ``sql_flavor`` value. Both surfaces derive +from here so the labels are written exactly once. + +This module is intentionally pure data (no I/O, no Streamlit, no SQLAlchemy) so it +can be imported from the UI, the MCP layer, and common services without cycles. + +(The per-flavor *connection-parameter* field labels and requirement rules are a +different, MCP-only concern and live in ``mcp/tools/common.py`` — they're hardcoded +in the JS form, not sourced from Python.) +""" + +from __future__ import annotations + +from enum import StrEnum + + +class SqlFlavorLabel(StrEnum): + """User-facing database-flavor labels (the ``sql_flavor`` value space).""" + + REDSHIFT = "Amazon Redshift" + REDSHIFT_SPECTRUM = "Amazon Redshift Spectrum" + AZURE_MSSQL = "Azure SQL Database" + SYNAPSE_MSSQL = "Azure Synapse Analytics" + DATABRICKS = "Databricks" + BIGQUERY = "Google BigQuery" + MSSQL = "Microsoft SQL Server" + ORACLE = "Oracle" + POSTGRESQL = "PostgreSQL" + SAP_HANA = "SAP HANA" + SALESFORCE_DATA360 = "Salesforce Data 360" + SNOWFLAKE = "Snowflake" + + +# code → label. Label strings are written once (on the enum); this only pairs them +# with flavor codes. +FLAVOR_CODE_TO_LABEL: dict[str, SqlFlavorLabel] = { + "redshift": SqlFlavorLabel.REDSHIFT, + "redshift_spectrum": SqlFlavorLabel.REDSHIFT_SPECTRUM, + "azure_mssql": SqlFlavorLabel.AZURE_MSSQL, + "synapse_mssql": SqlFlavorLabel.SYNAPSE_MSSQL, + "databricks": SqlFlavorLabel.DATABRICKS, + "bigquery": SqlFlavorLabel.BIGQUERY, + "mssql": SqlFlavorLabel.MSSQL, + "oracle": SqlFlavorLabel.ORACLE, + "postgresql": SqlFlavorLabel.POSTGRESQL, + "sap_hana": SqlFlavorLabel.SAP_HANA, + "salesforce_data360": SqlFlavorLabel.SALESFORCE_DATA360, + "snowflake": SqlFlavorLabel.SNOWFLAKE, +} + +# code → family (multiple codes can share one engine family, e.g. the two Azure +# variants both map to "mssql"). +FLAVOR_CODE_TO_FAMILY: dict[str, str] = { + "redshift": "redshift", + "redshift_spectrum": "redshift_spectrum", + "azure_mssql": "mssql", + "synapse_mssql": "mssql", + "databricks": "databricks", + "bigquery": "bigquery", + "mssql": "mssql", + "oracle": "oracle", + "postgresql": "postgresql", + "sap_hana": "sap_hana", + "salesforce_data360": "salesforce_data360", + "snowflake": "snowflake", +} diff --git a/testgen/common/freshness_service.py b/testgen/common/freshness_service.py index b87e5a93..2e74abe8 100644 --- a/testgen/common/freshness_service.py +++ b/testgen/common/freshness_service.py @@ -8,7 +8,8 @@ import numpy as np import pandas as pd -from testgen.common.time_series_service import NotEnoughData, get_holiday_dates +from testgen.common.holiday_service import get_holiday_dates +from testgen.common.time_series_service import NotEnoughData LOG = logging.getLogger("testgen") diff --git a/testgen/common/holiday_service.py b/testgen/common/holiday_service.py new file mode 100644 index 00000000..e72e0fdd --- /dev/null +++ b/testgen/common/holiday_service.py @@ -0,0 +1,64 @@ +"""Holiday calendar utilities: resolve and validate country / financial-market holiday codes. + +Holiday codes name the calendars excluded from prediction baselines — used as SARIMAX +exogenous regressors and in freshness business-time gap calculations. Each code resolves +via the ``holidays`` package as either a country (ISO code, e.g. ``US``) or a financial +market (e.g. ``NYSE``). +""" + +import logging +from datetime import datetime + +import holidays +import pandas as pd + +LOG = logging.getLogger("testgen") + + +def get_holiday_dates(holiday_codes: list[str], datetime_index: pd.DatetimeIndex) -> set[datetime]: + years = list(range(datetime_index.year.min(), datetime_index.year.max() + 1)) + + holiday_dates = set() + if holiday_codes: + for code in holiday_codes: + code = code.strip().upper() + found = False + + try: + country_holidays = holidays.country_holidays(code, years=years) + holiday_dates.update(country_holidays.keys()) + found = True + except NotImplementedError: + pass # Not a valid country code + + if not found: + try: + financial_holidays = holidays.financial_holidays(code, years=years) + holiday_dates.update(financial_holidays.keys()) + found = True + except NotImplementedError: + pass # Not a valid financial code + + if not found: + LOG.warning(f"Holiday code '{code}' could not be resolved as a country or financial market") + + return holiday_dates + + +def is_supported_holiday_code(code: str) -> bool: + """Whether a holiday code resolves to a country or financial-market calendar. + + Applies the same normalization as :func:`get_holiday_dates` (strip + upper), so a code + that passes here is one that resolver will actually honor. + """ + normalized = code.strip().upper() + if not normalized: + return False + for resolver in (holidays.country_holidays, holidays.financial_holidays): + try: + resolver(normalized) + except NotImplementedError: + continue + else: + return True + return False diff --git a/testgen/common/mixpanel_service.py b/testgen/common/mixpanel_service.py index 1af2ae4c..eb26e57e 100644 --- a/testgen/common/mixpanel_service.py +++ b/testgen/common/mixpanel_service.py @@ -1,7 +1,11 @@ +import atexit import functools import json import logging +import queue import ssl +import threading +import time import uuid from base64 import b64encode from functools import cached_property, wraps @@ -18,6 +22,13 @@ LOG = logging.getLogger("testgen") +_BATCH_SIZE = 50 # Mixpanel /track array cap +_FLUSH_INTERVAL_SEC = 10 # Time-based flush +_QUEUE_MAX_SIZE = 1000 # Memory cap before drop-on-overflow +_DRAIN_TIMEOUT_SEC = 5 # Bounded shutdown drain + +_SHUTDOWN = object() # sentinel enqueued by drain() to stop the worker + def safe_method(method): @wraps(method) @@ -33,6 +44,14 @@ def wrapped(*args, **kwargs): class MixpanelService(Singleton): + def __init__(self) -> None: + self._queue: queue.Queue | None = None + self._worker: threading.Thread | None = None + self._worker_lock = threading.Lock() + self._started = False + self._stopped = False + atexit.register(self.drain) + @cached_property @with_database_session def instance_id(self): @@ -54,33 +73,106 @@ def _hash_value(self, value: bytes | str, digest_size: int = 8) -> str: @safe_method def send_event(self, event_name, include_usage=False, **properties): - self._track(event_name, include_usage=include_usage, **properties) + self._enqueue(self._build_event(event_name, include_usage=include_usage, **properties)) def send_feedback(self, **properties): - # User-submitted feedback is content the user explicitly chose to share - # so it is not gated by the TG_ANALYTICS opt-out. + # User-submitted feedback is content the user explicitly chose to share, + # so it is not gated by the TG_ANALYTICS opt-out. It is a foreground action + # posted synchronously — never enqueued — so it never starts the worker. try: - self._track("feedback", **properties) + self.send_mp_request("track?ip=1", self._build_event("feedback", **properties)) except Exception: LOG.exception("Error sending feedback") - def _track(self, event_name, include_usage=False, **properties): + def _build_event(self, event_name, include_usage=False, **properties) -> dict: properties.setdefault("instance_id", self.instance_id) properties.setdefault("edition", settings.DOCKER_HUB_REPOSITORY) properties.setdefault("version", settings.VERSION) properties.setdefault("username", session.auth.user_display if session.auth else None) properties.setdefault("distinct_id", self.get_distinct_id(properties["username"])) + properties.setdefault("time", int(time.time())) if include_usage: properties.update(self.get_usage()) - track_payload = { + return { "event": event_name, "properties": { "token": settings.MIXPANEL_TOKEN, **properties, - } + }, } - self.send_mp_request("track?ip=1", track_payload) + + def _ensure_worker(self) -> None: + if self._started: + return + with self._worker_lock: + if self._started: + return + self._queue = queue.Queue(maxsize=_QUEUE_MAX_SIZE) + self._worker = threading.Thread(target=self._worker_loop, name="mixpanel-flush", daemon=True) + self._worker.start() + self._started = True + + def _enqueue(self, event: dict) -> None: + if self._stopped: + LOG.warning("analytics worker stopped; dropping event") + return + self._ensure_worker() + try: + self._queue.put_nowait(event) + except queue.Full: + LOG.warning("analytics queue full; dropping event") + + def _worker_loop(self) -> None: + while True: + batch, stop = self._next_batch() + if batch: + self._flush(batch) + if stop: + return + + def _next_batch(self) -> tuple[list[dict], bool]: + """Block up to _FLUSH_INTERVAL_SEC for the first event, then drain up to + _BATCH_SIZE without blocking. Returns (batch, stop).""" + try: + first = self._queue.get(timeout=_FLUSH_INTERVAL_SEC) + except queue.Empty: + return [], False + if first is _SHUTDOWN: + return [], True + batch = [first] + while len(batch) < _BATCH_SIZE: + try: + event = self._queue.get_nowait() + except queue.Empty: + break + if event is _SHUTDOWN: + return batch, True + batch.append(event) + return batch, False + + def _flush(self, events: list[dict]) -> None: + for start in range(0, len(events), _BATCH_SIZE): + chunk = events[start:start + _BATCH_SIZE] + try: + self.send_mp_request("track?ip=1", chunk) + except Exception: + LOG.exception("Failed to flush analytics batch") + + def drain(self) -> None: + """Flush queued events and stop the worker. Idempotent; bounded by + _DRAIN_TIMEOUT_SEC. Called by the atexit hook and the server lifespan.""" + if not self._started or self._stopped: + return + self._stopped = True + try: + self._queue.put_nowait(_SHUTDOWN) + except queue.Full: + LOG.warning("analytics queue full at shutdown; in-flight events may be dropped") + if self._worker is not None: + self._worker.join(timeout=_DRAIN_TIMEOUT_SEC) + if self._worker.is_alive(): + LOG.warning("analytics drain timed out after %ss", _DRAIN_TIMEOUT_SEC) def get_ssl_context(self): ssl_context = ssl.create_default_context() diff --git a/testgen/common/models/connection.py b/testgen/common/models/connection.py index a940edba..863efbb0 100644 --- a/testgen/common/models/connection.py +++ b/testgen/common/models/connection.py @@ -28,6 +28,11 @@ SQLFlavorCode = Literal["redshift", "redshift_spectrum", "snowflake", "mssql", "azure_mssql", "synapse_mssql", "postgresql", "databricks", "bigquery", "oracle", "sap_hana", "salesforce_data360"] +# Fallback when a connection row has a NULL max_query_chars (no DB default; older +# rows / non-UI insert paths may not seed it). Matches the UI's non-Salesforce +# default and the 0160 migration. +DEFAULT_MAX_QUERY_CHARS = 20000 + @dataclass class ConnectionMinimal(EntityMinimal): @@ -38,6 +43,17 @@ class ConnectionMinimal(EntityMinimal): connection_name: str +@dataclass +class ConnectionListItem(EntityMinimal): + connection_id: int + connection_name: str + project_code: str + sql_flavor_code: SQLFlavorCode + project_host: str | None + project_db: str | None + table_group_count: int + + class Connection(Entity): __tablename__ = "connections" @@ -88,6 +104,35 @@ def select_minimal_where( results = cls._select_columns_where(cls._minimal_columns, *clauses, order_by=order_by) return [ConnectionMinimal(**row) for row in results] + @classmethod + def list_for_project( + cls, project_code: str, *clauses, page: int = 1, limit: int = 20 + ) -> tuple[list[ConnectionListItem], int]: + """Paginated listing of connections in a project, with their table-group counts.""" + tg_count_subq = ( + select( + TableGroup.connection_id.label("connection_id"), + func.count().label("table_group_count"), + ) + .group_by(TableGroup.connection_id) + .subquery() + ) + query = ( + select( + cls.connection_id, + cls.connection_name, + cls.project_code, + cls.sql_flavor_code, + cls.project_host, + cls.project_db, + func.coalesce(tg_count_subq.c.table_group_count, 0).label("table_group_count"), + ) + .outerjoin(tg_count_subq, tg_count_subq.c.connection_id == cls.connection_id) + .where(cls.project_code == project_code, *clauses) + .order_by(*cls._default_order_by) + ) + return cls._paginate(query, page=page, limit=limit, data_class=ConnectionListItem) + @classmethod def is_in_use(cls, ids: list[str]) -> bool: table_groups = TableGroup.select_minimal_where(TableGroup.connection_id.in_(ids)) diff --git a/testgen/common/models/data_column.py b/testgen/common/models/data_column.py index cee7d088..12a7473b 100644 --- a/testgen/common/models/data_column.py +++ b/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 @@ -236,6 +237,13 @@ class ColumnSearchHit(EntityMinimal): column_name: str +@dataclass +class CreateScriptColumn: + column_name: str + db_data_type: str | None + datatype_suggestion: str | None + + class DataColumnChars(Entity): __tablename__ = "data_column_chars" @@ -253,6 +261,16 @@ class DataColumnChars(Entity): critical_data_element: bool | None = Column(Boolean) excluded_data_element: bool | None = Column(Boolean, nullable=True) pii_flag: str | None = Column(String(50), nullable=True) + description: str | None = Column(String(1000)) + data_source: str | None = Column(String(40)) + source_system: str | None = Column(String(40)) + source_process: str | None = Column(String(40)) + business_domain: str | None = Column(String(40)) + stakeholder_group: str | None = Column(String(40)) + transform_level: str | None = Column(String(40)) + aggregation_level: str | None = Column(String(40)) + data_product: str | None = Column(String(40)) + data_classification: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) @@ -260,14 +278,58 @@ class DataColumnChars(Entity): _default_order_by = (asc(ordinal_position), asc(column_name)) - # Unmapped columns: description, data_source, source_system, source_process, - # business_domain, stakeholder_group, transform_level, aggregation_level, - # data_product, add_date, last_mod_date, test_ct, last_test_date, + # Unmapped columns: add_date, last_mod_date, test_ct, last_test_date, # tests_last_run, tests_7_days_prior, tests_30_days_prior, fails_last_run, # fails_7_days_prior, fails_30_days_prior, warnings_last_run, # warnings_7_days_prior, warnings_30_days_prior, valid_profile_issue_ct, # valid_test_issue_ct + @classmethod + def list_for_create_script( + cls, table_groups_id: UUID, table_name: str, + ) -> tuple[str | None, list[CreateScriptColumn]]: + """Return ``(schema_name, columns)`` for a table's CREATE TABLE script. + + Columns are ordered by ordinal position and carry the profiling-derived type + suggestion from their latest complete profiling run. Returns ``(None, [])`` when + the table is not in the table group's profiled catalog. + """ + query = ( + select( + cls.schema_name, + cls.column_name, + cls.db_data_type, + ProfileResult.datatype_suggestion, + ) + .outerjoin( + ProfileResult, + and_( + ProfileResult.profile_run_id == cls.last_complete_profile_run_id, + ProfileResult.schema_name == cls.schema_name, + ProfileResult.table_name == cls.table_name, + ProfileResult.column_name == cls.column_name, + ), + ) + .where( + cls.table_groups_id == table_groups_id, + cls.table_name == table_name, + cls.drop_date.is_(None), + ) + .order_by(asc(cls.ordinal_position), asc(cls.column_name)) + ) + rows = get_current_session().execute(query).mappings().all() + if not rows: + return None, [] + columns = [ + CreateScriptColumn( + column_name=row["column_name"], + db_data_type=row["db_data_type"], + datatype_suggestion=row["datatype_suggestion"], + ) + for row in rows + ] + return rows[0]["schema_name"], columns + @classmethod def list_for_table_group( cls, @@ -499,11 +561,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( @@ -525,6 +587,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, diff --git a/testgen/common/models/data_structure_log.py b/testgen/common/models/data_structure_log.py new file mode 100644 index 00000000..552c932d --- /dev/null +++ b/testgen/common/models/data_structure_log.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass +from datetime import datetime +from uuid import UUID, uuid4 + +from sqlalchemy import Column, String, asc, desc, select +from sqlalchemy.dialects import postgresql + +from testgen.common.models import get_current_session +from testgen.common.models.entity import Entity, EntityMinimal + +# Schema-change codes stored in ``data_structure_log.change``. ``M`` only ever +# appears on column-level rows (column-name set, old/new data type populated). +SCHEMA_CHANGE_ADDED = "A" +SCHEMA_CHANGE_DROPPED = "D" +SCHEMA_CHANGE_MODIFIED = "M" + + +@dataclass +class DataStructureLogEntry(EntityMinimal): + """One schema-change event for a table or column in a table group. + + ``column_name`` is ``None`` for table-level events (add / drop of an entire + table); set for column-level events. ``change`` is the stored single-letter + code (``A`` / ``D`` / ``M``); callers translate to user-facing words. + """ + log_id: UUID + table_groups_id: UUID + table_name: str + column_name: str | None + change_date: datetime + change: str + old_data_type: str | None + new_data_type: str | None + + +class DataStructureLog(Entity): + __tablename__ = "data_structure_log" + + log_id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, default=uuid4) + table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True)) + table_id: UUID = Column(postgresql.UUID(as_uuid=True)) + column_id: UUID = Column(postgresql.UUID(as_uuid=True)) + table_name: str = Column(String) + column_name: str = Column(String) + change_date: datetime = Column(postgresql.TIMESTAMP) + change: str = Column(String) + old_data_type: str = Column(String) + new_data_type: str = Column(String) + + _get_by = "log_id" + _default_order_by = (desc(change_date), asc(table_name), asc(column_name)) + + @classmethod + def list_for_table_group( + cls, + table_group_id: str | UUID, + *clauses, + page: int = 1, + limit: int | None = 20, + ) -> tuple[list[DataStructureLogEntry], int]: + """Paginated schema-change audit log for one table group, newest first. + + Caller-supplied ``*clauses`` are WHERE expressions on this model — e.g. + ``cls.table_name == name`` (exact match; the audit log stores names + case-sensitively as the source emitted them) or a ``cls.change_date`` + bound. ``limit=None`` skips pagination — the caller gets every matching + row in one shot. + """ + query = select( + cls.log_id, + cls.table_groups_id, + cls.table_name, + cls.column_name, + cls.change_date, + cls.change, + cls.old_data_type, + cls.new_data_type, + ).where(cls.table_groups_id == table_group_id, *clauses) + query = query.order_by(*cls._default_order_by) + + if limit is None: + rows = get_current_session().execute(query).mappings().all() + entries = [DataStructureLogEntry(**row) for row in rows] + return entries, len(entries) + return cls._paginate(query, page=page, limit=limit, data_class=DataStructureLogEntry) diff --git a/testgen/common/models/data_table.py b/testgen/common/models/data_table.py index 21ea22b3..d3d3c7d0 100644 --- a/testgen/common/models/data_table.py +++ b/testgen/common/models/data_table.py @@ -17,6 +17,7 @@ ) from sqlalchemy.dialects import postgresql +from testgen.common.database.column_chars import ObjectType from testgen.common.models import get_current_session from testgen.common.models.data_column import DataColumnChars from testgen.common.models.entity import Entity @@ -42,6 +43,7 @@ class TableProfilingOverview: table_groups_id: UUID schema_name: str | None table_name: str + object_type: ObjectType | None record_ct: int | None column_ct: int | None dq_score_profiling: float | None @@ -61,18 +63,31 @@ class DataTable(Entity): table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("table_groups.id")) schema_name: str | None = Column(String) table_name: str = Column(String) + object_type: ObjectType | None = Column(String) column_ct: int | None = Column(BigInteger) record_ct: int | None = Column(BigInteger) approx_record_ct: int | None = Column(BigInteger) critical_data_element: bool | None = Column(Boolean) + description: str | None = Column(String(1000)) + data_source: str | None = Column(String(40)) + source_system: str | None = Column(String(40)) + source_process: str | None = Column(String(40)) + business_domain: str | None = Column(String(40)) + stakeholder_group: str | None = Column(String(40)) + transform_level: str | None = Column(String(40)) + aggregation_level: str | None = Column(String(40)) + data_product: str | None = Column(String(40)) + data_classification: str | None = Column(String(40)) drop_date: datetime | None = Column(postgresql.TIMESTAMP) last_complete_profile_run_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float | None = Column(Float) dq_score_testing: float | None = Column(Float) - # Unmapped columns: functional_table_type, description, data_source, - # source_system, source_process, business_domain, stakeholder_group, - # transform_level, aggregation_level, data_product, add_date, + # The inherited Entity default orders by the textual label "id", but this model's + # primary key column is "table_id" — order by a real column instead. + _default_order_by = (asc(table_name),) + + # Unmapped columns: functional_table_type, add_date, # last_refresh_date, last_profile_record_ct @classmethod @@ -110,6 +125,7 @@ def get_profiling_overview( cls.table_groups_id, cls.schema_name, cls.table_name, + cls.object_type, cls.record_ct, cls.column_ct, cls.dq_score_profiling, @@ -119,7 +135,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, diff --git a/testgen/common/models/hygiene_issue.py b/testgen/common/models/hygiene_issue.py index c7683479..9bccae6b 100644 --- a/testgen/common/models/hygiene_issue.py +++ b/testgen/common/models/hygiene_issue.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import aliased, relationship from sqlalchemy.sql.functions import func -from testgen.common.enums import Disposition +from testgen.common.enums import Disposition, IssueLikelihood, PiiRisk from testgen.common.models import Base, get_current_session from testgen.common.models.entity import Entity from testgen.common.models.job_execution import JobExecution @@ -23,12 +23,28 @@ @dataclass -class IssueLikelihoodCounts: - """Counts of hygiene issues by likelihood category, with dismissed/inactive separated.""" +class HygieneIssueCounts: + """Counts of active data-quality hygiene issues by likelihood category.""" definite: int = 0 likely: int = 0 possible: int = 0 + + +@dataclass +class PotentialPiiCounts: + """Counts of active Potential PII findings by risk level.""" + + high: int = 0 + moderate: int = 0 + + +@dataclass +class IssueCounts: + """Profiling-finding breakdown: active counts by kind, plus a single dismissed total.""" + + hygiene_issues: HygieneIssueCounts + potential_pii: PotentialPiiCounts dismissed: int = 0 @@ -202,19 +218,25 @@ def select_count_by_priority(cls, profiling_run_id: UUID) -> dict[str, IssueCoun return result @classmethod - def count_by_likelihood(cls, profile_run_id: UUID) -> IssueLikelihoodCounts: - """Count hygiene issues by likelihood category for a single profiling run.""" - dismissed = func.coalesce(cls.disposition, "Confirmed").in_(("Dismissed", "Inactive")) + def count_for_run(cls, profile_run_id: UUID) -> IssueCounts: + """Count profiling findings for a single run: active hygiene issues by likelihood, + active Potential PII findings by risk level, and a single dismissed total.""" + dismissed = func.coalesce(cls.disposition, Disposition.CONFIRMED).in_( + (Disposition.DISMISSED, Disposition.INACTIVE) + ) + is_pii = HygieneIssueType.likelihood == IssueLikelihood.POTENTIAL_PII - def _count_active(likelihood_values: tuple[str, ...]): - return func.sum(case((~dismissed & HygieneIssueType.likelihood.in_(likelihood_values), 1), else_=0)) + def _count(condition): + return func.coalesce(func.sum(case((condition, 1), else_=0)), 0) query = ( select( - _count_active(("Definite",)).label("definite"), - _count_active(("Likely",)).label("likely"), - _count_active(("Possible", "Potential PII")).label("possible"), - func.sum(case((dismissed, 1), else_=0)).label("dismissed"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.DEFINITE)).label("definite"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.LIKELY)).label("likely"), + _count(~dismissed & (HygieneIssueType.likelihood == IssueLikelihood.POSSIBLE)).label("possible"), + _count(~dismissed & is_pii & (cls.priority == PiiRisk.HIGH)).label("high"), + _count(~dismissed & is_pii & (cls.priority == PiiRisk.MODERATE)).label("moderate"), + _count(dismissed).label("dismissed"), ) .select_from(cls) .join(HygieneIssueType, HygieneIssueType.id == cls.type_id) @@ -222,7 +244,11 @@ def _count_active(likelihood_values: tuple[str, ...]): ) row = get_current_session().execute(query).first() - return IssueLikelihoodCounts(**{k: v for k, v in row._mapping.items() if v is not None}) + return IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=row.definite, likely=row.likely, possible=row.possible), + potential_pii=PotentialPiiCounts(high=row.high, moderate=row.moderate), + dismissed=row.dismissed, + ) @classmethod def _priority_order(cls): @@ -275,7 +301,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 +325,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 +340,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 +384,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 +394,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_( diff --git a/testgen/common/models/job_execution.py b/testgen/common/models/job_execution.py index 8d8160c2..bb6d8d83 100644 --- a/testgen/common/models/job_execution.py +++ b/testgen/common/models/job_execution.py @@ -5,22 +5,14 @@ from sqlalchemy import Column, String, Text, case, delete, func, select, text, update from sqlalchemy.dialects import postgresql +from sqlalchemy.sql.elements import ColumnElement -from testgen.common.enums import JobKey, JobSource, JobStatus +from testgen.common.enums import JobSource, JobStatus from testgen.common.models import Base, database_session, get_current_session LOG = logging.getLogger("testgen") -# Job kinds that are externally triggerable. Internal kinds (run-score-update, -# recalculate-project-scores, ...) are absent and filtered out by construction. -PUBLIC_JOB_KEYS: frozenset[JobKey] = frozenset({ - JobKey.run_profile, - JobKey.run_tests, - JobKey.run_test_generation, -}) - - _VALID_TRANSITIONS: dict[JobStatus, frozenset[JobStatus]] = { JobStatus.PENDING: frozenset({JobStatus.CLAIMED, JobStatus.CANCEL_REQUESTED}), JobStatus.CLAIMED: frozenset({JobStatus.RUNNING, JobStatus.ERROR, JobStatus.CANCEL_REQUESTED}), @@ -106,6 +98,7 @@ def claim_actionable(cls, limit: int = 5) -> list[Self]: @classmethod def select_active_by_kwargs( cls, + *clauses: ColumnElement[bool], project_code: str, job_key: str, kwargs_match: dict[str, str | list[str]], @@ -114,6 +107,7 @@ def select_active_by_kwargs( """Find JE rows whose ``kwargs`` JSONB matches the given (key, value) pairs. Values may be a single string or a list of strings (which becomes an ``IN`` filter). + Additional caller-supplied WHERE expressions may be passed as ``*clauses``. Defaults to active (non-terminal) statuses. """ statuses = statuses or cls._ACTIVE_STATUSES @@ -121,6 +115,7 @@ def select_active_by_kwargs( cls.project_code == project_code, cls.job_key == job_key, cls.status.in_(statuses), + *clauses, ) for k, v in kwargs_match.items(): if isinstance(v, list): diff --git a/testgen/common/models/oauth.py b/testgen/common/models/oauth.py new file mode 100644 index 00000000..f8f6856a --- /dev/null +++ b/testgen/common/models/oauth.py @@ -0,0 +1,90 @@ +import time +from enum import StrEnum + +from authlib.integrations.sqla_oauth2 import ( + OAuth2AuthorizationCodeMixin, + OAuth2ClientMixin, + OAuth2TokenMixin, +) +from sqlalchemy import Column, ForeignKey, String, case, func +from sqlalchemy.dialects import postgresql +from sqlalchemy.ext.hybrid import hybrid_property + +from testgen import settings +from testgen.common.models import Base + + +class OAuth2ClientType(StrEnum): + """How an OAuth2 client was provisioned. + + PAT clients are created by the authenticated personal-access-token path and owned by + a user; EXTERNAL clients are dynamically registered (MCP apps, automation scripts). + """ + + PAT = "pat" + EXTERNAL = "external" + + +class PersonalAccessTokenStatus(StrEnum): + """Display status of a personal access token. Revoked takes precedence over expired.""" + + ACTIVE = "active" + REVOKED = "revoked" + EXPIRED = "expired" + + +class OAuth2Client(Base, OAuth2ClientMixin): + __tablename__ = "oauth2_clients" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + client_type = Column(String(20), nullable=False, default=OAuth2ClientType.EXTERNAL) + + # Override to widen — JWTs can exceed 255 chars + # (the mixin defines client_id as VARCHAR(48) which is fine) + + +class OAuth2AuthorizationCode(Base, OAuth2AuthorizationCodeMixin): + __tablename__ = "oauth2_authorization_codes" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=False) + + +class OAuth2Token(Base, OAuth2TokenMixin): + __tablename__ = "oauth2_tokens" + + id = Column(postgresql.UUID(as_uuid=True), primary_key=True, server_default="gen_random_uuid()") + user_id = Column(postgresql.UUID(as_uuid=True), ForeignKey("auth_users.id", ondelete="CASCADE"), nullable=True) + + # Override to allow longer JWTs as access tokens + access_token = Column(String(2048), unique=True, nullable=False) + + # Set only for personal access tokens; NULL for tokens from the auth-code / MCP flows + name = Column(String(255), nullable=True) + + def is_refresh_token_active(self) -> bool: + if self.refresh_token_revoked_at: + return False + expires_at = self.issued_at + settings.REFRESH_TOKEN_EXPIRES_IN + return expires_at >= time.time() + + @hybrid_property + def status(self) -> PersonalAccessTokenStatus: + """Personal access token status, derived from revocation and expiry. + + Instance side reads the app clock; the SQL expression reads the DB clock — + both are wall-clock "now", evaluated independently. + """ + if self.access_token_revoked_at: + return PersonalAccessTokenStatus.REVOKED + if self.issued_at + self.expires_in < time.time(): + return PersonalAccessTokenStatus.EXPIRED + return PersonalAccessTokenStatus.ACTIVE + + @status.expression # type: ignore[no-redef] + def status(cls): + return case( + (cls.access_token_revoked_at > 0, PersonalAccessTokenStatus.REVOKED.value), + (cls.issued_at + cls.expires_in < func.extract("epoch", func.now()), PersonalAccessTokenStatus.EXPIRED.value), + else_=PersonalAccessTokenStatus.ACTIVE.value, + ) diff --git a/testgen/common/models/profiling_run.py b/testgen/common/models/profiling_run.py index 225c9942..028af6a6 100644 --- a/testgen/common/models/profiling_run.py +++ b/testgen/common/models/profiling_run.py @@ -1,16 +1,16 @@ from collections.abc import Iterable from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime from typing import ClassVar, Literal, NamedTuple, Self, TypedDict -from uuid import UUID, uuid4 +from uuid import UUID -from sqlalchemy import BigInteger, Column, Float, Integer, String, desc, func, select, text, update +from sqlalchemy import BigInteger, Column, Float, ForeignKey, String, delete, desc, func, select, text from sqlalchemy.dialects import postgresql -from sqlalchemy.orm import InstrumentedAttribute +from sqlalchemy.orm import InstrumentedAttribute, foreign, relationship from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case -from testgen.common.enums import Disposition, JobStatus +from testgen.common.enums import JOB_STATUS_LABEL, Disposition, JobStatus from testgen.common.models import database_session, get_current_session from testgen.common.models.connection import Connection from testgen.common.models.entity import Entity, EntityMinimal @@ -20,7 +20,6 @@ from testgen.common.models.table_group import TableGroup from testgen.utils import is_uuid4 -ProfilingRunStatus = Literal["Running", "Complete", "Error", "Cancelled"] ProgressKey = Literal["data_chars", "col_profiling", "freq_analysis", "hygiene_issues"] ProgressStatus = Literal["Pending", "Running", "Completed", "Warning"] @@ -57,8 +56,6 @@ class ProfilingRunSummary(EntityMinimal): progress: list[ProgressStep] table_groups_name: str | None table_group_schema: str | None - process_id: int | None - log_message: str | None table_ct: int | None column_ct: int | None record_ct: int | None @@ -71,15 +68,7 @@ class ProfilingRunSummary(EntityMinimal): dq_score_profiling: float | None total_count: int - STATUS_LABEL: ClassVar[dict[str, str]] = { - JobStatus.COMPLETED: "Completed", - JobStatus.CANCELED: "Canceled", - JobStatus.CANCEL_REQUESTED: "Canceling", - JobStatus.PENDING: "Pending", - JobStatus.CLAIMED: "Starting", - JobStatus.RUNNING: "Running", - JobStatus.ERROR: "Error", - } + STATUS_LABEL: ClassVar[dict[str, str]] = JOB_STATUS_LABEL @property def status_label(self) -> str: @@ -103,15 +92,16 @@ class LatestProfilingRun(NamedTuple): class ProfilingRun(Entity): __tablename__ = "profiling_runs" - id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, default=uuid4) + id: UUID = Column( + postgresql.UUID(as_uuid=True), + ForeignKey("job_executions.id", ondelete="CASCADE"), + primary_key=True, + ) project_code: str = Column(String, nullable=False) connection_id: str = Column(BigInteger, nullable=False) table_groups_id: UUID = Column(postgresql.UUID(as_uuid=True), nullable=False) profiling_starttime: datetime = Column(postgresql.TIMESTAMP) - profiling_endtime: datetime = Column(postgresql.TIMESTAMP) - status: ProfilingRunStatus = Column(String, default="Running") progress: list[ProgressStep] = Column(postgresql.JSONB, default=[]) - log_message: str = Column(String) table_ct: int = Column(BigInteger) column_ct: int = Column(BigInteger) record_ct: int = Column(BigInteger) @@ -122,8 +112,13 @@ class ProfilingRun(Entity): dq_affected_data_points: int = Column(BigInteger) dq_total_data_points: int = Column(BigInteger) dq_score_profiling: float = Column(Float) - process_id: int = Column(Integer) - job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) + + job_execution = relationship( + JobExecution, + primaryjoin=foreign(id) == JobExecution.id, + uselist=False, + viewonly=True, + ) _default_order_by = (desc(profiling_starttime),) _minimal_columns = ( @@ -140,12 +135,6 @@ class ProfilingRun(Entity): ).label("is_latest_run"), ) - @classmethod - def get_by_id_or_job(cls, identifier: UUID) -> Self | None: - """Look up a profiling run by its own ID or by job_execution_id.""" - query = select(cls).where((cls.id == identifier) | (cls.job_execution_id == identifier)) - return get_current_session().scalars(query).first() - @classmethod def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: if not is_uuid4(run_id): @@ -154,7 +143,7 @@ def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: query = ( select(*cls._minimal_columns) .join(TableGroup, cls.table_groups_id == TableGroup.id) - .where((cls.id == run_id) | (cls.job_execution_id == run_id)) + .where(cls.id == run_id) ) result = get_current_session().execute(query).mappings().first() return ProfilingRunMinimal(**result) if result else None @@ -163,7 +152,7 @@ def get_minimal(cls, run_id: str | UUID) -> ProfilingRunMinimal | None: def get_latest_run(cls, project_code: str) -> LatestProfilingRun | None: query = ( select(ProfilingRun.id, JobExecution.started_at.label("run_time")) - .join(JobExecution, ProfilingRun.job_execution_id == JobExecution.id) + .join(JobExecution, ProfilingRun.id == JobExecution.id) .where(ProfilingRun.project_code == project_code, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) .limit(1) @@ -175,15 +164,14 @@ def get_latest_run(cls, project_code: str) -> LatestProfilingRun | None: @classmethod def get_latest_complete_je_id_for_table_group(cls, table_groups_id: UUID) -> UUID | None: - """Return the ``job_execution_id`` of the latest completed profiling run for a table group. + """Return the id of the latest completed profiling run for a table group. - Computed live from ``profiling_runs`` joined to ``job_executions`` — does not read the - legacy ``table_groups.last_complete_profile_run_id`` cache, which points at the internal - run PK rather than the JE id. + Computed live from ``profiling_runs`` joined to ``job_executions`` rather than + reading the ``table_groups.last_complete_profile_run_id`` cache. """ query = ( - select(cls.job_execution_id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + select(cls.id) + .join(JobExecution, cls.id == JobExecution.id) .where(cls.table_groups_id == table_groups_id, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) .limit(1) @@ -209,6 +197,7 @@ def select_summary( project_code: str | None = None, table_group_id: str | UUID | None = None, job_execution_id: str | UUID | None = None, + schedule_id: str | None = None, statuses: list[JobStatus] | None = None, page: int = 1, page_size: int = 20, @@ -216,6 +205,7 @@ def select_summary( if ( (table_group_id and not is_uuid4(table_group_id)) or (job_execution_id and not is_uuid4(job_execution_id)) + or (schedule_id and not is_uuid4(schedule_id)) ): return [], 0 @@ -251,8 +241,6 @@ def select_summary( COALESCE(pr.progress, '[]'::jsonb) AS progress, tg.table_groups_name, tg.table_group_schema, - pr.process_id, - pr.log_message, pr.table_ct, pr.column_ct, pr.record_ct, @@ -265,13 +253,14 @@ def select_summary( pr.dq_score_profiling, COUNT(*) OVER() AS total_count FROM job_executions je - LEFT JOIN profiling_runs pr ON pr.job_execution_id = je.id + LEFT JOIN profiling_runs pr ON pr.id = je.id LEFT JOIN table_groups tg ON tg.id = pr.table_groups_id LEFT JOIN profile_anomalies pa ON pa.profile_run_id = pr.id WHERE je.job_key = 'run-profile' {" AND je.project_code = :project_code" if project_code else ""} {" AND tg.id = :table_group_id" if table_group_id else ""} {" AND je.id = :job_execution_id" if job_execution_id else ""} + {" AND je.job_schedule_id = :schedule_id" if schedule_id else ""} {" AND je.status IN :statuses" if statuses else ""} ORDER BY je.created_at DESC LIMIT :limit OFFSET :offset; @@ -280,6 +269,7 @@ def select_summary( "project_code": project_code, "table_group_id": str(table_group_id) if table_group_id else None, "job_execution_id": str(job_execution_id) if job_execution_id else None, + "schedule_id": schedule_id, "statuses": tuple(statuses) if statuses else (), "limit": page_size, "offset": (page - 1) * page_size, @@ -346,7 +336,7 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U """Check whether any active profiling job exists for the given entity or entities.""" query = ( select(func.count(cls.id)) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where(JobExecution.status.in_(cls._ACTIVE_JOB_STATUSES)) ) if entity_cls is cls: @@ -361,33 +351,15 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U raise ValueError(f"Unsupported entity: {entity_cls.__name__}") return get_current_session().execute(query).scalar() > 0 - @classmethod - def cancel_run(cls, run_id: str | UUID) -> None: - query = update(cls).where(cls.id == run_id).values(status="Cancelled", profiling_endtime=datetime.now(UTC)) - db_session = get_current_session() - db_session.execute(query) - @classmethod def cascade_delete(cls, ids: list[str]) -> None: - query = """ - DELETE FROM profile_pair_rules - WHERE profile_run_id IN :profiling_run_ids; + """Delete runs and their results by removing the parent job executions. - DELETE FROM profile_anomaly_results - WHERE profile_run_id IN :profiling_run_ids; - - DELETE FROM profile_results - WHERE profile_run_id IN :profiling_run_ids; - - DELETE FROM job_executions - WHERE id IN ( - SELECT job_execution_id FROM profiling_runs - WHERE id IN :profiling_run_ids AND job_execution_id IS NOT NULL - ); + The run id is the job execution id, and the foreign-key chain + (profile_results / profile_anomaly_results -> profiling_runs -> + job_executions) cascades on delete. """ - db_session = get_current_session() - db_session.execute(text(query), {"profiling_run_ids": tuple(ids)}) - cls.delete_where(cls.id.in_(ids)) + get_current_session().execute(delete(JobExecution).where(JobExecution.id.in_(ids))) @classmethod def find_latest_per_table_group(cls, project_code: str) -> set[UUID]: @@ -404,7 +376,7 @@ def find_latest_per_table_group(cls, project_code: str) -> set[UUID]: """ rows = get_current_session().scalars( select(cls.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( cls.project_code == project_code, JobExecution.status == JobStatus.COMPLETED, @@ -443,7 +415,7 @@ def delete_older_than( if protected_ids: where_clauses.append(cls.id.notin_(protected_ids)) - base_select = select(cls.id).join(JobExecution, cls.job_execution_id == JobExecution.id) + base_select = select(cls.id).join(JobExecution, cls.id == JobExecution.id) if dry_run: return get_current_session().scalar( @@ -460,18 +432,6 @@ def delete_older_than( total += len(ids) return total - @classmethod - def get_job_execution_ids(cls, profiling_run_ids: list[UUID]) -> dict[UUID, UUID | None]: - """Map profiling_run PKs to their job_execution_ids (batch lookup). - - Mirrors TestRun.get_job_execution_ids. - """ - if not profiling_run_ids: - return {} - query = select(cls.id, cls.job_execution_id).where(cls.id.in_(profiling_run_ids)) - rows = get_current_session().execute(query).all() - return {row.id: row.job_execution_id for row in rows} - def init_progress(self) -> None: self._progress = { "data_chars": {"label": "Refreshing data catalog"}, @@ -495,7 +455,7 @@ def set_progress(self, key: ProgressKey, status: ProgressStatus, detail: str | N def get_previous(self) -> Self | None: query = ( select(ProfilingRun) - .join(JobExecution, ProfilingRun.job_execution_id == JobExecution.id) + .join(JobExecution, ProfilingRun.id == JobExecution.id) .where( ProfilingRun.table_groups_id == self.table_groups_id, JobExecution.status == JobStatus.COMPLETED, @@ -511,7 +471,7 @@ def list_recent_complete(cls, table_groups_id: UUID, limit: int) -> list[Self]: """Return the most recent completed profiling runs for a table group, newest first.""" query = ( select(cls) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( cls.table_groups_id == table_groups_id, JobExecution.status == JobStatus.COMPLETED, diff --git a/testgen/common/models/project.py b/testgen/common/models/project.py index 40d04974..bb132802 100644 --- a/testgen/common/models/project.py +++ b/testgen/common/models/project.py @@ -66,19 +66,23 @@ def get_summary(cls, project_code: str) -> ProjectSummary | None: WHERE table_groups.project_code = :project_code ) AS profiling_run_count, ( - SELECT COUNT(*) FROM test_suites WHERE test_suites.project_code = :project_code + SELECT COUNT(*) FROM test_suites + WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_suite_count, ( SELECT COUNT(*) FROM test_definitions LEFT JOIN test_suites ON test_definitions.test_suite_id = test_suites.id WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_definition_count, ( SELECT COUNT(*) FROM test_runs LEFT JOIN test_suites ON test_runs.test_suite_id = test_suites.id WHERE test_suites.project_code = :project_code + AND test_suites.is_monitor IS NOT TRUE ) AS test_run_count, ( SELECT COALESCE(observability_api_key, '') <> '' diff --git a/testgen/common/models/scheduler.py b/testgen/common/models/scheduler.py index d1c03657..8b9e1e63 100644 --- a/testgen/common/models/scheduler.py +++ b/testgen/common/models/scheduler.py @@ -135,6 +135,14 @@ def select_active_by_kwargs( query = query.where(cls.kwargs[k].astext == str(v)) return list(get_current_session().scalars(query).all()) + @classmethod + def get_for_monitor_suite(cls, monitor_suite_id: str | UUID) -> Self | None: + """The run-monitors schedule for a monitor suite, active or paused.""" + return cls.get( + cls.key == RUN_MONITORS_JOB_KEY, + cls.kwargs["test_suite_id"].astext == str(monitor_suite_id), + ) + @classmethod def upsert_for_retention( cls, diff --git a/testgen/common/models/scores.py b/testgen/common/models/scores.py index 4dfb4c10..7b2e8d40 100644 --- a/testgen/common/models/scores.py +++ b/testgen/common/models/scores.py @@ -52,6 +52,7 @@ "stakeholder_group", "transform_level", "data_product", + "data_classification", ] Categories = Literal[ "column_name", @@ -68,6 +69,7 @@ "stakeholder_group", "transform_level", "data_product", + "data_classification", ] ScoreTypes = Literal["score", "cde_score"] @@ -88,6 +90,7 @@ class ScoreCategory(enum.Enum): dq_dimension = "dq_dimension" impact_dimension = "impact_dimension" data_product = "data_product" + data_classification = "data_classification" class ScoreDefinition(Base): @@ -826,6 +829,7 @@ class ScoreDefinitionBreakdownItem(Base): stakeholder_group: str = Column(String, nullable=True) transform_level: str = Column(String, nullable=True) data_product: str = Column(String, nullable=True) + data_classification: str = Column(String, nullable=True) impact: float = Column(Float) score: float = Column(Float) issue_ct: int = Column(Integer) diff --git a/testgen/common/models/table_group.py b/testgen/common/models/table_group.py index 251dcc46..095f54e3 100644 --- a/testgen/common/models/table_group.py +++ b/testgen/common/models/table_group.py @@ -47,6 +47,23 @@ class TableGroupStats(EntityMinimal): data_point_ct: int +@dataclass +class TableGroupListItem(EntityMinimal): + id: UUID + table_groups_name: str + table_group_schema: str + project_code: str + connection_name: str | None + table_count: int + column_count: int | None + row_count: int | None + last_profiled_date: datetime | None + last_tested_date: datetime | None + profiling_score: float | None + testing_score: float | None + quality_score: float | None + + @dataclass class TableGroupSummary(EntityMinimal): id: UUID @@ -90,6 +107,118 @@ class TableGroupSummary(EntityMinimal): total_count: int = 0 +@dataclass +class MonitorTableSummary: + """One row per monitored table in a table group's lookback window. + + Per-type ``*_anomalies`` count result_code = 0 results; ``*_is_training`` reflects the + latest run's training-mode state (result_code = -1); ``*_is_pending`` is True when no + result of that type exists in the window (monitor not yet configured / executed). + Schema is special-cased: ``*_is_pending`` only when no events at all, and ``table_state`` + captures whether the table was added / dropped / column-modified in the window. + """ + table_name: str + lookback: int + lookback_start: datetime | None + lookback_end: datetime | None + freshness_anomalies: int + volume_anomalies: int + schema_anomalies: int + metric_anomalies: int + freshness_is_training: bool | None + volume_is_training: bool | None + metric_is_training: bool | None + freshness_is_pending: bool + volume_is_pending: bool + schema_is_pending: bool + metric_is_pending: bool + freshness_error_message: str | None + volume_error_message: str | None + schema_error_message: str | None + metric_error_message: str | None + latest_update: datetime | None + row_count: int | None + previous_row_count: int | None + column_adds: int + column_drops: int + column_mods: int + table_state: str | None + + +@dataclass +class MonitorGroupSummary: + """Aggregated monitor health for a single table group across its lookback window. + + Booleans are group-wide: ``*_has_errors`` is True iff at least one monitored table errored; + ``*_is_training`` is True iff every monitored table is in training (and at least one is); + ``*_is_pending`` is True iff every monitored table is pending. + """ + lookback: int + lookback_start: datetime | None + lookback_end: datetime | None + total_monitored_tables: int + freshness_anomalies: int + volume_anomalies: int + schema_anomalies: int + metric_anomalies: int + freshness_has_errors: bool + volume_has_errors: bool + schema_has_errors: bool + metric_has_errors: bool + freshness_is_training: bool + volume_is_training: bool + metric_is_training: bool + freshness_is_pending: bool + volume_is_pending: bool + schema_is_pending: bool + metric_is_pending: bool + + +_ANOMALY_TYPE_TO_COLUMN: dict[str, str] = { + "Freshness_Trend": "freshness_anomalies", + "Volume_Trend": "volume_anomalies", + "Schema_Drift": "schema_anomalies", + "Metric_Trend": "metric_anomalies", +} + +_MONITOR_SORT_COLUMN: dict[str, str] = { + "table_name": "LOWER(monitor_tables.table_name)", + "freshness_anomalies": "monitor_tables.freshness_anomalies", + "volume_anomalies": "monitor_tables.volume_anomalies", + "schema_anomalies": "monitor_tables.schema_anomalies", + "metric_anomalies": "monitor_tables.metric_anomalies", + "total_anomalies": ( + "monitor_tables.freshness_anomalies + monitor_tables.volume_anomalies" + " + monitor_tables.schema_anomalies + monitor_tables.metric_anomalies" + ), + "latest_update": "monitor_tables.latest_update", + "row_count": "monitor_tables.row_count", + "row_count_change": "(monitor_tables.row_count - baseline_tables.previous_row_count)", +} + + +def _build_monitor_order_clause(sort_by: str | None, anomaly_type: str | None) -> str: + """Build an ORDER BY clause for the monitor-changes-by-tables query. + + ``sort_by`` is a field name with optional ``_desc`` suffix. When ``sort_by`` is + ``total_anomalies_desc`` and ``anomaly_type`` is set, the sort collapses to that + type's column so callers see "most anomalies of the filtered type first." Falls + back to ``LOWER(table_name) ASC`` for unknown / missing values. + """ + descending = False + field = sort_by + if field and field.endswith("_desc"): + descending = True + field = field[: -len("_desc")] + + if field == "total_anomalies" and anomaly_type is not None: + field = _ANOMALY_TYPE_TO_COLUMN[anomaly_type] + + column = _MONITOR_SORT_COLUMN.get(field or "table_name", _MONITOR_SORT_COLUMN["table_name"]) + direction = "DESC" if descending else "ASC" + return f"ORDER BY {column} {direction} NULLS LAST" + + class TableGroup(Entity): __tablename__ = "table_groups" @@ -132,6 +261,7 @@ class TableGroup(Entity): stakeholder_group: str = Column(NullIfEmptyString) transform_level: str = Column(NullIfEmptyString) data_product: str = Column(NullIfEmptyString) + data_classification: str = Column(NullIfEmptyString) last_complete_profile_run_id: UUID = Column(postgresql.UUID(as_uuid=True)) dq_score_profiling: float = Column(Float) dq_score_testing: float = Column(Float) @@ -149,6 +279,20 @@ class TableGroup(Entity): dq_score_testing, ) + @property + def quality_score(self) -> float: + """Overall quality score per ``utils.score``: profiling * testing when both + exist; the non-null one when only one ran; ``0.0`` when neither has run. + + Mirrors what Project Dashboard, data_catalog, and inventory_service display. + Render through ``utils.friendly_score`` to get the user-facing percentage form + (``"95.0"`` rather than ``"0.95"``). The list-side equivalent is computed in + SQL inside ``_list_with_activity`` — keep both in sync if the formula changes. + """ + from testgen.utils import score + + return score(self.dq_score_profiling, self.dq_score_testing) + @classmethod def get_minimal(cls, id_: str | UUID) -> TableGroupMinimal | None: result = cls._get_columns(id_, cls._minimal_columns) @@ -224,7 +368,6 @@ def select_summary( latest_profile AS ( SELECT latest_run.table_groups_id, latest_run.id, - latest_run.job_execution_id, MAX(latest_je.started_at) AS started_at, latest_run.anomaly_ct, SUM( @@ -259,7 +402,7 @@ def select_summary( groups.last_complete_profile_run_id = latest_run.id ) LEFT JOIN job_executions latest_je ON ( - latest_run.job_execution_id = latest_je.id + latest_run.id = latest_je.id ) LEFT JOIN profile_anomaly_results latest_anomalies ON ( latest_run.id = latest_anomalies.profile_run_id @@ -334,7 +477,7 @@ def select_summary( groups.dq_score_profiling, groups.dq_score_testing, latest_profile.id AS latest_profile_id, - latest_profile.job_execution_id AS latest_profile_job_execution_id, + latest_profile.id AS latest_profile_job_execution_id, latest_profile.started_at AS latest_profile_start, latest_profile.anomaly_ct AS latest_hygiene_issues_ct, latest_profile.definite_ct AS latest_hygiene_issues_definite_ct, @@ -385,6 +528,421 @@ def select_summary( total = items[0].total_count if items else 0 return items, total + @classmethod + def list_for_project( + cls, project_code: str, *, page: int = 1, limit: int = 20 + ) -> tuple[list[TableGroupListItem], int]: + """Config-focused paginated listing for a project, with table count and activity timestamps.""" + return cls._list_with_activity( + scope_sql="groups.project_code = :project_code", + scope_params={"project_code": project_code}, + page=page, + limit=limit, + ) + + @classmethod + def list_for_connection( + cls, connection_id: int, *, page: int = 1, limit: int = 20 + ) -> tuple[list[TableGroupListItem], int]: + """Config-focused paginated listing for a connection, with table count and activity timestamps.""" + return cls._list_with_activity( + scope_sql="groups.connection_id = :connection_id", + scope_params={"connection_id": connection_id}, + page=page, + limit=limit, + ) + + @classmethod + def _list_with_activity( + cls, + *, + scope_sql: str, + scope_params: dict, + page: int, + limit: int, + ) -> tuple[list[TableGroupListItem], int]: + session = get_current_session() + + # Separate COUNT(*) query — keeps `total` correct on out-of-range pages where + # the page rows are empty (a window function over zero rows would return 0). + total_query = f"SELECT COUNT(*) FROM table_groups AS groups WHERE {scope_sql};" + total = session.execute(text(total_query), scope_params).scalar() or 0 + if total == 0: + return [], 0 + + params: dict = {**scope_params, "limit": limit, "offset": (page - 1) * limit} + rows_query = f""" + WITH stats AS ( + SELECT table_groups_id, + COUNT(*) AS table_count, + SUM(column_ct) AS column_count, + SUM(COALESCE(record_ct, approx_record_ct)) AS row_count + FROM data_table_chars + GROUP BY table_groups_id + ), + latest_profile AS ( + SELECT pr.table_groups_id, MAX(je.started_at) AS started_at + FROM profiling_runs pr + LEFT JOIN job_executions je ON je.id = pr.id + GROUP BY pr.table_groups_id + ), + latest_test AS ( + SELECT ts.table_groups_id, MAX(tr.test_starttime) AS test_starttime + FROM test_runs tr + JOIN test_suites ts ON ts.id = tr.test_suite_id + WHERE ts.is_monitor IS NOT TRUE + GROUP BY ts.table_groups_id + ) + SELECT + groups.id, + groups.table_groups_name, + groups.table_group_schema, + groups.project_code, + connections.connection_name, + COALESCE(stats.table_count, 0) AS table_count, + stats.column_count, + stats.row_count, + latest_profile.started_at AS last_profiled_date, + latest_test.test_starttime AS last_tested_date, + groups.dq_score_profiling AS profiling_score, + groups.dq_score_testing AS testing_score, + -- Mirrors utils.score: product when both exist, else the non-null one, else NULL. + CASE + WHEN groups.dq_score_profiling IS NOT NULL AND groups.dq_score_testing IS NOT NULL + THEN groups.dq_score_profiling * groups.dq_score_testing + ELSE COALESCE(groups.dq_score_profiling, groups.dq_score_testing) + END AS quality_score + FROM table_groups AS groups + LEFT JOIN connections ON connections.connection_id = groups.connection_id + LEFT JOIN stats ON stats.table_groups_id = groups.id + LEFT JOIN latest_profile ON latest_profile.table_groups_id = groups.id + LEFT JOIN latest_test ON latest_test.table_groups_id = groups.id + WHERE {scope_sql} + ORDER BY LOWER(groups.table_groups_name) + LIMIT :limit OFFSET :offset; + """ + rows = session.execute(text(rows_query), params).mappings().all() + items = [TableGroupListItem(**row) for row in rows] + return items, total + + @classmethod + def list_monitor_table_summaries( + cls, + table_group_id: str | UUID, + *, + anomaly_types: list[str] | None = None, + sort_by: str | None = None, + lookback_override: int | None = None, + table_name_filter: str | None = None, + page: int = 1, + limit: int = 20, + ) -> tuple[list[MonitorTableSummary], int]: + """Per-monitored-table summary, paginated within the group's lookback window. + + ``anomaly_types`` are internal ``test_type`` values (``Freshness_Trend`` / + ``Volume_Trend`` / ``Schema_Drift`` / ``Metric_Trend``) — filters to tables with + at least one anomaly of any listed type in the window. ``lookback_override`` + replaces the suite-configured lookback for ad-hoc views. ``sort_by`` accepts + the dataclass field names (e.g. ``freshness_anomalies``, ``latest_update``, + ``row_count``, ``total_anomalies``) suffixed with ``_desc`` for descending + order; ``table_name`` is the default and sorts case-insensitively. + """ + # Run the CTE twice (rows + COUNT) so the total reflects every matching table, + # not just rows on the requested page — a window function would short-cut to 0 + # on out-of-range pages. + query, params = cls._monitor_changes_by_tables_query( + table_group_id, + anomaly_types=anomaly_types, + sort_by=sort_by, + lookback_override=lookback_override, + table_name_filter=table_name_filter, + limit=limit, + offset=(page - 1) * limit, + ) + count_query, count_params = cls._monitor_changes_by_tables_query( + table_group_id, + anomaly_types=anomaly_types, + lookback_override=lookback_override, + table_name_filter=table_name_filter, + ) + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + total = session.scalar( + text(f"SELECT COUNT(*) FROM ({count_query}) AS subquery"), count_params + ) or 0 + return [MonitorTableSummary(**row) for row in rows], int(total) + + @classmethod + def get_monitor_group_summary( + cls, + table_group_id: str | UUID, + *, + lookback_override: int | None = None, + ) -> MonitorGroupSummary: + """Group-level monitor health across the lookback window. + + Aggregates per-table results from ``list_monitor_table_summaries`` into a single + row. Returns zeroed counts and all-pending booleans when no tables are + monitored — callers detect "not monitored" upstream via the linked suite. + """ + inner_query, params = cls._monitor_changes_by_tables_query( + table_group_id, lookback_override=lookback_override, + ) + query = f""" + SELECT + COALESCE(MAX(lookback), 0)::INTEGER AS lookback, + MIN(lookback_start) AS lookback_start, + MAX(lookback_end) AS lookback_end, + COUNT(*)::INTEGER AS total_monitored_tables, + COALESCE(SUM(freshness_anomalies), 0)::INTEGER AS freshness_anomalies, + COALESCE(SUM(volume_anomalies), 0)::INTEGER AS volume_anomalies, + COALESCE(SUM(schema_anomalies), 0)::INTEGER AS schema_anomalies, + COALESCE(SUM(metric_anomalies), 0)::INTEGER AS metric_anomalies, + COALESCE(BOOL_OR(freshness_error_message IS NOT NULL), FALSE) AS freshness_has_errors, + COALESCE(BOOL_OR(volume_error_message IS NOT NULL), FALSE) AS volume_has_errors, + COALESCE(BOOL_OR(schema_error_message IS NOT NULL), FALSE) AS schema_has_errors, + COALESCE(BOOL_OR(metric_error_message IS NOT NULL), FALSE) AS metric_has_errors, + COALESCE( + BOOL_OR(freshness_is_training) AND BOOL_AND(freshness_is_training OR freshness_is_pending), + FALSE + ) AS freshness_is_training, + COALESCE( + BOOL_OR(volume_is_training) AND BOOL_AND(volume_is_training OR volume_is_pending), + FALSE + ) AS volume_is_training, + COALESCE( + BOOL_OR(metric_is_training) AND BOOL_AND(metric_is_training OR metric_is_pending), + FALSE + ) AS metric_is_training, + COALESCE(BOOL_AND(freshness_is_pending), TRUE) AS freshness_is_pending, + COALESCE(BOOL_AND(volume_is_pending), TRUE) AS volume_is_pending, + COALESCE(BOOL_AND(schema_is_pending), TRUE) AS schema_is_pending, + COALESCE(BOOL_AND(metric_is_pending), TRUE) AS metric_is_pending + FROM ({inner_query}) AS subquery + """ + # Outer query has no GROUP BY — aggregates over zero rows still yield one + # COALESCE'd row, so .first() never returns None here. ``lookback`` is 0 + # when the per-table CTE has no rows OR no runs against the monitor suite, + # so the dashboard / MCP can render the "no monitor runs yet" state. + row = get_current_session().execute(text(query), params).mappings().first() + return MonitorGroupSummary(**row) + + @classmethod + def _monitor_changes_by_tables_query( + cls, + table_group_id: str | UUID, + *, + anomaly_types: list[str] | None = None, + sort_by: str | None = None, + lookback_override: int | None = None, + table_name_filter: str | None = None, + limit: int | None = None, + offset: int | None = None, + ) -> tuple[str, dict]: + """Build the CTE that produces one ``MonitorTableSummary``-shaped row per table. + + Shared by ``list_monitor_table_summaries`` and ``get_monitor_group_summary`` + and by the monitors dashboard. ``anomaly_types`` filters the outer SELECT (so + the group summary can omit it and get all tables). ``sort_by`` is the + dataclass field name with optional ``_desc`` suffix — the caller (MCP / + dashboard) is responsible for validating it against a known set. + """ + lookback_expr = ( + ":lookback_override" if lookback_override is not None + else "COALESCE(test_suites.monitor_lookback, 1)" + ) + + anomaly_filter_clause = "" + if anomaly_types: + columns = [_ANOMALY_TYPE_TO_COLUMN[t] for t in anomaly_types] + anomaly_filter_clause = ( + "WHERE (" + + " OR ".join(f"monitor_tables.{col} > 0" for col in columns) + + ")" + ) + + sort_anomaly_type = anomaly_types[0] if anomaly_types and len(anomaly_types) == 1 else None + order_clause = _build_monitor_order_clause(sort_by, sort_anomaly_type) + + query = f""" + WITH ranked_test_runs AS ( + SELECT + test_runs.id, + test_runs.test_starttime, + {lookback_expr} AS lookback, + ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position + FROM table_groups + INNER JOIN test_runs + ON (test_runs.test_suite_id = table_groups.monitor_test_suite_id) + INNER JOIN test_suites + ON (table_groups.monitor_test_suite_id = test_suites.id) + WHERE table_groups.id = :table_group_id + ), + lookback_window AS ( + SELECT MIN(test_starttime) AS lookback_start + FROM ranked_test_runs + WHERE position <= lookback + ), + latest_tables AS ( + SELECT DISTINCT + table_chars.schema_name, + table_chars.table_name + FROM data_table_chars table_chars + CROSS JOIN lookback_window + WHERE table_chars.table_groups_id = :table_group_id + -- Include current tables and tables dropped within lookback window + AND (table_chars.drop_date IS NULL OR table_chars.drop_date >= lookback_window.lookback_start) + {"AND table_chars.table_name ILIKE :table_name_filter" if table_name_filter else ''} + ), + monitor_results AS ( + SELECT + latest_tables.table_name, + results.test_time, + results.test_type, + results.result_code, + ranked_test_runs.lookback, + ranked_test_runs.position, + ranked_test_runs.test_starttime, + -- result_code = -1 indicates training mode + CASE WHEN results.result_code = -1 THEN 1 ELSE 0 END AS is_training, + CASE WHEN results.test_type = 'Freshness_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS freshness_anomaly, + CASE WHEN results.test_type = 'Volume_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS volume_anomaly, + CASE WHEN results.test_type = 'Schema_Drift' AND results.result_code = 0 THEN 1 ELSE 0 END AS schema_anomaly, + CASE WHEN results.test_type = 'Metric_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS metric_anomaly, + CASE WHEN results.test_type = 'Freshness_Trend' THEN results.result_signal ELSE NULL END AS freshness_interval, + CASE WHEN results.test_type = 'Volume_Trend' THEN results.result_signal::BIGINT ELSE NULL END AS row_count, + CASE WHEN results.test_type = 'Schema_Drift' THEN SPLIT_PART(results.result_signal, '|', 1) ELSE NULL END AS table_change, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 2), '')::INT ELSE 0 END AS col_adds, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 3), '')::INT ELSE 0 END AS col_drops, + CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 4), '')::INT ELSE 0 END AS col_mods, + CASE WHEN results.result_status = 'Error' THEN results.result_message ELSE NULL END AS error_message + FROM latest_tables + LEFT JOIN ranked_test_runs ON TRUE + LEFT JOIN test_results AS results + ON results.test_run_id = ranked_test_runs.id + AND results.table_name = latest_tables.table_name + WHERE ranked_test_runs.position IS NULL + -- Also capture 1 run before the lookback to get baseline results + OR ranked_test_runs.position <= ranked_test_runs.lookback + 1 + ), + monitor_tables AS ( + SELECT + table_name, + MAX(lookback)::INTEGER AS lookback, + COALESCE(SUM(freshness_anomaly), 0)::INTEGER AS freshness_anomalies, + COALESCE(SUM(volume_anomaly), 0)::INTEGER AS volume_anomalies, + COALESCE(SUM(schema_anomaly), 0)::INTEGER AS schema_anomalies, + COALESCE(SUM(metric_anomaly), 0)::INTEGER AS metric_anomalies, + MAX(test_time - (COALESCE(NULLIF(freshness_interval, 'Unknown')::INTEGER, 0) * INTERVAL '1 minute')) + FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS latest_update, + MAX(row_count) FILTER (WHERE position = 1) AS row_count, + COALESCE(SUM(col_adds), 0)::INTEGER AS column_adds, + COALESCE(SUM(col_drops), 0)::INTEGER AS column_drops, + COALESCE(SUM(col_mods), 0)::INTEGER AS column_mods, + MAX(error_message) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Schema_Drift' AND position = 1) AS schema_error_message, + MAX(error_message) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_error_message, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_is_training, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_is_training, + BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_is_training, + BOOL_OR(test_type = 'Freshness_Trend') IS NOT TRUE AS freshness_is_pending, + BOOL_OR(test_type = 'Volume_Trend') IS NOT TRUE AS volume_is_pending, + -- Schema monitor only creates results on schema changes (Failed) + -- Mark it as pending only if there are no results of any test type + BOOL_OR(test_time IS NOT NULL) IS NOT TRUE AS schema_is_pending, + BOOL_OR(test_type = 'Metric_Trend') IS NOT TRUE AS metric_is_pending, + CASE + -- Mark as Dropped if latest Schema Drift result for the table indicates it was dropped + WHEN (ARRAY_AGG(table_change ORDER BY test_time DESC) FILTER (WHERE table_change IS NOT NULL))[1] = 'D' + THEN 'dropped' + -- Only mark as Added if latest change does not indicate a drop + WHEN MAX(CASE WHEN table_change = 'A' THEN 1 ELSE 0 END) = 1 + THEN 'added' + WHEN SUM(schema_anomaly) > 0 + THEN 'modified' + ELSE NULL + END AS table_state + FROM monitor_results + -- Only aggregate within lookback runs + WHERE position IS NULL OR position <= COALESCE(lookback, 1) + GROUP BY table_name + ), + table_bounds AS ( + SELECT + table_name, + MIN(position) AS min_position, + MAX(position) AS max_position + FROM monitor_results + WHERE position IS NOT NULL + GROUP BY table_name + ), + baseline_tables AS ( + SELECT + monitor_results.table_name, + MIN(monitor_results.test_starttime) FILTER ( + WHERE monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) + ) AS lookback_start, + MAX(monitor_results.test_starttime) FILTER ( + WHERE monitor_results.position = GREATEST(1, table_bounds.min_position) + ) AS lookback_end, + MAX(monitor_results.row_count) FILTER ( + WHERE monitor_results.test_type = 'Volume_Trend' + AND monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) + ) AS previous_row_count + FROM monitor_results + JOIN table_bounds ON monitor_results.table_name = table_bounds.table_name + GROUP BY monitor_results.table_name + ) + SELECT + monitor_tables.table_name, + monitor_tables.lookback, + baseline_tables.lookback_start, + baseline_tables.lookback_end, + monitor_tables.freshness_anomalies, + monitor_tables.volume_anomalies, + monitor_tables.schema_anomalies, + monitor_tables.metric_anomalies, + monitor_tables.freshness_is_training, + monitor_tables.volume_is_training, + monitor_tables.metric_is_training, + monitor_tables.freshness_is_pending, + monitor_tables.volume_is_pending, + monitor_tables.schema_is_pending, + monitor_tables.metric_is_pending, + monitor_tables.freshness_error_message, + monitor_tables.volume_error_message, + monitor_tables.schema_error_message, + monitor_tables.metric_error_message, + monitor_tables.latest_update, + monitor_tables.row_count, + baseline_tables.previous_row_count, + monitor_tables.column_adds, + monitor_tables.column_drops, + monitor_tables.column_mods, + monitor_tables.table_state + FROM monitor_tables + LEFT JOIN baseline_tables ON monitor_tables.table_name = baseline_tables.table_name + {anomaly_filter_clause} + {order_clause} + {"LIMIT :limit" if limit is not None else ""} + {"OFFSET :offset" if offset is not None else ""} + """ + + escaped_name_filter = ( + table_name_filter.replace("_", "\\_") if table_name_filter else None + ) + params: dict = {"table_group_id": str(table_group_id)} + if escaped_name_filter is not None: + params["table_name_filter"] = f"%{escaped_name_filter}%" + if lookback_override is not None: + params["lookback_override"] = lookback_override + if limit is not None: + params["limit"] = limit + if offset is not None: + params["offset"] = offset + return query, params + @classmethod def is_in_use(cls, ids: list[str]) -> bool: test_suites = TestSuite.select_minimal_where(TestSuite.table_groups_id.in_(ids)) @@ -417,8 +975,8 @@ def cascade_delete(cls, ids: list[str]) -> None: DELETE FROM job_executions WHERE id IN ( - SELECT pr.job_execution_id FROM profiling_runs pr - WHERE pr.table_groups_id IN :table_group_ids AND pr.job_execution_id IS NOT NULL + SELECT pr.id FROM profiling_runs pr + WHERE pr.table_groups_id IN :table_group_ids ); DELETE FROM profiling_runs pr diff --git a/testgen/common/models/test_definition.py b/testgen/common/models/test_definition.py index ec9ddde0..1479ef2e 100644 --- a/testgen/common/models/test_definition.py +++ b/testgen/common/models/test_definition.py @@ -1,3 +1,4 @@ +import json from collections.abc import Iterable from dataclasses import dataclass from datetime import UTC, datetime @@ -9,6 +10,7 @@ from sqlalchemy import ( Boolean, Column, + Enum, ForeignKey, String, Text, @@ -25,6 +27,7 @@ from sqlalchemy.orm import InstrumentedAttribute from sqlalchemy.sql.expression import case, literal +from testgen.common.enums import MonitorType from testgen.common.models import Base, get_current_session from testgen.common.models.custom_types import NullIfEmptyString, YNString, ZeroIfEmptyInteger from testgen.common.models.entity import Entity, EntityMinimal @@ -40,6 +43,87 @@ class Severity(StrEnum): WARNING = "Warning" +class TestAlgorithm(StrEnum): + """SQL-derived algorithm family for a test type (faceted picker axis).""" + + BOUNDARY_CHECK = "Boundary check" + COUNTING = "Counting" + PATTERN_REGEX = "Pattern / regex" + SET_LOOKUP = "Set / lookup" + STATISTICAL_DRIFT = "Statistical drift" + AGGREGATE_RECONCILIATION = "Aggregate reconciliation" + FRESHNESS_TIME = "Freshness / time" + SCHEMA_METADATA = "Schema / metadata" + CUSTOM_SQL = "Custom SQL" + + +class StatisticalTechnique(StrEnum): + """Named statistical technique a test type uses to evaluate its measure.""" + + COHENS_D = "Cohen's D" + COHENS_H = "Cohen's H" + OUTLIER_DETECTION = "Outlier Detection" + SD_SHIFT = "SD Shift" + JENSEN_SHANNON_DIVERGENCE = "Jensen-Shannon Divergence" + PREDICTIVE_MODEL = "Predictive Model" + + +class TestCriteria(StrEnum): + """What a test type needs to be set up (faceted picker axis). + + Derived, not stored: ``derive_test_criteria`` is the single source of truth, shared by the + UI lookup and MCP so the value never drifts between surfaces. + """ + + DEFINED_RULE = "Defined Rule" + DEFINED_THRESHOLD = "Defined Threshold" + DEFINED_VALUE = "Defined Value" + LIST_OF_VALUES = "List of Values" + REFERENCE_DATASET = "Reference Dataset" + CUSTOM_CRITERIA = "Custom Criteria" + + +# Predefined validity/integrity rules — the user just enables them; the rule itself is fixed +# (dedup/uniqueness and standard format validators). Not separable from structural attributes +# alone: e.g. Valid_Month is structurally identical to the Defined-Value test Pattern_Match +# (both column-scoped, Pattern / regex, with baseline_value + threshold_value params). +_DEFINED_RULE_TESTS = frozenset({ + "Dupe_Rows", "Unique", "Email_Format", "Street_Addr_Pattern", + "Valid_Characters", "Valid_Month", "Valid_US_Zip", "Valid_US_Zip3", +}) +# The user asserts the expected value/pattern the column should hold (a constant, a regex +# baseline, or simply that a value is present). Enumerated for the same reason as above. +_DEFINED_VALUE_TESTS = frozenset({ + "Constant", "Pattern_Match", "Required", +}) + + +def derive_test_criteria( + test_type: str, + test_scope: str | None, + algorithm: str | None, +) -> TestCriteria: + """Classify a test type by the kind of setup it requires. + + Single source of truth for the Criteria facet — call from both the UI lookup and MCP rather + than reproducing the rules. Scope and algorithm resolve the cleanly-typed buckets (referential + is checked before Set / lookup so Combo_Match stays Reference Dataset). Defined Rule, Defined + Value, and Defined Threshold can't be told apart from structural attributes, so the first two + are enumerated and Defined Threshold is the fallthrough. + """ + if test_scope == "custom": + return TestCriteria.CUSTOM_CRITERIA + if test_scope == "referential": + return TestCriteria.REFERENCE_DATASET + if algorithm == TestAlgorithm.SET_LOOKUP: + return TestCriteria.LIST_OF_VALUES + if test_type in _DEFINED_RULE_TESTS: + return TestCriteria.DEFINED_RULE + if test_type in _DEFINED_VALUE_TESTS: + return TestCriteria.DEFINED_VALUE + return TestCriteria.DEFINED_THRESHOLD + + class InvalidTestDefinitionFields(ValueError): """Aggregated field-level validation errors. ``errors``: ``dict[field_name, reason]``.""" @@ -53,6 +137,28 @@ def _is_blank(value: object) -> bool: return value is None or value == "" +CUSTOM_METADATA_MAX_KEYS = 50 +CUSTOM_METADATA_MAX_BYTES = 10_240 + + +def validate_custom_metadata(value: object) -> str | None: + """Return an error message if ``value`` is not a valid ``custom_metadata`` payload, else ``None``. + + ``custom_metadata`` must be a JSON object (key-value pairs), bounded in key count and serialized + size. Shared by every write path — ``TestDefinition.validate`` (UI/CLI/MCP) and the + ``TestDefinitionExport`` import schema — so the rule has a single definition. + """ + if value is None: + return None + if not isinstance(value, dict): + return "must be a JSON object of key-value pairs" + if len(value) > CUSTOM_METADATA_MAX_KEYS: + return f"must have at most {CUSTOM_METADATA_MAX_KEYS} keys" + if len(json.dumps(value)) > CUSTOM_METADATA_MAX_BYTES: + return f"must be at most {CUSTOM_METADATA_MAX_BYTES} bytes when serialized as JSON" + return None + + class ParamFieldsMixin: """Parsed access to default_parm_columns/prompts/help metadata. @@ -142,6 +248,8 @@ class TestDefinitionSummary(TestTypeSummary): prediction: dict[str, dict[str, float]] | None flagged: bool impact_dimension: str | None + external_url: str | None + custom_metadata: dict | None @property def display_name(self) -> str: @@ -163,6 +271,97 @@ class TestDefinitionMinimal(EntityMinimal): test_name_short: str +class ThresholdMode(StrEnum): + """How a monitor's bounds are determined — derived from which fields on the + definition are populated. See ``TestDefinition._derive_threshold_mode``.""" + PREDICTION = "Prediction Model" + HISTORICAL = "Historical Calculation" + STATIC = "Static" + NONE = "N/A" + + +@dataclass +class MonitorForecastPoint(EntityMinimal): + """One future forecast point read from a Prediction-Model monitor's stored + ``prediction`` JSONB. Each point is a ``(test_time, lower_bound, upper_bound)`` + triple at a specific upcoming timestamp; collectively they extend the + historical event series forward under the suite's active prediction + sensitivity. Surfaced as a separate forecast section on + ``list_monitor_events``, never as an event.""" + test_time: datetime + lower_bound: float | None + upper_bound: float | None + + +def forecast_points_from_prediction( + prediction: dict | None, + sensitivity: str, +) -> list[MonitorForecastPoint]: + """Extract forecast points for a sensitivity from a monitor's stored + ``prediction`` JSONB, sorted by time ascending (nearest future point + first). + + The JSONB is keyed as ``"lower_tolerance|" → {epoch_ms: value}`` + and ``"upper_tolerance|" → {epoch_ms: value}`` — matches the + format the dashboard reads at ``monitors_dashboard.py`` via + ``datetime.fromtimestamp(int(timestamp) / 1000.0, UTC)``. Returns ``[]`` + when the monitor isn't in Prediction Model mode (no JSONB) or when the + sensitivity has no stored series. Standalone so it works against either + the ``TestDefinition`` ORM row or the ``TestDefinitionSummary`` dataclass + — both carry the same ``prediction`` field shape. + """ + if not prediction: + return [] + lower_series = prediction.get(f"lower_tolerance|{sensitivity}") or {} + upper_series = prediction.get(f"upper_tolerance|{sensitivity}") or {} + if not lower_series and not upper_series: + return [] + + all_keys = sorted(set(lower_series) | set(upper_series), key=int) + points: list[MonitorForecastPoint] = [] + for k in all_keys: + ts = datetime.fromtimestamp(int(k) / 1000.0, UTC) + lower = lower_series.get(k) + upper = upper_series.get(k) + points.append(MonitorForecastPoint( + test_time=ts, + lower_bound=float(lower) if lower is not None else None, + upper_bound=float(upper) if upper is not None else None, + )) + return points + + +@dataclass +class MonitorConfig(EntityMinimal): + """One configured monitor — produced by ``TestDefinition.list_monitor_configs_for_table``. + + ``threshold_mode`` is derived from the underlying definition's + ``history_calculation``: ``"PREDICT"`` flags Prediction Model; any other + non-empty value flags Historical Calculation (not available for Freshness); + empty falls through to Static — the default for Freshness, Volume, and + Metric. Schema_Drift is presence-only and reports N/A. + + ``threshold_lower`` / ``threshold_upper`` carry the bounds active under + the current mode — static tolerances for Static, history-calc expressions + (e.g. ``"Minimum"`` / ``"Maximum"``) for Historical, ``None`` for + Prediction (the runtime bands live on each event, not the configuration). + For Freshness in Static mode only the upper bound applies. + + ``metric_name`` is the user-defined name for a Metric monitor (stored on + the underlying definition's ``column_name`` column, but it is the metric's + name rather than a column reference). ``custom_query`` is the metric's SQL + expression. Both are only set for ``Metric_Trend``. + """ + monitor_id: UUID + test_type: str + table_name: str + metric_name: str | None + threshold_mode: ThresholdMode + threshold_lower: str | None + threshold_upper: str | None + custom_query: str | None + + class QueryString(TypeDecorator): impl = String cache_ok = True @@ -175,6 +374,15 @@ def process_bind_param(self, value: str | None, _dialect) -> str | None: return value or None +def _enum_by_value(enum_cls: type[StrEnum]) -> Enum: + """Map a StrEnum to its VARCHAR column by member value (not name) so reads return enum members. + + These columns store the display value (e.g. ``"Boundary check"``), which differs from the enum + member name, so the default name-based mapping would not round-trip. + """ + return Enum(enum_cls, native_enum=False, values_callable=lambda cls: [member.value for member in cls]) + + class TestType(ParamFieldsMixin, Entity): __tablename__ = "test_types" @@ -204,12 +412,19 @@ class TestType(ParamFieldsMixin, Entity): dq_dimension: str = Column(String) impact_dimension: str = Column(String) health_dimension: str = Column(String) + algorithm: TestAlgorithm | None = Column(_enum_by_value(TestAlgorithm)) + statistical_technique: StatisticalTechnique | None = Column(_enum_by_value(StatisticalTechnique)) threshold_description: str = Column(String) usage_notes: str = Column(String) active: str = Column(String) # Unmapped columns: generation_template, result_visualization, result_visualization_params + @property + def criteria(self) -> TestCriteria: + """Setup-kind facet, derived via the shared classifier (see ``derive_test_criteria``).""" + return derive_test_criteria(self.test_type, self.test_scope, self.algorithm) + _summary_columns = ( *[key for key in TestTypeSummary.__annotations__.keys() if key not in ("default_test_description", "default_impact_dimension")], test_description.label("default_test_description"), @@ -299,6 +514,8 @@ class TestDefinition(Entity): flagged: bool = Column(Boolean, default=False, nullable=False) external_id: UUID | None = Column(postgresql.UUID(as_uuid=True)) impact_dimension: str | None = Column(String, nullable=True) + external_url: str | None = Column(NullIfEmptyString) + custom_metadata: dict | None = Column(postgresql.JSONB) _default_order_by = ( asc(func.lower(schema_name)), @@ -432,6 +649,95 @@ def list_for_suite( query = query.order_by(*cls._default_order_by) return cls._paginate(query, page=page, limit=limit, data_class=TestDefinitionSummary) + @classmethod + def get_singleton_monitor( + cls, + test_suite_id: str | UUID, + table_name: str, + test_type: str, + ) -> "TestDefinition | None": + """Return the single ``TestDefinition`` row for a singleton monitor + type (Freshness / Volume / Schema) on a given table. Metric is + multi-instance and should be looked up by ``id`` instead — this + helper would silently pick the first match. Returns ``None`` when no + monitor is configured.""" + session = get_current_session() + return session.execute( + select(cls) + .where(cls.test_suite_id == test_suite_id) + .where(cls.table_name == table_name) + .where(cls.test_type == test_type) + .limit(1) + ).scalars().first() + + @classmethod + def list_monitor_configs_for_table( + cls, + test_suite_id: str | UUID, + table_name: str, + ) -> list[MonitorConfig]: + """List configured monitors for a single table within a monitor suite. + + Returns one entry per ``test_definition`` row whose ``test_type`` is one + of the four monitor types. Typically 3-4 rows per table; Metric_Trend + contributes one entry per metric, so a table may have more. + """ + monitor_codes = [m.value for m in MonitorType] + defs = cls.select_where( + cls.test_suite_id == test_suite_id, + cls.table_name == table_name, + cls.test_type.in_(monitor_codes), + ) + return [cls._build_monitor_config(td) for td in defs] + + @classmethod + def _build_monitor_config(cls, td: "TestDefinition") -> MonitorConfig: + mode, threshold_lower, threshold_upper = cls._derive_threshold_mode(td) + return MonitorConfig( + monitor_id=td.id, + test_type=td.test_type, + table_name=td.table_name, + metric_name=td.column_name or None if td.test_type == MonitorType.METRIC.value else None, + threshold_mode=mode, + threshold_lower=threshold_lower, + threshold_upper=threshold_upper, + custom_query=td.custom_query if td.test_type == MonitorType.METRIC.value else None, + ) + + @classmethod + def _derive_threshold_mode( + cls, td: "TestDefinition", + ) -> tuple[ThresholdMode, str | None, str | None]: + """Pick a mode and the bounds tuple that applies under that mode. + + Detection mirrors the UI form (``test_definition_form.js``): a + ``history_calculation`` of exactly ``"PREDICT"`` flags Prediction + mode; any other non-empty value flags Historical (not available for + Freshness); empty falls through to Static — the default for + Freshness / Volume / Metric. Schema never has thresholds. + + Bounds returned per mode: + + * Prediction: ``None, None``. Runtime bounds live in the per-run + prediction JSONB; they are not configuration and do not surface + here. + * Historical: ``(history_calculation, history_calculation_upper)`` — + stored as expressions like ``"Minimum"`` / ``"Maximum"`` that the + execution layer evaluates against the lookback window. + * Static for Freshness: ``(None, upper_tolerance)``. Only an upper + bound applies. + * Static (Volume / Metric): ``(lower_tolerance, upper_tolerance)``. + """ + if td.test_type == MonitorType.SCHEMA.value: + return ThresholdMode.NONE, None, None + if td.history_calculation == "PREDICT": + return ThresholdMode.PREDICTION, None, None + if td.history_calculation and td.test_type != MonitorType.FRESHNESS.value: + return ThresholdMode.HISTORICAL, td.history_calculation, td.history_calculation_upper + if td.test_type == MonitorType.FRESHNESS.value: + return ThresholdMode.STATIC, None, td.upper_tolerance + return ThresholdMode.STATIC, td.lower_tolerance, td.upper_tolerance + @classmethod def select_page( cls, @@ -457,14 +763,16 @@ def select_page( # Fields editable on every test type regardless of param_columns. EDITABLE_BASE_FIELDS: ClassVar[frozenset[str]] = frozenset({ "test_active", "severity", "lock_refresh", "flagged", "test_description", + "external_url", "custom_metadata", }) def editable_fields(self, test_type: TestType) -> set[str]: """Fields a caller may set or change on this test definition under the given test type.""" fields = self.EDITABLE_BASE_FIELDS | test_type.param_columns - # column_name is meaningful for column-scoped tests (the column under test) and - # custom-scoped tests (a "Test Focus" label). Other scopes don't use it. - if test_type.test_scope in ("column", "custom"): + # column_name is meaningful for column-scoped tests (the column under test), + # custom-scoped tests (a "Test Focus" label), and referential tests (the aggregate + # expression or categorical column list under test). Table-scoped tests don't use it. + if test_type.test_scope in ("column", "custom", "referential"): fields = fields | {"column_name"} # impact_dimension is overridable only for user-defined-semantic scopes # (custom-scope = user-authored SQL; referential-scope = comparison-based tests). @@ -490,9 +798,10 @@ def validate(self, test_type: TestType) -> None: f"(got `{self.severity}`)" ) - # column_name applies to column-scoped tests (the column under test) and - # custom-scoped tests (a "Test Focus" label). Other scopes don't use it. - if test_type.test_scope not in ("column", "custom") and not _is_blank(self.column_name): + # column_name applies to column-scoped tests (the column under test), + # custom-scoped tests (a "Test Focus" label), and referential tests (the aggregate + # expression or categorical column list under test). Table-scoped tests don't use it. + if test_type.test_scope not in ("column", "custom", "referential") and not _is_blank(self.column_name): errors["column_name"] = ( f"test type `{test_type.test_type}` has scope `{test_type.test_scope}`; " f"column_name does not apply to this scope" @@ -503,6 +812,10 @@ def validate(self, test_type: TestType) -> None: f"test type `{test_type.test_type}` does not accept a custom query" ) + metadata_error = validate_custom_metadata(self.custom_metadata) + if metadata_error: + errors["custom_metadata"] = metadata_error + for required in _required_fields_for(test_type): if _is_blank(getattr(self, required, None)): errors[required] = f"required for test type `{test_type.test_type}`" diff --git a/testgen/common/models/test_result.py b/testgen/common/models/test_result.py index c1d25e73..a8ed2753 100644 --- a/testgen/common/models/test_result.py +++ b/testgen/common/models/test_result.py @@ -5,11 +5,12 @@ from typing import Self from uuid import UUID, uuid4 -from sqlalchemy import Boolean, Column, Enum, ForeignKey, Integer, String, desc, func, or_, select +from sqlalchemy import Boolean, Column, Enum, ForeignKey, Integer, String, desc, func, or_, select, text from sqlalchemy.dialects import postgresql from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import case +from testgen.common.enums import MonitorType from testgen.common.models import get_current_session from testgen.common.models.entity import Entity from testgen.common.models.test_definition import TestType @@ -64,6 +65,23 @@ class TestResultSearchRow: result_message: str | None +@dataclass +class TestRunResultRow: + """One individual result within a single test run for the API results endpoint.""" + + test_definition_id: UUID + test_type: str + schema_name: str + table_name: str | None + column_names: str | None + status: TestResultStatus | None + result_measure: str | None + threshold_value: str | None + message: str | None + test_time: datetime | None + disposition: str | None + + @dataclass class TrendBucket: """One time-bucket of failure aggregates for ``get_failure_trend``.""" @@ -101,6 +119,7 @@ class RunDiff: total_baseline: int total_target: int + stable_passes: int = 0 regressions: list[DiffRow] = field(default_factory=list) improvements: list[DiffRow] = field(default_factory=list) persistent_failures: list[DiffRow] = field(default_factory=list) @@ -108,6 +127,118 @@ class RunDiff: removed_tests: list[DiffRow] = field(default_factory=list) +@dataclass +class MonitorEvent: + """One monitor result for a (table, monitor_type) pair within the lookback window. + + Type-specific fields are populated by ``test_type``: + + * ``Freshness_Trend``: ``signal`` carries minutes-since-last-update (string + like ``"120"`` / ``"Unknown"``); ``message`` carries the human-readable + detection text. + * ``Volume_Trend``: ``signal`` carries the row count (string of integer); + ``lower_bound`` / ``upper_bound`` carry the tolerance bounds the run was + evaluated against (sourced from ``input_parameters`` at run time). + * ``Schema_Drift``: ``schema_change_kind`` is the table-level code + (``A`` / ``D`` / ``M``); ``column_adds`` / ``column_drops`` / + ``column_mods`` carry per-event column-change counts. + * ``Metric_Trend``: same as Volume_Trend plus ``metric_name`` (the + user-defined name for the metric — the underlying SQL expression lives + on ``MonitorConfig``, not on the event). + + Pending rows have no underlying ``test_results`` row — ``monitor_id``, + ``test_time``, and result flags are ``None``; ``is_pending`` is True. + Forecast points (future timestamps with predicted bounds) are NOT + events and surface separately via ``forecast_points_from_prediction``. + """ + monitor_id: UUID | None + test_type: str + test_time: datetime | None + is_anomaly: bool | None + is_training: bool | None + is_pending: bool + is_error: bool + message: str | None + signal: str | None + lower_bound: str | None = None + upper_bound: str | None = None + schema_change_kind: str | None = None + column_adds: int | None = None + column_drops: int | None = None + column_mods: int | None = None + metric_name: str | None = None + + +def _parse_kv_pairs(raw: str | None) -> dict[str, str]: + """Parse an ``input_parameters`` blob (``key=value; key=value; ...``) into a dict. + + ``input_parameters`` is built with ``"; ".join(...)`` in ``execute_tests_query`` + and read by the dashboard via ``dict_from_kv`` (default separator ``;``). + Tolerant of missing values, trailing/leading whitespace, empty strings. + Returns ``{}`` on empty input. + """ + if not raw: + return {} + pairs: dict[str, str] = {} + for entry in raw.split(";"): + if "=" not in entry: + continue + key, _, value = entry.partition("=") + pairs[key.strip()] = value.strip() + return pairs + + +def _build_monitor_event(row) -> MonitorEvent: + """Translate one CTE row into a ``MonitorEvent``, populating type-specific extras.""" + is_pending = row["result_id"] is None + result_code = row["result_code"] + is_anomaly = (result_code == 0) if result_code is not None else None + is_training = (result_code == -1) if result_code is not None else None + is_error = (row["result_status"] == "Error") + + event = MonitorEvent( + monitor_id=row["test_definition_id"], + test_type=row["test_type"], + test_time=row["test_time"] if not is_pending else None, + is_anomaly=is_anomaly, + is_training=is_training, + is_pending=is_pending, + is_error=is_error, + message=row["result_message"], + signal=row["result_signal"], + ) + + params = _parse_kv_pairs(row["input_parameters"]) + if event.test_type in (MonitorType.VOLUME.value, MonitorType.METRIC.value): + event.lower_bound = params.get("lower_tolerance") or None + event.upper_bound = params.get("upper_tolerance") or None + if event.test_type == MonitorType.METRIC.value: + event.metric_name = row["column_names"] or None + elif event.test_type == MonitorType.SCHEMA.value: + signal = row["result_signal"] + if signal: + parts = signal.split("|") + event.schema_change_kind = parts[0] or None + event.column_adds = _int_or_none(parts, 1) + event.column_drops = _int_or_none(parts, 2) + event.column_mods = _int_or_none(parts, 3) + + return event + + +def _int_or_none(parts: list[str], index: int) -> int | None: + try: + value = parts[index] + except IndexError: + return None + if not value: + return None + try: + return int(value) + except ValueError: + return None + + class TestResult(Entity): __tablename__ = "test_results" @@ -174,6 +305,38 @@ def select_results( query = query.order_by(cls.status, cls.table_name, cls.column_names).offset(offset).limit(limit) return get_current_session().scalars(query).all() + @classmethod + def list_for_run( + cls, + test_run_id: UUID, + *clauses, + page: int = 1, + limit: int = 20, + ) -> tuple[list[TestRunResultRow], int]: + """Paginated individual results for a single run, scoped by caller-supplied WHERE clauses. + + Monitor suites are always filtered out. + """ + query = ( + select( + cls.test_definition_id.label("test_definition_id"), + cls.test_type.label("test_type"), + cls.schema_name.label("schema_name"), + cls.table_name.label("table_name"), + cls.column_names.label("column_names"), + cls.status.label("status"), + cls.result_measure.label("result_measure"), + cls.threshold_value.label("threshold_value"), + cls.message.label("message"), + cls.test_time.label("test_time"), + cls.disposition.label("disposition"), + ) + .join(TestSuite, cls.test_suite_id == TestSuite.id) + .where(cls.test_run_id == test_run_id, TestSuite.is_monitor.isnot(True), *clauses) + .order_by(cls.status, cls.table_name, cls.column_names, cls.id) + ) + return cls._paginate(query, page=page, limit=limit, data_class=TestRunResultRow) + @classmethod def select_failures( cls, @@ -308,7 +471,7 @@ def search_results( select( cls.test_definition_id.label("test_definition_id"), cls.test_run_id.label("test_run_id"), - TestRun.job_execution_id.label("job_execution_id"), + TestRun.id.label("job_execution_id"), cls.test_time.label("test_time"), TestSuite.id.label("test_suite_id"), TestSuite.test_suite.label("test_suite_name"), @@ -471,12 +634,16 @@ def _row(tid: UUID, baseline_info: dict | None, target_info: dict | None) -> Dif for tid in baseline_results.keys() & target_results.keys(): baseline_info, target_info = baseline_results[tid], target_results[tid] + baseline_status, target_status = baseline_info["status"], target_info["status"] + if baseline_status == TestResultStatus.Passed and target_status == TestResultStatus.Passed: + diff.stable_passes += 1 + continue row = _row(tid, baseline_info, target_info) - if baseline_info["status"] == TestResultStatus.Passed and target_info["status"] in failing: + if baseline_status == TestResultStatus.Passed and target_status in failing: diff.regressions.append(row) - elif baseline_info["status"] in failing and target_info["status"] == TestResultStatus.Passed: + elif baseline_status in failing and target_status == TestResultStatus.Passed: diff.improvements.append(row) - elif baseline_info["status"] in failing and target_info["status"] in failing: + elif baseline_status in failing and target_status in failing: diff.persistent_failures.append(row) for tid in target_results.keys() - baseline_results.keys(): @@ -486,3 +653,153 @@ def _row(tid: UUID, baseline_info: dict | None, target_info: dict | None) -> Dif diff.removed_tests.append(_row(tid, baseline_results[tid], None)) return diff + + @classmethod + def list_monitor_events_for_table( + cls, + test_suite_id: str | UUID, + table_name: str, + *, + monitor_type: str | None = None, + lookback_multiplier: int = 1, + page: int = 1, + limit: int | None = None, + ) -> tuple[list[MonitorEvent], int]: + """Per-table monitor events within the suite's lookback window, newest first. + + ``monitor_type`` is the internal ``test_type`` value; when ``None`` events + for all four monitor types are returned (used by the dashboard, which + groups by type on the JS side). ``lookback_multiplier`` extends the + active window for the "show more history" toggle. ``limit=None`` + skips pagination entirely (the caller wants every event in the window). + + Forecast points for Prediction-Model monitors are NOT included here — + events are only past, observed runs. Read forecasts separately via + ``forecast_points_from_prediction(prediction, sensitivity)``. + """ + monitor_codes = ( + [monitor_type] if monitor_type is not None + else [m.value for m in MonitorType] + ) + type_filter_sql = "AND results.test_type = :monitor_type" if monitor_type else "" + + query = f""" + WITH ranked_test_runs AS ( + SELECT + test_runs.id, + test_runs.test_starttime, + COALESCE(test_suites.monitor_lookback, 1) * :lookback_multiplier AS lookback, + ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position + FROM test_suites + INNER JOIN test_runs ON (test_suites.id = test_runs.test_suite_id) + WHERE test_suites.id = :test_suite_id + ), + active_runs AS ( + SELECT id, test_starttime FROM ranked_test_runs WHERE position <= lookback + ), + target_types AS ( + SELECT UNNEST(CAST(:monitor_codes AS TEXT[])) AS test_type + ) + SELECT + COALESCE(results.test_time, active_runs.test_starttime) AS test_time, + tt.test_type, + results.id AS result_id, + results.test_definition_id, + results.result_code, + COALESCE(results.result_status, 'Log') AS result_status, + results.result_signal, + results.result_message, + COALESCE(results.input_parameters, '') AS input_parameters, + results.column_names + FROM active_runs + CROSS JOIN target_types tt + LEFT JOIN test_results AS results + ON results.test_run_id = active_runs.id + AND results.test_type = tt.test_type + AND results.table_name = :table_name + {type_filter_sql} + ORDER BY test_time DESC, tt.test_type, results.id NULLS LAST, active_runs.id + """ + + params: dict = { + "test_suite_id": str(test_suite_id), + "table_name": table_name, + "lookback_multiplier": lookback_multiplier, + "monitor_codes": monitor_codes, + } + if monitor_type is not None: + params["monitor_type"] = monitor_type + + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + events = [_build_monitor_event(row) for row in rows] + + # Paginate in Python — the CTE is bounded by lookback x |monitor_types| + # (typically <= ~120 rows for a single table). Revisit if either grows. + total = len(events) + if limit is not None: + start = (page - 1) * limit + events = events[start:start + limit] + return events, total + + @classmethod + def list_metric_monitor_events( + cls, + test_suite_id: str | UUID, + test_definition_id: str | UUID, + *, + lookback_multiplier: int = 1, + page: int = 1, + limit: int | None = None, + ) -> tuple[list[MonitorEvent], int]: + """Per-metric event history within the suite's lookback window, newest + first. Scoped to one ``test_definition_id`` since Metric_Trend is the + only multi-instance monitor type — table + type alone would interleave + events across every metric on the table. + + Distinct from ``list_monitor_events_for_table`` in two ways: (1) no + cross join with target_types — only one monitor type to query; (2) no + synthesized pending rows — a pending result for a specific + ``test_definition_id`` can't be distinguished from "no run yet" + without a results row to anchor on, so the model returns only rows + that actually ran. ``limit=None`` skips pagination. + """ + query = """ + WITH suite_window AS ( + SELECT COALESCE(monitor_lookback, 1) * :lookback_multiplier AS lookback + FROM test_suites + WHERE id = :test_suite_id + ) + SELECT + results.test_time, + results.test_type, + results.id AS result_id, + results.test_definition_id, + results.result_code, + COALESCE(results.result_status, 'Log') AS result_status, + results.result_signal, + results.result_message, + COALESCE(results.input_parameters, '') AS input_parameters, + results.column_names + FROM test_results AS results + WHERE results.test_suite_id = :test_suite_id + AND results.test_definition_id = :test_definition_id + ORDER BY results.test_time DESC, results.id + LIMIT (SELECT lookback FROM suite_window) + """ + + params: dict = { + "test_suite_id": str(test_suite_id), + "test_definition_id": str(test_definition_id), + "lookback_multiplier": lookback_multiplier, + } + + session = get_current_session() + rows = session.execute(text(query), params).mappings().all() + events = [_build_monitor_event(row) for row in rows] + + total = len(events) + if limit is not None: + start = (page - 1) * limit + events = events[start:start + limit] + return events, total diff --git a/testgen/common/models/test_run.py b/testgen/common/models/test_run.py index 2bb9bc51..e3db2fcc 100644 --- a/testgen/common/models/test_run.py +++ b/testgen/common/models/test_run.py @@ -1,14 +1,15 @@ from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import datetime from typing import ClassVar, Literal, NamedTuple, Self, TypedDict -from uuid import UUID, uuid4 +from uuid import UUID -from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, String, Text, desc, func, select, text, update +from sqlalchemy import BigInteger, Column, Float, ForeignKey, Integer, delete, desc, func, select, text from sqlalchemy.dialects import postgresql +from sqlalchemy.orm import foreign, relationship from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.sql.expression import case -from testgen.common.enums import JobStatus +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import database_session, get_current_session from testgen.common.models.connection import Connection from testgen.common.models.entity import Entity, EntityMinimal @@ -19,7 +20,6 @@ from testgen.common.models.test_suite import TestSuite from testgen.utils import is_uuid4 -TestRunStatus = Literal["Running", "Complete", "Error", "Cancelled"] ProgressKey = Literal["data_chars", "validation", "QUERY", "CAT", "METADATA"] ProgressStatus = Literal["Pending", "Running", "Completed", "Warning"] @@ -58,8 +58,6 @@ class TestRunSummary(EntityMinimal): test_suite: str | None project_code: str project_name: str - process_id: int | None - log_message: str | None test_ct: int | None passed_ct: int | None warning_ct: int | None @@ -70,15 +68,7 @@ class TestRunSummary(EntityMinimal): dq_score_testing: float | None total_count: int - STATUS_LABEL: ClassVar[dict[str, str]] = { - JobStatus.COMPLETED: "Completed", - JobStatus.CANCELED: "Canceled", - JobStatus.CANCEL_REQUESTED: "Canceling", - JobStatus.PENDING: "Pending", - JobStatus.CLAIMED: "Starting", - JobStatus.RUNNING: "Running", - JobStatus.ERROR: "Error", - } + STATUS_LABEL: ClassVar[dict[str, str]] = JOB_STATUS_LABEL @property def status_label(self) -> str: @@ -106,13 +96,14 @@ class LatestTestRun(NamedTuple): class TestRun(Entity): __tablename__ = "test_runs" - id: UUID = Column(postgresql.UUID(as_uuid=True), primary_key=True, nullable=False, default=uuid4) + id: UUID = Column( + postgresql.UUID(as_uuid=True), + ForeignKey("job_executions.id", ondelete="CASCADE"), + primary_key=True, + ) test_suite_id: UUID = Column(postgresql.UUID(as_uuid=True), ForeignKey("test_suites.id"), nullable=False) test_starttime: datetime = Column(postgresql.TIMESTAMP) - test_endtime: datetime = Column(postgresql.TIMESTAMP) - status: TestRunStatus = Column(String, default="Running") progress: list[ProgressStep] = Column(postgresql.JSONB, default=[]) - log_message: str = Column(Text) test_ct: int = Column(Integer) passed_ct: int = Column(Integer) failed_ct: int = Column(Integer) @@ -126,8 +117,13 @@ class TestRun(Entity): dq_affected_data_points: int = Column(BigInteger) dq_total_data_points: int = Column(BigInteger) dq_score_test_run: float = Column(Float) - process_id: int = Column(Integer) - job_execution_id: UUID | None = Column(postgresql.UUID(as_uuid=True), nullable=True) + + job_execution = relationship( + JobExecution, + primaryjoin=foreign(id) == JobExecution.id, + uselist=False, + viewonly=True, + ) _default_order_by = (desc(test_starttime),) _minimal_columns = ( @@ -144,21 +140,6 @@ class TestRun(Entity): ).label("is_latest_run"), ) - @classmethod - def get_by_id_or_job(cls, identifier: UUID) -> Self | None: - """Look up a test run by its own ID or by job_execution_id.""" - query = select(cls).where((cls.id == identifier) | (cls.job_execution_id == identifier)) - return get_current_session().scalars(query).first() - - @classmethod - def get_job_execution_ids(cls, test_run_ids: list[UUID]) -> dict[UUID, UUID | None]: - """Map test_run PKs to their job_execution_ids (batch lookup).""" - if not test_run_ids: - return {} - query = select(cls.id, cls.job_execution_id).where(cls.id.in_(test_run_ids)) - rows = get_current_session().execute(query).all() - return {row.id: row.job_execution_id for row in rows} - @classmethod def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: if not is_uuid4(run_id): @@ -167,7 +148,7 @@ def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: query = ( select(*cls._minimal_columns) .join(TestSuite) - .where((cls.id == run_id) | (cls.job_execution_id == run_id)) + .where(cls.id == run_id) ) result = get_current_session().execute(query).mappings().first() return TestRunMinimal(**result) if result else None @@ -176,7 +157,7 @@ def get_minimal(cls, run_id: str | UUID) -> TestRunMinimal | None: def get_latest_run(cls, project_code: str) -> LatestTestRun | None: query = ( select(TestRun.id, JobExecution.started_at.label("run_time")) - .join(JobExecution, TestRun.job_execution_id == JobExecution.id) + .join(JobExecution, TestRun.id == JobExecution.id) .join(TestSuite) .where(TestSuite.project_code == project_code, JobExecution.status == JobStatus.COMPLETED) .order_by(desc(JobExecution.started_at)) @@ -190,7 +171,7 @@ def get_latest_run(cls, project_code: str) -> LatestTestRun | None: def get_previous(self) -> Self | None: query = ( select(TestRun) - .join(JobExecution, TestRun.job_execution_id == JobExecution.id) + .join(JobExecution, TestRun.id == JobExecution.id) .where( TestRun.test_suite_id == self.test_suite_id, JobExecution.status == JobStatus.COMPLETED, @@ -219,6 +200,7 @@ def select_summary( test_suite_id: str | None = None, test_run_ids: list[str | UUID] | None = None, job_execution_id: str | UUID | None = None, + schedule_id: str | None = None, statuses: list[JobStatus] | None = None, page: int = 1, page_size: int = 20, @@ -228,6 +210,7 @@ def select_summary( or (test_suite_id and not is_uuid4(test_suite_id)) or (test_run_ids and not all(is_uuid4(run_id) for run_id in test_run_ids)) or (job_execution_id and not is_uuid4(job_execution_id)) + or (schedule_id and not is_uuid4(schedule_id)) ): return [], 0 @@ -266,8 +249,6 @@ def select_summary( ts.test_suite, je.project_code, p.project_name, - tr.process_id, - tr.log_message, tr.test_ct, rr.passed_ct, rr.warning_ct, @@ -278,7 +259,7 @@ def select_summary( tr.dq_score_test_run AS dq_score_testing, COUNT(*) OVER() AS total_count FROM job_executions je - LEFT JOIN test_runs tr ON tr.job_execution_id = je.id + LEFT JOIN test_runs tr ON tr.id = je.id LEFT JOIN test_suites ts ON ts.id = tr.test_suite_id LEFT JOIN table_groups tg ON tg.id = ts.table_groups_id LEFT JOIN projects p ON p.project_code = je.project_code @@ -290,6 +271,7 @@ def select_summary( {" AND ts.id = :test_suite_id" if test_suite_id else ""} {" AND tr.id IN :test_run_ids" if test_run_ids else ""} {" AND je.id = :job_execution_id" if job_execution_id else ""} + {" AND je.job_schedule_id = :schedule_id" if schedule_id else ""} {" AND je.status IN :statuses" if statuses else ""} ORDER BY je.created_at DESC LIMIT :limit OFFSET :offset; @@ -300,6 +282,7 @@ def select_summary( "test_suite_id": test_suite_id, "test_run_ids": tuple(test_run_ids or []), "job_execution_id": str(job_execution_id) if job_execution_id else None, + "schedule_id": schedule_id, "statuses": tuple(statuses) if statuses else (), "limit": page_size, "offset": (page - 1) * page_size, @@ -325,7 +308,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito )) projection = [ TestRun.id.label("test_run_id"), - TestRun.test_endtime, + JobExecution.completed_at.label("test_endtime"), TableGroup.id.label("table_group_id"), TableGroup.table_groups_name, Project.project_name, @@ -335,7 +318,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito ] group_by = [ TestRun.id, - TestRun.test_endtime, + JobExecution.completed_at, TableGroup.id, TableGroup.table_groups_name, Project.project_name, @@ -349,6 +332,7 @@ def get_monitoring_summary(self, table_name: str | None = None) -> TestRunMonito .join(TableGroup, TableGroup.monitor_test_suite_id == TestRun.test_suite_id) .join(Project, Project.project_code == TableGroup.project_code) .join(TestResult, TestResult.test_run_id == TestRun.id) + .join(JobExecution, JobExecution.id == TestRun.id) .where( TestRun.id == self.id, (TestResult.table_name == table_name) if table_name else True, @@ -365,7 +349,7 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U """Check whether any active test run job exists for the given entity or entities.""" query = ( select(func.count(cls.id)) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where(JobExecution.status.in_(cls._ACTIVE_JOB_STATUSES)) ) if entity_cls is cls: @@ -382,27 +366,14 @@ def has_active_job_for(cls, entity_cls: type[Entity], *entity_ids: str | int | U raise ValueError(f"Unsupported entity: {entity_cls.__name__}") return get_current_session().execute(query).scalar() > 0 - @classmethod - def cancel_run(cls, run_id: str | UUID) -> None: - query = update(cls).where(cls.id == run_id).values(status="Cancelled", test_endtime=datetime.now(UTC)) - db_session = get_current_session() - db_session.execute(query) - @classmethod def cascade_delete(cls, ids: list[str]) -> None: - query = """ - DELETE FROM test_results - WHERE test_run_id IN :test_run_ids; - - DELETE FROM job_executions - WHERE id IN ( - SELECT job_execution_id FROM test_runs - WHERE id IN :test_run_ids AND job_execution_id IS NOT NULL - ); + """Delete runs and their results by removing the parent job executions. + + The run id is the job execution id, and the foreign-key chain + (test_results -> test_runs -> job_executions) cascades on delete. """ - db_session = get_current_session() - db_session.execute(text(query), {"test_run_ids": tuple(ids)}) - cls.delete_where(cls.id.in_(ids)) + get_current_session().execute(delete(JobExecution).where(JobExecution.id.in_(ids))) @classmethod def find_latest_per_test_suite(cls, project_code: str) -> set[UUID]: @@ -417,7 +388,7 @@ def find_latest_per_test_suite(cls, project_code: str) -> set[UUID]: rows = get_current_session().scalars( select(cls.id) .join(TestSuite, cls.test_suite_id == TestSuite.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) .where( TestSuite.project_code == project_code, JobExecution.status == JobStatus.COMPLETED, @@ -459,7 +430,7 @@ def delete_older_than( base_select = ( select(cls.id) .join(TestSuite, cls.test_suite_id == TestSuite.id) - .join(JobExecution, cls.job_execution_id == JobExecution.id) + .join(JobExecution, cls.id == JobExecution.id) ) if dry_run: diff --git a/testgen/common/models/test_suite.py b/testgen/common/models/test_suite.py index ac29ddce..fab9720c 100644 --- a/testgen/common/models/test_suite.py +++ b/testgen/common/models/test_suite.py @@ -29,6 +29,15 @@ class TestSuiteMinimal(EntityMinimal): export_to_observability: str +@dataclass(frozen=True) +class TestDefinitionStats: + """Aggregate test-definition counts for a test suite.""" + + total: int + locked: int + counts_by_type: dict[str, int] + + @dataclass class TestSuiteSummary(EntityMinimal): id: UUID @@ -113,7 +122,6 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N WITH last_run AS ( SELECT test_runs.test_suite_id, test_runs.id, - test_runs.job_execution_id, test_runs.test_starttime, test_runs.test_ct, SUM( @@ -184,7 +192,7 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N test_defs.count AS test_ct, last_complete_profile_run_id, last_run.id AS latest_run_id, - last_run.job_execution_id AS latest_run_job_execution_id, + last_run.id AS latest_run_job_execution_id, last_run.test_starttime AS latest_run_start, last_run.test_ct AS last_run_test_ct, last_run.passed_ct AS last_run_passed_ct, @@ -213,6 +221,36 @@ def select_summary(cls, project_code: str, table_group_id: str | UUID | None = N results = db_session.execute(text(query), params).mappings().all() return [TestSuiteSummary(**row) for row in results] + @classmethod + def test_definition_stats(cls, test_suite_id: str | UUID) -> "TestDefinitionStats": + """Aggregate test-definition counts for a suite: total, locked, and per-test-type. + + The per-type bucket uses the user-facing ``test_name_short`` label from + ``test_types``, falling back to the raw ``test_type`` code when the lookup + row is missing (defensive — every shipping test type has a short name). + """ + query = """ + SELECT + COALESCE(tt.test_name_short, td.test_type) AS type_label, + COUNT(*) AS total, + SUM(CASE WHEN COALESCE(td.lock_refresh, 'N') = 'Y' THEN 1 ELSE 0 END) AS locked + FROM test_definitions td + LEFT JOIN test_types tt ON tt.test_type = td.test_type + WHERE td.test_suite_id = :test_suite_id + GROUP BY type_label + ORDER BY type_label; + """ + rows = ( + get_current_session() + .execute(text(query), {"test_suite_id": str(test_suite_id)}) + .mappings() + .all() + ) + counts_by_type = {row["type_label"]: int(row["total"]) for row in rows} + total = sum(counts_by_type.values()) + locked = sum(int(row["locked"]) for row in rows) + return TestDefinitionStats(total=total, locked=locked, counts_by_type=counts_by_type) + @classmethod def is_in_use(cls, ids: list[str]) -> bool: query = """ @@ -229,8 +267,8 @@ def cascade_delete(cls, ids: list[str]) -> None: query = """ DELETE FROM job_executions WHERE id IN ( - SELECT job_execution_id FROM test_runs - WHERE test_suite_id IN :test_suite_ids AND job_execution_id IS NOT NULL + SELECT id FROM test_runs + WHERE test_suite_id IN :test_suite_ids ); DELETE FROM test_runs @@ -248,4 +286,3 @@ def cascade_delete(cls, ids: list[str]) -> None: db_session = get_current_session() db_session.execute(text(query), {"test_suite_ids": tuple(ids)}) cls.delete_where(cls.id.in_(ids)) - diff --git a/testgen/common/monitor_forecast.py b/testgen/common/monitor_forecast.py new file mode 100644 index 00000000..a3c65a46 --- /dev/null +++ b/testgen/common/monitor_forecast.py @@ -0,0 +1,156 @@ +"""Monitor forecast computation shared by the monitors dashboard and the MCP +monitor tools. + +The forecast a monitor displays depends on its threshold mode and type: + +* Volume / Metric in Prediction Model mode: a value band extending forward. For + monitors coupled to a Freshness monitor the band holds at the baseline until + the next expected refresh, then steps to the forecast value (``gated_forecast_prediction``); + otherwise it is the raw per-step prediction band. +* Freshness in Prediction Model mode: a predicted next-update *time window* + rather than a value band (``next_update_window``). + +These functions are intentionally free of any UI dependency so both surfaces +render the same forecast from the same source of truth. +""" + +from datetime import UTC, date, datetime + +import pandas as pd + +from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates +from testgen.common.models.test_definition import MonitorForecastPoint, TestDefinition, TestDefinitionSummary +from testgen.common.models.test_suite import TestSuite + +# A monitor definition carrying the forecast-relevant fields — satisfied by both +# the ORM row and the read-only summary dataclass. +MonitorDefinition = TestDefinition | TestDefinitionSummary + + +def resolve_suite_holiday_dates(test_suite: TestSuite) -> set[date] | None: + """Holiday dates excluded from the suite's schedule, over a window spanning + recent history through the forecast horizon. ``None`` when no calendars are set.""" + if not test_suite.holiday_codes_list: + return None + now = pd.Timestamp.now("UTC") + idx = pd.DatetimeIndex([now - pd.Timedelta(days=7), now + pd.Timedelta(days=30)]) + return resolve_holiday_dates(test_suite.holiday_codes_list, idx) + + +def next_update_window( + freshness_definition: MonitorDefinition | None, + last_detection_time: datetime | None, + *, + test_suite: TestSuite, + cron_tz: str | None, +) -> dict | None: + """Predicted next-update window as ``{"start", "end"}`` epoch-ms, or ``None``. + + The schedule-derived business-time interval from the last detected update out + to the lower/upper staleness tolerance. ``start`` is ``None`` when only an + upper tolerance is configured (the update is expected by ``end``). Returns + ``None`` unless the Freshness monitor is in Prediction Model mode with a + trained schedule and at least one observed update. + """ + if ( + freshness_definition is None + or freshness_definition.history_calculation != "PREDICT" + or (freshness_definition.prediction and not freshness_definition.prediction.get("schedule_stage")) + or freshness_definition.upper_tolerance is None + or last_detection_time is None + ): + return None + + tz = cron_tz or "UTC" + exclude_weekends = test_suite.predict_exclude_weekends + holiday_dates = resolve_suite_holiday_dates(test_suite) + sched = get_schedule_params(freshness_definition.prediction) + + window_end = add_business_minutes( + pd.Timestamp(last_detection_time), + float(freshness_definition.upper_tolerance), + exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + window_start = None + if lower_minutes := (float(freshness_definition.lower_tolerance) if freshness_definition.lower_tolerance else None): + window_start = add_business_minutes( + pd.Timestamp(last_detection_time), + lower_minutes, + exclude_weekends, + holiday_dates, tz, + excluded_days=sched.excluded_days, + ) + + return { + "start": int(window_start.timestamp() * 1000) if window_start else None, + "end": int(window_end.timestamp() * 1000), + } + + +def gated_forecast_prediction( + definition: MonitorDefinition, + freshness_window: dict | None, + last_run_time: datetime | None, +) -> dict | None: + """Coupled forecast band for a Volume/Metric monitor whose value holds at its + baseline between refreshes. + + A flat baseline line from the latest run up to the predicted next-update + window, then a step to the forecast's next-refresh value with the band + opening to its tolerance. The anchor is never earlier than the latest run, so + the forecast extends forward rather than back over history. Returns ``None`` + when there is no usable forward window — no predicted window, the window has + already elapsed, or no stored baseline — keyed as ``lower_tolerance`` / + ``upper_tolerance`` (no sensitivity suffix); read with ``forecast_band_points``. + """ + baseline = definition.prediction.get("baseline_value") if definition.prediction else None + now_ms = int(pd.Timestamp(last_run_time).timestamp() * 1000) if last_run_time is not None else None + if freshness_window is None or now_ms is None or baseline is None: + return None + window_end = freshness_window.get("end") + if window_end is None or window_end <= now_ms: + return None + + forecast_means = (definition.prediction.get("mean") if definition.prediction else None) or {} + next_refresh_mean = forecast_means[min(forecast_means, key=lambda k: int(k))] if forecast_means else baseline + flat_anchor = max(freshness_window.get("start") or now_ms, now_ms) + # lower/upper_tolerance are VARCHAR columns — coerce to float so the band dicts are numerically + # typed throughout (baseline is already a float). + lower_tol = float(definition.lower_tolerance) if definition.lower_tolerance is not None else None + upper_tol = float(definition.upper_tolerance) if definition.upper_tolerance is not None else None + # The band carries the next-refresh tolerance across the whole forecast, flat anchor included. + # Anchoring the band to the tolerance — rather than collapsing it to the baseline at the flat + # anchor — keeps it a continuous width that connects to the historical band and avoids a + # zero-width pinch (which renders as an hourglass for a future window, or a band detached from + # history for an imminent one). The mean line still holds flat at baseline, then steps. + return { + "method": "predict", + "mean": {flat_anchor: baseline, window_end: next_refresh_mean}, + "lower_tolerance": {flat_anchor: lower_tol, window_end: lower_tol}, + "upper_tolerance": {flat_anchor: upper_tol, window_end: upper_tol}, + } + + +def forecast_band_points(prediction: dict | None) -> list[MonitorForecastPoint]: + """Convert a forecast band dict with plain ``lower_tolerance`` / ``upper_tolerance`` + epoch-ms series (as built by ``gated_forecast_prediction``) into time-ordered + forecast points. For the raw per-sensitivity prediction JSONB use + ``forecast_points_from_prediction`` instead.""" + if not prediction: + return [] + lower_series = prediction.get("lower_tolerance") or {} + upper_series = prediction.get("upper_tolerance") or {} + keys = sorted(set(lower_series) | set(upper_series), key=lambda k: int(k)) + points: list[MonitorForecastPoint] = [] + for k in keys: + ts = datetime.fromtimestamp(int(k) / 1000.0, UTC) + lower = lower_series.get(k) + upper = upper_series.get(k) + points.append(MonitorForecastPoint( + test_time=ts, + lower_bound=float(lower) if lower is not None else None, + upper_bound=float(upper) if upper is not None else None, + )) + return points diff --git a/testgen/common/monitor_service.py b/testgen/common/monitor_service.py new file mode 100644 index 00000000..ebacf645 --- /dev/null +++ b/testgen/common/monitor_service.py @@ -0,0 +1,161 @@ +"""Monitor lifecycle: enable, update, and disable monitoring for a table group. + +Shared by the MCP tools, the monitors dashboard, and the table-group creation +wizard so all three drive monitoring through one path. Functions mutate the +monitor ``TestSuite``, its ``JobSchedule``, and the table-group link, and raise +stdlib exceptions — the MCP and UI layers translate those into their own +user-facing errors. +""" + +import logging +from typing import Any + +from sqlalchemy import func, select + +from testgen.commands.test_generation import run_monitor_generation +from testgen.common.models import get_current_session +from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.table_group import TableGroup +from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_result import TestResult +from testgen.common.models.test_run import TestRun +from testgen.common.models.test_suite import TestSuite + +LOG = logging.getLogger("testgen") + +# Monitors generated when monitoring is first enabled. Freshness_Trend depends on +# per-table freshness fingerprinting and is generated separately, so it is not part +# of the initial bootstrap set. +INITIAL_MONITOR_TYPES = ["Volume_Trend", "Schema_Drift"] + +# Default monitor configuration applied at bootstrap, mirroring the UI's setup form. +_DEFAULT_SUITE_ATTRS: dict[str, Any] = { + "monitor_lookback": 14, + "monitor_regenerate_freshness": True, + "predict_min_lookback": 30, + "predict_sensitivity": "medium", + "predict_exclude_weekends": False, + "predict_holiday_codes": None, +} + +# The monitor-configuration columns callers may set via ``suite_attrs``. Used as a +# whitelist so a caller's dict (e.g. a UI form payload) can't write arbitrary suite columns. +_MONITOR_SETTING_COLUMNS = tuple(_DEFAULT_SUITE_ATTRS) + + +def enable_monitoring( + table_group: TableGroup, + cron_expr: str, + cron_tz: str = "UTC", + *, + suite_attrs: dict[str, Any] | None = None, + active: bool = True, +) -> tuple[TestSuite, int]: + """Bootstrap monitoring for a table group. + + Creates the monitor test suite, generates the initial monitors, creates the + run-monitors schedule, and links the suite to the table group. + + ``suite_attrs`` overrides the default monitor configuration. + + Returns ``(monitor_suite, monitor_count)``. Raises ``ValueError`` if the table + group already has monitoring enabled. + """ + if table_group.monitor_test_suite_id: + raise ValueError("Monitoring is already enabled for this table group.") + + provided = suite_attrs or {} + attrs = dict(_DEFAULT_SUITE_ATTRS) + for key in _MONITOR_SETTING_COLUMNS: + if (value := provided.get(key)) is not None: + attrs[key] = value + + monitor_suite = TestSuite( + project_code=table_group.project_code, + test_suite=f"{table_group.table_groups_name} Monitors", + connection_id=table_group.connection_id, + table_groups_id=table_group.id, + export_to_observability=False, + dq_score_exclude=True, + is_monitor=True, + **attrs, + ) + monitor_suite.save() + + JobSchedule( + project_code=table_group.project_code, + key=RUN_MONITORS_JOB_KEY, + kwargs={"test_suite_id": str(monitor_suite.id)}, + cron_expr=cron_expr, + cron_tz=cron_tz, + active=active, + ).save() + + table_group.monitor_test_suite_id = monitor_suite.id + table_group.save() + + # Commit needed to make the test suite visible to run_monitor_generation's separate DB connection. + get_current_session().commit() + run_monitor_generation(monitor_suite.id, INITIAL_MONITOR_TYPES) + + count = get_current_session().scalar( + select(func.count()).select_from(TestDefinition).where(TestDefinition.test_suite_id == monitor_suite.id) + ) + return monitor_suite, count or 0 + + +def update_monitoring( + monitor_suite: TestSuite, + schedule: JobSchedule, + *, + suite_attrs: dict[str, Any] | None = None, + cron_expr: str | None = None, + cron_tz: str | None = None, + active: bool | None = None, +) -> None: + """Apply a partial update to monitor settings and/or schedule. + + ``suite_attrs`` maps monitor ``TestSuite`` columns to new values; only the keys in + ``_MONITOR_SETTING_COLUMNS`` that are present are applied (a present ``None`` clears + the column). Supplied schedule fields are updated in place. Does not commit — the + caller's session lifecycle does. + """ + provided = suite_attrs or {} + for key in _MONITOR_SETTING_COLUMNS: + if key in provided: + setattr(monitor_suite, key, provided[key]) + monitor_suite.save() + + if cron_expr is not None: + schedule.cron_expr = cron_expr + if cron_tz is not None: + schedule.cron_tz = cron_tz + if active is not None: + schedule.active = active + schedule.save() + + +def disable_monitoring(monitor_suite: TestSuite) -> dict[str, int]: + """Remove a monitor suite and all its monitors, runs, and history. + + Counts what will be removed before cascading the delete. Returns counts keyed + ``monitors`` (test definitions), ``events`` (test results), and ``runs``. + """ + session = get_current_session() + suite_id = monitor_suite.id + counts = { + "monitors": session.scalar( + select(func.count()).select_from(TestDefinition).where(TestDefinition.test_suite_id == suite_id) + ) + or 0, + "events": session.scalar( + select(func.count()).select_from(TestResult).where(TestResult.test_suite_id == suite_id) + ) + or 0, + "runs": session.scalar( + select(func.count()).select_from(TestRun).where(TestRun.test_suite_id == suite_id) + ) + or 0, + } + TestSuite.cascade_delete([suite_id]) + return counts diff --git a/testgen/common/notifications/monitor_run.py b/testgen/common/notifications/monitor_run.py index 4967216a..9c9a55ec 100644 --- a/testgen/common/notifications/monitor_run.py +++ b/testgen/common/notifications/monitor_run.py @@ -236,7 +236,7 @@ def send_monitor_notifications(test_run: TestRun, result_list_ct=20): notification.recipients, { "summary": { - "test_endtime": test_run.test_endtime, + "test_endtime": test_run.job_execution.completed_at, "table_groups_name": table_group.table_groups_name, "project_name": project.project_name, "table_name": table_name, diff --git a/testgen/common/notifications/notifications.py b/testgen/common/notifications/notifications.py index 68e20732..52226950 100644 --- a/testgen/common/notifications/notifications.py +++ b/testgen/common/notifications/notifications.py @@ -1,8 +1,7 @@ import math from datetime import datetime -from testgen.common.models.profiling_run import ProfilingRunStatus -from testgen.common.models.test_definition import TestRunStatus +from testgen.common import date_service from testgen.common.notifications.base import BaseEmailTemplate from testgen.utils import friendly_score @@ -16,7 +15,7 @@ def format_number_helper(self, number: int) -> str: return "" if number is None else f"{number:,}" def format_dt_helper(self, dt: datetime) -> str: - return "" if dt is None else dt.strftime("%b %d, %-I:%M %p UTC") + return "" if dt is None else date_service.format_friendly_datetime(dt, "%b %d, %-I:%M %p UTC") def format_duration_helper(self, start_time: datetime, end_time: datetime) -> str: total_seconds = abs(end_time - start_time).total_seconds() @@ -29,12 +28,6 @@ def format_duration_helper(self, start_time: datetime, end_time: datetime) -> st formatted = " ".join([ f"{unit[0]}{unit[1]}" for unit in units if unit[0] ]) return formatted.strip() or "< 1s" - def format_status_helper(self, status: TestRunStatus | ProfilingRunStatus) -> str: - return { - "Complete": "Completed", - "Cancelled": "Canceled", - }.get(status, status) - def format_score_helper(self, score: float) -> str: return friendly_score(score) diff --git a/testgen/common/notifications/profiling_run.py b/testgen/common/notifications/profiling_run.py index 2a4c254c..1ceec969 100644 --- a/testgen/common/notifications/profiling_run.py +++ b/testgen/common/notifications/profiling_run.py @@ -4,6 +4,7 @@ from sqlalchemy import select from testgen import settings +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.hygiene_issue import HygieneIssue from testgen.common.models.notification_settings import ( @@ -23,7 +24,7 @@ class ProfilingRunEmailTemplate(BaseNotificationTemplate): def get_subject_template(self) -> str: return ( - "[TestGen] Profiling Run {{format_status profiling_run.status}}: {{table_groups_name}}" + "[TestGen] Profiling Run {{profiling_run.status_label}}: {{table_groups_name}}" "{{#if issue_count}}" ' | {{format_number issue_count}} hygiene {{pluralize issue_count "issue" "issues"}}' "{{/if}}" @@ -32,9 +33,9 @@ def get_subject_template(self) -> str: def get_title_template(self): return """ TestGen Profiling Run - {{format_status profiling_run.status}} + {{#if (eq profiling_run.status 'error')}} text-red {{/if}} + {{#if (eq profiling_run.status 'canceled')}} text-purple {{/if}} + ">{{profiling_run.status_label}} """ def get_main_content_template(self): @@ -85,7 +86,7 @@ def get_main_content_template(self): - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} {{#if (eq notification_trigger 'on_changes')}} Profiling run detected new hygiene issues. {{/if}} @@ -95,15 +96,15 @@ def get_main_content_template(self): {{/if}} {{/if}} {{/if}} - {{#if (eq profiling_run.status 'Error')}} + {{#if (eq profiling_run.status 'error')}} Profiling encountered an error. {{/if}} - {{#if (eq profiling_run.status 'Cancelled')}} + {{#if (eq profiling_run.status 'canceled')}} Profiling run was canceled. {{/if}} - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} @@ -131,12 +132,12 @@ def get_main_content_template(self): {{/if}} - {{#if (eq profiling_run.status 'Error')}} + {{#if (eq profiling_run.status 'error')}} {{/if}} - {{#if (eq profiling_run.status 'Complete')}} + {{#if (eq profiling_run.status 'completed')}} - + {{/if}} {{#if (eq test_run.status 'completed')}} @@ -261,7 +262,7 @@ def send_test_run_notifications(test_run: TestRun, result_list_ct=20, result_sta changed_td_id_list.extend(td_id_list) triggers = {TestRunNotificationTrigger.always} - if test_run.status in ("Error", "Cancelled"): + if test_run.job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED): triggers.update(TestRunNotificationTrigger) else: if test_run.error_ct + test_run.failed_ct: @@ -329,7 +330,7 @@ def send_test_run_notifications(test_run: TestRun, result_list_ct=20, result_sta "/test-runs:results?project_code=", str(tr_summary.project_code), "&run_id=", - str(test_run.job_execution_id), + str(test_run.id), "&source=email" ) ) diff --git a/testgen/common/source_data_service.py b/testgen/common/source_data_service.py index 52cf2723..941679f1 100644 --- a/testgen/common/source_data_service.py +++ b/testgen/common/source_data_service.py @@ -11,7 +11,7 @@ import pandas as pd from sqlalchemy import text -from testgen.common.clean_sql import concat_columns +from testgen.common.clean_sql import concat_columns, null_if_empty from testgen.common.database.database_service import get_flavor_service, replace_params from testgen.common.date_service import parse_fuzzy_date from testgen.common.models import get_current_session @@ -207,12 +207,14 @@ def _build_query_standard(issue_data: dict, limit: int) -> str | None: if (parsed_test_date := parse_fuzzy_date(issue_data["test_date"])) else None, "CUSTOM_QUERY": test_definition.custom_query, + # BASELINE_VALUE varies (quoted literal / number / IN-list) by test type — templated as-is. + # The numeric baselines render empty as NULL to avoid invalid SQL like CAST( AS FLOAT). "BASELINE_VALUE": test_definition.baseline_value, - "BASELINE_CT": test_definition.baseline_ct, - "BASELINE_AVG": test_definition.baseline_avg, - "BASELINE_SD": test_definition.baseline_sd, - "LOWER_TOLERANCE": "NULL" if test_definition.lower_tolerance in (None, "") else test_definition.lower_tolerance, - "UPPER_TOLERANCE": "NULL" if test_definition.upper_tolerance in (None, "") else test_definition.upper_tolerance, + "BASELINE_CT": null_if_empty(test_definition.baseline_ct), + "BASELINE_AVG": null_if_empty(test_definition.baseline_avg), + "BASELINE_SD": null_if_empty(test_definition.baseline_sd), + "LOWER_TOLERANCE": null_if_empty(test_definition.lower_tolerance), + "UPPER_TOLERANCE": null_if_empty(test_definition.upper_tolerance), "THRESHOLD_VALUE": test_definition.threshold_value or 0, # SUBSET_CONDITION should be replaced after CUSTOM_QUERY # since the latter may contain the former diff --git a/testgen/api/test_definition_service.py b/testgen/common/test_definition_export_import_service.py similarity index 64% rename from testgen/api/test_definition_service.py rename to testgen/common/test_definition_export_import_service.py index e8b2199f..fb9a4e57 100644 --- a/testgen/api/test_definition_service.py +++ b/testgen/common/test_definition_export_import_service.py @@ -1,39 +1,231 @@ -"""Business logic for test definition export/import.""" +"""Test definition export/import — portable document schemas and import/export logic.""" from dataclasses import dataclass from datetime import UTC, datetime +from enum import StrEnum from typing import Any from uuid import UUID +from pydantic import BaseModel, field_validator from sqlalchemy import delete, func, select, update from sqlalchemy.engine import Row from testgen import settings -from testgen.api.deps import api_error -from testgen.api.schemas import ( - ExportDocument, - ExportSource, - ImportAction, - ImportConfig, - ImportItem, - ImportItemTD, - ImportMode, - ImportPayload, - ImportReason, - ImportResponse, - ImportSummary, - OnAbsence, - OnMatch, - OnNew, - Origin, - TestDefinitionExport, -) from testgen.common.models import get_current_session from testgen.common.models.data_table import DataTable from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_definition import TestDefinition, TestType +from testgen.common.models.test_definition import TestDefinition, TestType, validate_custom_metadata from testgen.common.models.test_suite import TestSuite +EXPORT_FORMAT_VERSION = 1 + + +class Origin(StrEnum): + manual = "manual" + auto = "auto" + both = "both" + + +class ImportMode(StrEnum): + preview = "preview" + apply = "apply" + apply_strict = "apply_strict" + + +class OnMatch(StrEnum): + overwrite_all = "overwrite_all" + overwrite_unlocked = "overwrite_unlocked" + skip = "skip" + + +class OnNew(StrEnum): + skip = "skip" + create = "create" + create_and_lock = "create_and_lock" + + +class OnAbsence(StrEnum): + do_nothing = "do_nothing" + delete_all = "delete_all" + delete_unlocked = "delete_unlocked" + + +class ImportAction(StrEnum): + create = "create" + update = "update" + skip = "skip" + delete = "delete" + + +class ImportReason(StrEnum): + matched = "matched" + no_match = "no_match" + policy = "policy" + locked = "locked" + invalid_test_type = "invalid_test_type" + invalid_table = "invalid_table" + missing_external_id = "missing_external_id" + absent = "absent" + + +# Non-None defaults must match the ORM column defaults in TestDefinition: +# test_active=True (YNString default="Y"), lock_refresh=False (YNString default="N"), +# skip_errors=0 (ZeroIfEmptyInteger), window_days=0 (ZeroIfEmptyInteger), +# history_lookback=0 (Column default=0). +# On export, fields equal to these defaults are omitted to keep the file compact — any +# serialization of the document must pass ``exclude_defaults=True`` (the REST route does this +# via ``response_model_exclude_defaults=True``). On import, model_fields_set distinguishes +# explicit from defaulted, so a non-compact file would force-write every defaulted field. +class TestDefinitionExport(BaseModel): + """Test definition fields included in the export/import file.""" + + model_config = {"from_attributes": True} + + # Matching / identity + test_type: str + external_id: UUID | None = None + last_auto_gen_date: datetime | None = None + + # Definition fields + table_name: str | None = None + column_name: str | None = None + test_description: str | None = None + test_active: bool = True + severity: str | None = None + lock_refresh: bool = False + export_to_observability: bool | None = None + skip_errors: int = 0 + + # Calibration fields + baseline_ct: str | None = None + baseline_unique_ct: str | None = None + baseline_value: str | None = None + baseline_value_ct: str | None = None + threshold_value: str | None = None + baseline_sum: str | None = None + baseline_avg: str | None = None + baseline_sd: str | None = None + lower_tolerance: str | None = None + upper_tolerance: str | None = None + + # Subset / grouping + subset_condition: str | None = None + groupby_names: str | None = None + having_condition: str | None = None + window_date_column: str | None = None + window_days: int = 0 + + # Referential + match_schema_name: str | None = None + match_table_name: str | None = None + match_column_names: str | None = None + match_subset_condition: str | None = None + match_groupby_names: str | None = None + match_having_condition: str | None = None + + # Query / history + custom_query: str | None = None + history_calculation: str | None = None + history_calculation_upper: str | None = None + history_lookback: int = 0 + + # URL / metadata + external_url: str | None = None + custom_metadata: dict[str, Any] | None = None + + @field_validator("skip_errors", "window_days", "history_lookback", mode="before") + @classmethod + def _coerce_none_to_zero(cls, v: int | None) -> int: + return v if v is not None else 0 + + @field_validator("custom_metadata") + @classmethod + def _check_custom_metadata(cls, v: dict[str, Any] | None) -> dict[str, Any] | None: + # The import path does not run TestDefinition.validate(), so this is the only enforcement + # of the custom_metadata shape/size bound on import. Shares validate_custom_metadata with + # the model so the rule has a single definition. + error = validate_custom_metadata(v) + if error: + raise ValueError(f"custom_metadata {error}") + return v + + +class ExportSource(BaseModel): + project_code: str + test_suite: str + table_group: str + table_group_schema: str + exported_at: datetime + testgen_version: str | None = None + + +class ExportDocument(BaseModel): + # No default: a defaulted version is stripped by default-omitting serialization, + # leaving exported documents without their schema-version marker. + version: int + source: ExportSource + definitions: list[TestDefinitionExport] + + +class ImportConfig(BaseModel): + mode: ImportMode + on_match: OnMatch + on_new: OnNew + on_absence: OnAbsence + + +class ImportPayload(BaseModel): + """Import payload — same structure as an export document, but definitions are typed.""" + + version: int = EXPORT_FORMAT_VERSION + source: ExportSource | None = None + definitions: list[TestDefinitionExport] + + +class ImportItemTD(BaseModel): + idx: int | None = None + target_id: UUID | None = None + + +class ImportItem(BaseModel): + action: ImportAction + reason: ImportReason + tds: list[ImportItemTD] + + +class ImportSummary(BaseModel): + created: int = 0 + updated: int = 0 + skipped: int = 0 + deleted: int = 0 + + +class ImportResponse(BaseModel): + summary: ImportSummary + items: list[ImportItem] + + +class ImportErrorCode(StrEnum): + duplicate_natural_key = "duplicate_natural_key" + unsupported_version = "unsupported_version" + + +class InvalidImportPayload(ValueError): + """Import payload rejected before any matching or persistence.""" + + def __init__(self, code: ImportErrorCode, detail: str): + super().__init__(detail) + self.code = code + + +class ImportStrictViolation(Exception): + """``apply_strict`` import with test definitions that would be skipped; nothing was applied.""" + + def __init__(self, result: ImportResponse): + super().__init__(f"{result.summary.skipped} test definitions would be skipped") + self.result = result + + # Fields that must never be written from the import payload on update. # These are either identity fields (set once on create) or determined by matching logic. _UPDATE_EXCLUDE_FIELDS = frozenset({"test_type", "last_auto_gen_date", "external_id"}) @@ -87,6 +279,7 @@ def export_definitions( definitions = [TestDefinitionExport.model_validate(td, from_attributes=True) for td in tds] return ExportDocument( + version=EXPORT_FORMAT_VERSION, source=ExportSource( project_code=test_suite.project_code, test_suite=test_suite.test_suite, @@ -110,6 +303,11 @@ def import_definitions( profiled_tables = set(DataTable.select_table_names(test_suite.table_groups_id, limit=None)) # --- Phase 1: Upfront validation --- + if payload.version != EXPORT_FORMAT_VERSION: + raise InvalidImportPayload( + ImportErrorCode.unsupported_version, + f"Unsupported export document version {payload.version}; supported version: {EXPORT_FORMAT_VERSION}", + ) _check_duplicate_keys(incoming) # --- Phase 2: Matching --- @@ -175,10 +373,13 @@ def import_definitions( # Locked TDs surviving delete_unlocked are omitted entirely (per design) # --- Phase 3: Apply --- - should_apply = config.mode in (ImportMode.apply, ImportMode.apply_strict) has_skips = any(a.action == ImportAction.skip for a in actions) + if config.mode == ImportMode.apply_strict and has_skips: + # Built before apply, so created TDs have no target_id — correct, nothing was persisted. + # The post-apply build below is NOT redundant: _apply_actions fills planned.target for creates. + raise ImportStrictViolation(_build_response(actions)) - if should_apply and not (config.mode == ImportMode.apply_strict and has_skips): + if config.mode in (ImportMode.apply, ImportMode.apply_strict): _apply_actions(actions, test_suite, table_group, config) return _build_response(actions) @@ -218,9 +419,8 @@ def _check_duplicate_keys(incoming: list[TestDefinitionExport]) -> None: if td.last_auto_gen_date is not None: key = (td.test_type, td.table_name, td.column_name) if key in auto_keys: - raise api_error( - 400, - "duplicate_natural_key", + raise InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, f"Duplicate auto-gen key at index {idx}: ({td.test_type}, {td.table_name}, {td.column_name})", ) auto_keys.add(key) @@ -228,9 +428,8 @@ def _check_duplicate_keys(incoming: list[TestDefinitionExport]) -> None: if td.external_id is None: continue if td.external_id in manual_keys: - raise api_error( - 400, - "duplicate_natural_key", + raise InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, f"Duplicate external_id at index {idx}: {td.external_id}", ) manual_keys.add(td.external_id) diff --git a/testgen/common/test_result_disposition_service.py b/testgen/common/test_result_disposition_service.py new file mode 100644 index 00000000..4e1c2033 --- /dev/null +++ b/testgen/common/test_result_disposition_service.py @@ -0,0 +1,101 @@ +"""Shared test-result disposition service. + +Sets the disposition on test results and keeps the parent test definition's +active/lock state coupled: a "Muted" (``Disposition.INACTIVE``) disposition +deactivates the test definition and locks it against auto-regeneration; any +other value — or clearing the disposition — reactivates and unlocks it. Passed +test results are never dispositioned. Used by both the Streamlit UI +(test_results page) and the MCP tools. Must not import Streamlit. +""" +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from uuid import UUID + +from sqlalchemy import func, select, update + +from testgen.common.enums import Disposition +from testgen.common.models import get_current_session +from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_result import TestResult, TestResultStatus + + +@dataclass(frozen=True) +class DispositionUpdate: + """Outcome of a disposition write. + + ``matched`` is the number of non-Passed results updated; ``passed_skipped`` is + the number of Passed results that matched but were left unchanged. + """ + + matched: int + passed_skipped: int + + +def coupled_test_definition_state(disposition: Disposition | None) -> tuple[bool, bool]: + """Return the parent test definition's ``(test_active, lock_refresh)`` for a disposition. + + Muted (``INACTIVE``) → ``(False, True)``: deactivate and lock against + auto-regeneration. Any other value, or cleared (``None``) → ``(True, False)``. + """ + deactivate = disposition == Disposition.INACTIVE + return (not deactivate, deactivate) + + +def coerce_ui_disposition(value: str | None) -> Disposition | None: + """Map a UI disposition string to the stored value. + + The UI passes ``Confirmed`` / ``Dismissed`` / ``Inactive``, or ``No Decision`` / + empty / ``None`` to clear. Unknown values raise ``ValueError`` (caller's bug). + """ + if value in (None, "", "No Decision"): + return None + return Disposition(value) + + +def set_test_results_disposition( + test_result_ids: Sequence[str | UUID], + disposition: Disposition | None, +) -> DispositionUpdate: + """Set ``disposition`` on the given results and couple their parent test definitions. + + Passed results are excluded from the write. ``disposition=None`` clears it (NULL). + Returns the matched (non-Passed, updated) and passed-skipped counts. + """ + ids = [UUID(str(rid)) for rid in test_result_ids] + if not ids: + return DispositionUpdate(matched=0, passed_skipped=0) + + session = get_current_session() + + passed_skipped = session.scalar( + select(func.count()) + .select_from(TestResult) + .where(TestResult.id.in_(ids), TestResult.status == TestResultStatus.Passed) + ) or 0 + + # NULL result_status rows (e.g. training-mode results) are excluded by `!= Passed`, + # matching the prior UI behavior — disposition only applies to evaluated results. + non_passed = (TestResult.id.in_(ids), TestResult.status != TestResultStatus.Passed) + + tr_stmt = ( + update(TestResult) + .where(*non_passed) + .values(disposition=disposition.value if disposition is not None else None) + ) + matched = session.execute(tr_stmt).rowcount + + test_active, lock_refresh = coupled_test_definition_state(disposition) + affected_td_ids = select(TestResult.test_definition_id).where(*non_passed) + td_stmt = ( + update(TestDefinition) + .where(TestDefinition.id.in_(affected_td_ids)) + .values( + test_active=test_active, + lock_refresh=lock_refresh, + last_manual_update=datetime.now(UTC).replace(tzinfo=None), + ) + ) + session.execute(td_stmt) + + return DispositionUpdate(matched=matched, passed_skipped=passed_skipped) diff --git a/testgen/common/time_series_service.py b/testgen/common/time_series_service.py index aeabb180..ddf329c0 100644 --- a/testgen/common/time_series_service.py +++ b/testgen/common/time_series_service.py @@ -1,11 +1,11 @@ import logging -from datetime import datetime -import holidays import numpy as np import pandas as pd from statsmodels.tsa.statespace.sarimax import SARIMAX +from testgen.common.holiday_service import get_holiday_dates + LOG = logging.getLogger("testgen") # This is a heuristic minimum to get a reasonable prediction @@ -23,6 +23,7 @@ def get_sarimax_forecast( exclude_weekends: bool = False, holiday_codes: list[str] | None = None, tz: str | None = None, + event_space: bool = False, ) -> pd.DataFrame: """ # Parameters @@ -33,6 +34,14 @@ def get_sarimax_forecast( :param exclude_weekends: Whether weekends should be considered exogenous when training the model and forecasting. :param holiday_codes: List of country or financial market codes defining holidays to be considered exogenous when training the model and forecasting. :param tz: IANA timezone (e.g. "America/New_York") for day-of-week/holiday checks. Naive timestamps are treated as UTC and converted to this timezone before determining weekday/holiday status. + :param event_space: When False (default), resample the series onto a regular calendar grid and + linearly interpolate gaps before fitting — appropriate for series sampled at a + steady cadence, and where weekend/holiday exogenous flags apply. When True, fit + in event-space: one model step per observation, no interpolation, no calendar + exog. Required for irregularly-spaced series (e.g. the freshness-update points of + a refresh-driven table) where interpolation would smooth multi-period jumps into + small uniform increments and collapse the very jump variance the forecast must + capture. # Return value Returns a Pandas dataframe with forecast DatetimeIndex, "mean" column, and "se" (standard error) column. @@ -40,23 +49,33 @@ def get_sarimax_forecast( if len(history) < MIN_TRAIN_VALUES: raise NotEnoughData("Not enough data points in history.") - # statsmodels requires DatetimeIndex with a regular frequency - # Resample the data to get a regular time series - datetimes = history.index.to_series() - frequency = infer_frequency(datetimes) - resampled_history = history.resample(frequency).mean().interpolate(method="linear") - - if len(resampled_history) < MIN_TRAIN_VALUES: - raise NotEnoughData("Not enough data points after resampling.") + if event_space: + # Event-space: one step per observation. Map the values onto a synthetic regular grid + # stepped by the median observed interval, so forecast timestamps remain realistic while + # the model sees each observation as a single step (no interpolated points between them). + median_step = history.index.to_series().diff().median() + if pd.isna(median_step) or median_step <= pd.Timedelta(0): + median_step = pd.Timedelta(days=1) + step = median_step + synthetic_index = pd.date_range(end=history.index[-1], periods=len(history), freq=step) + train_history = pd.DataFrame(history.iloc[:, 0].values, index=synthetic_index, columns=[history.columns[0]]) + else: + # statsmodels requires DatetimeIndex with a regular frequency + # Resample the data to get a regular time series + datetimes = history.index.to_series() + frequency = infer_frequency(datetimes) + train_history = history.resample(frequency).mean().interpolate(method="linear") + if len(train_history) < MIN_TRAIN_VALUES: + raise NotEnoughData("Not enough data points after resampling.") + step = pd.to_timedelta(frequency) # Generate DatetimeIndex with future dates - forecast_start = resampled_history.index[-1] + pd.to_timedelta(frequency) - forecast_index = pd.date_range(start=forecast_start, periods=num_forecast, freq=frequency) + forecast_index = pd.date_range(start=train_history.index[-1] + step, periods=num_forecast, freq=step) - # Detect holidays in entire date range + # Detect holidays in entire date range (calendar-aware path only) holiday_dates = None - if holiday_codes: - all_dates_index = resampled_history.index.append(forecast_index) + if not event_space and holiday_codes: + all_dates_index = train_history.index.append(forecast_index) holiday_dates = get_holiday_dates(holiday_codes, all_dates_index) def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: @@ -71,11 +90,13 @@ def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: exog.loc[pd.Index(check_index.date).isin(holiday_dates), "is_excluded"] = 1 return exog - exog_train = get_exog_flags(resampled_history.index) + # Calendar exogenous flags only apply when fitting against real calendar time. + exog_train = None if event_space else get_exog_flags(train_history.index) + exog_forecast = None if event_space else get_exog_flags(forecast_index) # When seasonal_order is not specified, this is effectively the ARIMAX model model = SARIMAX( - resampled_history.iloc[:, 0], + train_history.iloc[:, 0], exog=exog_train, # This is a good starting point according to Gemini - tune if needed order=(1, 1, 1), @@ -85,12 +106,6 @@ def get_exog_flags(index: pd.DatetimeIndex) -> pd.DataFrame: ) fitted_model = model.fit(disp=False) - forecast_index = pd.date_range( - start=resampled_history.index[-1] + pd.to_timedelta(frequency), - periods=num_forecast, - freq=frequency - ) - exog_forecast = get_exog_flags(forecast_index) forecast = fitted_model.get_forecast(steps=num_forecast, exog=exog_forecast) results = pd.DataFrame(index=forecast_index) @@ -137,31 +152,3 @@ def infer_frequency(datetime_series: pd.Series) -> str: return frequency if frequency != "0min" else f"{int(total_seconds)}S" -def get_holiday_dates(holiday_codes: list[str], datetime_index: pd.DatetimeIndex) -> set[datetime]: - years = list(range(datetime_index.year.min(), datetime_index.year.max() + 1)) - - holiday_dates = set() - if holiday_codes: - for code in holiday_codes: - code = code.strip().upper() - found = False - - try: - country_holidays = holidays.country_holidays(code, years=years) - holiday_dates.update(country_holidays.keys()) - found = True - except NotImplementedError: - pass # Not a valid country code - - if not found: - try: - financial_holidays = holidays.financial_holidays(code, years=years) - holiday_dates.update(financial_holidays.keys()) - found = True - except NotImplementedError: - pass # Not a valid financial code - - if not found: - LOG.warning(f"Holiday code '{code}' could not be resolved as a country or financial market") - - return holiday_dates diff --git a/testgen/mcp/exceptions.py b/testgen/mcp/exceptions.py index fc89f98a..713f5b40 100644 --- a/testgen/mcp/exceptions.py +++ b/testgen/mcp/exceptions.py @@ -20,6 +20,14 @@ class MCPUserError(Exception): """ +class MCPAuthenticationError(MCPUserError): + """Token authenticated at the transport but failed authorization at the tool boundary. + + Raised when the token's user no longer exists or the token was revoked — distinct from + per-project permission denial. Carries a uniform, re-authenticate message for the client. + """ + + class MCPPermissionDenied(MCPUserError): """Raised when access is denied due to insufficient project permissions.""" diff --git a/testgen/mcp/permissions.py b/testgen/mcp/permissions.py index 0850c753..6042d5dd 100644 --- a/testgen/mcp/permissions.py +++ b/testgen/mcp/permissions.py @@ -2,14 +2,17 @@ import contextvars import functools +import logging from collections.abc import Callable from dataclasses import dataclass from testgen.common.models.project_membership import ProjectMembership from testgen.common.models.user import User -from testgen.mcp.exceptions import MCPPermissionDenied +from testgen.mcp.exceptions import MCPAuthenticationError, MCPPermissionDenied from testgen.utils.plugins import PluginHook +LOG = logging.getLogger("testgen") + _NOT_SET = object() _mcp_username: contextvars.ContextVar[str | None] = contextvars.ContextVar("mcp_username", default=None) @@ -77,13 +80,18 @@ def set_mcp_token(token: str | None) -> None: _mcp_token.set(token) +def get_mcp_username() -> str | None: + """Return the authenticated username for the current MCP request (or None).""" + return _mcp_username.get() + + def get_authorized_mcp_user() -> User: """Get the authenticated and authorized User for the current MCP request. Checks user existence and token revocation status. Must be called within @with_database_session scope. """ - from testgen.common.auth import authorize_token + from testgen.common.auth import AuthError, authorize_token from testgen.common.models import get_current_session username = _mcp_username.get() @@ -92,7 +100,13 @@ def get_authorized_mcp_user() -> User: token_str = _mcp_token.get() session = get_current_session() - return authorize_token(token_str or "", username, session) + try: + return authorize_token(token_str or "", username, session) + except AuthError as err: + LOG.warning("MCP token authorization failed: %s", err) + raise MCPAuthenticationError( + "Authentication failed: your access token is no longer valid. Please sign in again." + ) from err def _compute_project_permissions(user: User, permission: str) -> ProjectPermissions: @@ -129,17 +143,29 @@ def mcp_permission(permission: str) -> Callable: Raises ``MCPPermissionDenied`` if the user has no projects with the required permission. Other ``MCPPermissionDenied`` exceptions from tool code propagate through — the ``safe_tool`` error boundary handles conversion to text. + + ``global_admin`` is a user-level flag, not a per-project role: the check + consults ``User.is_global_admin`` and the resulting ``ProjectPermissions`` + has empty ``memberships``. Tools gated on ``global_admin`` operate above + project scope and should not call ``get_project_permissions()``. """ def decorator(fn: Callable) -> Callable: @functools.wraps(fn) def wrapper(*args, **kwargs): user = get_authorized_mcp_user() - perms = _compute_project_permissions(user, permission) - if not perms.allowed_codes: - raise MCPPermissionDenied( - "Your role does not include the necessary permission for this operation on any project." - ) + if permission == "global_admin": + if not user.is_global_admin: + raise MCPPermissionDenied( + "Your role does not include the necessary permission for this operation." + ) + perms = ProjectPermissions(memberships={}, permission=permission, username=user.username) + else: + perms = _compute_project_permissions(user, permission) + if not perms.allowed_codes: + raise MCPPermissionDenied( + "Your role does not include the necessary permission for this operation on any project." + ) tok = _mcp_project_permissions.set(perms) try: return fn(*args, **kwargs) diff --git a/testgen/mcp/server.py b/testgen/mcp/server.py index 5e91b2f1..3d5c4287 100644 --- a/testgen/mcp/server.py +++ b/testgen/mcp/server.py @@ -1,4 +1,7 @@ +import functools import logging +import time +from enum import StrEnum from urllib.parse import urlparse from mcp.server.auth.provider import AccessToken @@ -9,11 +12,27 @@ from starlette.applications import Starlette from testgen import settings -from testgen.common.auth import decode_jwt_token -from testgen.mcp.permissions import set_mcp_token, set_mcp_username +from testgen.common.auth import AuthError, decode_jwt_token +from testgen.common.mixpanel_service import MixpanelService +from testgen.mcp.exceptions import MCPPermissionDenied, MCPUserError +from testgen.mcp.permissions import get_mcp_username, set_mcp_token, set_mcp_username LOG = logging.getLogger("testgen") + +class MCPCallStatus(StrEnum): + SUCCESS = "success" + ERROR = "error" + PERMISSION_DENIED = "permission_denied" + USER_ERROR = "user_error" + + +class HandlerKind(StrEnum): + TOOL = "tool" + RESOURCE = "resource" + PROMPT = "prompt" + + SERVER_INSTRUCTIONS = """\ TestGen is a data quality platform that profiles databases, generates tests, and monitors tables. @@ -65,7 +84,7 @@ async def verify_token(self, token: str) -> AccessToken | None: scopes=[], expires_at=int(payload["exp"]), ) - except (ValueError, KeyError): + except (AuthError, KeyError): return None @@ -93,8 +112,12 @@ def _build_transport_security() -> TransportSecuritySettings: netloc = parsed.netloc scheme = parsed.scheme or "http" + # base_host (bare) is allowed alongside netloc: when BASE_URL carries a non-default + # port, netloc and the :* wildcard only match a host that includes a port, but some + # clients (e.g. hosted MCP gateways) send a port-less Host/Origin for their own host. allowed_hosts: set[str] = { netloc, + base_host, f"{base_host}:*", "127.0.0.1:*", "localhost:*", @@ -102,6 +125,7 @@ def _build_transport_security() -> TransportSecuritySettings: } allowed_origins: set[str] = { f"{scheme}://{netloc}", + f"{scheme}://{base_host}", "http://127.0.0.1:*", "https://127.0.0.1:*", "http://localhost:*", "https://localhost:*", "http://[::1]:*", "https://[::1]:*", @@ -110,6 +134,9 @@ def _build_transport_security() -> TransportSecuritySettings: host_pattern = host if ":" in host else f"{host}:*" allowed_hosts.add(host_pattern) allowed_origins.update({f"http://{host_pattern}", f"https://{host_pattern}"}) + bare = host.split(":", 1)[0] + allowed_hosts.add(bare) + allowed_origins.update({f"http://{bare}", f"https://{bare}"}) return TransportSecuritySettings( enable_dns_rebinding_protection=True, @@ -118,6 +145,42 @@ def _build_transport_security() -> TransportSecuritySettings: ) +def _instrument(fn, kind: HandlerKind): + """Wrap an MCP handler to emit one ``mcp-call`` event per invocation. + + Sits inside ``mcp_error_handler`` so it observes the raised exception type + before it is converted to a text response. ``kind`` distinguishes tools, + resources, and prompts in the emitted event. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + start = time.monotonic() + status = MCPCallStatus.SUCCESS + try: + return fn(*args, **kwargs) + except MCPPermissionDenied: + status = MCPCallStatus.PERMISSION_DENIED + raise + except MCPUserError: + status = MCPCallStatus.USER_ERROR + raise + except Exception: + status = MCPCallStatus.ERROR + raise + finally: + MixpanelService().send_event( + "mcp-call", + kind=kind, + handler_name=fn.__name__, + username=get_mcp_username(), + latency_ms=int((time.monotonic() - start) * 1000), + status=status, + ) + + return wrapper + + def build_mcp_server( api_base_url: str, server_url: str | None = None, @@ -137,7 +200,20 @@ def build_mcp_server( profiling_overview, table_health, ) - from testgen.mcp.tools.discovery import get_data_inventory, list_projects, list_tables, list_test_suites + from testgen.mcp.tools.connections import ( + get_connection, + list_connections, + test_connection, + ) + from testgen.mcp.tools.data_catalog import update_catalog_metadata + from testgen.mcp.tools.discovery import ( + get_data_inventory, + get_project, + get_test_suite, + list_projects, + list_tables, + list_test_suites, + ) from testgen.mcp.tools.execution import ( cancel_profiling_run, cancel_test_run, @@ -151,6 +227,17 @@ def build_mcp_server( search_hygiene_issues, update_hygiene_issue, ) + from testgen.mcp.tools.monitors import ( + disable_monitors, + enable_monitors, + get_monitor_settings, + get_monitor_summary, + list_monitor_events, + list_monitor_schema_changes, + list_monitored_tables, + list_monitors, + update_monitor_settings, + ) from testgen.mcp.tools.notifications import ( create_notification, delete_notification, @@ -164,6 +251,7 @@ def build_mcp_server( get_schema_history, ) from testgen.mcp.tools.profiling import ( + generate_create_table_script, get_column_frequent_values, get_column_patterns, get_column_profile_detail, @@ -174,6 +262,7 @@ def build_mcp_server( list_profiling_summaries, search_columns, ) + from testgen.mcp.tools.projects import update_project from testgen.mcp.tools.quality_scores import ( create_scorecard, delete_scorecard, @@ -184,6 +273,8 @@ def build_mcp_server( ) from testgen.mcp.tools.reference import ( column_profile_fields_resource, + connection_parameters_index_resource, + connection_parameters_resource, get_test_type, glossary_resource, hygiene_issue_types_resource, @@ -197,13 +288,22 @@ def build_mcp_server( list_schedules, update_schedule, ) - from testgen.mcp.tools.source_data import get_source_data, get_source_data_query + from testgen.mcp.tools.source_data import get_source_data, get_source_data_query, get_table_sample + from testgen.mcp.tools.table_groups import ( + create_table_group, + get_table_group, + list_table_groups, + preview_table_group, + update_table_group, + ) from testgen.mcp.tools.test_definitions import ( bulk_update_tests, create_test, create_test_note, delete_test_note, + export_tests, get_test, + import_tests, list_test_notes, list_test_types, list_tests, @@ -212,14 +312,17 @@ def build_mcp_server( validate_custom_test, ) from testgen.mcp.tools.test_results import ( + bulk_update_test_results, compare_test_runs, get_failure_summary, get_failure_trend, list_test_result_history, list_test_results, search_test_results, + update_test_result, ) from testgen.mcp.tools.test_runs import get_test_run, list_test_runs + from testgen.mcp.tools.test_suites import create_test_suite, update_test_suite if server_url is None: server_url = f"{api_base_url}/mcp" @@ -236,86 +339,117 @@ def build_mcp_server( ) _configure_mcp_logging() - def safe_tool(fn): - mcp.tool()(mcp_error_handler(fn)) + def tool_wrapper(fn): + mcp.tool()(mcp_error_handler(_instrument(fn, HandlerKind.TOOL))) def safe_resource(uri, fn): - mcp.resource(uri)(mcp_error_handler(fn)) + mcp.resource(uri)(mcp_error_handler(_instrument(fn, HandlerKind.RESOURCE))) def safe_prompt(fn): - mcp.prompt()(mcp_error_handler(fn)) + mcp.prompt()(mcp_error_handler(_instrument(fn, HandlerKind.PROMPT))) # Tools - safe_tool(get_data_inventory) - safe_tool(list_projects) - safe_tool(list_tables) - safe_tool(list_test_suites) - safe_tool(list_test_runs) - safe_tool(get_test_run) - safe_tool(list_test_results) - safe_tool(list_test_result_history) - safe_tool(get_failure_summary) - safe_tool(search_test_results) - safe_tool(get_failure_trend) - safe_tool(compare_test_runs) - safe_tool(get_test_type) - safe_tool(get_source_data) - safe_tool(get_source_data_query) - safe_tool(list_tests) - safe_tool(get_test) - safe_tool(list_test_notes) - safe_tool(list_test_types) - safe_tool(get_table) - safe_tool(list_column_profiles) - safe_tool(list_profiling_summaries) - safe_tool(list_profiling_runs) - safe_tool(get_profiling_run) - safe_tool(get_column_profile_detail) - safe_tool(get_column_frequent_values) - safe_tool(get_column_patterns) - safe_tool(search_columns) - safe_tool(compare_profiling_runs) - safe_tool(get_profiling_trends) - safe_tool(get_schema_history) - safe_tool(run_tests) - safe_tool(run_profiling) - safe_tool(cancel_test_run) - safe_tool(cancel_profiling_run) - safe_tool(generate_tests) - safe_tool(create_test) - safe_tool(update_test) - safe_tool(validate_custom_test) - safe_tool(bulk_update_tests) - safe_tool(create_test_note) - safe_tool(update_test_note) - safe_tool(delete_test_note) - safe_tool(list_hygiene_issues) - safe_tool(get_hygiene_issue) - safe_tool(search_hygiene_issues) - safe_tool(update_hygiene_issue) - safe_tool(create_profiling_schedule) - safe_tool(create_test_run_schedule) - safe_tool(list_schedules) - safe_tool(get_schedule) - safe_tool(update_schedule) - safe_tool(delete_schedule) - safe_tool(get_quality_scores) - safe_tool(list_scorecards) - safe_tool(get_scorecard) - safe_tool(create_scorecard) - safe_tool(update_scorecard) - safe_tool(delete_scorecard) - safe_tool(list_notifications) - safe_tool(get_notification) - safe_tool(create_notification) - safe_tool(update_notification) - safe_tool(delete_notification) + tool_wrapper(get_data_inventory) + tool_wrapper(list_projects) + tool_wrapper(get_project) + tool_wrapper(list_tables) + tool_wrapper(list_test_suites) + tool_wrapper(get_test_suite) + tool_wrapper(list_test_runs) + tool_wrapper(get_test_run) + tool_wrapper(list_test_results) + tool_wrapper(list_test_result_history) + tool_wrapper(get_failure_summary) + tool_wrapper(search_test_results) + tool_wrapper(get_failure_trend) + tool_wrapper(compare_test_runs) + tool_wrapper(update_test_result) + tool_wrapper(bulk_update_test_results) + tool_wrapper(get_test_type) + tool_wrapper(get_source_data) + tool_wrapper(get_source_data_query) + tool_wrapper(list_tests) + tool_wrapper(get_test) + tool_wrapper(list_test_notes) + tool_wrapper(list_test_types) + tool_wrapper(get_table) + tool_wrapper(generate_create_table_script) + tool_wrapper(get_table_sample) + tool_wrapper(list_column_profiles) + tool_wrapper(list_profiling_summaries) + tool_wrapper(list_profiling_runs) + tool_wrapper(get_profiling_run) + tool_wrapper(get_column_profile_detail) + tool_wrapper(get_column_frequent_values) + tool_wrapper(get_column_patterns) + tool_wrapper(search_columns) + tool_wrapper(compare_profiling_runs) + tool_wrapper(get_profiling_trends) + tool_wrapper(get_schema_history) + tool_wrapper(get_monitor_summary) + tool_wrapper(list_monitored_tables) + tool_wrapper(list_monitor_events) + tool_wrapper(list_monitors) + tool_wrapper(list_monitor_schema_changes) + tool_wrapper(enable_monitors) + tool_wrapper(get_monitor_settings) + tool_wrapper(update_monitor_settings) + tool_wrapper(disable_monitors) + tool_wrapper(run_tests) + tool_wrapper(run_profiling) + tool_wrapper(cancel_test_run) + tool_wrapper(cancel_profiling_run) + tool_wrapper(generate_tests) + tool_wrapper(create_test) + tool_wrapper(update_test) + tool_wrapper(validate_custom_test) + tool_wrapper(bulk_update_tests) + tool_wrapper(export_tests) + tool_wrapper(import_tests) + tool_wrapper(create_test_note) + tool_wrapper(update_test_note) + tool_wrapper(delete_test_note) + tool_wrapper(list_hygiene_issues) + tool_wrapper(get_hygiene_issue) + tool_wrapper(search_hygiene_issues) + tool_wrapper(update_hygiene_issue) + tool_wrapper(create_profiling_schedule) + tool_wrapper(create_test_run_schedule) + tool_wrapper(list_schedules) + tool_wrapper(get_schedule) + tool_wrapper(update_schedule) + tool_wrapper(delete_schedule) + tool_wrapper(get_quality_scores) + tool_wrapper(list_scorecards) + tool_wrapper(get_scorecard) + tool_wrapper(create_scorecard) + tool_wrapper(update_scorecard) + tool_wrapper(delete_scorecard) + tool_wrapper(list_notifications) + tool_wrapper(get_notification) + tool_wrapper(create_notification) + tool_wrapper(update_notification) + tool_wrapper(delete_notification) + tool_wrapper(list_connections) + tool_wrapper(get_connection) + tool_wrapper(test_connection) + tool_wrapper(list_table_groups) + tool_wrapper(get_table_group) + tool_wrapper(create_table_group) + tool_wrapper(update_table_group) + tool_wrapper(preview_table_group) + tool_wrapper(update_project) + tool_wrapper(create_test_suite) + tool_wrapper(update_test_suite) + tool_wrapper(update_catalog_metadata) # Resources safe_resource("testgen://test-types", test_types_resource) safe_resource("testgen://hygiene-issue-types", hygiene_issue_types_resource) safe_resource("testgen://column-profile-fields", column_profile_fields_resource) safe_resource("testgen://glossary", glossary_resource) + safe_resource("testgen://connection-parameters", connection_parameters_index_resource) + safe_resource("testgen://connection-parameters/{flavor}", connection_parameters_resource) # Prompts safe_prompt(health_check) @@ -325,6 +459,17 @@ def safe_prompt(fn): safe_prompt(profiling_overview) safe_prompt(hygiene_triage) + # Register plugin-provided tools through the same tool_wrapper as the core tools above. + from testgen.utils.plugins import discover + + for plugin in discover(): + try: + spec = plugin.load() + for tool in spec.get_mcp_tools(): + tool_wrapper(tool) + except Exception: + LOG.warning("Plugin %s failed to load; skipping its MCP tools", plugin.package) + return mcp diff --git a/testgen/mcp/services/inventory_service.py b/testgen/mcp/services/inventory_service.py index cee1c7a6..6d504ed3 100644 --- a/testgen/mcp/services/inventory_service.py +++ b/testgen/mcp/services/inventory_service.py @@ -117,9 +117,9 @@ def get_inventory( lines.append("_No table groups._\n") continue - for _conn_id, conn in proj["connections"].items(): + for conn_id, conn in proj["connections"].items(): if can_view: - lines.append(f"### Connection: {conn['name']}\n") + lines.append(f"### Connection: {conn['name']} (id: `{conn_id}`)\n") if not conn["groups"]: if can_view: diff --git a/testgen/mcp/tools/common.py b/testgen/mcp/tools/common.py index e08ca861..25b302fa 100644 --- a/testgen/mcp/tools/common.py +++ b/testgen/mcp/tools/common.py @@ -1,12 +1,25 @@ +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field from datetime import date, datetime from enum import StrEnum +from typing import NoReturn, TypeVar from uuid import UUID from sqlalchemy import select from testgen.common.date_service import parse_since -from testgen.common.enums import Disposition, ImpactDimension, IssueLikelihood, JobStatus, PiiRisk, QualityDimension +from testgen.common.enums import ( + Disposition, + ImpactDimension, + IssueLikelihood, + JobStatus, + MonitorType, + PiiRisk, + QualityDimension, +) +from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL, SqlFlavorLabel from testgen.common.models import get_current_session +from testgen.common.models.connection import Connection from testgen.common.models.data_column import ( GENERAL_TYPE_TO_CODE, ColumnOrderBy, @@ -14,7 +27,7 @@ ProfileMetric, SuggestedDataType, ) -from testgen.common.models.hygiene_issue import HygieneIssueType +from testgen.common.models.hygiene_issue import HygieneIssue, HygieneIssueType from testgen.common.models.notification_settings import ( MonitorNotificationTrigger, NotificationEvent, @@ -23,14 +36,16 @@ TestRunNotificationTrigger, ) from testgen.common.models.profiling_run import ProfilingRun +from testgen.common.models.project import Project from testgen.common.models.scheduler import SCHEDULABLE_JOB_KEYS, JobSchedule from testgen.common.models.scores import ScoreCategory, ScoreDefinition from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_definition import TestDefinition, TestDefinitionNote, TestType -from testgen.common.models.test_result import TestResultStatus +from testgen.common.models.test_definition import Severity, TestDefinition, TestDefinitionNote, TestType +from testgen.common.models.test_result import TestResult, TestResultStatus from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions +from testgen.mcp.tools.markdown import MdDoc # User-facing label for ``Disposition.INACTIVE`` is "Muted" — accept that label on input. _DISPOSITION_USER_TO_DB: dict[str, Disposition] = { @@ -54,6 +69,7 @@ class DocGroup(StrEnum): DISCOVER = "Discover what TestGen knows about" INVESTIGATE = "Investigate quality issues" BROWSE_PROFILING = "Browse profiling results" + MONITORS = "Browse monitor health and events" TRIGGER = "Trigger profiling, tests, and test generation" SCORING = "Track data quality scores" MANAGE = "Manage TestGen configuration" @@ -66,6 +82,17 @@ def parse_uuid(value: str, label: str = "ID") -> UUID: raise MCPUserError(f"Invalid {label}: `{value}` is not a valid UUID.") from err +_ParsedEnum = TypeVar("_ParsedEnum", bound=StrEnum) + + +def parse_enum(value: str, enum_cls: type[_ParsedEnum], label: str) -> _ParsedEnum: + try: + return enum_cls(value) + except ValueError as err: + valid = ", ".join(member.value for member in enum_cls) + raise MCPUserError(f"Invalid {label} `{value}`. Valid values: {valid}") from err + + def parse_result_status(value: str) -> TestResultStatus: try: return TestResultStatus(value) @@ -107,6 +134,15 @@ def parse_quality_dimension(value: str) -> QualityDimension: raise MCPUserError(f"Invalid quality_dimension `{value}`. Valid values: {valid}") from err +def parse_severity(value: str) -> Severity: + """Validate a test-suite default severity. Accepts ``Fail`` or ``Warning``.""" + try: + return Severity(value) + except ValueError as err: + valid = ", ".join(s.value for s in Severity) + raise MCPUserError(f"Invalid severity `{value}`. Valid values: {valid}") from err + + class ScoreGroupBy(StrEnum): """User-facing values accepted for the ``group_by`` argument on quality-score rollups.""" @@ -122,6 +158,7 @@ class ScoreGroupBy(StrEnum): STAKEHOLDER_GROUP = "Stakeholder Group" TRANSFORM_LEVEL = "Transform Level" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" # Translates the user-facing label to the internal DB column name used by @@ -139,6 +176,7 @@ class ScoreGroupBy(StrEnum): ScoreGroupBy.STAKEHOLDER_GROUP: "stakeholder_group", ScoreGroupBy.TRANSFORM_LEVEL: "transform_level", ScoreGroupBy.DATA_PRODUCT: "data_product", + ScoreGroupBy.DATA_CLASSIFICATION: "data_classification", } @@ -161,6 +199,7 @@ class ScoreFilterField(StrEnum): STAKEHOLDER_GROUP = "Stakeholder Group" TRANSFORM_LEVEL = "Transform Level" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" SCORE_FILTER_FIELD_TO_COLUMN: dict[ScoreFilterField, str] = { @@ -174,6 +213,7 @@ class ScoreFilterField(StrEnum): ScoreFilterField.STAKEHOLDER_GROUP: "stakeholder_group", ScoreFilterField.TRANSFORM_LEVEL: "transform_level", ScoreFilterField.DATA_PRODUCT: "data_product", + ScoreFilterField.DATA_CLASSIFICATION: "data_classification", } @@ -197,6 +237,7 @@ class ScoreCategoryArg(StrEnum): QUALITY_DIMENSION = "Quality Dimension" IMPACT_DIMENSION = "Impact Dimension" DATA_PRODUCT = "Data Product" + DATA_CLASSIFICATION = "Data Classification" SCORE_CATEGORY_ARG_TO_COLUMN: dict[ScoreCategoryArg, str] = { @@ -211,6 +252,7 @@ class ScoreCategoryArg(StrEnum): ScoreCategoryArg.QUALITY_DIMENSION: "dq_dimension", ScoreCategoryArg.IMPACT_DIMENSION: "impact_dimension", ScoreCategoryArg.DATA_PRODUCT: "data_product", + ScoreCategoryArg.DATA_CLASSIFICATION: "data_classification", } @@ -297,6 +339,22 @@ def parse_run_status_filter(value: str) -> list[JobStatus]: return statuses +class FailureGroupBy(StrEnum): + """User-facing values accepted for the ``group_by`` argument on ``get_failure_summary``.""" + + TEST_TYPE = "test_type" + TABLE = "table" + COLUMN = "column" + + +def parse_failure_group_by(value: str) -> FailureGroupBy: + try: + return FailureGroupBy(value) + except ValueError as err: + valid = ", ".join(g.value for g in FailureGroupBy) + raise MCPUserError(f"Invalid group_by `{value}`. Valid values: {valid}") from err + + def format_run_duration(started_at: datetime | None, completed_at: datetime | None) -> str | None: """Render an elapsed duration as ``Xs`` / ``Xm Ys`` / ``Xh Ym``. Returns ``None`` if either bound is missing.""" if not started_at or not completed_at: @@ -321,6 +379,53 @@ def next_scheduled_run( return min(s.get_sample_triggering_timestamps(1)[0] for s in schedules) +# Monitor type — internal ``test_type`` value ↔ user-facing short label (the form +# callers pass and the form rendered in output). +_MONITOR_TYPE_USER_TO_DB: dict[str, MonitorType] = { + "freshness": MonitorType.FRESHNESS, + "volume": MonitorType.VOLUME, + "schema": MonitorType.SCHEMA, + "metric": MonitorType.METRIC, +} + + +def parse_monitor_type(value: str, label: str = "monitor_type") -> MonitorType: + """Validate a user-facing monitor type label and return the stored ``MonitorType``. + + Accepts ``Freshness`` / ``Volume`` / ``Schema`` / ``Metric`` (Title Case + as rendered in tool output) or the equivalent lowercase forms. ``label`` + names the caller's argument in the error message — pass + ``"anomaly_type"`` when the public arg is named differently from + ``monitor_type``. + """ + db_value = _MONITOR_TYPE_USER_TO_DB.get(value.lower()) if isinstance(value, str) else None + if db_value is None: + valid = ", ".join(_MONITOR_TYPE_USER_TO_DB) + raise MCPUserError(f"Invalid {label} `{value}`. Valid values: {valid}") + return db_value + + +class MonitorTableSort(StrEnum): + """User-facing values accepted for the ``sort_by`` argument on ``list_monitored_tables``. + + When an ``anomaly_type`` filter is set, ``anomaly_count_desc`` sorts by that type's + count; with no filter, it sorts by total anomalies across all types. + """ + + TABLE_NAME = "table_name" + ANOMALY_COUNT_DESC = "anomaly_count_desc" + LATEST_UPDATE_DESC = "latest_update_desc" + ROW_COUNT_CHANGE_DESC = "row_count_change_desc" + + +def parse_monitor_table_sort(value: str) -> MonitorTableSort: + try: + return MonitorTableSort(value) + except ValueError as err: + valid = ", ".join(s.value for s in MonitorTableSort) + raise MCPUserError(f"Invalid sort_by `{value}`. Valid values: {valid}") from err + + def parse_disposition(value: str) -> Disposition: """Validate a user-facing disposition label and return the stored ``Disposition``. @@ -334,6 +439,25 @@ def parse_disposition(value: str) -> Disposition: return db_value +_NO_DECISION = "No Decision" + + +def parse_test_result_disposition(value: str) -> Disposition | None: + """Validate a user-facing test-result disposition and return the stored value. + + Accepts ``Confirmed``, ``Dismissed``, ``Muted``, and ``No Decision``. ``Muted`` + maps to ``Disposition.INACTIVE``; ``No Decision`` clears the disposition (returns + ``None`` → NULL). + """ + if value == _NO_DECISION: + return None + db_value = _DISPOSITION_USER_TO_DB.get(value) + if db_value is None: + valid = ", ".join([*_DISPOSITION_USER_TO_DB, _NO_DECISION]) + raise MCPUserError(f"Invalid disposition `{value}`. Valid values: {valid}") + return db_value + + def format_disposition(value: Disposition | str) -> str: """Map a stored disposition to its user-facing label (``INACTIVE`` → "Muted").""" try: @@ -491,10 +615,17 @@ def resolve_issue_type(name: str) -> str: def format_page_info(total: int, page: int, limit: int) -> str: - """Shared pagination summary line for MCP tool output.""" + """Shared pagination summary line for MCP tool output. + + Returns empty for a zero total *or* when ``page`` is past the last page \u2014 + the caller's own "no rows on page N" message is more useful than a + ``Showing 6\u20135 of 5`` nonsense range. + """ if total == 0: return "" start = (page - 1) * limit + 1 + if start > total: + return "" end = min(start + limit - 1, total) return f"Showing {start}\u2013{end} of {total} (page {page})." @@ -507,10 +638,101 @@ def format_page_footer(total: int, page: int, limit: int) -> str: return f"_Page {page} of {total_pages}. Use `page={page + 1}` for more._" +def _default_render_diff_value(value: object) -> str | None: + """Default formatter for before/after cells in :func:`render_diff_table`. + + * ``bool`` → ``"Yes"`` / ``"No"`` + * ``None`` or ``""`` → ``None`` (rendered as em-dash by the table cell) + * else → ``str(value)`` + """ + if isinstance(value, bool): + return "Yes" if value else "No" + if value is None or value == "": + return None + return str(value) + + +def render_diff_table( + doc: MdDoc, + before: Mapping[str, object], + after: Mapping[str, object], + *, + attrs: Sequence[str], + labels: Mapping[str, str], + secret_attrs: frozenset[str] = frozenset(), + value_renderer: Callable[[object], str | None] | None = None, +) -> bool: + """Emit a ``Field / Before / After`` table for attrs whose values changed. + + Update tools across the MCP surface share this shape: snapshot before, apply + field-by-field, snapshot after, render the delta. The shared helper keeps the + ordering, label lookup, and secret redaction consistent. + + Args: + doc: ``MdDoc`` to append the table to. + before / after: same-shape dicts of the snapshot keyed by attr name. + Only entries in ``attrs`` are inspected. + attrs: ordered tuple of attrs to consider; controls row order. + labels: attr → user-facing label for the first column. + secret_attrs: attrs whose values must not be echoed (API keys, passwords). + Rendered as ``[secret]`` when present, em-dash when absent — distinct + from a "rotated" cue. Tools that need to communicate rotation + specifically (e.g. ``update_connection``) keep their own renderer. + value_renderer: override for the non-secret cells. Defaults to + :func:`_default_render_diff_value`. + + Returns ``True`` if at least one row was rendered, ``False`` when nothing + changed (lets the caller emit a "No fields changed" message instead). + """ + changed = {attr for attr in attrs if before.get(attr) != after.get(attr)} + if not changed: + return False + render = value_renderer or _default_render_diff_value + rows: list[list[object]] = [] + for attr in attrs: + if attr not in changed: + continue + if attr in secret_attrs: + b = "[secret]" if before.get(attr) else None + a = "[secret]" if after.get(attr) else None + else: + b = render(before.get(attr)) + a = render(after.get(attr)) + rows.append([labels.get(attr, attr), b, a]) + doc.table(["Field", "Before", "After"], rows, code=[0]) + return True + + # Entity resolution helpers — see mcp-roadmap.md "Entity Resolution Helpers" guideline. # Extract a new resolve_ here when a second caller needs the same parse-uuid + # perm-scoped lookup + collapsed-error pattern. + +def resolve_project(project_code: str) -> Project: + """Resolve a project code to a Project, scoped to the user's allowed projects. + + Collapses missing-or-inaccessible into a single ``MCPResourceNotAccessible`` so + callers can't enumerate whether a project they can't see exists. + """ + perms = get_project_permissions() + project = Project.get(project_code, Project.project_code.in_(perms.allowed_codes)) + if project is None: + raise MCPResourceNotAccessible("Project", project_code) + return project + + +def resolve_connection(connection_id: int) -> Connection: + """Resolve a connection ID, collapsing missing-or-inaccessible into one error path.""" + perms = get_project_permissions() + conn = Connection.get( + connection_id, + Connection.project_code.in_(perms.allowed_codes), + ) + if conn is None: + raise MCPResourceNotAccessible("Connection", str(connection_id)) + return conn + + def resolve_table_group(table_group_id: str) -> TableGroup: """Resolve a TG ID, collapsing missing-or-inaccessible into one error path.""" tg_uuid = parse_uuid(table_group_id, "table_group_id") @@ -521,6 +743,16 @@ def resolve_table_group(table_group_id: str) -> TableGroup: return tg +def resolve_hygiene_issue(issue_id: str) -> HygieneIssue: + """Resolve a hygiene issue ID, collapsing missing-or-inaccessible into one error path.""" + issue_uuid = parse_uuid(issue_id, "issue_id") + perms = get_project_permissions() + issue = HygieneIssue.get(issue_uuid, HygieneIssue.project_code.in_(perms.allowed_codes)) + if issue is None: + raise MCPResourceNotAccessible("Hygiene issue", issue_id) + return issue + + def resolve_test_suite(test_suite_id: str) -> TestSuite: """Resolve a regular (non-monitor) test suite ID, collapsing missing-or-inaccessible into one error path.""" suite_uuid = parse_uuid(test_suite_id, "test_suite_id") @@ -535,6 +767,23 @@ def resolve_test_suite(test_suite_id: str) -> TestSuite: return suite +def resolve_monitored_table_group(table_group_id: str) -> tuple[TableGroup, TestSuite | None]: + """Resolve a table group ID and look up its linked monitor suite. + + Returns ``(table_group, monitor_suite)``. ``monitor_suite`` is ``None`` when the + table group has no ``monitor_test_suite_id`` set, or that pointer doesn't resolve + to an ``is_monitor=True`` suite — callers render the "not monitored" output in + that case rather than raising an inaccessible error. Raises + ``MCPResourceNotAccessible`` when the table group itself is missing or out of + scope. + """ + tg = resolve_table_group(table_group_id) + if tg.monitor_test_suite_id is None: + return tg, None + suite = TestSuite.get(tg.monitor_test_suite_id, TestSuite.is_monitor.is_(True)) + return tg, suite + + def resolve_profiling_run(job_execution_id: str) -> ProfilingRun: """Resolve a profiling run by id-or-JE-id, scoped to allowed projects. @@ -542,7 +791,7 @@ def resolve_profiling_run(job_execution_id: str) -> ProfilingRun: so callers don't leak existence of runs they shouldn't see. """ run_uuid = parse_uuid(job_execution_id, "job_execution_id") - run = ProfilingRun.get_by_id_or_job(run_uuid) + run = ProfilingRun.get(run_uuid) perms = get_project_permissions() if run is None or not perms.has_access(run.project_code): raise MCPResourceNotAccessible("Profiling run", job_execution_id) @@ -582,6 +831,28 @@ def resolve_test_definition(test_definition_id: str) -> TestDefinition: return td +def resolve_test_result(test_result_id: str) -> TestResult: + """Resolve a test result ID to the live ORM model, collapsing missing-or-inaccessible. + + Filters monitor suites and project access via the result's parent test suite. + """ + result_uuid = parse_uuid(test_result_id, "test_result_id") + perms = get_project_permissions() + query = ( + select(TestResult) + .join(TestSuite, TestResult.test_suite_id == TestSuite.id) + .where( + TestResult.id == result_uuid, + TestSuite.is_monitor.isnot(True), + TestSuite.project_code.in_(perms.allowed_codes), + ) + ) + result = get_current_session().scalars(query).first() + if result is None: + raise MCPResourceNotAccessible("Test result", test_result_id) + return result + + def resolve_test_note(test_note_id: str) -> TestDefinitionNote: """Resolve a test note ID to the live ORM model, collapsing missing-or-inaccessible. @@ -636,6 +907,52 @@ def resolve_notification(notification_id: str) -> NotificationSettings: return notif +def resolve_aggregate_scope( + project_code: str | None, + test_suite_id: str | None = None, + table_group_id: str | None = None, +) -> list[str]: + """Validate optional project / test-suite / table-group scope for a cross-run aggregation. + + Resolves any supplied suite or table group (existence + access via the + ``resolve_*`` helpers), and — when ``project_code`` is also given — requires the + resolved entity to belong to it, raising a clear error on a cross-project mismatch + so the query never silently returns empty. Returns the project codes to scope the + aggregation to. + """ + perms = get_project_permissions() + if project_code: + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + scoped_projects: set[str] = set() + if test_suite_id: + suite = resolve_test_suite(test_suite_id) + if project_code and suite.project_code != project_code: + raise MCPUserError( + f"Test suite `{test_suite_id}` belongs to project `{suite.project_code}`, not `{project_code}`." + ) + scoped_projects.add(suite.project_code) + if table_group_id: + table_group = resolve_table_group(table_group_id) + if project_code and table_group.project_code != project_code: + raise MCPUserError( + f"Table group `{table_group_id}` belongs to project `{table_group.project_code}`, not `{project_code}`." + ) + scoped_projects.add(table_group.project_code) + + if len(scoped_projects) > 1: + # Suite and table group resolve to different projects (only reachable when no + # project_code pins the scope). The two filters would AND to an empty result — + # reject rather than silently return nothing. + raise MCPUserError("The test suite and table group belong to different projects — narrow to one scope.") + + if project_code: + return [project_code] + if scoped_projects: + return list(scoped_projects) + return perms.allowed_codes + + # Notification event-type labels. class NotificationEventLabel(StrEnum): @@ -724,3 +1041,498 @@ def format_notification_trigger(event: NotificationEvent | str, settings: dict | if event_enum is NotificationEvent.monitor_run: return _MONITOR_TRIGGER_INTERNAL_TO_LABEL[MonitorNotificationTrigger(raw)].value return None + + +# Flavor display labels are the single source of truth in ``common/flavors.py`` +# (shared with the UI page). These maps just re-shape them for the MCP layer. +SQL_FLAVOR_CODE_TO_LABEL: dict[str, SqlFlavorLabel] = dict(FLAVOR_CODE_TO_LABEL) +SQL_FLAVOR_LABEL_TO_CODE: dict[SqlFlavorLabel, str] = { + label: code for code, label in FLAVOR_CODE_TO_LABEL.items() +} +SQL_FLAVOR_CODE_TO_FAMILY: dict[str, str] = dict(FLAVOR_CODE_TO_FAMILY) + + +def parse_sql_flavor(value: str) -> tuple[SqlFlavorLabel, str, str]: + """Validate a user-facing ``sql_flavor`` value and return ``(label, code, family)``.""" + try: + label = SqlFlavorLabel(value) + except ValueError as err: + valid = ", ".join(f.value for f in SqlFlavorLabel) + raise MCPUserError(f"Invalid sql_flavor `{value}`. Valid values: {valid}") from err + code = SQL_FLAVOR_LABEL_TO_CODE[label] + return label, code, SQL_FLAVOR_CODE_TO_FAMILY[code] + + +def format_flavor_label(sql_flavor_code: str | None) -> str: + """Map a stored ``sql_flavor_code`` to its user-facing display label. + + Returns the raw code as a fallback when the code is not in the registry — defensive + against a never-shipping-but-still-stored value rather than letting an LLM see ``None``. + """ + if sql_flavor_code is None: + return "" + label = SQL_FLAVOR_CODE_TO_LABEL.get(sql_flavor_code) + return label.value if label else sql_flavor_code + + +# =========================================================================== +# Connection-parameter contract (MCP input vocabulary) +# +# The per-flavor connection shape: which auth modes exist and which fields each +# needs, keyed by their UI labels (which double as ``connection_params`` keys). +# This is MCP-only input vocabulary + parsing — it mirrors the per-flavor JS +# forms in ``ui/static/js/components/connection_form.js`` (NOT sourced from +# Python), and drives both the connection tools and the +# ``testgen://connection-parameters/{flavor}`` resource. Field labels map to the +# (often leaky) ``Connection`` columns here so the tool's arg surface speaks the +# target-DB vocabulary. +# =========================================================================== + + +class ConnectionMode(StrEnum): + PASSWORD = "Password" # noqa: S105 — auth-mode label, not a credential + KEY_PAIR = "Key-Pair" + MANAGED_IDENTITY = "Managed Identity" + ACCESS_TOKEN = "Access Token" # noqa: S105 — auth-mode label, not a credential + SERVICE_PRINCIPAL = "Service Principal (OAuth)" + JWT_BEARER = "JWT Bearer Flow" + CLIENT_CREDENTIALS = "Client Credentials Flow" + SERVICE_ACCOUNT = "Service Account Key" + + +class Req(StrEnum): + REQUIRED = "required" # always required in this mode + REQUIRED_UNLESS_URL = "required_unless_url" # required only in host mode + OPTIONAL = "optional" + + +@dataclass(frozen=True) +class ConnField: + label: str # exact UI label == connection_params key + column: str # Connection ORM attribute it maps to + requirement: Req + secret: bool = False + + +@dataclass(frozen=True) +class FlavorMode: + mode: ConnectionMode | None # None for single-auth flavors + fields: tuple[ConnField, ...] # excludes the URL field + supports_url: bool # whether the URL alternative is offered + # Columns forced regardless of supplied params (auth flags, Databricks PAT user). + sets: dict[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True) +class FlavorSchema: + code: str + label: str + modes: tuple[FlavorMode, ...] + url_field: ConnField | None # the shared URL field when any mode supports_url + + +# -- reusable field definitions --------------------------------------------- + +_HOST = ConnField("Host", "project_host", Req.REQUIRED_UNLESS_URL) +_PORT = ConnField("Port", "project_port", Req.REQUIRED_UNLESS_URL) +_DATABASE = ConnField("Database", "project_db", Req.REQUIRED_UNLESS_URL) +_SERVICE_NAME = ConnField("Service Name", "project_db", Req.REQUIRED_UNLESS_URL) +_USERNAME = ConnField("Username", "project_user", Req.REQUIRED) +_PASSWORD_OPT = ConnField("Password", "project_pw_encrypted", Req.OPTIONAL, secret=True) +_PASSWORD_REQ = ConnField("Password", "project_pw_encrypted", Req.REQUIRED, secret=True) +_WAREHOUSE = ConnField("Warehouse", "warehouse", Req.OPTIONAL) +_PRIVATE_KEY = ConnField("Private Key", "private_key", Req.REQUIRED, secret=True) +_PRIVATE_KEY_PASSPHRASE = ConnField("Private Key Passphrase", "private_key_passphrase", Req.OPTIONAL, secret=True) +_URL = ConnField("URL", "url", Req.REQUIRED) + +# Databricks +_HTTP_PATH_RU = ConnField("HTTP Path", "http_path", Req.REQUIRED_UNLESS_URL) +_HTTP_PATH_REQ = ConnField("HTTP Path", "http_path", Req.REQUIRED) +_HOST_REQ = ConnField("Host", "project_host", Req.REQUIRED) +_PORT_REQ = ConnField("Port", "project_port", Req.REQUIRED) +_CATALOG = ConnField("Catalog", "project_db", Req.REQUIRED_UNLESS_URL) +_CATALOG_REQ = ConnField("Catalog", "project_db", Req.REQUIRED) +_ACCESS_TOKEN = ConnField("Access Token", "project_pw_encrypted", Req.REQUIRED, secret=True) +_CLIENT_ID = ConnField("Client ID", "project_user", Req.REQUIRED) +_CLIENT_SECRET = ConnField("Client Secret", "project_pw_encrypted", Req.REQUIRED, secret=True) + +# BigQuery / Salesforce +_SERVICE_ACCOUNT_KEY = ConnField("Service Account Key", "service_account_key", Req.REQUIRED, secret=True) +_LOGIN_URL = ConnField("Login URL", "project_host", Req.REQUIRED) +_CONSUMER_KEY = ConnField("Consumer Key", "project_user", Req.REQUIRED) +_SF_USERNAME = ConnField("Username", "project_db", Req.REQUIRED) +_CONSUMER_SECRET = ConnField("Consumer Secret", "project_pw_encrypted", Req.REQUIRED, secret=True) + + +def _host_auth_schema(code: str, *, db_field: ConnField = _DATABASE) -> FlavorSchema: + """Single-mode host/URL flavors (PostgreSQL, Redshift, MSSQL, Oracle, SAP HANA).""" + return FlavorSchema( + code=code, + label=FLAVOR_CODE_TO_LABEL[code], + modes=( + FlavorMode(mode=None, fields=(_HOST, _PORT, db_field, _USERNAME, _PASSWORD_OPT), supports_url=True), + ), + url_field=_URL, + ) + + +def _azure_schema(code: str) -> FlavorSchema: + return FlavorSchema( + code=code, + label=FLAVOR_CODE_TO_LABEL[code], + modes=( + FlavorMode( + mode=ConnectionMode.PASSWORD, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _PASSWORD_OPT), + supports_url=True, + sets={"connect_with_identity": False}, + ), + FlavorMode( + mode=ConnectionMode.MANAGED_IDENTITY, + fields=(_HOST, _PORT, _DATABASE), + supports_url=True, + sets={"connect_with_identity": True}, + ), + ), + url_field=_URL, + ) + + +FLAVOR_CONNECTION_SCHEMA: dict[str, FlavorSchema] = { + "postgresql": _host_auth_schema("postgresql"), + "redshift": _host_auth_schema("redshift"), + "redshift_spectrum": _host_auth_schema("redshift_spectrum"), + "mssql": _host_auth_schema("mssql"), + "oracle": _host_auth_schema("oracle", db_field=_SERVICE_NAME), + "sap_hana": _host_auth_schema("sap_hana"), + "azure_mssql": _azure_schema("azure_mssql"), + "synapse_mssql": _azure_schema("synapse_mssql"), + "snowflake": FlavorSchema( + code="snowflake", + label=FLAVOR_CODE_TO_LABEL["snowflake"], + modes=( + FlavorMode( + mode=ConnectionMode.KEY_PAIR, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _WAREHOUSE, _PRIVATE_KEY, _PRIVATE_KEY_PASSPHRASE), + supports_url=True, + sets={"connect_by_key": True}, + ), + FlavorMode( + mode=ConnectionMode.PASSWORD, + fields=(_HOST, _PORT, _DATABASE, _USERNAME, _WAREHOUSE, _PASSWORD_REQ), + supports_url=True, + sets={"connect_by_key": False}, + ), + ), + url_field=_URL, + ), + "databricks": FlavorSchema( + code="databricks", + label=FLAVOR_CODE_TO_LABEL["databricks"], + modes=( + FlavorMode( + mode=ConnectionMode.ACCESS_TOKEN, + fields=(_HOST, _PORT, _HTTP_PATH_RU, _CATALOG, _ACCESS_TOKEN), + supports_url=True, + # PAT auth: the username is always the literal 'token'. + sets={"connect_by_key": False, "project_user": "token"}, + ), + FlavorMode( + mode=ConnectionMode.SERVICE_PRINCIPAL, + fields=(_HOST_REQ, _PORT_REQ, _HTTP_PATH_REQ, _CATALOG_REQ, _CLIENT_ID, _CLIENT_SECRET), + supports_url=False, + sets={"connect_by_key": True}, + ), + ), + url_field=_URL, + ), + "bigquery": FlavorSchema( + code="bigquery", + label=FLAVOR_CODE_TO_LABEL["bigquery"], + modes=(FlavorMode(mode=None, fields=(_SERVICE_ACCOUNT_KEY,), supports_url=False),), + url_field=None, + ), + "salesforce_data360": FlavorSchema( + code="salesforce_data360", + label=FLAVOR_CODE_TO_LABEL["salesforce_data360"], + modes=( + FlavorMode( + mode=ConnectionMode.JWT_BEARER, + fields=(_LOGIN_URL, _CONSUMER_KEY, _SF_USERNAME, _PRIVATE_KEY), + supports_url=False, + sets={"connect_by_key": True}, + ), + FlavorMode( + mode=ConnectionMode.CLIENT_CREDENTIALS, + fields=(_LOGIN_URL, _CONSUMER_KEY, _CONSUMER_SECRET), + supports_url=False, + sets={"connect_by_key": False}, + ), + ), + url_field=None, + ), +} + + +def schema_for(code: str) -> FlavorSchema: + """Return the connection schema for a flavor code. Raises ``KeyError`` if unknown.""" + return FLAVOR_CONNECTION_SCHEMA[code] + + +def resolve_mode(code: str, mode_label: str | None) -> FlavorMode: + """Resolve the ``FlavorMode`` for a flavor + supplied ``connection_mode`` label. + + Single-mode flavors take no ``connection_mode`` (passing one is an error). + Multi-mode flavors require a valid one. + """ + schema = schema_for(code) + if len(schema.modes) == 1 and schema.modes[0].mode is None: + if mode_label is not None: + raise MCPUserError(f"{schema.label} does not take a connection_mode.") + return schema.modes[0] + + valid = [str(m.mode) for m in schema.modes if m.mode is not None] + if mode_label is None: + raise MCPUserError(f"{schema.label} requires a connection_mode. Valid values: {', '.join(valid)}.") + for fmode in schema.modes: + if fmode.mode is not None and str(fmode.mode) == mode_label: + return fmode + raise MCPUserError( + f"Invalid connection_mode `{mode_label}` for {schema.label}. Valid values: {', '.join(valid)}." + ) + + +def infer_mode(connection: Connection) -> ConnectionMode | None: + """Reverse of a mode's ``sets`` flags — derive the active mode from a stored + connection so update/validation can pick the right field set without the + caller re-supplying ``connection_mode``. + """ + code = connection.sql_flavor_code + schema = FLAVOR_CONNECTION_SCHEMA.get(code) + if schema is None or (len(schema.modes) == 1 and schema.modes[0].mode is None): + return None + + if code in {"azure_mssql", "synapse_mssql"}: + return ConnectionMode.MANAGED_IDENTITY if connection.connect_with_identity else ConnectionMode.PASSWORD + if code == "snowflake": + return ConnectionMode.KEY_PAIR if connection.connect_by_key else ConnectionMode.PASSWORD + if code == "databricks": + return ConnectionMode.SERVICE_PRINCIPAL if connection.connect_by_key else ConnectionMode.ACCESS_TOKEN + if code == "salesforce_data360": + return ConnectionMode.JWT_BEARER if connection.connect_by_key else ConnectionMode.CLIENT_CREDENTIALS + return None + + +def _mode_for_connection(connection: Connection) -> FlavorMode: + """Pick the FlavorMode matching a connection's current flags (for validation).""" + schema = schema_for(connection.sql_flavor_code) + if len(schema.modes) == 1: + return schema.modes[0] + active = infer_mode(connection) + for fmode in schema.modes: + if fmode.mode == active: + return fmode + return schema.modes[0] + + +def connection_display_fields(connection: Connection) -> list[ConnField]: + """Active-mode fields (plus the URL field when in URL mode), in schema order. + + For rendering a connection back to the user with the flavor-specific UI label + for each populated column (e.g. ``Catalog`` for Databricks, ``Login URL`` for + Salesforce). Callers skip secrets and empty values. + """ + schema = schema_for(connection.sql_flavor_code) + fields = list(_mode_for_connection(connection).fields) + if schema.url_field is not None and getattr(connection, "connect_by_url", False): + fields.append(schema.url_field) + return fields + + +def connection_field_labels(connection: Connection) -> dict[str, str]: + """Map each ``Connection`` column to its flavor/mode-specific UI label. + + Used to label diff output. Columns outside the active mode fall back to the + caller's generic labels. + """ + schema = schema_for(connection.sql_flavor_code) + labels = {fld.column: fld.label for fld in _mode_for_connection(connection).fields} + if schema.url_field is not None: + labels[schema.url_field.column] = schema.url_field.label + return labels + + +def apply_connection_params( + connection: Connection, + code: str, + mode_label: str | None, + params: dict[str, object], +) -> None: + """Map a label-keyed ``connection_params`` dict onto a ``Connection``. + + * Resolves the mode and applies its forced ``sets`` columns (auth flags, + Databricks PAT ``project_user='token'``). + * Maps each supplied label to its ``Connection`` column (casting ``Port``). + * Toggles ``connect_by_url`` from the presence of ``URL`` vs the host-group + fields, rejecting an ambiguous mix. + + Raises ``MCPUserError`` on unknown keys, an unsupported / conflicting ``URL``, + or an invalid / missing mode. + """ + fmode = resolve_mode(code, mode_label) + fields_by_label = {f.label: f for f in fmode.fields} + url_fields = {f.label for f in fmode.fields if f.requirement is Req.REQUIRED_UNLESS_URL} + + valid_keys = set(fields_by_label) + if fmode.supports_url: + valid_keys.add(_URL.label) + unknown = [key for key in params if key not in valid_keys] + if unknown: + raise MCPUserError( + f"Unknown connection_params for {schema_for(code).label}: {', '.join(sorted(unknown))}. " + f"Allowed: {', '.join(sorted(valid_keys))}." + ) + + has_url = _URL.label in params + provided_url_fields = [key for key in params if key in url_fields] + if has_url and not fmode.supports_url: + raise MCPUserError(f"{schema_for(code).label} does not support URL connections in this mode.") + if has_url and provided_url_fields: + raise MCPUserError( + f"Provide either a `URL` or host fields ({', '.join(sorted(url_fields))}), not both." + ) + + # Forced columns first so explicit params can't be clobbered by sets. + for attr, value in fmode.sets.items(): + setattr(connection, attr, value) + + for label, value in params.items(): + if label == _URL.label: + continue + column = fields_by_label[label].column + setattr(connection, column, str(value) if column == "project_port" and value is not None else value) + + if has_url: + connection.connect_by_url = True + connection.url = str(params[_URL.label]) + elif provided_url_fields: + connection.connect_by_url = False + + +# -- field-requirement validation ------------------------------------------- + +_CONNECTION_NAME_MIN = 3 +_CONNECTION_NAME_MAX = 40 +_MAX_THREADS_MIN = 1 +_MAX_THREADS_MAX = 8 +_MAX_QUERY_CHARS_MIN = 500 +_MAX_QUERY_CHARS_MAX = 50000 + + +def _missing(value: object) -> bool: + return value is None or (isinstance(value, str) and not value.strip()) + + +def _required_fields_for(connection: Connection) -> list[ConnField]: + """Schema fields that must be non-empty for this connection's flavor + active + mode, accounting for the URL alternative. + """ + fmode = _mode_for_connection(connection) + using_url = fmode.supports_url and bool(connection.connect_by_url) + + required: list[ConnField] = [] + for fld in fmode.fields: + if fld.requirement is Req.REQUIRED: + required.append(fld) + elif fld.requirement is Req.REQUIRED_UNLESS_URL and not using_url: + required.append(fld) + + schema = schema_for(connection.sql_flavor_code) + if using_url and schema.url_field is not None: + required.append(schema.url_field) + return required + + +def validate_connection_fields(connection: Connection) -> list[str]: + """Return every validation error (empty list = valid). + + Mirrors the per-flavor JS form validators in + ``ui/static/js/components/connection_form.js``. The connection tools call this + and raise their user-facing error containing the bullets. Field errors use the + UI label (e.g. ``Host``) so the LLM sees the same wording as a UI user. + """ + errors: list[str] = [] + flavor_label = FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code, connection.sql_flavor_code) + + name = connection.connection_name + if _missing(name): + errors.append("`connection_name` is required.") + elif not (_CONNECTION_NAME_MIN <= len(name.strip()) <= _CONNECTION_NAME_MAX): + errors.append( + f"`connection_name` must be between {_CONNECTION_NAME_MIN} and {_CONNECTION_NAME_MAX} characters." + ) + + for fld in _required_fields_for(connection): + if _missing(getattr(connection, fld.column, None)): + errors.append(f"`{fld.label}` is required for {flavor_label}.") + + threads = connection.max_threads + if threads is not None and not (_MAX_THREADS_MIN <= threads <= _MAX_THREADS_MAX): + errors.append(f"`max_threads` must be between {_MAX_THREADS_MIN} and {_MAX_THREADS_MAX}.") + + query_chars = connection.max_query_chars + if query_chars is not None and not (_MAX_QUERY_CHARS_MIN <= query_chars <= _MAX_QUERY_CHARS_MAX): + errors.append(f"`max_query_chars` must be between {_MAX_QUERY_CHARS_MIN} and {_MAX_QUERY_CHARS_MAX}.") + + return errors + + +def effective_mode(connection: Connection, connection_mode: str | None) -> str | None: + """Mode label to apply: the explicit override, else the connection's current mode.""" + if connection_mode is not None: + return connection_mode + inferred = infer_mode(connection) + return str(inferred) if inferred is not None else None + + +def raise_validation_error(errors: list[str], header: str) -> NoReturn: + bullets = "\n".join(f"- {err}" for err in errors) + raise MCPUserError(f"{header}\n\n{bullets}") + + +def render_connection_body(doc: MdDoc, connection: Connection) -> None: + """Render every non-secret connection field below the heading. + + Encrypted columns are filtered out via ``ConnField.secret``. + """ + doc.field("ID", connection.connection_id, code=True) + doc.field("Project", connection.project_code, code=True) + doc.field("Type", format_flavor_label(connection.sql_flavor_code)) + + # Each populated, non-secret field under its flavor-specific label + # (e.g. "Catalog" for Databricks, "Login URL" for Salesforce). + for fld in connection_display_fields(connection): + if fld.secret: + continue + value = getattr(connection, fld.column, None) + if value in (None, ""): + continue + doc.field(fld.label, value, code=fld.column != "project_port") + + doc.field("Authentication", authentication_label(connection)) + if connection.max_threads is not None: + doc.field("Max Threads", connection.max_threads) + if connection.max_query_chars is not None: + doc.field("Max Expression Length", connection.max_query_chars) + + +def authentication_label(connection: Connection) -> str: + """The connection's auth method: the active connection mode for multi-mode + flavors, else the implicit method (service account key, else password). + """ + mode = infer_mode(connection) + if mode is not None: + return str(mode) + if connection.service_account_key: + return "Service Account Key" + return "Password" diff --git a/testgen/mcp/tools/connections.py b/testgen/mcp/tools/connections.py new file mode 100644 index 00000000..55423144 --- /dev/null +++ b/testgen/mcp/tools/connections.py @@ -0,0 +1,187 @@ +"""MCP tools for database connections — list, get, and test database connections. + +The per-flavor connection shape (which auth modes exist and which ``connection_params`` keys +each needs) lives in ``testgen.common.database.connection_service`` and is exposed to the +model through the ``testgen://connection-parameters/{flavor}`` resource. Validation and +auth-path normalization are delegated to that same module so the rules stay in one place. +""" + +from __future__ import annotations + +from testgen.common.database.connection_service import ( + apply_connection_defaults, + normalize_auth_fields, + test_connection_status, +) +from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import get_project_permissions, mcp_permission +from testgen.mcp.tools.common import ( + SQL_FLAVOR_CODE_TO_LABEL, + DocGroup, + apply_connection_params, + effective_mode, + format_flavor_label, + format_page_footer, + format_page_info, + parse_sql_flavor, + raise_validation_error, + render_connection_body, + resolve_connection, + validate_connection_fields, + validate_limit, + validate_page, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + + +@with_database_session +@mcp_permission("view") +def list_connections(project_code: str, page: int = 1, limit: int = 20) -> str: + """List database connections in a project with their database type and table-group counts. + + Use this before changing or referencing a connection to confirm its ID, name, and host. + Credentials are never returned. + + Args: + project_code: The project code to list connections for. + page: Page number starting at 1 (default 1). + limit: Page size (default 20, max 100). + """ + validate_page(page) + validate_limit(limit, 100) + + perms = get_project_permissions() + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + rows, total = Connection.list_for_project(project_code, page=page, limit=limit) + + if not rows: + if page > 1: + return f"No connections on page {page} (total: {total})." + return f"No connections found for project `{project_code}`." + + doc = MdDoc() + doc.heading(1, f"Connections for `{project_code}`") + doc.text(format_page_info(total, page, limit)) + table_rows: list[list[object]] = [] + for row in rows: + table_rows.append( + [ + row.connection_id, + row.connection_name, + format_flavor_label(row.sql_flavor_code), + row.project_host, + row.project_db, + row.table_group_count, + ] + ) + doc.table( + ["ID", "Name", "Type", "Host", "Database", "Table groups"], + table_rows, + code=[0, 3, 4], + ) + if footer := format_page_footer(total, page, limit): + doc.text(footer) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_connection(connection_id: int) -> str: + """Get a connection's configuration, including database type, host, and authentication mode. + + Credentials (password, private key, service-account key) are never returned. + Use this before editing a connection or creating a table group on it. + + Args: + connection_id: Bigint connection ID returned by `list_connections`. + """ + connection = resolve_connection(connection_id) + doc = MdDoc() + doc.heading(1, f"Connection `{connection.connection_name}`") + render_connection_body(doc, connection) + return doc.render() + + +@with_database_session +@mcp_permission("administer") +def test_connection( + connection_id: int | None = None, + sql_flavor: str | None = None, + connection_params: dict | None = None, + connection_mode: str | None = None, +) -> str: + """Test connectivity against a stored or inline-supplied connection. + + Two call shapes: pass ``connection_id`` to test a stored connection + (``connection_params`` / ``connection_mode`` override the stored values for + the test only, nothing is saved); or pass ``sql_flavor`` plus + ``connection_params`` (and ``connection_mode`` for multi-mode flavors) to + test inline without persisting. See the + ``testgen://connection-parameters/{flavor}`` resource for each flavor's + ``connection_mode`` values and ``connection_params`` keys. + + Args: + connection_id: Stored connection ID, e.g. from ``get_data_inventory``. + Omit for inline tests. + sql_flavor: SQL Database flavor. Required for inline tests; not accepted + when ``connection_id`` is set. + connection_params: Connection field values. For a stored connection, + omitted fields keep their stored value. + connection_mode: Authentication mode for multi-mode flavors. + """ + if connection_id is None and sql_flavor is None: + raise MCPUserError( + "Provide `connection_id` to test an existing connection, " + "or `sql_flavor` plus `connection_params` to test inline." + ) + + connection: Connection | None = None + if connection_id is not None: + if sql_flavor is not None: + raise MCPUserError( + "`sql_flavor` cannot be overridden when `connection_id` is set " + "— database type is immutable on an existing connection." + ) + connection = resolve_connection(connection_id) + + inline = connection is None + if connection is None: + _, code, family = parse_sql_flavor(sql_flavor) # type: ignore[arg-type] + connection = Connection(sql_flavor=family, sql_flavor_code=code) + + if connection_params is not None or connection_mode is not None: + mode = connection_mode if inline else effective_mode(connection, connection_mode) + apply_connection_params(connection, connection.sql_flavor_code, mode, connection_params or {}) + + normalize_auth_fields(connection) + # No-op for stored connections; keeps the inline-test contract identical to create. + apply_connection_defaults(connection) + + errors = validate_connection_fields(connection) + if errors: + raise_validation_error(errors, "Cannot test connection. Required fields missing or invalid.") + + status = test_connection_status(connection) + + doc = MdDoc() + heading = "Connection test succeeded" if status.successful else "Connection test failed" + doc.heading(1, heading) + + if connection_id is not None: + doc.field("ID", connection.connection_id, code=True) + if connection.connection_name: + doc.field("Name", connection.connection_name, code=True) + label = SQL_FLAVOR_CODE_TO_LABEL.get(connection.sql_flavor_code) + doc.field("Type", label.value if label else connection.sql_flavor_code) + if connection.project_host: + doc.field("Host", connection.project_host, code=True) + + doc.text(status.message) + if status.details: + doc.code_block(status.details) + return doc.render() diff --git a/testgen/mcp/tools/data_catalog.py b/testgen/mcp/tools/data_catalog.py new file mode 100644 index 00000000..645d160a --- /dev/null +++ b/testgen/mcp/tools/data_catalog.py @@ -0,0 +1,222 @@ +from collections.abc import Mapping + +from testgen.common.data_catalog_service import ( + TAG_FIELDS, + apply_column_metadata, + apply_table_metadata, + disable_autoflags, + validate_metadata_fields, +) +from testgen.common.models import with_database_session +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import get_project_permissions, mcp_permission +from testgen.mcp.tools.common import DocGroup, resolve_table_group +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +_BOOLEAN_ARGS = ("cde", "xde", "pii") + +# MCP argument name -> data_*_chars column name. +_ARG_TO_COLUMN = {"cde": "critical_data_element", "xde": "excluded_data_element", "pii": "pii_flag"} +_COLUMN_TO_ARG = {v: k for k, v in _ARG_TO_COLUMN.items()} + + +@with_database_session +@mcp_permission("disposition") +def update_catalog_metadata(updates: list[dict]) -> str: + """Apply metadata updates to tables and columns within table groups. + + Each update targets one table, or one column within a table, and sets one or + more metadata fields. Rows are processed independently: a failure on one row + does not stop the others, and the response reports the outcome of every row. + + Omit a field to leave it unchanged, pass null to clear it, or pass a value to set it. + + Args: + updates: List of per-row update specs. Each spec accepts: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Name of the target table. + column_name: Name of the target column. Omit for a table-level update. + description: Free-text description (max 1000 characters). + cde: Whether this is a critical data element (true/false). Set at the + table level to apply to all its columns unless a column overrides it. + xde: Whether to exclude the column from future profiling and test + generation (true/false). Columns only. + pii: Whether the column contains personally identifiable information + (true/false). Columns only. + data_source, source_system, source_process, business_domain, + stakeholder_group, transform_level, aggregation_level, data_product, + data_classification: + Catalog tags (max 40 characters each). + """ + if not updates: + raise MCPUserError("Provide at least one update.") + + perms = get_project_permissions() + tg_cache: dict[str, object] = {} + flag_writes: dict = {} # tg.id -> {"tg": tg, "cde": bool, "pii": bool} + inheritance_notices: list[str] = [] + exclusion_notices: list[str] = [] + rows: list[tuple[str, str, str]] = [] + + for spec in updates: + target = _target_label(spec) + try: + outcome, detail = _apply_update(spec, perms, tg_cache, flag_writes, inheritance_notices, exclusion_notices) + rows.append((target, outcome, detail)) + except MCPUserError as err: + rows.append((target, "Failed", str(err))) + + autoflag_notices: list[str] = [] + for entry in flag_writes.values(): + disabled = disable_autoflags(entry["tg"], wrote_cde=entry["cde"], wrote_pii=entry["pii"]) + for flag in disabled: + autoflag_notices.append( + f"Auto-disabled {flag} on table group `{entry['tg'].table_groups_name}` to preserve manual marks." + ) + + return _render(rows, autoflag_notices + inheritance_notices + exclusion_notices) + + +def _target_label(spec: Mapping) -> str: + table = spec.get("table_name") or "?" + column = spec.get("column_name") + return f"{table}.{column}" if column else table + + +def _apply_update(spec, perms, tg_cache, flag_writes, inheritance_notices, exclusion_notices) -> tuple[str, str]: + if not isinstance(spec, Mapping): + raise MCPUserError("Each update must be an object.") + + table_group_id = spec.get("table_group_id") + table_name = spec.get("table_name") + if not table_group_id: + raise MCPUserError("table_group_id is required.") + if not table_name: + raise MCPUserError("table_name is required.") + column_name = spec.get("column_name") + is_column = column_name is not None + + for arg in _BOOLEAN_ARGS: + if arg in spec and spec[arg] is not None and not isinstance(spec[arg], bool): + raise MCPUserError(f"{arg} must be true or false.") + + if not is_column: + if "xde" in spec: + raise MCPUserError("xde applies to columns only; omit it for table-level updates.") + if "pii" in spec: + raise MCPUserError("pii applies to columns only; omit it for table-level updates.") + + tg = tg_cache.get(table_group_id) + if tg is None: + tg = resolve_table_group(table_group_id) + tg_cache[table_group_id] = tg + + if "pii" in spec and not perms.has_permission("view_pii", tg.project_code): + raise MCPUserError(f"Setting pii requires permission to view PII on project `{tg.project_code}`.") + + fields = _build_fields(spec) + errors = validate_metadata_fields(fields) + if errors: + raise MCPUserError("; ".join(errors)) + + if not fields: + return "Skipped", "no metadata fields provided" + + if is_column: + target = _resolve_column(tg.id, table_name, column_name) + if target is None: + raise MCPUserError(f"Column `{column_name}` not found in table `{table_name}`.") + else: + target = _resolve_table(tg.id, table_name) + if target is None: + raise MCPUserError(f"Table `{table_name}` not found in this table group.") + + # Disable a table group's auto-detect flag only on a real change, matching the UI's confirm-to-disable + # behavior. A no-op write (e.g. cde: false on an already-false column) must not silently turn it off. + cde_changed = "cde" in spec and target.critical_data_element != spec["cde"] + pii_changed = "pii" in spec and target.pii_flag != ("MANUAL" if spec["pii"] else None) + + if is_column: + apply_column_metadata(target, fields) + else: + apply_table_metadata(target, fields) + + entry = flag_writes.setdefault(tg.id, {"tg": tg, "cde": False, "pii": False}) + if cde_changed: + entry["cde"] = True + if pii_changed: + entry["pii"] = True + + if not is_column and spec.get("cde") is True: + inheritance_notices.append( + f"Table-level CDE set; affects all columns of `{table_name}` unless explicitly overridden." + ) + if is_column and spec.get("xde") is True: + exclusion_notices.append( + f"Column `{table_name}.{column_name}` excluded; next profiling run and test generation will skip it." + ) + + return "Updated", _change_summary(fields) + + +def _build_fields(spec: Mapping) -> dict: + """Translate a spec into a {data_*_chars column: value} dict (pii bool -> MANUAL/None).""" + fields: dict = {} + if "description" in spec: + fields["description"] = spec["description"] + for tag in TAG_FIELDS: + if tag in spec: + fields[tag] = spec[tag] + if "cde" in spec: + fields["critical_data_element"] = spec["cde"] + if "xde" in spec: + fields["excluded_data_element"] = spec["xde"] + if "pii" in spec: + fields["pii_flag"] = "MANUAL" if spec["pii"] else None + return fields + + +def _resolve_table(table_groups_id, table_name: str): + matches = list(DataTable.select_where( + DataTable.table_groups_id == table_groups_id, + DataTable.table_name == table_name, + )) + return matches[0] if matches else None + + +def _resolve_column(table_groups_id, table_name: str, column_name: str): + matches = list(DataColumnChars.select_where( + DataColumnChars.table_groups_id == table_groups_id, + DataColumnChars.table_name == table_name, + DataColumnChars.column_name == column_name, + )) + return matches[0] if matches else None + + +def _change_summary(fields: dict) -> str: + parts = [] + for column, value in fields.items(): + label = _COLUMN_TO_ARG.get(column, column) + parts.append(f"{label} cleared" if value is None else f"{label} set") + return ", ".join(parts) + + +def _render(rows: list[tuple[str, str, str]], notices: list[str]) -> str: + succeeded = sum(1 for _, outcome, _ in rows if outcome == "Updated") + skipped = sum(1 for _, outcome, _ in rows if outcome == "Skipped") + failed = sum(1 for _, outcome, _ in rows if outcome == "Failed") + + doc = MdDoc() + doc.heading(1, "Catalog metadata update") + doc.field("Rows attempted", len(rows)) + doc.field("Updated", succeeded) + doc.field("Skipped", skipped) + doc.field("Failed", failed) + doc.table(["Target", "Outcome", "Details"], [list(row) for row in rows], code=[0]) + for notice in notices: + doc.text(notice) + return doc.render() diff --git a/testgen/mcp/tools/discovery.py b/testgen/mcp/tools/discovery.py index 05b7ab5b..84eb03a1 100644 --- a/testgen/mcp/tools/discovery.py +++ b/testgen/mcp/tools/discovery.py @@ -1,11 +1,18 @@ -from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session from testgen.common.models.data_table import DataTable from testgen.common.models.project import Project -from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite +from testgen.mcp.exceptions import MCPResourceNotAccessible from testgen.mcp.permissions import get_project_permissions, mcp_permission -from testgen.mcp.tools.common import DocGroup, resolve_table_group, validate_limit, validate_page +from testgen.mcp.tools.common import ( + DocGroup, + format_flavor_label, + resolve_connection, + resolve_table_group, + resolve_test_suite, + validate_limit, + validate_page, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.DISCOVER @@ -23,7 +30,6 @@ def get_data_inventory() -> str: from testgen.mcp.services.inventory_service import get_inventory perms = get_project_permissions() - MixpanelService().send_event("mcp-get-data-inventory", username=perms.username) return get_inventory( project_codes=perms.allowed_codes, view_project_codes=perms.codes_allowed_to("view"), @@ -51,6 +57,53 @@ def list_projects() -> str: return doc.render() +@with_database_session +@mcp_permission("view") +def get_project(project_code: str) -> str: + """Get a project's configuration and configuration counts. + + Returns the project name, observability and data-retention settings, and counts + of connections, table groups, test suites, test definitions, profiling runs, and + test runs scoped to the project. Use this before configuration changes to confirm + a project's current shape. + + Args: + project_code: The project code, e.g. from `list_projects`. + """ + perms = get_project_permissions() + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + + project = Project.get(project_code) + summary = Project.get_summary(project_code) + if project is None or summary is None: + raise MCPResourceNotAccessible("Project", project_code) + + doc = MdDoc() + doc.heading(1, f"Project `{project_code}`") + doc.field("Name", project.project_name) + doc.field("Connections", summary.connection_count) + doc.field("Table groups", summary.table_group_count) + doc.field("Test suites", summary.test_suite_count) + doc.field("Test definitions", summary.test_definition_count) + doc.field("Profiling runs", summary.profiling_run_count) + doc.field("Test runs", summary.test_run_count) + + doc.heading(2, "Configuration") + doc.field("Weighted data quality scoring", project.use_dq_score_weights) + + doc.heading(2, "Observability Integration") + doc.field("Configured", summary.can_export_to_observability) + if project.observability_api_url: + doc.field("API URL", project.observability_api_url, code=True) + + doc.heading(2, "Data Retention") + doc.field("Automatically delete old profiling and test history", project.data_retention_enabled) + if project.data_retention_enabled and project.data_retention_days is not None: + doc.field("Delete history older than (days)", project.data_retention_days) + + return doc.render() + + @with_database_session @mcp_permission("view") def list_test_suites(project_code: str) -> str: @@ -63,17 +116,13 @@ def list_test_suites(project_code: str) -> str: return "Missing required parameter `project_code`." perms = get_project_permissions() - perms.verify_access(project_code, not_found=f"No test suites found for project `{project_code}`.") + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) summaries = TestSuite.select_summary(project_code) if not summaries: return f"No test suites found for project `{project_code}`." - # Batch-lookup job_execution_ids for latest runs - run_ids = [s.latest_run_id for s in summaries if s.latest_run_id] - job_exec_map = TestRun.get_job_execution_ids(run_ids) if run_ids else {} - doc = MdDoc() doc.heading(1, f"Test Suites for `{project_code}`") for s in summaries: @@ -85,8 +134,7 @@ def list_test_suites(project_code: str) -> str: doc.field("Test definitions", s.test_ct or 0) if s.latest_run_id: - run_id = job_exec_map.get(s.latest_run_id) or s.latest_run_id - doc.field("Latest run", f"`{run_id}` ({s.latest_run_start})") + doc.field("Latest run", f"`{s.latest_run_id}` ({s.latest_run_start})") results_summary = ( f"{s.last_run_test_ct or 0} tests: " f"{s.last_run_passed_ct or 0} passed, " @@ -103,6 +151,58 @@ def list_test_suites(project_code: str) -> str: return doc.render() +@with_database_session +@mcp_permission("view") +def get_test_suite(test_suite_id: str) -> str: + """Get a test suite's configuration: connection, table group, default severity, and per-test-type counts. + + Returns the test suite's identity and configuration along with a breakdown of how many test + definitions it contains by type and how many are locked (excluded from regeneration). + Use this before changing a suite's tests to understand what will be affected. + + Args: + test_suite_id: The test suite UUID, e.g. from `list_test_suites`. + """ + suite = resolve_test_suite(test_suite_id) + # Defense in depth: resolve via perm-filtered helpers rather than `Model.get(...)`. + # FK constraints guarantee same-project today; the resolvers are the established wrapper + # for project-scoped lookups and keep us aligned if those guarantees ever change. + connection = resolve_connection(suite.connection_id) if suite.connection_id else None + table_group = resolve_table_group(str(suite.table_groups_id)) if suite.table_groups_id else None + stats = TestSuite.test_definition_stats(suite.id) + + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}`") + doc.field("ID", str(suite.id), code=True) + doc.field("Project", suite.project_code, code=True) + + if connection is not None: + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`, {format_flavor_label(connection.sql_flavor_code)})", + ) + if table_group is not None: + doc.field( + "Table group", + f"{table_group.table_groups_name} (`{table_group.id}`)", + ) + + if suite.test_suite_description: + doc.field("Description", suite.test_suite_description) + doc.field("Default severity", suite.severity or "Inherit from test type") + doc.field("Export to observability", suite.export_to_observability) + + doc.field("Total tests", stats.total) + doc.field("Locked tests", stats.locked) + + if stats.counts_by_type: + doc.heading(2, "Tests by type") + rows = [[type_label, count] for type_label, count in stats.counts_by_type.items()] + doc.table(["Test", "Count"], rows) + + return doc.render() + + @with_database_session @mcp_permission("catalog") def list_tables(table_group_id: str, limit: int = 200, page: int = 1) -> str: diff --git a/testgen/mcp/tools/execution.py b/testgen/mcp/tools/execution.py index 3ed02e0e..2847e77a 100644 --- a/testgen/mcp/tools/execution.py +++ b/testgen/mcp/tools/execution.py @@ -133,7 +133,7 @@ def _render_submission( ) -> str: doc = MdDoc() doc.heading(1, f"{kind} submitted for `{scope_name}`") - doc.field("Job ID", job.id, code=True) + doc.field(kind.title(), job.id, code=True) doc.field(scope_label, scope_name) doc.field("Status", "Pending") doc.text(f"Use `{poll_tool}` {poll_hint}.") @@ -147,7 +147,7 @@ def _render_cancel(job: JobExecution, kind: str, poll_tool: str) -> str: ) doc = MdDoc() doc.heading(1, f"{kind} cancellation requested") - doc.field("Job ID", job.id, code=True) + doc.field(kind.title(), job.id, code=True) doc.field("Status", job.status) doc.text(f"Use `{poll_tool}` to confirm cancellation.") return doc.render() diff --git a/testgen/mcp/tools/hygiene_issues.py b/testgen/mcp/tools/hygiene_issues.py index 19042588..ac7846cf 100644 --- a/testgen/mcp/tools/hygiene_issues.py +++ b/testgen/mcp/tools/hygiene_issues.py @@ -25,6 +25,7 @@ parse_quality_dimension, parse_since_arg, parse_uuid, + resolve_hygiene_issue, resolve_issue_type, resolve_table_group, validate_limit, @@ -86,12 +87,12 @@ def _resolve_profile_run_je_id( return je_uuid job_uuid = parse_uuid(job_execution_id, "job_execution_id") - run = ProfilingRun.get_by_id_or_job(job_uuid) + run = ProfilingRun.get(job_uuid) perms = get_project_permissions() tg = TableGroup.get(run.table_groups_id) if run else None if run is None or tg is None or not perms.has_access(tg.project_code): raise MCPResourceNotAccessible("Profiling run", job_execution_id) - return run.job_execution_id + return run.id @with_database_session @@ -211,17 +212,9 @@ def update_hygiene_issue(*, issue_id: str, disposition: str) -> str: issue_id: UUID of the hygiene issue. disposition: New disposition. Valid values: 'Confirmed', 'Dismissed', 'Muted'. """ - issue_uuid = parse_uuid(issue_id, "issue_id") db_disposition = parse_disposition(disposition) - perms = get_project_permissions() - - updated = HygieneIssue.update_disposition( - issue_uuid, - db_disposition, - HygieneIssue.project_code.in_(perms.allowed_codes), - ) - if not updated: - raise MCPResourceNotAccessible("Hygiene issue", issue_id) + issue = resolve_hygiene_issue(issue_id) + issue.disposition = db_disposition doc = MdDoc() doc.text(f"Updated hygiene issue {MdDoc.code(issue_id)} disposition to **{disposition}**.") diff --git a/testgen/mcp/tools/monitors.py b/testgen/mcp/tools/monitors.py new file mode 100644 index 00000000..426190be --- /dev/null +++ b/testgen/mcp/tools/monitors.py @@ -0,0 +1,1098 @@ +import re +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Any, cast +from uuid import UUID + +from sqlalchemy import select + +from testgen.common.cron_service import describe_cron, get_cron_sample +from testgen.common.enums import MonitorType +from testgen.common.holiday_service import is_supported_holiday_code +from testgen.common.models import get_current_session, with_database_session +from testgen.common.models.data_structure_log import ( + SCHEMA_CHANGE_ADDED, + SCHEMA_CHANGE_DROPPED, + SCHEMA_CHANGE_MODIFIED, + DataStructureLog, +) +from testgen.common.models.job_execution import JobExecution +from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.table_group import MonitorTableSummary, TableGroup +from testgen.common.models.test_definition import ( + MonitorForecastPoint, + TestDefinition, + TestDefinitionSummary, + forecast_points_from_prediction, +) +from testgen.common.models.test_result import MonitorEvent, TestResult +from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_forecast import ( + forecast_band_points, + gated_forecast_prediction, + next_update_window, +) +from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + MonitorTableSort, + format_page_footer, + format_page_info, + next_scheduled_run, + parse_monitor_table_sort, + parse_monitor_type, + parse_since_arg, + parse_uuid, + resolve_monitored_table_group, + validate_limit, + validate_page, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MONITORS + +_NOT_MONITORED_OUTPUT = "This table group is not monitored." + +_FORECAST_PENDING_NOTE = ( + "_No forecast yet — monitors in Prediction Model mode produce a forecast once they " + "have trained on enough history._" +) + +# Used where a forecast can be absent for reasons other than training (e.g. the +# predicted next-update window has already passed, or no baseline is stored), so +# the message must not claim the model still needs to train. +_FORECAST_UNAVAILABLE_NOTE = "_No forecast available for this monitor right now._" + +_MONITOR_LABEL: dict[MonitorType, str] = { + MonitorType.FRESHNESS: "Freshness", + MonitorType.VOLUME: "Volume", + MonitorType.SCHEMA: "Schema", + MonitorType.METRIC: "Metric", +} + +# Maps a ``MonitorTableSort`` value to the model-method ``sort_by`` argument. The +# anomaly-count sort uses the ``total_anomalies`` field — the model layer collapses +# it to the filtered type's column when exactly one ``anomaly_type`` is set. +_SORT_TO_MODEL_FIELD: dict[MonitorTableSort, str] = { + MonitorTableSort.TABLE_NAME: "table_name", + MonitorTableSort.ANOMALY_COUNT_DESC: "total_anomalies_desc", + MonitorTableSort.LATEST_UPDATE_DESC: "latest_update_desc", + MonitorTableSort.ROW_COUNT_CHANGE_DESC: "row_count_change_desc", +} + + +@with_database_session +@mcp_permission("view") +def get_monitor_summary(table_group_id: str, lookback: int | None = None) -> str: + """Get monitor health for a table group: per-type anomaly counts, error / training / pending status, and the active lookback window. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + lookback: Number of monitor runs to summarize. Omit to use the lookback runs configured for the table group. + """ + if lookback is not None and not 1 <= lookback <= 365: + raise MCPUserError(f"Invalid lookback `{lookback}`: must be between 1 and 365.") + + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + summary = TableGroup.get_monitor_group_summary(tg.id, lookback_override=lookback) + next_run = next_scheduled_run( + RUN_MONITORS_JOB_KEY, {"test_suite_id": str(monitor_suite.id)}, tg.project_code + ) + + doc = MdDoc() + doc.heading(1, f"Monitor summary for `{tg.table_groups_name}`") + doc.field("Project", tg.project_code, code=True) + doc.field("Monitored tables", summary.total_monitored_tables) + lookback_label = f"{summary.lookback} runs" + if lookback is not None: + lookback_label += " (override)" + doc.field("Lookback", lookback_label) + if summary.lookback_start: + doc.field("Window start", summary.lookback_start) + if summary.lookback_end: + doc.field("Window end", summary.lookback_end) + if next_run: + # Next-firing timestamps from JobSchedule are tz-aware in the schedule's + # configured cron_tz. MdDoc renders any datetime with a blind "UTC" suffix, + # so convert here to avoid claiming a local-tz timestamp is UTC. + doc.field("Next scheduled run", next_run.astimezone(UTC)) + else: + doc.field("Next scheduled run", "not scheduled") + + doc.heading(2, "Anomalies by type") + doc.table( + ["Monitor", "Anomalies", "Status"], + [ + [ + _MONITOR_LABEL[MonitorType.FRESHNESS], + summary.freshness_anomalies, + _summary_status( + is_pending=summary.freshness_is_pending, + is_training=summary.freshness_is_training, + has_errors=summary.freshness_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.VOLUME], + summary.volume_anomalies, + _summary_status( + is_pending=summary.volume_is_pending, + is_training=summary.volume_is_training, + has_errors=summary.volume_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.SCHEMA], + summary.schema_anomalies, + _summary_status( + is_pending=summary.schema_is_pending, + is_training=False, + has_errors=summary.schema_has_errors, + ), + ], + [ + _MONITOR_LABEL[MonitorType.METRIC], + summary.metric_anomalies, + _summary_status( + is_pending=summary.metric_is_pending, + is_training=summary.metric_is_training, + has_errors=summary.metric_has_errors, + ), + ], + ], + ) + + return doc.render() + + +@with_database_session +@mcp_permission("view") +def list_monitored_tables( + table_group_id: str, + anomaly_type: str | None = None, + sort_by: str | None = None, + limit: int = 20, + page: int = 1, +) -> str: + """List monitored tables in a table group with per-type anomaly counts, training / pending / error status, latest update timestamp, and row count change. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + anomaly_type: Filter to tables with at least one anomaly of this type. One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. + sort_by: Sort order. One of ``table_name`` (default, case-insensitive ascending), ``anomaly_count_desc`` (sorts by the filtered type when ``anomaly_type`` is set, total anomalies otherwise), ``latest_update_desc``, ``row_count_change_desc``. + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + + monitor_type = parse_monitor_type(anomaly_type, "anomaly_type") if anomaly_type is not None else None + sort = parse_monitor_table_sort(sort_by) if sort_by is not None else None + + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + rows, total = TableGroup.list_monitor_table_summaries( + tg.id, + anomaly_types=[monitor_type.value] if monitor_type is not None else None, + sort_by=_SORT_TO_MODEL_FIELD[sort] if sort is not None else None, + page=page, + limit=limit, + ) + + doc = MdDoc() + scope = f" — anomaly type `{anomaly_type}`" if anomaly_type else "" + doc.heading(1, f"Monitored tables in `{tg.table_groups_name}`{scope}") + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not rows: + if page > 1: + doc.text(f"No tables on page {page} (total: {total}).") + else: + doc.text("_No monitored tables match this filter._") + return doc.render() + + doc.table( + [ + "Table", + "Freshness", + "Volume", + "Schema", + "Schema change", + "Metric", + "Latest update", + "Row count change", + ], + [ + [ + row.table_name, + _format_monitor_cell( + count=row.freshness_anomalies, + is_pending=row.freshness_is_pending, + is_training=row.freshness_is_training, + has_error=row.freshness_error_message is not None, + ), + _format_monitor_cell( + count=row.volume_anomalies, + is_pending=row.volume_is_pending, + is_training=row.volume_is_training, + has_error=row.volume_error_message is not None, + ), + _format_schema_cell(row), + _format_schema_change(row), + _format_monitor_cell( + count=row.metric_anomalies, + is_pending=row.metric_is_pending, + is_training=row.metric_is_training, + has_error=row.metric_error_message is not None, + ), + row.latest_update, + _format_row_count_change(row), + ] + for row in rows + ], + code=[0], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + return doc.render() + + +def _summary_status(*, is_pending: bool, is_training: bool, has_errors: bool) -> str: + """Single status word for the group-level summary cell.""" + if has_errors: + return "Error" + if is_pending: + return "No results yet or not configured" + if is_training: + return "Training" + return "Ok" + + +def _format_monitor_cell( + *, + count: int, + is_pending: bool, + is_training: bool | None, + has_error: bool, +) -> str: + """Render a per-table, per-monitor cell (non-schema). + + Precedence: error > positive count > pending > training > zero. A positive + count wins over training/pending so anomalies stay visible when a monitor is + still learning, and so a row that surfaces via ``anomaly_type=X`` actually + shows the X count in its column. + """ + if has_error: + return "Error" + if count > 0: + return str(count) + if is_pending: + return "Pending" + if is_training: + return "Training" + return "0" + + +def _format_schema_cell(row: MonitorTableSummary) -> str: + """Schema anomaly count, or pending / error status. Verbose detail is in the + sibling ``Schema change`` column.""" + if row.schema_error_message is not None: + return "Error" + if row.schema_is_pending: + return "Pending" + return str(row.schema_anomalies) + + +def _format_schema_change(row: MonitorTableSummary) -> str | None: + """Verbose description of schema events in the lookback window. + + Mirrors the dashboard tooltip wording: "Table added with N columns.", + "Table dropped with N columns.", or a per-kind breakdown when the table was + modified ("1 column added. 2 columns dropped."). Empty when there are no + schema events or the monitor is in an unfinished state. + """ + if row.schema_is_pending or row.schema_error_message is not None: + return None + state = row.table_state + if state == "added": + return f"Table added with {row.column_adds} columns." + if state == "dropped": + return f"Table dropped with {row.column_drops} columns." + if state == "modified": + parts: list[str] = [] + if row.column_adds: + parts.append(f"{row.column_adds} column{'' if row.column_adds == 1 else 's'} added") + if row.column_drops: + parts.append(f"{row.column_drops} column{'' if row.column_drops == 1 else 's'} dropped") + if row.column_mods: + parts.append(f"{row.column_mods} column{'' if row.column_mods == 1 else 's'} modified") + return ". ".join(parts) + "." if parts else None + return None + + +def _format_row_count_change(row: MonitorTableSummary) -> str | None: + """Signed delta between the latest and pre-window row count. + + ``+1,234`` / ``-1,234`` / ``0`` when both endpoints are known; ``None`` (em-dash) + when either is missing (e.g. first run with no baseline). Sign reflects net + change across the window, not run-to-run variance. + """ + current = row.row_count + previous = row.previous_row_count + if current is None or previous is None: + return None + delta = current - previous + if delta == 0: + return "0" + return f"{delta:+,}" + + +# --------------------------------------------------------------------------- +# Lifecycle & settings tools +# --------------------------------------------------------------------------- + +_LOOKBACK_RUNS_RANGE = (1, 200) +_MIN_TRAINING_LOOKBACK_RANGE = (20, 1000) + + +@with_database_session +@mcp_permission("edit") +def enable_monitors(table_group_id: str, cron_expression: str, cron_tz: str = "UTC") -> str: + """Turn on monitoring for a table group: create its monitors and schedule them to run. + + Sets up the initial Volume and Schema monitors with default settings; adjust them afterward + with ``update_monitor_settings``. Freshness monitors are added automatically once the table + group has profiling data. Fails if monitoring is already on for the group. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + cron_expression: Cron expression for the monitor schedule (required — monitors only run on a schedule), e.g. ``0 6 * * *`` for 6 AM daily. + cron_tz: IANA timezone for the schedule. Defaults to ``UTC``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is not None: + raise MCPUserError("Monitoring is already enabled for this table group.") + _validate_cron_verbatim(cron_expression, cron_tz) + + _, count = enable_monitoring(tg, cron_expression, cron_tz) + + doc = MdDoc() + doc.heading(1, f"Monitoring enabled for `{tg.table_groups_name}`") + doc.field("Initial monitors created", count) + doc.field("Cron expression", cron_expression, code=True) + if (readable := describe_cron(cron_expression)) is not None: + doc.field("Cron description", readable) + doc.field("Timezone", cron_tz) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_monitor_settings(table_group_id: str) -> str: + """Get a table group's monitor configuration and schedule. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + # A monitored table group always has a run-monitors schedule (enable_monitors creates it). + schedule = cast(JobSchedule, JobSchedule.get_for_monitor_suite(monitor_suite.id)) + + doc = MdDoc() + doc.heading(1, f"Monitor settings for `{tg.table_groups_name}`") + doc.field("Project", tg.project_code, code=True) + _render_monitor_settings(doc, monitor_suite, schedule) + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def update_monitor_settings( + table_group_id: str, + sensitivity: str | None = None, + lookback_runs: int | None = None, + min_training_lookback: int | None = None, + exclude_weekends: bool | None = None, + holiday_codes: list[str] | None = None, + regenerate_freshness: bool | None = None, + cron_expression: str | None = None, + cron_tz: str | None = None, + active: bool | None = None, +) -> str: + """Update a table group's monitor configuration and/or schedule. Partial — omitted fields are left unchanged. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + sensitivity: How readily monitors flag a deviation. ``high`` flags smaller deviations (more alerts), ``low`` only large ones (fewer alerts), ``medium`` is balanced. + lookback_runs: Monitor runs aggregated for dashboard summaries. Display only — does not affect detection (1-200). + min_training_lookback: Minimum monitor runs required to train the prediction model (20-1000). + exclude_weekends: Whether to exclude weekends from the model's training data. + holiday_codes: Holiday calendars to exclude from the model's training data — ISO country codes (e.g. ``US``, ``GB``) or financial-market codes (e.g. ``NYSE``, ``ECB``); see https://holidays.readthedocs.io/en/latest/#available-countries. Pass an empty list to clear. + regenerate_freshness: Whether to automatically reconfigure Freshness monitors with new fingerprints after each profiling run. + cron_expression: New cron expression for the schedule, e.g. ``0 6 * * *``. + cron_tz: New IANA timezone for the schedule, e.g. ``America/New_York``. + active: ``True`` to resume the schedule, ``False`` to pause it. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + return _NOT_MONITORED_OUTPUT + + if all( + value is None + for value in ( + sensitivity, + lookback_runs, + min_training_lookback, + exclude_weekends, + holiday_codes, + regenerate_freshness, + cron_expression, + cron_tz, + active, + ) + ): + raise MCPUserError("No fields supplied to update.") + + # A monitored table group always has a run-monitors schedule (enable_monitors creates it). + schedule = cast(JobSchedule, JobSchedule.get_for_monitor_suite(monitor_suite.id)) + + suite_attrs: dict[str, Any] = {} + if sensitivity is not None: + suite_attrs["predict_sensitivity"] = _parse_sensitivity(sensitivity) + if lookback_runs is not None: + _validate_range(lookback_runs, "lookback_runs", *_LOOKBACK_RUNS_RANGE) + suite_attrs["monitor_lookback"] = lookback_runs + if min_training_lookback is not None: + _validate_range(min_training_lookback, "min_training_lookback", *_MIN_TRAINING_LOOKBACK_RANGE) + suite_attrs["predict_min_lookback"] = min_training_lookback + if exclude_weekends is not None: + suite_attrs["predict_exclude_weekends"] = exclude_weekends + if holiday_codes is not None: + cleaned = [code.strip() for code in holiday_codes if code.strip()] + invalid = [code for code in cleaned if not is_supported_holiday_code(code)] + if invalid: + raise MCPUserError( + f"Unknown holiday codes: {', '.join(invalid)}. Use ISO country codes (e.g. `US`, `GB`) " + "or financial-market codes (e.g. `NYSE`, `ECB`); see " + "https://holidays.readthedocs.io/en/latest/#available-countries." + ) + suite_attrs["predict_holiday_codes"] = ",".join(cleaned) if cleaned else None + if regenerate_freshness is not None: + suite_attrs["monitor_regenerate_freshness"] = regenerate_freshness + + if cron_expression is not None or cron_tz is not None: + effective_expr = cron_expression if cron_expression is not None else schedule.cron_expr + effective_tz = cron_tz if cron_tz is not None else schedule.cron_tz + _validate_cron_verbatim(effective_expr, effective_tz) + + update_monitoring( + monitor_suite, + schedule, + suite_attrs=suite_attrs, + cron_expr=cron_expression, + cron_tz=cron_tz, + active=active, + ) + + doc = MdDoc() + doc.heading(1, f"Monitor settings updated for `{tg.table_groups_name}`") + _render_monitor_settings(doc, monitor_suite, schedule) + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def disable_monitors(table_group_id: str) -> str: + """Turn off monitoring for a table group, removing all its monitors, the schedule, and monitor history. + + This is irreversible — the monitors and their accumulated history are deleted. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + """ + tg, monitor_suite = resolve_monitored_table_group(table_group_id) + if monitor_suite is None: + raise MCPUserError("Monitoring is not enabled for this table group.") + + counts = disable_monitoring(monitor_suite) + + doc = MdDoc() + doc.heading(1, f"Monitoring disabled for `{tg.table_groups_name}`") + doc.text("Removed all monitors, the schedule, and monitor history.") + doc.field("Monitors removed", counts["monitors"]) + doc.field("Events removed", counts["events"]) + doc.field("Runs removed", counts["runs"]) + return doc.render() + + +def _parse_sensitivity(value: str) -> PredictSensitivity: + try: + return PredictSensitivity(value) + except ValueError as err: + valid = ", ".join(s.value for s in PredictSensitivity) + raise MCPUserError(f"Invalid sensitivity `{value}`. Valid values: {valid}") from err + + +def _validate_range(value: int, label: str, low: int, high: int) -> None: + if not low <= value <= high: + raise MCPUserError(f"Invalid {label} `{value}`: must be between {low} and {high}.") + + +def _validate_cron_verbatim(cron_expr: str, cron_tz: str) -> None: + """Validate a cron expression + timezone, surfacing the parser's message verbatim.""" + sample = get_cron_sample(cron_expr, cron_tz, sample_count=1) + if "error" in sample: + raise MCPUserError(sample["error"]) + + +def _last_monitor_run(schedule_id: UUID) -> datetime | None: + je = ( + get_current_session() + .scalars( + select(JobExecution) + .where(JobExecution.job_schedule_id == schedule_id) + .order_by(JobExecution.created_at.desc()) + .limit(1) + ) + .first() + ) + if je is None: + return None + return je.completed_at or je.started_at + + +def _render_monitor_settings(doc: MdDoc, monitor_suite: TestSuite, schedule: JobSchedule) -> None: + doc.field("Lookback runs", monitor_suite.monitor_lookback) + doc.field("Regenerate freshness", monitor_suite.monitor_regenerate_freshness) + + doc.heading(2, "Prediction Model") + sensitivity = monitor_suite.predict_sensitivity.value if monitor_suite.predict_sensitivity else None + doc.field("Sensitivity", sensitivity) + doc.field("Min training lookback", monitor_suite.predict_min_lookback) + doc.field("Exclude weekends", monitor_suite.predict_exclude_weekends) + holidays = monitor_suite.holiday_codes_list + doc.field("Holiday codes", ", ".join(holidays) if holidays else None) + + doc.heading(2, "Schedule") + doc.field("Cron expression", schedule.cron_expr, code=True) + if (readable := describe_cron(schedule.cron_expr)) is not None: + doc.field("Cron description", readable) + doc.field("Timezone", schedule.cron_tz) + doc.field("Status", "Active" if schedule.active else "Paused") + if schedule.active: + try: + next_runs = schedule.get_sample_triggering_timestamps(1) + except Exception: + next_runs = [] + if next_runs: + # Cron samples are tz-aware in the schedule's timezone; convert to UTC so the + # rendered " UTC" suffix is accurate. + doc.field("Next run", next_runs[0].astimezone(UTC)) + if (last_run := _last_monitor_run(schedule.id)) is not None: + # job_executions timestamps are tz-aware UTC; convert so the rendered " UTC" suffix is accurate. + doc.field("Last run", last_run.astimezone(UTC)) +# list_monitor_events +# --------------------------------------------------------------------------- + + +@with_database_session +@mcp_permission("view") +def list_monitor_events( + table_group_id: str, + table_name: str, + monitor_type: str, + monitor_id: str | None = None, + include_predictions: bool = False, + limit: int = 20, + page: int = 1, +) -> str: + """List per-table monitor events for one monitor type within the lookback window, newest first. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Table name exactly as stored in TestGen (case-sensitive). + monitor_type: One of ``freshness`` / ``volume`` / ``schema`` / ``metric``. + monitor_id: Required for ``monitor_type="metric"``, rejected for other types. Get it from ``list_monitors``. + include_predictions: When True, append the monitor's forecast, if available (shown on the first page only). + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + parsed_type = parse_monitor_type(monitor_type) + + if parsed_type == MonitorType.METRIC and monitor_id is None: + raise MCPUserError( + '`monitor_id` is required for `monitor_type="metric"`. ' + "Metric monitors are user-defined and many can apply to one table — " + "use `list_monitors(table_group_id, table_name)` to find the metric's " + "`monitor_id` first, then call this tool with it." + ) + if monitor_id is not None and parsed_type != MonitorType.METRIC: + raise MCPUserError( + '`monitor_id` only applies when `monitor_type="metric"`. ' + "For singleton monitor types (`freshness`, `volume`, `schema`), " + "the monitor is uniquely identified by table + type." + ) + + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + monitor_def: TestDefinition | TestDefinitionSummary | None + if parsed_type == MonitorType.METRIC: + monitor_uuid = parse_uuid(monitor_id, "monitor_id") + # Scope the lookup to this suite + table + Metric_Trend so a monitor_id + # for a metric on another table (or another suite) can't be rendered + # under this table's heading. + monitor_def = next( + iter( + TestDefinition.select_where( + TestDefinition.id == monitor_uuid, + TestDefinition.test_suite_id == suite.id, + TestDefinition.table_name == table_name, + TestDefinition.test_type == MonitorType.METRIC.value, + ) + ), + None, + ) + if monitor_def is None: + raise MCPUserError( + f"No metric monitor `{monitor_id}` found on table `{table_name}`. " + "Use `list_monitors(table_group_id, table_name)` to find a metric's `monitor_id`." + ) + events, total = TestResult.list_metric_monitor_events( + suite.id, + monitor_uuid, + page=page, + limit=limit, + ) + metric_name = monitor_def.column_name or None + else: + events, total = TestResult.list_monitor_events_for_table( + suite.id, + table_name, + monitor_type=parsed_type.value, + page=page, + limit=limit, + ) + # Singleton lookup only needed when we will render a forecast — Schema + # short-circuits to "not applicable" without touching the definition. + monitor_def = ( + TestDefinition.get_singleton_monitor(suite.id, table_name, parsed_type.value) + if include_predictions and parsed_type != MonitorType.SCHEMA + else None + ) + metric_name = None + + doc = MdDoc() + if parsed_type == MonitorType.METRIC and metric_name: + doc.heading( + 1, + f"Monitor events: Metric `{metric_name}` on `{table_name}` in `{tg.table_groups_name}`", + ) + else: + doc.heading( + 1, + f"Monitor events: `{table_name}` — `{monitor_type}` in `{tg.table_groups_name}`", + ) + + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not events: + if page > 1: + doc.text(f"No events on page {page} (total: {total}).") + else: + doc.text("_No monitor events in the active lookback window._") + if include_predictions and page == 1: + _render_forecast_section(doc, _compute_forecast(suite, table_name, parsed_type, monitor_def, events)) + return doc.render() + + if parsed_type == MonitorType.SCHEMA: + doc.table( + ["Time", "Status", "Table change", "Columns added", "Columns dropped", "Columns modified"], + [ + [ + event.test_time, + _event_status(event), + _format_schema_change_kind(event.schema_change_kind), + event.column_adds, + event.column_drops, + event.column_mods, + ] + for event in events + ], + ) + elif parsed_type == MonitorType.FRESHNESS: + doc.table( + ["Time", "Status", "Update detected", "Detail"], + [ + [event.test_time, _event_status(event), *_parse_freshness_message(event.message)] + for event in events + ], + ) + elif parsed_type == MonitorType.METRIC: + doc.table( + ["Time", "Status", "Value", "Lower bound", "Upper bound"], + [ + [ + event.test_time, + _event_status(event), + event.signal, + event.lower_bound, + event.upper_bound, + ] + for event in events + ], + ) + else: # VOLUME + doc.table( + ["Time", "Status", "Row count", "Lower bound", "Upper bound"], + [ + [ + event.test_time, + _event_status(event), + event.signal, + event.lower_bound, + event.upper_bound, + ] + for event in events + ], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + if include_predictions: + _render_forecast_section(doc, _compute_forecast(suite, table_name, parsed_type, monitor_def, events)) + + return doc.render() + + +@dataclass +class _Forecast: + """The forecast to render under a monitor's events. Exactly one of + ``window`` / ``points`` / ``note`` is meaningful — see ``_compute_forecast``.""" + note: str | None = None + sensitivity: str | None = None + points: list[MonitorForecastPoint] = field(default_factory=list) + window: dict | None = None # {"start": epoch_ms | None, "end": epoch_ms} + + +def _compute_forecast( + suite: TestSuite, + table_name: str, + parsed_type: MonitorType, + monitor_def: TestDefinition | TestDefinitionSummary | None, + events: list[MonitorEvent], +) -> _Forecast: + """Compute the same forecast the monitors dashboard plots for this monitor. + + Volume / Metric in Prediction Model mode get a forward value band — the + coupled baseline-then-refresh band when the monitor is tied to a Freshness + monitor, otherwise the per-step prediction band. Freshness gets its predicted + next-update window. Everything else gets an explanatory note.""" + if parsed_type == MonitorType.SCHEMA: + return _Forecast( + note="_Predictions not applicable to Schema monitors._" + ) + if monitor_def is None: + return _Forecast(note="_No monitor configured for this table._") + if monitor_def.history_calculation != "PREDICT": + return _Forecast( + note="_Predictions not available — this monitor's threshold mode is " + "Static or Historical Calculation, not Prediction Model._" + ) + + if parsed_type == MonitorType.FRESHNESS: + window = _next_update_window_for_table(suite, monitor_def, events) + if window is None: + return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(window=window) + + last_run_time = max((e.test_time for e in events if e.test_time is not None), default=None) + + # A monitor coupled to a Freshness monitor holds at its baseline until the + # next expected refresh, so its band is the coupled baseline-then-refresh + # shape keyed off the freshness next-update window — not the raw per-step + # prediction series. The tolerance precondition mirrors the dashboard: a + # coupled monitor with no configured tolerance has no band on either surface. + if monitor_def.prediction and monitor_def.prediction.get("freshness_gated"): + if monitor_def.lower_tolerance is None and monitor_def.upper_tolerance is None: + return _Forecast(note=_FORECAST_UNAVAILABLE_NOTE) + freshness_def = TestDefinition.get_singleton_monitor( + suite.id, table_name, MonitorType.FRESHNESS.value + ) + freshness_events, _ = TestResult.list_monitor_events_for_table( + suite.id, table_name, monitor_type=MonitorType.FRESHNESS.value, limit=None + ) + window = _next_update_window_for_table(suite, freshness_def, freshness_events) + points = forecast_band_points(gated_forecast_prediction(monitor_def, window, last_run_time)) + if not points: + return _Forecast(note=_FORECAST_UNAVAILABLE_NOTE) + return _Forecast(points=points) + + sensitivity = suite.predict_sensitivity.value if suite.predict_sensitivity is not None else "medium" + points = forecast_points_from_prediction(monitor_def.prediction, sensitivity) + if not points: + return _Forecast(note=_FORECAST_PENDING_NOTE) + return _Forecast(sensitivity=sensitivity, points=points) + + +def _next_update_window_for_table( + suite: TestSuite, + freshness_def: TestDefinition | TestDefinitionSummary | None, + freshness_events: list[MonitorEvent], +) -> dict | None: + """Predicted next-update window from the table's Freshness monitor, computed + the same way the dashboard does (last detected update + business-time tolerance).""" + last_detection = max( + ( + e.test_time for e in freshness_events + if e.test_time is not None and not e.is_training and not e.is_pending and not e.is_error + and _parse_freshness_message(e.message)[0] == "Yes" + ), + default=None, + ) + schedule = JobSchedule.get_for_monitor_suite(suite.id) + return next_update_window( + freshness_def, + last_detection, + test_suite=suite, + cron_tz=schedule.cron_tz if schedule else None, + ) + + +def _render_forecast_section(doc: MdDoc, forecast: _Forecast) -> None: + """Append the forecast under its own `## Forecast` heading. Predictions are + NEVER mixed into the events table so the consumer can tell observed-past from + predicted-future at a glance.""" + doc.heading(2, "Forecast") + + if forecast.window is not None: + end = datetime.fromtimestamp(forecast.window["end"] / 1000.0, UTC) + start_ms = forecast.window.get("start") + if start_ms is not None: + start = datetime.fromtimestamp(start_ms / 1000.0, UTC) + doc.field("Next update expected", f"{start:%Y-%m-%d %H:%M} to {end:%Y-%m-%d %H:%M} UTC") + else: + doc.field("Next update expected by", f"{end:%Y-%m-%d %H:%M} UTC") + return + + if forecast.points: + if forecast.sensitivity: + doc.field("Sensitivity", forecast.sensitivity) + doc.table( + ["Time", "Predicted lower", "Predicted upper"], + [[p.test_time, p.lower_bound, p.upper_bound] for p in forecast.points], + ) + return + + doc.text(forecast.note or "_No forecast available._") + + +def _event_status(event: MonitorEvent) -> str: + """Single status word for a monitor-event row.""" + if event.is_error: + return "Error" + if event.is_pending: + return "Pending" + if event.is_training: + return "Training" + if event.is_anomaly: + return "Anomaly" + return "Ok" + + +# Freshness events carry a structured message of the form +# ``"Table update detected: {Yes|No}[. {Detail}]"`` — see the SQL templates at +# ``template/dbsetup_test_types/test_types_Freshness_Trend.yaml``. Parse it into +# separate columns so the LLM doesn't have to re-derive the same fields. +_FRESHNESS_MESSAGE_RE = re.compile( + r"^Table update detected:\s*(Yes|No)(?:\.\s*(.+?))?\.?\s*$", + re.IGNORECASE, +) + + +def _parse_freshness_message(message: str | None) -> tuple[str, str]: + """Return ``(update_detected, detail)``: ``"Yes"``/``"No"``/``"—"`` and the + descriptive tail (``"On time"`` / ``"Earlier than expected"`` / + ``"Later than expected"`` / ``"Late"`` / ``"—"``). Falls back to the raw + message text when the format doesn't match (e.g. error rows).""" + if not message: + return "—", "—" + match = _FRESHNESS_MESSAGE_RE.match(message.strip()) + if not match: + return "—", message + detected = match.group(1).capitalize() + detail = match.group(2) or "—" + return detected, detail + + +# --------------------------------------------------------------------------- +# list_monitors +# --------------------------------------------------------------------------- + + +_SCHEMA_CHANGE_LABEL: dict[str, str] = { + SCHEMA_CHANGE_ADDED: "added", + SCHEMA_CHANGE_DROPPED: "dropped", + SCHEMA_CHANGE_MODIFIED: "modified", +} + + +def _format_schema_change_kind(code: str | None) -> str | None: + """Map an audit-log change code (``A`` / ``D`` / ``M``) to its user-facing word.""" + if code is None: + return None + return _SCHEMA_CHANGE_LABEL.get(code, code) + + +@with_database_session +@mcp_permission("view") +def list_monitors(table_group_id: str, table_name: str) -> str: + """List configured monitors for a table: monitor IDs, types, threshold modes, and bounds. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Table name exactly as stored in TestGen (case-sensitive). + """ + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + configs = TestDefinition.list_monitor_configs_for_table(suite.id, table_name) + + doc = MdDoc() + doc.heading(1, f"Monitors on `{table_name}` in `{tg.table_groups_name}`") + sensitivity = ( + suite.predict_sensitivity.value + if suite.predict_sensitivity is not None + else "medium" + ) + doc.field("Prediction model sensitivity", sensitivity) + + if not configs: + doc.text("_No monitors configured for this table._") + return doc.render() + + doc.table( + ["Monitor ID", "Type", "Metric name", "Threshold mode", "Lower", "Upper", "Metric expression"], + [ + [ + str(c.monitor_id), + _MONITOR_LABEL[MonitorType(c.test_type)], + c.metric_name, + c.threshold_mode, + c.threshold_lower, + c.threshold_upper, + c.custom_query, + ] + for c in configs + ], + code=[0, 6], + ) + + return doc.render() + + +# --------------------------------------------------------------------------- +# list_monitor_schema_changes +# --------------------------------------------------------------------------- + + +@with_database_session +@mcp_permission("view") +def list_monitor_schema_changes( + table_group_id: str, + table_name: str | None = None, + since: str | None = None, + limit: int = 20, + page: int = 1, +) -> str: + """List schema changes detected for a table group, newest first. + + Args: + table_group_id: UUID of the table group, e.g. from ``list_table_groups``. + table_name: Filter to one table (exact, case-sensitive). Omit to list changes across every table in the group. + since: Lower-bound date (e.g. ``'7 days'``, ``'2 weeks'``, ``'2026-04-01'``). Omit to include the entire stored history. + limit: Page size (default 20, max 100). + page: Page number starting at 1 (default 1). + """ + validate_page(page) + validate_limit(limit, 100) + since_date = parse_since_arg(since) if since is not None else None + + tg, suite = resolve_monitored_table_group(table_group_id) + if suite is None: + return _NOT_MONITORED_OUTPUT + + clauses = [] + if table_name is not None: + clauses.append(DataStructureLog.table_name == table_name) + if since_date is not None: + clauses.append(DataStructureLog.change_date >= since_date) + + entries, total = DataStructureLog.list_for_table_group( + tg.id, + *clauses, + page=page, + limit=limit, + ) + + doc = MdDoc() + scope_parts: list[str] = [] + if table_name: + scope_parts.append(f"table `{table_name}`") + if since: + scope_parts.append(f"since `{since}`") + scope = f" — {' — '.join(scope_parts)}" if scope_parts else "" + doc.heading(1, f"Schema changes in `{tg.table_groups_name}`{scope}") + + page_info = format_page_info(total, page, limit) + if page_info: + doc.text(page_info) + + if not entries: + if page > 1: + doc.text(f"No schema changes on page {page} (total: {total}).") + else: + doc.text("_No schema changes recorded for this scope._") + return doc.render() + + doc.table( + ["Time", "Change", "Table", "Column", "Old type", "New type"], + [ + [ + entry.change_date, + _format_schema_change_kind(entry.change), + entry.table_name, + entry.column_name, + entry.old_data_type, + entry.new_data_type, + ] + for entry in entries + ], + code=[2, 3], + ) + + if footer := format_page_footer(total, page, limit): + doc.text(footer) + + return doc.render() diff --git a/testgen/mcp/tools/profile_history.py b/testgen/mcp/tools/profile_history.py index 17839798..62537cc4 100644 --- a/testgen/mcp/tools/profile_history.py +++ b/testgen/mcp/tools/profile_history.py @@ -203,7 +203,7 @@ def _delta_cell(metric: ProfileMetric, baseline: object | None, target: object | def _require_completed(run: ProfilingRun, label: str) -> None: """Raise if the run's job execution isn't completed.""" - je = get_current_session().get(JobExecution, run.job_execution_id) + je = get_current_session().get(JobExecution, run.id) if je.status != JobStatus.COMPLETED: status_label = ProfilingRunSummary.STATUS_LABEL.get(je.status, je.status) raise MCPUserError( @@ -475,8 +475,8 @@ def _render_run_comparison( ["", "Target", "Baseline"], [ ["Profiling Run", - MdDoc.code(str(target_run.job_execution_id)), - MdDoc.code(str(baseline_run.job_execution_id))], + MdDoc.code(str(target_run.id)), + MdDoc.code(str(baseline_run.id))], ["Started", target_run.profiling_starttime, baseline_run.profiling_starttime], ], ) @@ -792,7 +792,7 @@ def _render_schema_history( baseline_snap=snapshots.get(baseline.id, {}), ) doc.heading(2, f"Run started {_format_run_label(target)}") - doc.field("Profiling Run", target.job_execution_id, code=True) + doc.field("Profiling Run", target.id, code=True) if section_lines: doc.bullets(section_lines) else: diff --git a/testgen/mcp/tools/profiling.py b/testgen/mcp/tools/profiling.py index 79ed638c..5ef11af9 100644 --- a/testgen/mcp/tools/profiling.py +++ b/testgen/mcp/tools/profiling.py @@ -3,6 +3,8 @@ from sqlalchemy import func, or_ +from testgen.common.data_catalog_service import build_create_table_script +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models import with_database_session from testgen.common.models.data_column import ( SUGGESTED_DATA_TYPE_TO_PREFIX, @@ -12,6 +14,7 @@ DataColumnChars, ) from testgen.common.models.data_table import DataTable +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, ProfilingRunSummary @@ -413,6 +416,7 @@ def _render_column_profile_row(c: ColumnProfileSummary) -> list: @mcp_permission("catalog") def list_profiling_runs( table_group_id: str, + schedule_id: str | None = None, status: str | None = None, limit: int = 10, page: int = 1, @@ -422,6 +426,8 @@ def list_profiling_runs( Args: table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + schedule_id: Optional UUID of a schedule, e.g. from `list_schedules`. Returns only runs + triggered by that schedule. status: Optional run status filter. One of: Pending, Running, Completed, Canceled, Error. limit: Page size (default 10, max 100). page: Page number starting at 1 (default 1). @@ -430,11 +436,14 @@ def list_profiling_runs( validate_page(page) statuses = parse_run_status_filter(status) if status else None + if schedule_id: + parse_uuid(schedule_id, "schedule_id") tg = resolve_table_group(table_group_id) summaries, total = ProfilingRun.select_summary( project_code=tg.project_code, table_group_id=tg.id, + schedule_id=schedule_id, statuses=statuses, page=page, page_size=limit, @@ -444,7 +453,9 @@ def list_profiling_runs( # joined-run queries. Surface them as a separate "Pending" section on page 1. pending_jes: list[JobExecution] = [] if page == 1: + clauses = [JobExecution.job_schedule_id == schedule_id] if schedule_id else [] pending_jes = JobExecution.select_active_by_kwargs( + *clauses, project_code=tg.project_code, job_key=RUN_PROFILE_JOB_KEY, kwargs_match={"table_group_id": str(tg.id)}, @@ -452,7 +463,12 @@ def list_profiling_runs( ) doc = MdDoc() - scope = f" — status `{status}`" if status else "" + scope_parts = [] + if schedule_id: + scope_parts.append(f"schedule `{schedule_id}`") + if status: + scope_parts.append(f"status `{status}`") + scope = f" — {', '.join(scope_parts)}" if scope_parts else "" doc.heading(1, f"Profiling runs for `{tg.table_groups_name}`{scope}") next_run = next_scheduled_run( @@ -509,7 +525,7 @@ def get_profiling_run(job_execution_id: str) -> str: doc = MdDoc() tg_label = summary.table_groups_name or "—" doc.heading(1, f"Profiling run: {tg_label}") - doc.field("Job ID", summary.job_execution_id, code=True) + doc.field("Profiling Run", summary.job_execution_id, code=True) if summary.table_groups_name: doc.field("Table group", summary.table_groups_name) if summary.table_group_schema: @@ -528,13 +544,19 @@ def get_profiling_run(job_execution_id: str) -> str: doc.field("Columns profiled", summary.column_ct or 0) if summary.record_ct is not None: doc.field("Records", summary.record_ct) - doc.field( - "Hygiene issues (confirmed)", - f"{(summary.anomalies_definite_ct or 0) + (summary.anomalies_likely_ct or 0) + (summary.anomalies_possible_ct or 0)} total " - f"— {summary.anomalies_definite_ct or 0} definite, " - f"{summary.anomalies_likely_ct or 0} likely, " - f"{summary.anomalies_possible_ct or 0} possible", - ) + if summary.profiling_run_id: + # Count from the canonical source so likelihood buckets and Potential PII + # stay separate (matches the REST profiling-run issue_counts). + counts = HygieneIssue.count_for_run(summary.profiling_run_id) + hygiene = counts.hygiene_issues + doc.field( + "Hygiene issues (confirmed)", + f"{hygiene.definite + hygiene.likely + hygiene.possible} total " + f"— {hygiene.definite} definite, {hygiene.likely} likely, {hygiene.possible} possible", + ) + pii = counts.potential_pii + if pii.high or pii.moderate: + doc.field("Potential PII", f"{pii.high} high, {pii.moderate} moderate") if summary.dq_score_profiling is not None: doc.field("Profiling Score", friendly_score(summary.dq_score_profiling)) @@ -561,7 +583,7 @@ def get_profiling_run(job_execution_id: str) -> str: def _render_pending_profiling_je(doc: MdDoc, je: JobExecution, label: str) -> None: status_label = ProfilingRunSummary.STATUS_LABEL.get(je.status, je.status) doc.heading(3, f"{label} — {status_label}") - doc.field("Job ID", je.id, code=True) + doc.field("Profiling Run", je.id, code=True) if je.job_schedule_id is not None: doc.field("Schedule", je.job_schedule_id, code=True) doc.field("Submitted", je.created_at) @@ -572,7 +594,7 @@ def _render_pending_profiling_je(doc: MdDoc, je: JobExecution, label: str) -> No def _render_profiling_run_section(doc: MdDoc, run: ProfilingRunSummary) -> None: title = run.table_groups_name or run.profiling_run_id or run.job_execution_id doc.heading(2, f"{title} — {run.status_label}") - doc.field("Job ID", run.job_execution_id, code=True) + doc.field("Profiling Run", run.job_execution_id, code=True) if run.job_schedule_id is not None: doc.field("Schedule", run.job_schedule_id, code=True) doc.field("Submitted", run.created_at) @@ -743,7 +765,7 @@ def get_column_profile_detail( "Run profiling for the table group first." ) - if detail.profile_run_status in ("Running", "Error", "Cancelled"): + if detail.profile_run_status in (JobStatus.RUNNING, JobStatus.ERROR, JobStatus.CANCELED): _raise_run_not_ready(detail) payload = dataclasses.asdict(detail) @@ -765,11 +787,12 @@ def _raise_run_not_ready(detail: ColumnProfileDetail) -> None: ended = detail.profile_run_ended_at started_label = started.strftime("%Y-%m-%d %H:%M UTC") if started else "—" ended_label = ended.strftime("%Y-%m-%d %H:%M UTC") if ended else "—" + status_label = JOB_STATUS_LABEL.get(status, status) lines = [ - f"Profiling run `{je}` is in `{status}` state — no profile detail available.", + f"Profiling run `{je}` is in `{status_label}` state — no profile detail available.", f"Started: {started_label}. Ended: {ended_label}.", ] - if status == "Error" and detail.profile_run_log_message: + if status == JobStatus.ERROR and detail.profile_run_log_message: lines.append(f"Error: {detail.profile_run_log_message}") raise MCPUserError("\n".join(lines)) @@ -952,7 +975,7 @@ def get_column_frequent_values( doc = MdDoc() doc.heading(1, f"Frequent values: {table_name}.{column_name}") doc.field("Table group", tg.id, code=True) - doc.field("Profiling Run", profiling_run.job_execution_id, code=True) + doc.field("Profiling Run", profiling_run.id, code=True) doc.field("Row Count", profile.record_ct) doc.field("Distinct values", profile.distinct_value_ct) if pii_flag: @@ -1007,7 +1030,7 @@ def get_column_patterns( doc = MdDoc() doc.heading(1, f"Character patterns: {table_name}.{column_name}") doc.field("Table group", tg.id, code=True) - doc.field("Profiling Run", profiling_run.job_execution_id, code=True) + doc.field("Profiling Run", profiling_run.id, code=True) doc.field("Row Count", profile.record_ct) doc.field("Distinct values", profile.distinct_value_ct) @@ -1136,3 +1159,21 @@ def search_columns( if footer: doc.text(footer) return doc.render() + + +@with_database_session +@mcp_permission("catalog") +def generate_create_table_script(table_group_id: str, table_name: str) -> str: + """Generate a CREATE TABLE script for a profiled table from its columns and suggested data types. + + Args: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Table name exactly as stored in TestGen (case-sensitive). + """ + tg = resolve_table_group(table_group_id) + + script = build_create_table_script(tg.id, table_name) + if script is None: + raise MCPResourceNotAccessible("Table", table_name) + + return MdDoc().code_block(script, language="sql").render() diff --git a/testgen/mcp/tools/projects.py b/testgen/mcp/tools/projects.py new file mode 100644 index 00000000..be075233 --- /dev/null +++ b/testgen/mcp/tools/projects.py @@ -0,0 +1,249 @@ +"""MCP write tools for projects. + +Project creation and deletion live in the enterprise project-management plugin +(gated on ``global_admin``); editing a project's own settings is a per-project +``administer`` operation that ships in core. This module holds the latter. +""" + +from __future__ import annotations + +from typing import Any + +from testgen.common.enums import JobKey, JobSource +from testgen.common.models import with_database_session +from testgen.common.models.job_execution import JobExecution +from testgen.common.models.project import Project +from testgen.common.models.scheduler import ( + DEFAULT_DATA_CLEANUP_CRON, + DEFAULT_RETENTION_CRON_TZ, + JobSchedule, +) +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + raise_validation_error, + render_diff_table, + resolve_project, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +DEFAULT_RETENTION_DAYS = 180 + + +def _empty_to_none(value: str | None) -> str | None: + return value if value else None + + +@with_database_session +@mcp_permission("administer") +def update_project( + project_code: str, + *, + project_name: str | None = None, + use_dq_score_weights: bool | None = None, + observability_api_url: str | None = None, + observability_api_key: str | None = None, + data_retention_enabled: bool | None = None, + data_retention_days: int | None = None, + retention_cron_expr: str | None = None, + retention_cron_tz: str | None = None, +) -> str: + """Update a project's settings. + + Args: + project_code: The project code, e.g. from `list_projects`. + project_name: New display name. + use_dq_score_weights: Whether quality scoring uses the configured weights. + Changing this re-runs project-wide score recalculation in the background. + observability_api_url: DataOps Observability API URL. Pass an empty string to clear. + observability_api_key: DataOps Observability API key. Pass an empty string to clear. + Never echoed back in tool output. + data_retention_enabled: Whether old profiling and test history is automatically deleted. + data_retention_days: How many days of history to keep (only when retention is enabled). + Defaults to 180 when retention is enabled without specifying days. + retention_cron_expr: Cron expression for the retention cleanup job + (only when retention is enabled). Defaults to daily at 01:00. + retention_cron_tz: Timezone for the retention cleanup cron + (only when retention is enabled). Defaults to UTC. + """ + supplied = { + "project_name": project_name, + "use_dq_score_weights": use_dq_score_weights, + "observability_api_url": observability_api_url, + "observability_api_key": observability_api_key, + "data_retention_enabled": data_retention_enabled, + "data_retention_days": data_retention_days, + "retention_cron_expr": retention_cron_expr, + "retention_cron_tz": retention_cron_tz, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + project = resolve_project(project_code) + + schedule = JobSchedule.get( + JobSchedule.project_code == project_code, + JobSchedule.key == JobKey.run_data_cleanup, + ) + + # Compute effective retention state from supplied args + current state. Errors below. + effective_enabled = ( + data_retention_enabled if data_retention_enabled is not None else project.data_retention_enabled + ) + effective_days = ( + data_retention_days if data_retention_days is not None else project.data_retention_days + ) + # When retention is being newly enabled without explicit days, fall back to the system default. + if effective_enabled and effective_days is None: + effective_days = DEFAULT_RETENTION_DAYS + + # Validate field-by-field; surface all errors at once. + errors: list[str] = [] + cleaned_name: str | None = None + if project_name is not None: + cleaned_name = project_name.strip() + if not cleaned_name: + errors.append("project_name: must not be empty.") + + if data_retention_days is not None and data_retention_days < 1: + errors.append("data_retention_days: must be a positive integer.") + + # Setting retention schedule args only makes sense when retention will be enabled. + if not effective_enabled: + if data_retention_days is not None: + errors.append("data_retention_days: cannot be set when data_retention_enabled is False.") + if retention_cron_expr is not None: + errors.append("retention_cron_expr: cannot be set when data_retention_enabled is False.") + if retention_cron_tz is not None: + errors.append("retention_cron_tz: cannot be set when data_retention_enabled is False.") + + if errors: + raise_validation_error(errors, "Update rejected. No changes saved.") + + weights_were = project.use_dq_score_weights + + # Snapshot the editable surface (including schedule cron + tz, which live on JobSchedule). + before = _snapshot(project, schedule) + + # Apply project changes. + if cleaned_name is not None: + project.project_name = cleaned_name + if use_dq_score_weights is not None: + project.use_dq_score_weights = use_dq_score_weights + if observability_api_url is not None: + project.observability_api_url = _empty_to_none(observability_api_url) + if observability_api_key is not None: + project.observability_api_key = _empty_to_none(observability_api_key) + project.data_retention_enabled = effective_enabled + project.data_retention_days = effective_days if effective_enabled else None + + # Compute the effective cron values now so the snapshot reflects what the schedule will be. + schedule_supplied = any( + v is not None + for v in (data_retention_enabled, data_retention_days, retention_cron_expr, retention_cron_tz) + ) + effective_cron_expr: str | None + effective_cron_tz: str | None + if effective_enabled: + effective_cron_expr = ( + retention_cron_expr or (schedule.cron_expr if schedule else None) or DEFAULT_DATA_CLEANUP_CRON + ) + effective_cron_tz = ( + retention_cron_tz or (schedule.cron_tz if schedule else None) or DEFAULT_RETENTION_CRON_TZ + ) + else: + effective_cron_expr = None + effective_cron_tz = None + + after = _snapshot( + project, + _ScheduleSnapshot(cron_expr=effective_cron_expr, cron_tz=effective_cron_tz) + if effective_enabled + else None, + ) + + doc = MdDoc() + doc.heading(1, f"Project `{project_code}` updated") + + rendered = render_diff_table( + doc, before, after, + attrs=_DIFF_ORDER, labels=_DIFF_LABELS, secret_attrs=_SECRET_ATTRS, + ) + if not rendered: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + project.save() + + # Schedule side effects: only touch JobSchedule when a retention-related arg was supplied. + if schedule_supplied: + if effective_enabled: + JobSchedule.upsert_for_retention( + project_code=project_code, + retention_days=effective_days, # type: ignore[arg-type] + cron_expr=effective_cron_expr, # type: ignore[arg-type] + cron_tz=effective_cron_tz, # type: ignore[arg-type] + ) + else: + JobSchedule.delete_for_retention(project_code) + + # Weights side effect: submit a background recalculation job, same as the UI. + if use_dq_score_weights is not None and use_dq_score_weights != weights_were: + JobExecution.submit( + job_key=JobKey.recalculate_project_scores, + kwargs={"project_code": project_code}, + source=JobSource.mcp, + project_code=project_code, + ) + + return doc.render() + + +class _ScheduleSnapshot: + """Stand-in for a JobSchedule row, used while computing the post-update snapshot + before any DB write has happened.""" + + def __init__(self, cron_expr: str | None, cron_tz: str | None) -> None: + self.cron_expr = cron_expr + self.cron_tz = cron_tz + + +def _snapshot(project: Project, schedule: Any) -> dict[str, Any]: + return { + "project_name": project.project_name, + "use_dq_score_weights": project.use_dq_score_weights, + "observability_api_url": project.observability_api_url, + "observability_api_key": project.observability_api_key, + "data_retention_enabled": project.data_retention_enabled, + "data_retention_days": project.data_retention_days, + "retention_cron_expr": schedule.cron_expr if schedule is not None else None, + "retention_cron_tz": schedule.cron_tz if schedule is not None else None, + } + + +_DIFF_ORDER: tuple[str, ...] = ( + "project_name", + "use_dq_score_weights", + "observability_api_url", + "observability_api_key", + "data_retention_enabled", + "data_retention_days", + "retention_cron_expr", + "retention_cron_tz", +) + +_DIFF_LABELS: dict[str, str] = { + "project_name": "Name", + "use_dq_score_weights": "Weighted quality scoring", + "observability_api_url": "DataOps Observability API URL", + "observability_api_key": "DataOps Observability API key", + "data_retention_enabled": "Data retention", + "data_retention_days": "Retention days", + "retention_cron_expr": "Retention cron expression", + "retention_cron_tz": "Retention timezone", +} + +_SECRET_ATTRS: frozenset[str] = frozenset({"observability_api_key"}) diff --git a/testgen/mcp/tools/quality_scores.py b/testgen/mcp/tools/quality_scores.py index 1545d7f4..692b9995 100644 --- a/testgen/mcp/tools/quality_scores.py +++ b/testgen/mcp/tools/quality_scores.py @@ -87,7 +87,7 @@ def get_quality_scores( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. ``"Impact Dimension"`` and ``"Quality Dimension"`` are valid as ``group_by`` only, not as filter @@ -106,7 +106,7 @@ def get_quality_scores( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, - ``"Data Product"``. + ``"Data Product"``, ``"Data Classification"``. score_type: Narrow returned scores. Omit to show all four (Total, CDE, Profiling, Testing); pass ``"Total"`` for Total + Profiling + Testing, or ``"CDE"`` for CDE alone. @@ -551,7 +551,7 @@ def create_scorecard( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. @@ -563,7 +563,8 @@ def create_scorecard( ``"Quality Dimension"``, ``"Impact Dimension"``, ``"Data Source"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Table Group"``, ``"Transform Level"``, ``"Data Location"``, - ``"Source System"``, ``"Source Process"``, ``"Data Product"``. + ``"Source System"``, ``"Source Process"``, ``"Data Product"``, + ``"Data Classification"``. show_total_score: Whether the scorecard exposes the Total Score. show_cde_score: Whether the scorecard exposes the CDE Score. """ @@ -624,7 +625,7 @@ def update_scorecard( ``"Table Group"``, ``"Data Location"``, ``"Data Source"``, ``"Source System"``, ``"Source Process"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Transform Level"``, ``"Semantic Data Type"``, - ``"Data Product"``. To target specific tables or columns, chain a + ``"Data Product"``, ``"Data Classification"``. To target specific tables or columns, chain a ``"Table Group"`` filter via ``others`` into ``"Table"`` (optionally then ``"Column"``); sibling chains OR. @@ -638,7 +639,8 @@ def update_scorecard( ``"Quality Dimension"``, ``"Impact Dimension"``, ``"Data Source"``, ``"Business Domain"``, ``"Stakeholder Group"``, ``"Table Group"``, ``"Transform Level"``, ``"Data Location"``, - ``"Source System"``, ``"Source Process"``, ``"Data Product"``. + ``"Source System"``, ``"Source Process"``, ``"Data Product"``, + ``"Data Classification"``. Pass ``""`` to clear an existing category. filters: List of filter entries. See **Filters** above for shape. """ diff --git a/testgen/mcp/tools/reference.py b/testgen/mcp/tools/reference.py index 210e4b89..94ada81c 100644 --- a/testgen/mcp/tools/reference.py +++ b/testgen/mcp/tools/reference.py @@ -1,7 +1,16 @@ from testgen.common.models import with_database_session from testgen.common.models.hygiene_issue import HygieneIssueType from testgen.common.models.test_definition import TestType -from testgen.mcp.tools.common import DocGroup +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.tools.common import ( + FLAVOR_CONNECTION_SCHEMA, + ConnField, + DocGroup, + FlavorMode, + FlavorSchema, + Req, + schema_for, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.DISCOVER @@ -326,3 +335,126 @@ def glossary_resource() -> str: - **referential** — Tests relationships between tables (e.g., foreign key match). - **custom** — User-defined SQL tests. """ + + +# Doc-only hints for the connection-parameters resource: value type / format +# guidance that helps the model but isn't part of the validation/mapping schema. +_FIELD_NOTES: dict[str, str] = { + "Port": "Integer.", + "Service Account Key": "Service-account key JSON, passed as a parsed object.", + "Private Key": "Private key as PEM text.", + "Private Key Passphrase": "Passphrase for the encrypted private key. Omit if the key is unencrypted.", + "Login URL": "My Domain URL of the Salesforce org.", + "Consumer Key": "Consumer key from the Salesforce external client app.", + "Consumer Secret": "Consumer secret from the Salesforce external client app.", + "Access Token": "Databricks personal access token.", + "Client ID": "Service-principal client ID.", + "Client Secret": "Service-principal OAuth secret.", + "HTTP Path": "Databricks SQL warehouse HTTP path.", + "Catalog": "Databricks catalog (the database/namespace).", + "Service Name": "Oracle service name.", + "Warehouse": "Snowflake warehouse name.", + "URL": "Full connection URL. Provide instead of the host fields.", +} + +# Doc-only: the conventional default port per flavor, so the model can fill the +# Port when the user doesn't specify one. +_DEFAULT_PORTS: dict[str, str] = { + "redshift": "5439", + "redshift_spectrum": "5439", + "azure_mssql": "1433", + "synapse_mssql": "1433", + "mssql": "1433", + "postgresql": "5432", + "snowflake": "443", + "databricks": "443", + "oracle": "1521", + "sap_hana": "39015", +} + + +def _requirement_label(field: ConnField) -> str: + if field.requirement is Req.REQUIRED: + return "Required" + if field.requirement is Req.REQUIRED_UNLESS_URL: + return "Required (host mode)" + return "Optional" + + +def _field_note(field: ConnField, schema: FlavorSchema) -> str | None: + parts = [] + if field.secret: + parts.append("Secret — encrypted at rest, never echoed back.") + if note := _FIELD_NOTES.get(field.label): + parts.append(note) + if field.column == "project_port" and (port := _DEFAULT_PORTS.get(schema.code)): + parts.append(f"Default for {schema.label} is {port} — use it unless the user specifies otherwise.") + return " ".join(parts) or None + + +def _append_mode(doc: MdDoc, mode: FlavorMode, schema: FlavorSchema, *, url_offered: bool) -> None: + if mode.mode is not None: + doc.heading(2, f"Mode: {mode.mode}") + doc.text(f'Pass `connection_mode="{mode.mode}"`.') + doc.table( + headers=["Field", "Required", "Notes"], + rows=[[field.label, _requirement_label(field), _field_note(field, schema)] for field in mode.fields], + code=[0], + ) + if mode.supports_url and url_offered: + doc.text( + "Alternatively, provide `URL` instead of the host fields " + "(the fields marked _Required (host mode)_) to connect by URL." + ) + + +def connection_parameters_resource(flavor: str) -> str: + """Per-flavor connection parameter shapes: the auth modes and the exact + ``connection_params`` keys (with required/optional + secret notes) for the flavor. + + Args: + flavor: Flavor code, e.g. ``snowflake``, ``azure_mssql``, ``salesforce_data360``. + """ + if flavor not in FLAVOR_CONNECTION_SCHEMA: + valid = ", ".join(sorted(FLAVOR_CONNECTION_SCHEMA)) + raise MCPUserError(f"Unknown flavor `{flavor}`. Valid flavor codes: {valid}.") + + schema = schema_for(flavor) + doc = MdDoc() + doc.heading(1, f"{schema.label} Connection Parameters") + doc.text( + f'Create with `sql_flavor="{schema.label}"` and a `connection_params` dict keyed by the ' + "field labels below. Secrets are encrypted at rest and never returned." + ) + + multi_mode = len([m for m in schema.modes if m.mode is not None]) > 1 + if multi_mode: + modes = ", ".join(f"`{m.mode}`" for m in schema.modes if m.mode is not None) + doc.text(f"Set `connection_mode` to one of: {modes}.") + + for mode in schema.modes: + _append_mode(doc, mode, schema, url_offered=schema.url_field is not None) + + return doc.render() + + +def connection_parameters_index_resource() -> str: + """Supported database flavors for the connection tools: the accepted + ``sql_flavor`` values and, for each, the resource that documents its + connection modes and fields. + """ + doc = MdDoc() + doc.heading(1, "Connection Flavors") + doc.text( + "Accepted `sql_flavor` values. Read the per-flavor resource for a flavor's " + "connection modes and the fields each needs." + ) + doc.table( + headers=["sql_flavor", "Parameters resource"], + rows=[ + [schema.label, f"testgen://connection-parameters/{code}"] + for code, schema in FLAVOR_CONNECTION_SCHEMA.items() + ], + code=[1], + ) + return doc.render() diff --git a/testgen/mcp/tools/schedules.py b/testgen/mcp/tools/schedules.py index 9c1ed1d0..d1db736e 100644 --- a/testgen/mcp/tools/schedules.py +++ b/testgen/mcp/tools/schedules.py @@ -6,12 +6,11 @@ from sqlalchemy import select from testgen.common.cron_service import describe_cron, get_cron_sample -from testgen.common.enums import JobKey +from testgen.common.enums import JOB_STATUS_LABEL, JobKey from testgen.common.models import get_current_session, with_database_session from testgen.common.models.job_execution import JobExecution from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup -from testgen.common.models.test_run import TestRunSummary # STATUS_LABEL is shared with ProfilingRunSummary from testgen.common.models.test_suite import TestSuite from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission @@ -206,6 +205,7 @@ def create_profiling_schedule( active=active, ) sched.save() + get_current_session().flush() # populate the default-generated id before rendering doc = MdDoc() doc.heading(1, f"Profiling schedule created for `{table_group.table_groups_name}`") @@ -240,6 +240,7 @@ def create_test_run_schedule( active=active, ) sched.save() + get_current_session().flush() # populate the default-generated id before rendering doc = MdDoc() doc.heading(1, f"Test run schedule created for `{suite.test_suite}`") @@ -421,13 +422,13 @@ def get_schedule(schedule_id: str) -> str: for je in history: rows.append([ je.id, - TestRunSummary.STATUS_LABEL.get(je.status, je.status), + JOB_STATUS_LABEL.get(je.status, je.status), je.started_at, je.completed_at, format_run_duration(je.started_at, je.completed_at), ]) doc.table( - ["Job ID", "Status", "Started", "Completed", "Duration"], + [_kind_display(sched.key), "Status", "Started", "Completed", "Duration"], rows, code=[0], ) diff --git a/testgen/mcp/tools/source_data.py b/testgen/mcp/tools/source_data.py index 1b3fb0bb..91321860 100644 --- a/testgen/mcp/tools/source_data.py +++ b/testgen/mcp/tools/source_data.py @@ -1,21 +1,43 @@ from datetime import datetime +from testgen.common.data_catalog_service import fetch_table_sample from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.profiling_run import ProfilingRun from testgen.common.models.test_definition import TestDefinition from testgen.common.source_data_service import ( SourceDataResult, + build_hygiene_query, build_test_result_query, + fetch_hygiene_source_data, fetch_test_result_source_data, ) from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission -from testgen.mcp.tools.common import DocGroup, parse_uuid, validate_limit +from testgen.mcp.tools.common import ( + DocGroup, + parse_uuid, + resolve_hygiene_issue, + resolve_table_group, + validate_limit, +) from testgen.mcp.tools.markdown import MdDoc _DOC_GROUP = DocGroup.INVESTIGATE -def _resolve_context(test_definition_id: str, reference_date: str | None) -> dict: +def _validate_source_args(test_definition_id: str | None, issue_id: str | None, reference_date: str | None) -> None: + """Enforce 'exactly one entity' and the reference_date-only-with-test-definition rule.""" + if bool(test_definition_id) == bool(issue_id): + raise MCPUserError("Provide exactly one of test_definition_id or issue_id.") + if issue_id and reference_date: + raise MCPUserError( + "reference_date applies only to test_definition_id; omit it when looking up a hygiene issue." + ) + + +def _resolve_test_definition_context(test_definition_id: str, reference_date: str | None) -> dict: """Look up the test definition context and validate permissions.""" td_uuid = parse_uuid(test_definition_id, "test_definition_id") perms = get_project_permissions() @@ -40,40 +62,86 @@ def _resolve_context(test_definition_id: str, reference_date: str | None) -> dic return context +def _resolve_hygiene_context(issue_id: str) -> dict: + """Resolve a hygiene issue (permission-scoped) into the lookup context the service expects. + + The source profiling run is intrinsic to the issue, so ``profiling_starttime`` comes from the + issue's ``ProfilingRun`` — there is no caller-supplied reference date. + """ + issue = resolve_hygiene_issue(issue_id) + run = ProfilingRun.get(issue.profile_run_id) + return { + "table_groups_id": issue.table_groups_id, + "anomaly_id": issue.type_id, + "detail": issue.detail, + "schema_name": issue.schema_name, + "table_name": issue.table_name, + "column_name": issue.column_name, + "profiling_starttime": run.profiling_starttime if run else None, + "project_code": issue.project_code, + } + + +def _render_header_fields(doc: MdDoc, context: dict) -> None: + """Render the entity-neutral location fields shared by both tools.""" + if context.get("test_type"): + doc.field("Test type", context.get("test_type"), code=True) + doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True) + column = context.get("column_names") or context.get("column_name") + if column: + doc.field("Column", column, code=True) + + @with_database_session @mcp_permission("view") def get_source_data_query( - test_definition_id: str, + test_definition_id: str | None = None, + issue_id: str | None = None, reference_date: str | None = None, limit: int = 100, ) -> str: - """Get the SQL query that would be used to look up source data for a test definition, without executing it. + """Get the SQL query that would be used to look up source data, without executing it. - Builds a lookup query using current test definition parameters (thresholds, conditions). + Builds a lookup query using the current criteria of a test definition or a hygiene issue. The query targets the connected database. Some test types (e.g. Freshness Trend, Schema Drift) do not have source data lookups. + Provide exactly one of ``test_definition_id`` or ``issue_id``. + Args: test_definition_id: UUID of a test definition, e.g. from ``list_test_results``. - reference_date: ISO 8601 date used as the test reference point (default: now). + issue_id: UUID of a hygiene issue, e.g. from ``list_hygiene_issues``. Mutually exclusive + with ``test_definition_id``. + reference_date: ISO 8601 date used as the test reference point (default: now). Applies only + to ``test_definition_id``. limit: Maximum rows the query would return (default 100, max 500). """ + _validate_source_args(test_definition_id, issue_id, reference_date) validate_limit(limit, 500) - context = _resolve_context(test_definition_id, reference_date) - query = build_test_result_query(context, limit) + if test_definition_id: + context = _resolve_test_definition_context(test_definition_id, reference_date) + entity_label, entity_id = "Test Definition", test_definition_id + query = build_test_result_query(context, limit) + else: + context = _resolve_hygiene_context(issue_id) + entity_label, entity_id = "Hygiene Issue", issue_id + query = build_hygiene_query(context, limit) + if not query: + if test_definition_id: + return ( + f"Source data lookup is not available for test type `{context.get('test_type', 'unknown')}`.\n\n" + "This test type does not have a defined lookup query." + ) return ( - f"Source data lookup is not available for test type `{context.get('test_type', 'unknown')}`.\n\n" - "This test type does not have a defined lookup query." + "Source data lookup is not available for this hygiene issue.\n\n" + "This hygiene issue type does not have a defined lookup query." ) doc = MdDoc() - doc.heading(1, f"Source Data Query for Test Definition `{test_definition_id}`") - doc.field("Test type", context.get("test_type"), code=True) - doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True) - if context.get("column_names"): - doc.field("Column", context["column_names"], code=True) + doc.heading(1, f"Source Data Query for {entity_label} `{entity_id}`") + _render_header_fields(doc, context) doc.field("Limit", limit) doc.code_block(query, language="sql") @@ -83,34 +151,46 @@ def get_source_data_query( @with_database_session @mcp_permission("view") def get_source_data( - test_definition_id: str, + test_definition_id: str | None = None, + issue_id: str | None = None, reference_date: str | None = None, limit: int = 100, ) -> str: - """Look up rows from the connected database that match or violate a test definition's criteria. + """Look up rows from the connected database that match or violate a test or hygiene issue's criteria. Executes the source data query against the connected database and returns matching rows. - Shows CURRENT data — rows may have changed since the test last ran. + Shows CURRENT data — rows may have changed since the test or profiling run. Some test types (e.g. Freshness Trend, Schema Drift) do not have source data lookups. + Provide exactly one of ``test_definition_id`` or ``issue_id``. + Args: test_definition_id: UUID of a test definition, e.g. from ``list_test_results``. - reference_date: ISO 8601 date used as the test reference point (default: now). + issue_id: UUID of a hygiene issue, e.g. from ``list_hygiene_issues``. Mutually exclusive + with ``test_definition_id``. + reference_date: ISO 8601 date used as the test reference point (default: now). Applies only + to ``test_definition_id``. limit: Maximum rows to return (default 100, max 500). """ + _validate_source_args(test_definition_id, issue_id, reference_date) validate_limit(limit, 500) - context = _resolve_context(test_definition_id, reference_date) + + if test_definition_id: + context = _resolve_test_definition_context(test_definition_id, reference_date) + entity_label, entity_id = "Test Definition", test_definition_id + fetch = fetch_test_result_source_data + else: + context = _resolve_hygiene_context(issue_id) + entity_label, entity_id = "Hygiene Issue", issue_id + fetch = fetch_hygiene_source_data mask_pii = not get_project_permissions().has_permission("view_pii", context.get("project_code")) - result: SourceDataResult = fetch_test_result_source_data(context, limit, mask_pii) + result: SourceDataResult = fetch(context, limit, mask_pii) doc = MdDoc() - doc.heading(1, f"Source Data for Test Definition `{test_definition_id}`") - doc.field("Test type", context.get("test_type"), code=True) - doc.field("Table", f"{context.get('schema_name')}.{context.get('table_name')}", code=True) - if context.get("column_names"): - doc.field("Column", context["column_names"], code=True) + doc.heading(1, f"Source Data for {entity_label} `{entity_id}`") + _render_header_fields(doc, context) if result.status == "OK": row_count = len(result.df) if result.df is not None else 0 @@ -135,3 +215,46 @@ def get_source_data( doc.code_block(result.query, language="sql") return doc.render() + + +@with_database_session +@mcp_permission("catalog") +def get_table_sample(table_group_id: str, table_name: str, limit: int = 100) -> str: + """Fetch sample rows from a source table for inspection. + + Args: + table_group_id: UUID of the table group, e.g. from `get_data_inventory`. + table_name: Table name exactly as stored in TestGen (case-sensitive). + limit: Maximum rows to return (default 100, max 500). + """ + validate_limit(limit, 500) + tg = resolve_table_group(table_group_id) + + schema_name, _ = DataColumnChars.list_for_create_script(tg.id, table_name) + if schema_name is None: + raise MCPResourceNotAccessible("Table", table_name) + + connection = Connection.get_by_table_group(tg.id) + if connection is None: + raise MCPResourceNotAccessible("Table", table_name) + + mask_pii = not get_project_permissions().has_permission("view_pii", tg.project_code) + result = fetch_table_sample( + connection, tg.id, schema_name, table_name, limit=limit, mask_pii=mask_pii, + ) + + if result.status == "ERR": + raise MCPUserError( + f"Could not read from the source database for connection `{connection.connection_name}`." + ) + + doc = MdDoc() + if result.status == "ND": + return doc.text("Table has no rows.").render() + + row_count = len(result.df) if result.df is not None else 0 + doc.field("Rows returned", row_count) + if result.pii_redacted: + doc.text("_PII columns have been redacted._") + doc.table_from_dataframe(result.df) + return doc.render() diff --git a/testgen/mcp/tools/table_groups.py b/testgen/mcp/tools/table_groups.py new file mode 100644 index 00000000..9f4fa259 --- /dev/null +++ b/testgen/mcp/tools/table_groups.py @@ -0,0 +1,778 @@ +"""MCP tools for table groups — create, update, append tables, preview. + +Each tool gates on the ``edit`` permission. Validation and target-DB +introspection are delegated to ``testgen.common.database.table_group_service`` +so the rules and SQL paths stay in one place. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy.exc import IntegrityError + +from testgen.common.database.table_group_service import ( + preview_table_group as preview_table_group_service, +) +from testgen.common.database.table_group_service import validate_table_group_fields +from testgen.common.models import with_database_session +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import get_project_permissions, mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + format_flavor_label, + format_page_footer, + format_page_info, + render_diff_table, + resolve_connection, + resolve_table_group, + validate_limit, + validate_page, +) +from testgen.mcp.tools.markdown import MdDoc +from testgen.utils import friendly_score + +_DOC_GROUP = DocGroup.MANAGE + +_DUPLICATE_NAME_MESSAGE = "A Table Group with the same name already exists." +_PII_FLAG_DENIED_MESSAGE = ( + "Changing PII detection requires permission to view PII. " + "Leave this setting unchanged or contact your administrator." +) +_SCHEMA_LOCKED_MESSAGE = ( + "Schema cannot be changed once the table group has been used. " + "Delete and recreate the table group to use a different schema." +) + + +@with_database_session +@mcp_permission("view") +def list_table_groups( + project_code: str | None = None, + connection_id: int | None = None, + page: int = 1, + limit: int = 20, +) -> str: + """List table groups in a project or on a specific connection. + + Pass exactly one of `project_code` or `connection_id`. Returns each group's + table count, last profile / test timestamps, and current quality score. + + Args: + project_code: List groups in a project. + connection_id: List groups on a specific connection. + page: Page number starting at 1 (default 1). + limit: Page size (default 20, max 100). + """ + if (project_code is None) == (connection_id is None): + raise MCPUserError("Pass either `project_code` or `connection_id`, not both.") + validate_page(page) + validate_limit(limit, 100) + + perms = get_project_permissions() + if project_code is not None: + perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) + rows, total = TableGroup.list_for_project(project_code, page=page, limit=limit) + heading = f"Table groups for project `{project_code}`" + else: + connection = resolve_connection(connection_id) + rows, total = TableGroup.list_for_connection(connection.connection_id, page=page, limit=limit) + heading = f"Table groups on connection `{connection.connection_name}` (`{connection.connection_id}`)" + + if not rows: + if page > 1: + return f"No table groups on page {page} (total: {total})." + return f"{heading} — none found." + + doc = MdDoc() + doc.heading(1, heading) + doc.text(format_page_info(total, page, limit)) + table_rows: list[list[object]] = [] + for row in rows: + table_rows.append( + [ + str(row.id), + row.table_groups_name, + row.connection_name, + row.table_group_schema, + row.table_count, + row.column_count, + row.row_count, + row.last_profiled_date, + row.last_tested_date, + friendly_score(row.quality_score), + ] + ) + doc.table( + ["ID", "Name", "Connection", "Schema", "Tables", "Columns", "Rows", "Last profiled", "Last tested", "Quality Score"], + table_rows, + code=[0, 3], + ) + if footer := format_page_footer(total, page, limit): + doc.text(footer) + return doc.render() + + +@with_database_session +@mcp_permission("view") +def get_table_group(table_group_id: str) -> str: + """Get a table group's full configuration: filters, sampling, profiling flags, catalog tags, and recent activity. + + Use this before editing a table group or generating tests. + + Args: + table_group_id: The table group UUID, e.g. from `list_table_groups` or `get_data_inventory`. + """ + table_group = resolve_table_group(table_group_id) + # Defense in depth: route through resolve_connection (perm-scoped) rather than Connection.get. + connection = resolve_connection(table_group.connection_id) if table_group.connection_id else None + + doc = MdDoc() + doc.heading(1, f"Table group `{table_group.table_groups_name}`") + doc.field("ID", str(table_group.id), code=True) + doc.field("Project", table_group.project_code, code=True) + if connection is not None: + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`, {format_flavor_label(connection.sql_flavor_code)})", + ) + doc.field("Schema", table_group.table_group_schema, code=True) + if table_group.description: + doc.field("Description", table_group.description) + + doc.heading(2, "Criteria") + if table_group.profiling_table_set: + doc.field(_DIFF_LABELS["profiling_table_set"], table_group.profiling_table_set, code=True) + if table_group.profiling_include_mask: + doc.field(_DIFF_LABELS["profiling_include_mask"], table_group.profiling_include_mask, code=True) + if table_group.profiling_exclude_mask: + doc.field(_DIFF_LABELS["profiling_exclude_mask"], table_group.profiling_exclude_mask, code=True) + doc.field(_DIFF_LABELS["profile_id_column_mask"], table_group.profile_id_column_mask, code=True) + doc.field(_DIFF_LABELS["profile_sk_column_mask"], table_group.profile_sk_column_mask, code=True) + + doc.heading(2, "Settings") + doc.field(_DIFF_LABELS["profile_flag_cdes"], table_group.profile_flag_cdes) + doc.field(_DIFF_LABELS["profile_flag_pii"], table_group.profile_flag_pii) + doc.field(_DIFF_LABELS["profile_exclude_xde"], table_group.profile_exclude_xde) + doc.field(_DIFF_LABELS["include_in_dashboard"], table_group.include_in_dashboard) + doc.field(_DIFF_LABELS["profiling_delay_days"], table_group.profiling_delay_days) + + doc.heading(2, "Sampling parameters") + doc.field(_DIFF_LABELS["profile_use_sampling"], table_group.profile_use_sampling) + if table_group.profile_use_sampling: + doc.field(_DIFF_LABELS["profile_sample_percent"], table_group.profile_sample_percent) + doc.field(_DIFF_LABELS["profile_sample_min_count"], table_group.profile_sample_min_count) + + if any(getattr(table_group, attr, None) for attr in _CATALOG_ATTRS): + doc.heading(2, "Catalog tags") + for attr in _CATALOG_ATTRS: + value = getattr(table_group, attr, None) + if value: + doc.field(_DIFF_LABELS[attr], value) + + if table_group.dq_score_testing is not None or table_group.dq_score_profiling is not None: + doc.heading(2, "Latest activity") + if (profiling := friendly_score(table_group.dq_score_profiling)) is not None: + doc.field("Profiling Score", profiling) + if (testing := friendly_score(table_group.dq_score_testing)) is not None: + doc.field("Testing Score", testing) + if (quality := friendly_score(table_group.quality_score)) is not None: + doc.field("Quality Score", quality) + + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def create_table_group( + connection_id: int, + table_group_name: str, + schema: str, + *, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + add_scorecard: bool = True, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, + data_classification: str | None = None, +) -> str: + """Create a table group on an existing connection. + + The table group inherits its project from the connection. ``include_mask`` + and ``exclude_mask`` are SQL ``LIKE`` patterns (e.g. ``fact_%,dim_%``); + ``table_set`` is an explicit list. All filters compose with ``AND``. + + Args: + connection_id: Bigint connection ID, e.g. from ``get_data_inventory``. + table_group_name: 3-40 character display name. Must be unique within the project. + schema: Schema name on the target database, e.g. ``public``. For Salesforce Data 360 + connections, use the data space name. + description: Optional free-text description. + table_set: Explicit list of table names. Combined with masks if also set. + include_mask: Comma-separated SQL LIKE patterns to include (e.g. ``fact_%,dim_%``). + exclude_mask: Comma-separated SQL LIKE patterns to exclude. + profile_id_column_mask: SQL LIKE pattern marking ID columns. Default ``%id``. + profile_sk_column_mask: SQL LIKE pattern marking surrogate-key columns. Default ``%_sk``. + profile_use_sampling: Whether to sample large tables during profiling. + profile_sample_percent: Sample size as a percent (1-100). + profile_sample_min_count: Minimum row count when sampling. + profiling_delay_days: Number of days to wait before new profiling will be available + to generate tests. + profile_flag_cdes: Whether profiling flags Critical Data Elements. + profile_flag_pii: Whether profiling flags Personally Identifiable Information. + profile_exclude_xde: Whether profiling excludes columns flagged as excluded data elements. + include_in_dashboard: Whether the table group appears on the project dashboard. + add_scorecard: Whether to add a scorecard for the table group to the Quality Dashboard. + data_source: Catalog tag — original source of the dataset. + source_system: Catalog tag — enterprise system source for the dataset. + source_process: Catalog tag — process, program, or data flow that produced the dataset. + data_location: Catalog tag — physical or virtual location of the dataset + (e.g. ``Headquarters``, ``Cloud``). + business_domain: Catalog tag — business division responsible for the dataset + (e.g. ``Finance``, ``Sales``, ``Manufacturing``). + stakeholder_group: Catalog tag — data owners or stakeholders responsible for the dataset. + transform_level: Catalog tag — data warehouse processing stage (e.g. ``Raw``, + ``Conformed``, ``Processed``, ``Reporting``) or Medallion level + (``bronze``, ``silver``, ``gold``). + data_product: Catalog tag — data domain that comprises the dataset. + data_classification: Catalog tag — information classification level of the dataset + (e.g. ``Public``, ``Internal``, ``Confidential``, ``Restricted``). + """ + connection = resolve_connection(connection_id) + + table_group = TableGroup( + project_code=connection.project_code, + connection_id=connection.connection_id, + table_groups_name=table_group_name, + table_group_schema=schema, + **_model_defaults(), + ) + _apply_args_to_table_group( + table_group, + description=description, + table_set=table_set, + include_mask=include_mask, + exclude_mask=exclude_mask, + profile_id_column_mask=profile_id_column_mask, + profile_sk_column_mask=profile_sk_column_mask, + profile_use_sampling=profile_use_sampling, + profile_sample_percent=profile_sample_percent, + profile_sample_min_count=profile_sample_min_count, + profiling_delay_days=profiling_delay_days, + profile_flag_cdes=profile_flag_cdes, + profile_flag_pii=profile_flag_pii, + profile_exclude_xde=profile_exclude_xde, + include_in_dashboard=include_in_dashboard, + data_source=data_source, + source_system=source_system, + source_process=source_process, + data_location=data_location, + business_domain=business_domain, + stakeholder_group=stakeholder_group, + transform_level=transform_level, + data_product=data_product, + data_classification=data_classification, + ) + + errors = validate_table_group_fields(table_group) + if errors: + _raise_validation_error(errors, "Table group creation rejected. No changes saved.") + + try: + table_group.save(add_scorecard_definition=add_scorecard) + except IntegrityError as err: + _maybe_raise_duplicate_name(err) + raise + + return _render_created_table_group(table_group, connection) + + +@with_database_session +@mcp_permission("edit") +def update_table_group( + table_group_id: str, + *, + table_group_name: str | None = None, + schema: str | None = None, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, + data_classification: str | None = None, +) -> str: + """Update fields on an existing table group. Atomic — no partial save. + + Connection and project are immutable — delete and recreate the table group + to re-parent it. ``schema`` is also immutable once the table group has been + used (profiled or has test suites); supply a different ``schema`` only on + unused table groups. + + Args: + table_group_id: UUID of the table group to update. + table_group_name: New display name (3-40 chars). Must be unique within the project. + schema: New target-DB schema. Rejected if the table group has been used. + description: Free-text description. + table_set: Replacement explicit table list (full replacement of the current list). + include_mask: Comma-separated SQL LIKE patterns to include. + exclude_mask: Comma-separated SQL LIKE patterns to exclude. + profile_id_column_mask: SQL LIKE pattern marking ID columns. + profile_sk_column_mask: SQL LIKE pattern marking surrogate-key columns. + profile_use_sampling: Whether to sample large tables during profiling. + profile_sample_percent: Sample size as a percent (1-100). + profile_sample_min_count: Minimum row count when sampling. + profiling_delay_days: Number of days to wait before new profiling will be available + to generate tests. + profile_flag_cdes: Whether profiling flags CDEs. + profile_flag_pii: Whether profiling flags PII. + profile_exclude_xde: Whether profiling excludes XDE columns. + include_in_dashboard: Whether the table group appears on the project dashboard. + data_source: Catalog tag — original source of the dataset. + source_system: Catalog tag — enterprise system source for the dataset. + source_process: Catalog tag — process, program, or data flow that produced the dataset. + data_location: Catalog tag — physical or virtual location of the dataset + (e.g. ``Headquarters``, ``Cloud``). + business_domain: Catalog tag — business division responsible for the dataset + (e.g. ``Finance``, ``Sales``, ``Manufacturing``). + stakeholder_group: Catalog tag — data owners or stakeholders responsible for the dataset. + transform_level: Catalog tag — data warehouse processing stage (e.g. ``Raw``, + ``Conformed``, ``Processed``, ``Reporting``) or Medallion level + (``bronze``, ``silver``, ``gold``). + data_product: Catalog tag — data domain that comprises the dataset. + data_classification: Catalog tag — information classification level of the dataset + (e.g. ``Public``, ``Internal``, ``Confidential``, ``Restricted``). + """ + supplied = { + "table_group_name": table_group_name, + "schema": schema, + "description": description, + "table_set": table_set, + "include_mask": include_mask, + "exclude_mask": exclude_mask, + "profile_id_column_mask": profile_id_column_mask, + "profile_sk_column_mask": profile_sk_column_mask, + "profile_use_sampling": profile_use_sampling, + "profile_sample_percent": profile_sample_percent, + "profile_sample_min_count": profile_sample_min_count, + "profiling_delay_days": profiling_delay_days, + "profile_flag_cdes": profile_flag_cdes, + "profile_flag_pii": profile_flag_pii, + "profile_exclude_xde": profile_exclude_xde, + "include_in_dashboard": include_in_dashboard, + "data_source": data_source, + "source_system": source_system, + "source_process": source_process, + "data_location": data_location, + "business_domain": business_domain, + "stakeholder_group": stakeholder_group, + "transform_level": transform_level, + "data_product": data_product, + "data_classification": data_classification, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + table_group = resolve_table_group(table_group_id) + + if ( + schema is not None + and schema != table_group.table_group_schema + and TableGroup.is_in_use([table_group.id]) + ): + raise MCPUserError(_SCHEMA_LOCKED_MESSAGE) + + if ( + profile_flag_pii is not None + and profile_flag_pii != table_group.profile_flag_pii + and not get_project_permissions().has_permission("view_pii", table_group.project_code) + ): + raise MCPPermissionDenied(_PII_FLAG_DENIED_MESSAGE) + + before = _snapshot(table_group) + _apply_args_to_table_group(table_group, **supplied) + + errors = validate_table_group_fields(table_group) + if errors: + _raise_validation_error(errors, "Update rejected. No changes saved.") + + after = _snapshot(table_group) + + doc = MdDoc() + doc.heading(1, f"Table Group `{table_group.table_groups_name}` updated") + doc.field("ID", str(table_group.id), code=True) + + rendered = render_diff_table(doc, before, after, attrs=_DIFF_ATTRS, labels=_DIFF_LABELS) + if not rendered: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + try: + table_group.save() + except IntegrityError as err: + _maybe_raise_duplicate_name(err) + raise + + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def preview_table_group( + table_group_id: str, + verify_access: bool = False, +) -> str: + """Probe the target database for tables matching a table group's filters. + + Returns counts plus a per-table breakdown. Does not save anything to the + application database. + + Args: + table_group_id: UUID of the table group. + verify_access: When True, probe read access on every matched table. + """ + table_group = resolve_table_group(table_group_id) + connection = Connection.get_by_table_group(table_group.id) + if connection is None: + raise MCPUserError("Cannot preview — the table group's connection is unavailable.") + + preview, _data_chars, _sql_generator = preview_table_group_service( + table_group, connection=connection, verify_access=verify_access, + ) + stats = preview["stats"] + name = stats.get("table_groups_name") or table_group.table_groups_name + + doc = MdDoc() + + if not preview["success"]: + if preview.get("message", "").startswith("No tables found matching the criteria"): + doc.heading(1, f"Preview for table group `{name}` returned no tables") + doc.text(preview["message"]) + else: + doc.heading(1, f"Preview failed for table group `{name}`") + doc.text(preview.get("message") or "Preview failed for an unknown reason.") + return doc.render() + + doc.heading(1, f"Preview for table group `{name}`") + doc.field("Table Group ID", str(table_group.id), code=True) + doc.field("Schema", stats.get("table_group_schema"), code=True) + doc.field("Tables matched", stats.get("table_ct") or 0) + doc.field("Total columns", stats.get("column_ct") or 0) + if stats.get("approx_record_ct") is not None: + doc.field("Approx rows", stats.get("approx_record_ct")) + if stats.get("approx_data_point_ct") is not None: + doc.field("Approx data points", stats.get("approx_data_point_ct")) + + headers = ["Table", "Columns", "Approx Rows", "Approx Data Points"] + if verify_access: + headers.append("Read Access") + + rows: list[list[object]] = [] + for table_name, info in preview["tables"].items(): + row: list[object] = [ + table_name, + info.get("column_ct"), + info.get("approx_record_ct"), + info.get("approx_data_point_ct"), + ] + if verify_access: + access = info.get("can_access") + row.append("Yes" if access is True else "No" if access is False else "Unknown") + rows.append(row) + doc.table(headers, rows, code=[0]) + + if verify_access and preview.get("message"): + doc.text(preview["message"]) + + return doc.render() + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +_DIFF_ATTRS: tuple[str, ...] = ( + "table_groups_name", + "table_group_schema", + "profiling_table_set", + "profiling_include_mask", + "profiling_exclude_mask", + "profile_id_column_mask", + "profile_sk_column_mask", + "profile_use_sampling", + "profile_sample_percent", + "profile_sample_min_count", + "profiling_delay_days", + "profile_flag_cdes", + "profile_flag_pii", + "profile_exclude_xde", + "include_in_dashboard", + "description", + "data_source", + "source_system", + "source_process", + "data_location", + "business_domain", + "stakeholder_group", + "transform_level", + "data_product", + "data_classification", +) + +_DIFF_LABELS: dict[str, str] = { + "table_groups_name": "Name", + "table_group_schema": "Schema", + "profiling_table_set": "Table set", + "profiling_include_mask": "Include mask", + "profiling_exclude_mask": "Exclude mask", + "profile_id_column_mask": "ID column mask", + "profile_sk_column_mask": "SK column mask", + "profile_use_sampling": "Sampling", + "profile_sample_percent": "Sample %", + "profile_sample_min_count": "Sample min rows", + "profiling_delay_days": "Min profiling age (days)", + "profile_flag_cdes": "Flag CDEs", + "profile_flag_pii": "Flag PII", + "profile_exclude_xde": "Exclude XDE", + "include_in_dashboard": "Include in dashboard", + "description": "Description", + "data_source": "Data source", + "source_system": "Source system", + "source_process": "Source process", + "data_location": "Data location", + "business_domain": "Business domain", + "stakeholder_group": "Stakeholder group", + "transform_level": "Transform level", + "data_product": "Data product", + "data_classification": "Data classification", +} + +_CATALOG_ATTRS: tuple[str, ...] = ( + "data_source", + "source_system", + "source_process", + "data_location", + "business_domain", + "stakeholder_group", + "transform_level", + "data_product", + "data_classification", +) + + +# Mirror SQLAlchemy ``Column(default=...)`` values into the constructor kwargs. +# Column defaults only fire at flush time, but ``validate_table_group_fields`` +# runs *before* flush. Without seeding these, every create call that omits +# e.g. ``profile_sample_percent`` would fail validation. Computing once at +# import time keeps the values in sync with the model and survives tests that +# patch the ``TableGroup`` class itself. +# +# ``YNString`` columns store raw "Y"/"N" strings as their default but expose +# the attribute as ``bool``; normalize so the in-memory render path doesn't +# treat "N" as truthy before the row is reloaded. +def _normalize_default(column, raw: Any) -> Any: + from testgen.common.models.custom_types import YNString + + if isinstance(column.type, YNString) and isinstance(raw, str): + return raw == "Y" + return raw + + +_MODEL_DEFAULTS: dict[str, Any] = { + column.name: _normalize_default(column, column.default.arg) + for column in TableGroup.__table__.columns + if ( + column.default is not None + and not column.primary_key + and getattr(column.default, "is_scalar", False) + ) +} + + +def _model_defaults() -> dict[str, Any]: + return dict(_MODEL_DEFAULTS) + + +def _apply_args_to_table_group( + table_group: TableGroup, + *, + table_group_name: str | None = None, + schema: str | None = None, + description: str | None = None, + table_set: list[str] | None = None, + include_mask: str | None = None, + exclude_mask: str | None = None, + profile_id_column_mask: str | None = None, + profile_sk_column_mask: str | None = None, + profile_use_sampling: bool | None = None, + profile_sample_percent: int | None = None, + profile_sample_min_count: int | None = None, + profiling_delay_days: int | None = None, + profile_flag_cdes: bool | None = None, + profile_flag_pii: bool | None = None, + profile_exclude_xde: bool | None = None, + include_in_dashboard: bool | None = None, + data_source: str | None = None, + source_system: str | None = None, + source_process: str | None = None, + data_location: str | None = None, + business_domain: str | None = None, + stakeholder_group: str | None = None, + transform_level: str | None = None, + data_product: str | None = None, + data_classification: str | None = None, +) -> None: + """Apply every non-None arg to its model field. + + Casts ``table_set: list[str]`` to comma-joined string, and casts + integer args for columns the model stores as strings. + """ + if table_group_name is not None: + table_group.table_groups_name = table_group_name + if schema is not None: + table_group.table_group_schema = schema + if description is not None: + table_group.description = description + if table_set is not None: + table_group.profiling_table_set = ",".join(table_set) + if include_mask is not None: + table_group.profiling_include_mask = include_mask + if exclude_mask is not None: + table_group.profiling_exclude_mask = exclude_mask + if profile_id_column_mask is not None: + table_group.profile_id_column_mask = profile_id_column_mask + if profile_sk_column_mask is not None: + table_group.profile_sk_column_mask = profile_sk_column_mask + if profile_use_sampling is not None: + table_group.profile_use_sampling = profile_use_sampling + if profile_sample_percent is not None: + table_group.profile_sample_percent = str(profile_sample_percent) + if profile_sample_min_count is not None: + table_group.profile_sample_min_count = profile_sample_min_count + if profiling_delay_days is not None: + table_group.profiling_delay_days = str(profiling_delay_days) + if profile_flag_cdes is not None: + table_group.profile_flag_cdes = profile_flag_cdes + if profile_flag_pii is not None: + table_group.profile_flag_pii = profile_flag_pii + if profile_exclude_xde is not None: + table_group.profile_exclude_xde = profile_exclude_xde + if include_in_dashboard is not None: + table_group.include_in_dashboard = include_in_dashboard + if data_source is not None: + table_group.data_source = data_source + if source_system is not None: + table_group.source_system = source_system + if source_process is not None: + table_group.source_process = source_process + if data_location is not None: + table_group.data_location = data_location + if business_domain is not None: + table_group.business_domain = business_domain + if stakeholder_group is not None: + table_group.stakeholder_group = stakeholder_group + if transform_level is not None: + table_group.transform_level = transform_level + if data_product is not None: + table_group.data_product = data_product + if data_classification is not None: + table_group.data_classification = data_classification + + +def _raise_validation_error(errors: list[str], header: str) -> None: + bullets = "\n".join(f"- {err}" for err in errors) + raise MCPUserError(f"{header}\n\n{bullets}") + + +def _maybe_raise_duplicate_name(err: IntegrityError) -> None: + if "table_groups_name_unique" in str(err.orig): + raise MCPUserError(_DUPLICATE_NAME_MESSAGE) from err + + +def _snapshot(table_group: TableGroup) -> dict[str, Any]: + return {attr: getattr(table_group, attr, None) for attr in _DIFF_ATTRS} + + +def _render_created_table_group(table_group: TableGroup, connection: Connection) -> str: + doc = MdDoc() + doc.heading(1, f"Table Group `{table_group.table_groups_name}` created") + doc.field("ID", str(table_group.id), code=True) + doc.field("Project", table_group.project_code, code=True) + doc.field( + "Connection", + f"{connection.connection_name} (`{connection.connection_id}`)", + ) + doc.field("Schema", table_group.table_group_schema, code=True) + if table_group.description: + doc.field("Description", table_group.description) + + if table_group.profiling_table_set: + doc.field("Table set", table_group.profiling_table_set, code=True) + if table_group.profiling_include_mask: + doc.field("Include mask", table_group.profiling_include_mask, code=True) + if table_group.profiling_exclude_mask: + doc.field("Exclude mask", table_group.profiling_exclude_mask, code=True) + + doc.field("Profile sampling", "Yes" if table_group.profile_use_sampling else "No") + if table_group.profile_use_sampling: + doc.field("Sample %", table_group.profile_sample_percent) + doc.field("Sample min rows", table_group.profile_sample_min_count) + + profile_flags = [] + if table_group.profile_flag_cdes: + profile_flags.append("Flag CDEs") + if table_group.profile_flag_pii: + profile_flags.append("Flag PII") + if table_group.profile_exclude_xde: + profile_flags.append("Exclude XDE") + if profile_flags: + doc.field("Profiling flags", ", ".join(profile_flags)) + + doc.field("Min profiling age (days)", table_group.profiling_delay_days) + doc.field("Include in dashboard", "Yes" if table_group.include_in_dashboard else "No") + + if any(getattr(table_group, attr, None) for attr in _CATALOG_ATTRS): + doc.heading(2, "Catalog") + for attr in _CATALOG_ATTRS: + value = getattr(table_group, attr, None) + if value: + doc.field(_DIFF_LABELS[attr], value) + + return doc.render() diff --git a/testgen/mcp/tools/test_definitions.py b/testgen/mcp/tools/test_definitions.py index 67197f20..c33c1edd 100644 --- a/testgen/mcp/tools/test_definitions.py +++ b/testgen/mcp/tools/test_definitions.py @@ -1,7 +1,9 @@ +import json from datetime import UTC, datetime from enum import StrEnum from typing import NoReturn +from pydantic import ValidationError from sqlalchemy import update from testgen.common.custom_test_validation import validate_custom_query @@ -17,15 +19,32 @@ TestType, ) from testgen.common.models.test_result import TestResult +from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportConfig, + ImportMode, + ImportPayload, + ImportResponse, + ImportStrictViolation, + InvalidImportPayload, + OnAbsence, + OnMatch, + OnNew, + Origin, + export_definitions, + import_definitions, +) from testgen.mcp.exceptions import MCPUserError from testgen.mcp.permissions import get_authorized_mcp_user, get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( DocGroup, format_page_footer, format_page_info, + parse_enum, parse_impact_dimension, parse_quality_dimension, parse_uuid, + raise_validation_error, resolve_test_definition, resolve_test_note, resolve_test_suite, @@ -224,6 +243,14 @@ def _append_td_summary(doc: MdDoc, td: TestDefinitionSummary) -> None: elif td.last_auto_gen_date: doc.field("Last Updated", f"{td.last_auto_gen_date} (auto-generated)") + # URL and metadata + if td.external_url: + doc.field("External URL", td.external_url) + if td.custom_metadata: + doc.heading(2, "Custom Metadata") + for key, value in td.custom_metadata.items(): + doc.field(MdDoc.escape(key), value) + # Parameters (editable fields from test type metadata) _append_parameters_section(doc, td) @@ -488,6 +515,25 @@ def _raise_validation_errors(err: InvalidTestDefinitionFields, header: str) -> N raise MCPUserError(f"{header}\n\n{bullets}") from err +def _coerce_custom_metadata(fields: dict) -> None: + """Accept ``custom_metadata`` as a JSON string for convenience, parsing it to an object in place. + + A blank string becomes ``None`` (clears the field). Non-object JSON is left for model + validation to reject with a consistent message. + """ + value = fields.get("custom_metadata") + if not isinstance(value, str): + return + stripped = value.strip() + if not stripped: + fields["custom_metadata"] = None + return + try: + fields["custom_metadata"] = json.loads(stripped) + except json.JSONDecodeError as e: + raise MCPUserError("`custom_metadata` must be a JSON object of key-value pairs.") from e + + @with_database_session @mcp_permission("edit") def create_test( @@ -529,6 +575,7 @@ def create_test( ) fields = fields or {} + _coerce_custom_metadata(fields) accepted = td.editable_fields(tt) rejected = sorted(set(fields) - accepted) if rejected: @@ -573,6 +620,7 @@ def update_test(test_definition_id: str, fields: dict) -> str: if not fields: raise MCPUserError("No fields supplied to update.") + _coerce_custom_metadata(fields) accepted = td.editable_fields(tt) rejected = sorted(set(fields) - accepted) if rejected: @@ -698,11 +746,7 @@ def bulk_update_tests( table_name: Optional table-name filter. Case-sensitive. test_type: Optional test type name (e.g. ``Alpha Truncation``). """ - try: - bulk_action = BulkAction(action) - except ValueError as err: - valid = ", ".join(f"`{a.value}`" for a in BulkAction) - raise MCPUserError(f"`action` must be one of: {valid}.") from err + bulk_action = parse_enum(action, BulkAction, "action") suite = resolve_test_suite(test_suite_id) tt_code = resolve_test_type(test_type) if test_type else None @@ -742,3 +786,152 @@ def bulk_update_tests( doc.heading(1, f"{verb} {count} test(s) in suite `{suite.test_suite}`") doc.field("Filter", filter_str) return doc.render() + + +# --------------------------------------------------------------------------- +# Export / import +# +# Thin wrappers over ``testgen.common.test_definition_export_import_service`` — the +# same export document and import semantics as the REST API. +# --------------------------------------------------------------------------- + + +def _parse_import_payload(payload: str) -> ImportPayload: + try: + return ImportPayload.model_validate_json(payload) + except ValidationError as err: + errors = err.errors() + bullets = [f"`{'.'.join(str(part) for part in e['loc']) or 'payload'}`: {e['msg']}" for e in errors[:5]] + if len(errors) > 5: + bullets.append(f"…and {len(errors) - 5} more") + raise_validation_error(bullets, "`payload` is not a valid export document.") + + +def _append_import_result(doc: MdDoc, result: ImportResponse, preview: bool) -> None: + verb_suffix = " (projected)" if preview else "" + doc.field(f"Created{verb_suffix}", result.summary.created) + doc.field(f"Updated{verb_suffix}", result.summary.updated) + doc.field(f"Skipped{verb_suffix}", result.summary.skipped) + doc.field(f"Deleted{verb_suffix}", result.summary.deleted) + + rows_by_action: dict[ImportAction, list[list]] = {} + for item in result.items: + for td in item.tds: + rows_by_action.setdefault(item.action, []).append( + [item.reason.value, td.idx, td.target_id] + ) + + for action in ImportAction: + rows = rows_by_action.get(action) + if not rows: + continue + doc.heading(2, f"{action.value.capitalize()} ({len(rows)})") + doc.table(["Reason", "Index", "Test Definition ID"], rows, code=[2]) + + +@with_database_session +@mcp_permission("view") +def export_tests( + test_suite_id: str, + origin: str = "both", + table_name: str | None = None, + test_type: str | None = None, +) -> str: + """Export test definitions from a test suite as a portable JSON document. + + The returned JSON payload can be applied to another suite (or back to the same + suite) with ``import_tests`` — as-is for promotion or copying, or after editing + it for bulk changes such as threshold calibration. + + Args: + test_suite_id: UUID of the test suite to export from. + origin: Which tests to include: ``both`` (default), ``manual`` (only manually + created tests), or ``auto`` (only auto-generated tests). + table_name: Optional table-name filter. Case-sensitive. + test_type: Optional test type name filter, e.g. ``Alpha Truncation``. + """ + origin_enum = parse_enum(origin, Origin, "origin") + suite = resolve_test_suite(test_suite_id) + table_name = table_name or None + test_type_code = resolve_test_type(test_type) if test_type else None + + export = export_definitions(suite, origin_enum, table_name, test_type_code) + + filters = [] + if origin_enum is not Origin.both: + filters.append(f"{origin_enum.value} tests only") + if table_name: + filters.append(f"table `{table_name}`") + if test_type: + filters.append(f"type `{test_type}`") + + doc = MdDoc() + doc.heading(1, f"Test Definition Export from suite `{suite.test_suite}`") + doc.field("Project", export.source.project_code, code=True) + doc.field("Table Group", export.source.table_group) + doc.field("Schema", export.source.table_group_schema, code=True) + doc.field("Exported At", export.source.exported_at) + if filters: + doc.field("Filters", ", ".join(filters)) + doc.field("Tests Exported", len(export.definitions)) + doc.code_block(export.model_dump_json(indent=2, exclude_defaults=True), language="json") + return doc.render() + + +@with_database_session +@mcp_permission("edit") +def import_tests( + test_suite_id: str, + payload: str, + mode: str = "preview", + on_match: str = "overwrite_unlocked", + on_new: str = "create", + on_absence: str = "do_nothing", +) -> str: + """Apply an export document from ``export_tests`` to a test suite. + + Args: + test_suite_id: UUID of the target test suite. + payload: The export document as a JSON string (the fenced JSON returned by + ``export_tests``). + mode: ``preview`` (default) reports the projected changes without persisting; + ``apply`` persists, skipping entries that can't be applied; ``apply_strict`` + persists only if nothing would be skipped, otherwise applies nothing. + on_match: What to do when an incoming test matches an existing one: + ``overwrite_unlocked`` (default), ``overwrite_all``, or ``skip``. + on_new: What to do when an incoming test has no match: ``create`` (default), + ``create_and_lock``, or ``skip``. + on_absence: What to do with existing tests that are absent from the payload: + ``do_nothing`` (default), ``delete_all``, or ``delete_unlocked``. + """ + config = ImportConfig( + mode=parse_enum(mode, ImportMode, "mode"), + on_match=parse_enum(on_match, OnMatch, "on_match"), + on_new=parse_enum(on_new, OnNew, "on_new"), + on_absence=parse_enum(on_absence, OnAbsence, "on_absence"), + ) + suite = resolve_test_suite(test_suite_id) + parsed_payload = _parse_import_payload(payload) + + try: + result = import_definitions(suite, config, parsed_payload) + except InvalidImportPayload as err: + raise MCPUserError(str(err)) from err + except ImportStrictViolation as err: + breakdown = MdDoc() + _append_import_result(breakdown, err.result, preview=True) + raise MCPUserError( + f"Strict import failed: {err.result.summary.skipped} tests would be skipped. " + f"Nothing was applied.\n\n{breakdown.render()}\n\n" + "Retry with `mode='apply'` to import while skipping these, or adjust the payload or policies." + ) from err + + preview = config.mode is ImportMode.preview + doc = MdDoc() + if preview: + doc.heading(1, f"Test Definition Import Preview for suite `{suite.test_suite}`") + doc.text("Preview only — no changes were persisted. Re-run with `mode='apply'` to persist.") + else: + doc.heading(1, f"Test Definition Import into suite `{suite.test_suite}`") + _append_import_result(doc, result, preview=preview) + return doc.render() diff --git a/testgen/mcp/tools/test_results.py b/testgen/mcp/tools/test_results.py index d35d3f0a..bb571954 100644 --- a/testgen/mcp/tools/test_results.py +++ b/testgen/mcp/tools/test_results.py @@ -1,22 +1,34 @@ from datetime import UTC, datetime, timedelta from uuid import UUID +from sqlalchemy import select + from testgen.common.enums import JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.job_execution import JobExecution -from testgen.common.models.test_definition import TestType +from testgen.common.models.test_definition import TestDefinition, TestType from testgen.common.models.test_result import BucketInterval, TestResult, TestResultStatus from testgen.common.models.test_run import TestRun, TestRunSummary from testgen.common.models.test_suite import TestSuite +from testgen.common.test_result_disposition_service import ( + DispositionUpdate, + set_test_results_disposition, +) from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import get_project_permissions, mcp_permission from testgen.mcp.tools.common import ( DocGroup, + FailureGroupBy, format_page_footer, format_page_info, + parse_failure_group_by, parse_result_status, parse_since_arg, + parse_test_result_disposition, parse_uuid, + resolve_aggregate_scope, + resolve_test_result, + resolve_test_suite, resolve_test_type, validate_limit, validate_page, @@ -27,6 +39,12 @@ _DEFAULT_SEARCH_STATUSES = [TestResultStatus.Failed, TestResultStatus.Warning] +_MODEL_GROUP_COLUMN = { + FailureGroupBy.TEST_TYPE: "test_type", + FailureGroupBy.TABLE: "table_name", + FailureGroupBy.COLUMN: "column_names", +} + @with_database_session @mcp_permission("view") @@ -72,14 +90,14 @@ def list_test_results( raise MCPResourceNotAccessible("Test suite", test_suite_id) if suite.last_complete_test_run_id is None: raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") - test_run = TestRun.get_by_id_or_job(suite.last_complete_test_run_id) + test_run = TestRun.get(suite.last_complete_test_run_id) if test_run is None: raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") resolved_via_suite = True - run_id_label = str(test_run.job_execution_id) + run_id_label = str(test_run.id) else: job_uuid = parse_uuid(job_execution_id, "job_execution_id") - test_run = TestRun.get_by_id_or_job(job_uuid) + test_run = TestRun.get(job_uuid) suite = TestSuite.get_regular(test_run.test_suite_id) if test_run else None if test_run is None or suite is None or not perms.has_access(suite.project_code): raise MCPResourceNotAccessible("Test run", job_execution_id) @@ -112,6 +130,12 @@ def list_test_results( type_names = {tt.test_type: tt.test_name_short for tt in TestType.select_where(TestType.active == "Y")} + td_ids = {r.test_definition_id for r in results if r.test_definition_id} + source_by_td = { + td.id: (td.external_url, td.custom_metadata) + for td in (TestDefinition.select_where(TestDefinition.id.in_(td_ids)) if td_ids else []) + } + doc = MdDoc() doc.heading(1, f"Test Results for run `{run_id_label}`") if resolved_via_suite: @@ -125,6 +149,7 @@ def list_test_results( doc.heading(2, f"[{status_str}] {test_name} on `{r.column_names}` in `{r.table_name}`") else: doc.heading(2, f"[{status_str}] {test_name} on `{r.table_name}`") + doc.field("Test result", r.id, code=True) doc.field("Test definition", r.test_definition_id, code=True) if r.column_names: doc.field("Column", r.column_names, code=True) @@ -134,6 +159,13 @@ def list_test_results( doc.field("Threshold", r.threshold_value) if r.message: doc.field("Message", r.message) + external_url, custom_metadata = source_by_td.get(r.test_definition_id, (None, None)) + if external_url: + doc.field("External URL", external_url) + if custom_metadata: + doc.heading(3, "Custom Metadata") + for key, value in custom_metadata.items(): + doc.field(MdDoc.escape(key), value) return doc.render() @@ -167,20 +199,20 @@ def get_failure_summary( group_by: Group failures by 'test_type', 'table', or 'column' (default: 'test_type'). """ perms = get_project_permissions() + group = parse_failure_group_by(group_by) if not any((job_execution_id, test_suite_id, since)): raise MCPUserError( "Provide 'job_execution_id' for a single run, or 'test_suite_id' or 'project_code' " "to aggregate across runs. 'since' is required when 'test_suite_id' is not provided." ) - if group_by in ("table", "column") and not (job_execution_id or test_suite_id): + if group in (FailureGroupBy.TABLE, FailureGroupBy.COLUMN) and not (job_execution_id or test_suite_id): raise MCPUserError( - f"'{group_by}' grouping requires a single-suite scope. " + f"'{group}' grouping requires a single-suite scope. " "Provide 'job_execution_id' or 'test_suite_id'." ) - model_group_map = {"table": "table_name", "column": "column_names"} - model_group_by = model_group_map.get(group_by, group_by) + model_group_by = _MODEL_GROUP_COLUMN[group] scope_label: str test_run_id = None @@ -189,7 +221,7 @@ def get_failure_summary( if job_execution_id: job_uuid = parse_uuid(job_execution_id, "job_execution_id") - test_run = TestRun.get_by_id_or_job(job_uuid) + test_run = TestRun.get(job_uuid) suite = TestSuite.get_regular(test_run.test_suite_id) if test_run else None if test_run is None or suite is None or not perms.has_access(suite.project_code): raise MCPResourceNotAccessible("Test run", job_execution_id) @@ -197,15 +229,7 @@ def get_failure_summary( scope_label = f"run `{job_execution_id}`" project_codes = perms.allowed_codes else: - if project_code: - perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) - project_codes = [project_code] - else: - project_codes = perms.allowed_codes - if test_suite_uuid is not None: - suite = TestSuite.get_regular(test_suite_uuid) - if suite is None or not perms.has_access(suite.project_code): - raise MCPResourceNotAccessible("Test suite", test_suite_id) + project_codes = resolve_aggregate_scope(project_code, test_suite_id=test_suite_id) scope_parts = [] if project_code: scope_parts.append(f"project `{project_code}`") @@ -227,14 +251,14 @@ def get_failure_summary( return f"No confirmed failures found for {scope_label}." total = sum(row[-1] for row in failures) - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: type_names = {tt.test_type: tt.test_name_short for tt in TestType.select_where(TestType.active == "Y")} doc = MdDoc() doc.heading(1, f"Failure Summary — {scope_label}") doc.text(f"**Total confirmed failures (Failed + Warning):** {total}") - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: headers = ["Test Type", "Severity", "Count"] rows = [] for row in failures: @@ -242,7 +266,7 @@ def get_failure_summary( name = type_names.get(code, code) severity = status.value if status else "Unknown" rows.append([name, severity, count]) - elif group_by == "column": + elif group is FailureGroupBy.COLUMN: headers = ["Column", "Count"] rows = [] for row in failures: @@ -253,9 +277,9 @@ def get_failure_summary( headers = ["Table Name", "Count"] rows = [[row[0], row[-1]] for row in failures] - doc.table(headers, rows, code=[0] if group_by == "table" else None) + doc.table(headers, rows, code=[0] if group is FailureGroupBy.TABLE else None) - if group_by == "test_type": + if group is FailureGroupBy.TEST_TYPE: doc.text( "Check `testgen://test-types` to understand what each test type checks " "and `get_test_type(test_type='...')` to fetch more details." @@ -450,12 +474,9 @@ def get_failure_trend( valid = ", ".join(v.value for v in BucketInterval) raise MCPUserError(f"Invalid `bucket`: `{bucket}`. Valid values: {valid}") from err - perms = get_project_permissions() - if project_code: - perms.verify_access(project_code, not_found=MCPResourceNotAccessible("Project", project_code)) - project_codes = [project_code] - else: - project_codes = perms.allowed_codes + project_codes = resolve_aggregate_scope( + project_code, test_suite_id=test_suite_id, table_group_id=table_group_id + ) anchor_today = datetime.now(UTC).date() if exclude_today: @@ -540,7 +561,7 @@ def compare_test_runs( perms = get_project_permissions() def _resolve_accessible(je_id_str: str, je_uuid: UUID) -> TestRun: - run = TestRun.get_by_id_or_job(je_uuid) + run = TestRun.get(je_uuid) if run is None: raise MCPResourceNotAccessible("Test run", je_id_str) suite = TestSuite.get_regular(run.test_suite_id) @@ -549,7 +570,7 @@ def _resolve_accessible(je_id_str: str, je_uuid: UUID) -> TestRun: return run def _require_completed(run: TestRun, label: str) -> None: - je = get_current_session().get(JobExecution, run.job_execution_id) + je = get_current_session().get(JobExecution, run.id) if je.status != JobStatus.COMPLETED: status_label = TestRunSummary.STATUS_LABEL.get(je.status, je.status) raise MCPUserError( @@ -587,14 +608,15 @@ def _require_completed(run: TestRun, label: str) -> None: ["", "Target", "Baseline"], [ ["Test Run", - MdDoc.code(str(target_run.job_execution_id)), - MdDoc.code(str(baseline_run.job_execution_id))], + MdDoc.code(str(target_run.id)), + MdDoc.code(str(baseline_run.id))], ["Started", target_run.test_starttime, baseline_run.test_starttime], ], ) doc.table( headers=["Category", "Count"], rows=[ + ["Stable passes (Baseline passed → Target passed)", diff.stable_passes], ["Regressions (Baseline passed → Target failed/warning)", len(diff.regressions)], ["Improvements (Baseline failed/warning → Target passed)", len(diff.improvements)], ["Persistent failures", len(diff.persistent_failures)], @@ -637,3 +659,116 @@ def _section(title: str, rows: list) -> None: _section("Removed Tests", diff.removed_tests) return doc.render() + + +@with_database_session +@mcp_permission("disposition") +def update_test_result(test_result_id: str, disposition: str) -> str: + """Set the disposition on a single test result (confirm, dismiss, mute, or clear). + + Args: + test_result_id: UUID of the test result, e.g. from ``list_test_results``. + disposition: New disposition. One of 'Confirmed', 'Dismissed', 'Muted', + 'No Decision' (clears it). 'Muted' deactivates the parent test and locks + it against auto-regeneration; any other value reactivates and unlocks it. + """ + result = resolve_test_result(test_result_id) + db_disposition = parse_test_result_disposition(disposition) + + update: DispositionUpdate = set_test_results_disposition([result.id], db_disposition) + + doc = MdDoc() + if update.matched == 0: + doc.text( + f"Test result {MdDoc.code(test_result_id)} was not dispositioned — disposition does " + f"not apply to passed results. No change made." + ) + return doc.render() + + doc.text(f"Updated test result {MdDoc.code(test_result_id)} disposition to **{disposition}**.") + return doc.render() + + +@with_database_session +@mcp_permission("disposition") +def bulk_update_test_results( + test_suite_id: str, + disposition: str, + job_execution_id: str | None = None, + table_name: str | None = None, + test_type: str | None = None, + status: str | None = None, + test_definition_id: str | None = None, +) -> str: + """Set the disposition on every matching test result in a suite (confirm, dismiss, mute, clear). + + Args: + test_suite_id: UUID of the test suite, e.g. from ``list_test_suites``. + disposition: New disposition. One of 'Confirmed', 'Dismissed', 'Muted', + 'No Decision' (clears it). 'Muted' deactivates the parent tests and locks + them against auto-regeneration; any other value reactivates and unlocks. + job_execution_id: UUID of a test run within the suite. Defaults to the suite's + latest completed run when omitted. + table_name: Optional table-name filter. Case-sensitive. + test_type: Optional test type name (e.g. 'Alpha Truncation'). + status: Optional result-status filter (Passed, Failed, Warning, Error, Log). + test_definition_id: Optional single test-definition filter. + """ + suite = resolve_test_suite(test_suite_id) + db_disposition = parse_test_result_disposition(disposition) + + if job_execution_id: + run = TestRun.get(parse_uuid(job_execution_id, "job_execution_id")) + if run is None or run.test_suite_id != suite.id: + raise MCPResourceNotAccessible("Test run", job_execution_id) + else: + run = ( + TestRun.get(suite.last_complete_test_run_id) + if suite.last_complete_test_run_id + else None + ) + if run is None: + raise MCPUserError(f"No completed test runs found for test suite `{test_suite_id}`.") + + clauses = [TestResult.test_suite_id == suite.id, TestResult.test_run_id == run.id] + if status: + clauses.append(TestResult.status == parse_result_status(status)) + if table_name: + clauses.append(TestResult.table_name == table_name) + if test_type: + clauses.append(TestResult.test_type == resolve_test_type(test_type)) + if test_definition_id: + clauses.append(TestResult.test_definition_id == parse_uuid(test_definition_id, "test_definition_id")) + + result_ids = list(get_current_session().scalars(select(TestResult.id).where(*clauses)).all()) + update = set_test_results_disposition(result_ids, db_disposition) + + filters = [] + if table_name: + filters.append(f"table_name=`{table_name}`") + if test_type: + filters.append(f"test_type=`{test_type}`") + if status: + filters.append(f"status=`{status}`") + if test_definition_id: + filters.append(f"test_definition_id=`{test_definition_id}`") + filter_str = ", ".join(filters) if filters else "no filter" + + doc = MdDoc() + if update.matched == 0 and update.passed_skipped == 0: + doc.heading(1, "No test results matched") + doc.text( + f"No test results in suite `{suite.test_suite}` matched the filter ({filter_str}). Nothing changed." + ) + return doc.render() + + doc.heading(1, f"Updated {update.matched} test results in suite `{suite.test_suite}`") + doc.field("Disposition", disposition) + doc.field("Run", run.id, code=True) + doc.field("Filter", filter_str) + if update.passed_skipped: + doc.text( + f"Left {update.passed_skipped} passed results unchanged — disposition is not " + f"applied to passed results." + ) + return doc.render() diff --git a/testgen/mcp/tools/test_runs.py b/testgen/mcp/tools/test_runs.py index 415571e6..8f92d86a 100644 --- a/testgen/mcp/tools/test_runs.py +++ b/testgen/mcp/tools/test_runs.py @@ -30,6 +30,7 @@ def list_test_runs( project_code: str | None = None, test_suite: str | None = None, table_group_id: str | None = None, + schedule_id: str | None = None, status: str | None = None, limit: int = 10, page: int = 1, @@ -43,6 +44,8 @@ def list_test_runs( test_suite: Optional test suite name to filter by (case-sensitive). table_group_id: Optional UUID of a table group, e.g. from `get_data_inventory`. Returns runs for any suite in the group. + schedule_id: Optional UUID of a schedule, e.g. from `list_schedules`. Returns only runs + triggered by that schedule. status: Optional run status filter. One of: Pending, Running, Completed, Canceled, Error. limit: Page size (default 10, max 100). page: Page number starting at 1 (default 1). @@ -51,6 +54,8 @@ def list_test_runs( validate_page(page) statuses = parse_run_status_filter(status) if status else None + if schedule_id: + parse_uuid(schedule_id, "schedule_id") if not project_code and not table_group_id: raise MCPUserError("Provide either `project_code` or `table_group_id`.") @@ -86,6 +91,7 @@ def list_test_runs( project_code=project_code, table_group_id=str(table_group.id) if table_group else None, test_suite_id=test_suite_id, + schedule_id=schedule_id, statuses=statuses, page=page, page_size=limit, @@ -99,10 +105,11 @@ def list_test_runs( project_code=project_code, test_suite_id=test_suite_id, table_group_id=str(table_group.id) if table_group else None, + schedule_id=schedule_id, statuses=statuses, ) - scope_descriptor = _scope_descriptor(project_code, test_suite, table_group_id, status) + scope_descriptor = _scope_descriptor(project_code, test_suite, table_group_id, schedule_id, status) doc = MdDoc() doc.heading(1, f"Test runs{scope_descriptor}") @@ -160,7 +167,7 @@ def get_test_run(job_execution_id: str) -> str: doc = MdDoc() suite_label = summary.test_suite or "—" doc.heading(1, f"Test run: {suite_label}") - doc.field("Job ID", summary.job_execution_id, code=True) + doc.field("Test Run", summary.job_execution_id, code=True) doc.field("Test suite", suite_label) if summary.table_groups_name: doc.field("Table group", summary.table_groups_name) @@ -199,6 +206,7 @@ def _scope_descriptor( project_code: str | None, test_suite: str | None, table_group_id: str | None, + schedule_id: str | None, status: str | None, ) -> str: parts: list[str] = [] @@ -208,6 +216,8 @@ def _scope_descriptor( parts.append(f"suite `{test_suite}`") if table_group_id: parts.append(f"table group `{table_group_id}`") + if schedule_id: + parts.append(f"schedule `{schedule_id}`") if status: parts.append(f"status `{status}`") return f" — {', '.join(parts)}" if parts else "" @@ -246,6 +256,7 @@ def _select_pending_test_jes( project_code: str, test_suite_id: str | None, table_group_id: str | None, + schedule_id: str | None, statuses, ) -> list[JobExecution]: """Find queued/in-flight test-run JEs for a given suite or table group scope. For a @@ -267,7 +278,9 @@ def _select_pending_test_jes( return [] else: return [] + clauses = [JobExecution.job_schedule_id == schedule_id] if schedule_id else [] return JobExecution.select_active_by_kwargs( + *clauses, project_code=project_code, job_key=RUN_TESTS_JOB_KEY, kwargs_match={"test_suite_id": suite_ids}, @@ -278,7 +291,7 @@ def _select_pending_test_jes( def _render_pending_je(doc: MdDoc, je: JobExecution, label: str) -> None: status_label = TestRunSummary.STATUS_LABEL.get(je.status, je.status) doc.heading(3, f"{label} — {status_label}") - doc.field("Job ID", je.id, code=True) + doc.field("Test Run", je.id, code=True) if je.job_schedule_id is not None: doc.field("Schedule", je.job_schedule_id, code=True) doc.field("Submitted", je.created_at) @@ -289,7 +302,7 @@ def _render_pending_je(doc: MdDoc, je: JobExecution, label: str) -> None: def _render_test_run_section(doc: MdDoc, run: TestRunSummary) -> None: title = run.test_suite or run.project_code doc.heading(2, f"{title} — {run.status_label}") - doc.field("Job ID", run.job_execution_id, code=True) + doc.field("Test Run", run.job_execution_id, code=True) if run.job_schedule_id is not None: doc.field("Schedule", run.job_schedule_id, code=True) if run.test_suite: diff --git a/testgen/mcp/tools/test_suites.py b/testgen/mcp/tools/test_suites.py new file mode 100644 index 00000000..dc1ab7bf --- /dev/null +++ b/testgen/mcp/tools/test_suites.py @@ -0,0 +1,266 @@ +"""MCP write tools for test suites — create and update regular (non-monitor) suites. + +Monitor suites (``is_monitor=True``) are managed exclusively through the dedicated +monitors surface. ``resolve_test_suite`` already filters them out, so an +``update_test_suite`` call against a monitor-suite ID surfaces the unified +``not found or not accessible`` error. +""" + +from __future__ import annotations + +from typing import Any + +from testgen.common.models import get_current_session, with_database_session +from testgen.common.models.test_definition import Severity +from testgen.common.models.test_suite import TestSuite +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.permissions import mcp_permission +from testgen.mcp.tools.common import ( + DocGroup, + parse_severity, + raise_validation_error, + render_diff_table, + resolve_table_group, + resolve_test_suite, +) +from testgen.mcp.tools.markdown import MdDoc + +_DOC_GROUP = DocGroup.MANAGE + +# Free-text + component fields are stored as ``NullIfEmptyString`` columns: writing ``""`` +# to one becomes ``NULL`` in the DB. Normalize on the way in so the diff-renderer doesn't +# flag a phantom change (``before=None`` → ``after=""``) on a value that the DB will read +# back as identical to before. +_OPTIONAL_TEXT_ATTRS = ( + "test_suite_description", + "component_key", + "component_type", + "component_name", +) + + +def _empty_to_none(value: str | None) -> str | None: + return value if value else None + + +@with_database_session +@mcp_permission("edit") +def create_test_suite( + table_group_id: str, + test_suite_name: str, + *, + description: str | None = None, + severity_default: str | None = None, + dq_score_exclude: bool = False, + export_to_observability: bool = False, + component_key: str | None = None, + component_type: str | None = "dataset", + component_name: str | None = None, +) -> str: + """Create a test suite under a table group. + + The new suite inherits the table group's project and connection. + + Args: + table_group_id: UUID of the table group, e.g. from `list_table_groups`. + test_suite_name: Display name for the suite. + description: Optional free-text description. + severity_default: Optional default severity applied to tests that do + not set their own. Accepts `Fail` or `Warning`. + dq_score_exclude: Whether to exclude this suite's results from data + quality scoring. Defaults to False. + export_to_observability: Whether to export test results to the + configured DataOps Observability API. Defaults to False. + component_key: Component identifier in DataOps Observability. + component_type: Component type in DataOps Observability + (e.g. `dataset`). Defaults to `dataset`. + component_name: Component display name in DataOps Observability. + """ + name = test_suite_name.strip() if test_suite_name else "" + errors: list[str] = [] + if not name: + errors.append("test_suite_name: must not be empty.") + parsed_severity: Severity | None = _parse_severity_field(severity_default, errors) + if errors: + raise_validation_error(errors, "Test suite creation rejected. No changes saved.") + + table_group = resolve_table_group(table_group_id) + + suite = TestSuite( + project_code=table_group.project_code, + connection_id=table_group.connection_id, + table_groups_id=table_group.id, + test_suite=name, + test_suite_description=_empty_to_none(description), + severity=parsed_severity.value if parsed_severity is not None else None, + dq_score_exclude=dq_score_exclude, + export_to_observability=export_to_observability, + component_key=_empty_to_none(component_key), + component_type=_empty_to_none(component_type), + component_name=_empty_to_none(component_name), + is_monitor=False, + ) + session = get_current_session() + session.add(suite) + # Flush so the autogenerated ``id`` is available for the rendered response. + session.flush() + + return _render_created_suite(suite, table_group_name=table_group.table_groups_name) + + +@with_database_session +@mcp_permission("edit") +def update_test_suite( + test_suite_id: str, + *, + test_suite_name: str | None = None, + description: str | None = None, + severity_default: str | None = None, + dq_score_exclude: bool | None = None, + export_to_observability: bool | None = None, + component_key: str | None = None, + component_type: str | None = None, + component_name: str | None = None, +) -> str: + """Update fields on a test suite. Atomic — nothing saved unless every supplied value is valid. + + Args: + test_suite_id: UUID of the test suite, e.g. from `list_test_suites`. + test_suite_name: New display name. + description: New free-text description. Pass an empty string to clear. + severity_default: New default severity for tests that do not set their + own. Accepts `Fail` or `Warning`. + dq_score_exclude: Whether this suite's results are excluded from data + quality scoring. + export_to_observability: Whether test results are exported to the + configured DataOps Observability API. + component_key: Component identifier in DataOps Observability. + Pass an empty string to clear. + component_type: Component type in DataOps Observability + (e.g. `dataset`). Pass an empty string to clear. + component_name: Component display name in DataOps Observability. + Pass an empty string to clear. + """ + supplied = { + "test_suite_name": test_suite_name, + "description": description, + "severity_default": severity_default, + "dq_score_exclude": dq_score_exclude, + "export_to_observability": export_to_observability, + "component_key": component_key, + "component_type": component_type, + "component_name": component_name, + } + if all(value is None for value in supplied.values()): + raise MCPUserError("No fields supplied to update.") + + suite = resolve_test_suite(test_suite_id) + + # Validate every supplied field first; surface all errors at once. + errors: list[str] = [] + updates: dict[str, Any] = {} + + if test_suite_name is not None: + cleaned = test_suite_name.strip() + if not cleaned: + errors.append("test_suite_name: must not be empty.") + else: + updates["test_suite"] = cleaned + + if description is not None: + updates["test_suite_description"] = _empty_to_none(description) + + if severity_default is not None: + parsed = _parse_severity_field(severity_default, errors) + if parsed is not None: + updates["severity"] = parsed.value + + if dq_score_exclude is not None: + updates["dq_score_exclude"] = dq_score_exclude + if export_to_observability is not None: + updates["export_to_observability"] = export_to_observability + if component_key is not None: + updates["component_key"] = _empty_to_none(component_key) + if component_type is not None: + updates["component_type"] = _empty_to_none(component_type) + if component_name is not None: + updates["component_name"] = _empty_to_none(component_name) + + if errors: + raise_validation_error(errors, "Update rejected. No changes saved.") + + # Diff scoped to the attributes actually touched by this call. + before = {attr: getattr(suite, attr, None) for attr in updates} + for attr, value in updates.items(): + setattr(suite, attr, value) + after = {attr: getattr(suite, attr, None) for attr in updates} + + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}` updated") + doc.field("ID", str(suite.id), code=True) + + rendered = render_diff_table(doc, before, after, attrs=_DIFF_ORDER, labels=_DIFF_LABELS) + if not rendered: + doc.text("No fields changed — supplied values matched the current state.") + return doc.render() + + suite.save() + return doc.render() + + +def _parse_severity_field(value: str | None, errors: list[str]) -> Severity | None: + """Validate a ``severity_default`` arg; append a field-scoped bullet on failure.""" + if value is None: + return None + try: + return parse_severity(value) + except MCPUserError: + valid = ", ".join(s.value for s in Severity) + errors.append(f"severity_default: must be one of {valid} (got `{value}`).") + return None + + +def _render_created_suite(suite: TestSuite, *, table_group_name: str) -> str: + doc = MdDoc() + doc.heading(1, f"Test Suite `{suite.test_suite}` created") + doc.field("ID", str(suite.id), code=True) + doc.field("Project", suite.project_code, code=True) + doc.field("Table Group", f"{table_group_name} (`{suite.table_groups_id}`)") + if suite.test_suite_description: + doc.field("Description", suite.test_suite_description) + if suite.severity: + doc.field("Default severity", suite.severity) + doc.field("Exclude from quality scoring", suite.dq_score_exclude) + doc.field("Export to Observability", suite.export_to_observability) + if suite.component_key: + doc.field("Component key", suite.component_key) + if suite.component_type: + doc.field("Component type", suite.component_type) + if suite.component_name: + doc.field("Component name", suite.component_name) + return doc.render() + + +# Order drives the diff-table rendering. Mirrors the update_test_suite arg order: +# observability + components grouped together, dq_score_exclude separate. +_DIFF_ORDER: tuple[str, ...] = ( + "test_suite", + "test_suite_description", + "severity", + "dq_score_exclude", + "export_to_observability", + "component_key", + "component_type", + "component_name", +) + +_DIFF_LABELS: dict[str, str] = { + "test_suite": "Name", + "test_suite_description": "Description", + "severity": "Default severity", + "dq_score_exclude": "Exclude from quality scoring", + "export_to_observability": "Export to Observability", + "component_key": "Component key", + "component_type": "Component type", + "component_name": "Component name", +} diff --git a/testgen/server/__init__.py b/testgen/server/__init__.py index d0867e9d..ba8dd5b3 100644 --- a/testgen/server/__init__.py +++ b/testgen/server/__init__.py @@ -23,6 +23,7 @@ from testgen.api.oauth.routes import router as oauth_router from testgen.api.oauth.server import create_authorization_server from testgen.common import version_service +from testgen.common.mixpanel_service import MixpanelService from testgen.common.models import with_database_session from testgen.server.middleware import BodySizeLimitMiddleware, SecurityHeadersMiddleware @@ -44,6 +45,11 @@ def _patch_openapi_schema(app: FastAPI) -> None: - Top-level component schema titles (shown in Redoc sidebar) - Path/query parameter schemas + It also strips ``description`` from enum component schemas. A ``StrEnum`` used + as a field type surfaces its class docstring as the schema description; enums + are internal primitives free to carry technical docstrings, so the public + schema must not echo them. Per-field API text comes from ``Field(description=...)``. + FastAPI caches the schema after the first call, so the patching runs once. """ _original = app.openapi @@ -56,6 +62,8 @@ def patched_openapi() -> dict: for branch in prop.get("anyOf", []): branch.pop("title", None) model_schema.pop("title", None) + if "enum" in model_schema: + model_schema.pop("description", None) for methods in schema.get("paths", {}).values(): for details in methods.values(): if isinstance(details, dict): @@ -78,11 +86,14 @@ def create_app(version: str | None = None) -> FastAPI: @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: - if mcp_session_manager is not None: - async with mcp_session_manager.run(): + try: + if mcp_session_manager is not None: + async with mcp_session_manager.run(): + yield + else: yield - else: - yield + finally: + MixpanelService().drain() tags_metadata = [ {"name": "Jobs", "description": "Submit, poll, cancel, and list job executions (profiling, tests, generation)."}, @@ -173,6 +184,17 @@ def run_server() -> None: if settings.API_TLS_ENABLED: ssl_kwargs["ssl_certfile"] = settings.SSL_CERT_FILE ssl_kwargs["ssl_keyfile"] = settings.SSL_KEY_FILE + # uvicorn defaults ssl_ciphers to "TLSv1", which restricts TLS 1.2 to legacy + # CBC/SHA-1 suites with no AEAD. Hardened TLS clients (e.g. the Databricks + # Unity Catalog HTTP-connection gateway) negotiate TLS 1.2 with an AEAD-only + # policy and find no shared cipher, so the handshake fails before any request + # is sent. Offer a modern AEAD cipher set (Mozilla intermediate). + ssl_kwargs["ssl_ciphers"] = ( + "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:" + "ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:" + "ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:" + "DHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES256-GCM-SHA384" + ) LOG.info( "Starting server on %s:%s (TLS: %s, MCP: %s)", diff --git a/testgen/settings.py b/testgen/settings.py index 69e922dc..cf760cba 100644 --- a/testgen/settings.py +++ b/testgen/settings.py @@ -533,6 +533,17 @@ def _ssl_files_present() -> bool: Lifetime of OAuth access and refresh tokens. """ +PAT_DEFAULT_LIFETIME_SECONDS: int = 31_536_000 # 365 days +""" +Default maximum lifetime a personal access token can be created with. +""" + +PAT_MAX_LIFETIME_SECONDS: int = 63_072_000 # 2 years (730 days) +""" +Absolute ceiling on the admin-configurable personal access token maximum lifetime. +The configured maximum cannot exceed this. +""" + JWT_HASHING_KEY_B64: str = getenv("TG_JWT_HASHING_KEY") """ Random key used to sign/verify the authentication token @@ -634,7 +645,8 @@ def _default_ui_base_url() -> str: Extra Host header values accepted by MCP DNS rebinding protection (comma-separated). BASE_URL's hostname and loopback are always allowed; this adds more for multi-domain deployments or reverse proxies that rewrite Host. Entries without a port (`tg.example.com`) -get an automatic `:*` wildcard; entries with a port are matched literally +are allowed both with an automatic `:*` port wildcard and bare, so gateways that send an +Origin with no port are accepted; entries with a port are matched literally (`tg.example.com:8080`) or with explicit wildcard (`tg.example.com:*`). Only affects MCP routes — the parent FastAPI app does not validate Host headers. diff --git a/testgen/template/data_chars/data_chars_update.sql b/testgen/template/data_chars/data_chars_update.sql index 38281dad..2084b587 100644 --- a/testgen/template/data_chars/data_chars_update.sql +++ b/testgen/template/data_chars/data_chars_update.sql @@ -8,6 +8,7 @@ WITH new_chars AS ( schema_name, table_name, run_date, + MAX(object_type) AS object_type, MAX(approx_record_ct) AS approx_record_ct, MAX(record_ct) AS record_ct, COUNT(*) AS column_ct @@ -20,7 +21,8 @@ WITH new_chars AS ( ), updated_records AS ( UPDATE data_table_chars - SET approx_record_ct = n.approx_record_ct, + SET object_type = n.object_type, + approx_record_ct = n.approx_record_ct, record_ct = n.record_ct, column_ct = n.column_ct, last_refresh_date = n.run_date, @@ -55,6 +57,7 @@ WITH new_chars AS ( schema_name, table_name, run_date, + MAX(object_type) AS object_type, MAX(approx_record_ct) AS approx_record_ct, MAX(record_ct) AS record_ct, COUNT(*) AS column_ct @@ -72,6 +75,7 @@ inserted_records AS ( table_name, add_date, last_refresh_date, + object_type, approx_record_ct, record_ct, column_ct @@ -81,6 +85,7 @@ inserted_records AS ( n.table_name, n.run_date, n.run_date, + n.object_type, n.approx_record_ct, n.record_ct, n.column_ct diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql index 66e4db8b..32e6b569 100644 --- a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql +++ b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql @@ -35,6 +35,7 @@ CREATE TABLE stg_data_chars_updates ( general_type VARCHAR(1), column_type VARCHAR(50), db_data_type VARCHAR(50), + object_type VARCHAR(20), approx_record_ct BIGINT, record_ct BIGINT ); @@ -130,6 +131,7 @@ CREATE TABLE table_groups stakeholder_group VARCHAR(40), transform_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), last_complete_profile_run_id UUID, dq_score_profiling FLOAT, dq_score_testing FLOAT @@ -143,10 +145,7 @@ CREATE TABLE profiling_runs ( connection_id BIGINT NOT NULL, table_groups_id UUID NOT NULL, profiling_starttime TIMESTAMP, - profiling_endtime TIMESTAMP, - status VARCHAR(100) DEFAULT 'Running', progress JSONB, - log_message VARCHAR, table_ct BIGINT, column_ct BIGINT, record_ct BIGINT, @@ -156,9 +155,7 @@ CREATE TABLE profiling_runs ( anomaly_column_ct BIGINT, dq_affected_data_points BIGINT, dq_total_data_points BIGINT, - dq_score_profiling FLOAT, - process_id INTEGER, - job_execution_id UUID + dq_score_profiling FLOAT ); CREATE TABLE test_suites ( @@ -245,6 +242,8 @@ CREATE TABLE test_definitions ( flagged BOOLEAN DEFAULT FALSE NOT NULL, external_id UUID, impact_dimension VARCHAR(20), + external_url VARCHAR, + custom_metadata JSONB, CONSTRAINT test_definitions_test_suites_test_suite_id_fk FOREIGN KEY (test_suite_id) REFERENCES test_suites ); @@ -419,6 +418,7 @@ CREATE TABLE data_table_chars ( table_groups_id UUID, schema_name VARCHAR(50), table_name VARCHAR(120), + object_type VARCHAR(20), functional_table_type VARCHAR(50), description VARCHAR(1000), critical_data_element BOOLEAN, @@ -430,6 +430,7 @@ CREATE TABLE data_table_chars ( transform_level VARCHAR(40), aggregation_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), add_date TIMESTAMP, drop_date TIMESTAMP, last_refresh_date TIMESTAMP, @@ -469,6 +470,7 @@ CREATE TABLE data_column_chars ( transform_level VARCHAR(40), aggregation_level VARCHAR(40), data_product VARCHAR(40), + data_classification VARCHAR(40), add_date TIMESTAMP, last_mod_date TIMESTAMP, drop_date TIMESTAMP, @@ -579,6 +581,8 @@ CREATE TABLE test_types ( dq_dimension VARCHAR(50), impact_dimension VARCHAR(20), health_dimension VARCHAR(50), + algorithm VARCHAR(64), + statistical_technique VARCHAR(64), threshold_description VARCHAR(200), result_visualization VARCHAR(50) DEFAULT 'line_chart', result_visualization_params TEXT DEFAULT NULL, @@ -610,10 +614,7 @@ CREATE TABLE test_runs ( PRIMARY KEY, test_suite_id UUID NOT NULL, test_starttime TIMESTAMP, - test_endtime TIMESTAMP, - status VARCHAR(100) DEFAULT 'Running', progress JSONB, - log_message TEXT, test_ct INTEGER, passed_ct INTEGER, failed_ct INTEGER, @@ -627,8 +628,6 @@ CREATE TABLE test_runs ( dq_affected_data_points BIGINT, dq_total_data_points BIGINT, dq_score_test_run FLOAT, - process_id INTEGER, - job_execution_id UUID, CONSTRAINT test_runs_test_suites_fk FOREIGN KEY (test_suite_id) REFERENCES test_suites ); @@ -746,12 +745,12 @@ CREATE INDEX ix_pm_role ON project_memberships(role); CREATE TABLE oauth2_clients ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - user_id UUID REFERENCES auth_users(id) ON DELETE SET NULL, client_id VARCHAR(48) NOT NULL UNIQUE, client_secret VARCHAR(120), client_id_issued_at INTEGER NOT NULL DEFAULT 0, client_secret_expires_at INTEGER NOT NULL DEFAULT 0, - client_metadata TEXT NOT NULL DEFAULT '{}' + client_metadata TEXT NOT NULL DEFAULT '{}', + client_type VARCHAR(20) NOT NULL DEFAULT 'external' ); CREATE INDEX idx_oauth2_clients_client_id ON oauth2_clients(client_id); @@ -782,7 +781,8 @@ CREATE TABLE oauth2_tokens ( issued_at INTEGER NOT NULL DEFAULT EXTRACT(EPOCH FROM NOW())::INTEGER, access_token_revoked_at INTEGER NOT NULL DEFAULT 0, refresh_token_revoked_at INTEGER NOT NULL DEFAULT 0, - expires_in INTEGER NOT NULL DEFAULT 0 + expires_in INTEGER NOT NULL DEFAULT 0, + name VARCHAR(255) ); CREATE INDEX idx_oauth2_tokens_refresh_token ON oauth2_tokens(refresh_token); @@ -854,6 +854,7 @@ CREATE TABLE IF NOT EXISTS score_definition_results_breakdown ( stakeholder_group TEXT DEFAULT NULL, transform_level TEXT DEFAULT NULL, data_product TEXT DEFAULT NULL, + data_classification TEXT DEFAULT NULL, impact DOUBLE PRECISION DEFAULT 0.0, score DOUBLE PRECISION DEFAULT 0.0, issue_ct INTEGER DEFAULT 0 @@ -1139,3 +1140,26 @@ CREATE TABLE notification_settings ( REFERENCES score_definitions (id) ON DELETE CASCADE ); + +-- Run identity = job execution id, with a cascading delete chain: +-- deleting a job_executions row removes its run and the run's results. +-- Declared here because job_executions is created after the run tables. +ALTER TABLE profiling_runs + ADD CONSTRAINT profiling_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +ALTER TABLE test_runs + ADD CONSTRAINT test_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +ALTER TABLE test_results + ADD CONSTRAINT test_results_test_runs_id_fk + FOREIGN KEY (test_run_id) REFERENCES test_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_results + ADD CONSTRAINT profile_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_anomaly_results + ADD CONSTRAINT profile_anomaly_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; diff --git a/testgen/template/dbsetup/060_create_standard_views.sql b/testgen/template/dbsetup/060_create_standard_views.sql index fffaa445..5b99ae67 100644 --- a/testgen/template/dbsetup/060_create_standard_views.sql +++ b/testgen/template/dbsetup/060_create_standard_views.sql @@ -122,6 +122,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, dtc.table_name, dcc.column_name, pr.profiling_starttime as profiling_run_date, @@ -162,6 +163,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, t.dq_dimension, pr.table_name, @@ -205,6 +207,7 @@ GROUP BY pr.profile_run_id, pr.table_groups_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, t.dq_dimension, pr.run_date, tg.project_code ; @@ -232,6 +235,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.test_time, r.table_name, r.column_names as column_name, COUNT(*) as test_ct, @@ -271,6 +275,7 @@ GROUP BY r.table_groups_id, r.table_name, r.column_names, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.test_time, tg.project_code; @@ -312,6 +317,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.dq_dimension, r.test_time, r.table_name, dcc.column_name, @@ -347,6 +353,7 @@ GROUP BY r.table_groups_id, r.test_run_id, r.test_suite_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.dq_dimension, r.test_time, r.table_name, dcc.column_name, tg.project_code; @@ -368,6 +375,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, t.impact_dimension, pr.table_name, @@ -411,6 +419,7 @@ GROUP BY pr.profile_run_id, pr.table_groups_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, t.impact_dimension, pr.run_date, tg.project_code; @@ -450,6 +459,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.impact_dimension, r.test_time, r.table_name, dcc.column_name, @@ -485,6 +495,7 @@ GROUP BY r.table_groups_id, r.test_run_id, r.test_suite_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.impact_dimension, r.test_time, r.table_name, dcc.column_name, tg.project_code; @@ -511,6 +522,7 @@ SELECT tg.project_code, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, pr.table_name, pr.column_name, @@ -557,6 +569,7 @@ GROUP BY pr.profile_run_id, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level), COALESCE(dcc.critical_data_element, dtc.critical_data_element), COALESCE(dcc.data_product, dtc.data_product, tg.data_product), + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification), dcc.functional_data_type, pr.run_date, tg.project_code ; @@ -581,6 +594,7 @@ SELECT COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.critical_data_element, dtc.critical_data_element) as critical_data_element, COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification, dcc.functional_data_type as semantic_data_type, r.test_time, r.table_name, r.column_names as column_name, COUNT(*) as test_ct, @@ -623,5 +637,6 @@ GROUP BY sr.definition_id, tg.stakeholder_group, dcc.transform_level, dtc.transform_level, tg.transform_level, dcc.critical_data_element, dtc.critical_data_element, dcc.data_product, dtc.data_product, tg.data_product, + dcc.data_classification, dtc.data_classification, tg.data_classification, dcc.functional_data_type, r.test_time, tg.project_code; diff --git a/testgen/template/dbsetup/075_grant_role_rights.sql b/testgen/template/dbsetup/075_grant_role_rights.sql index a18f20b0..287663a8 100644 --- a/testgen/template/dbsetup/075_grant_role_rights.sql +++ b/testgen/template/dbsetup/075_grant_role_rights.sql @@ -44,6 +44,7 @@ GRANT SELECT, INSERT, DELETE, UPDATE ON {SCHEMA_NAME}.job_executions, {SCHEMA_NAME}.settings, {SCHEMA_NAME}.notification_settings, + {SCHEMA_NAME}.project_memberships, {SCHEMA_NAME}.test_definition_notes, {SCHEMA_NAME}.oauth2_clients, {SCHEMA_NAME}.oauth2_authorization_codes, diff --git a/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml b/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml index 723b3862..a20bc8df 100644 --- a/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml +++ b/testgen/template/dbsetup_anomaly_types/profile_anomaly_types_Invalid_Zip3_USA.yaml @@ -16,7 +16,7 @@ profile_anomaly_types: suggested_action: |- Review your source data, ingestion process, and any processing steps that update this column. dq_score_prevalence_formula: |- - (NULLIF(p.record_ct, 0)::INT - NULLIF(SPLIT_PART(p.top_patterns, ' | ', 1), '')::BIGINT)::FLOAT/NULLIF(p.record_ct, 0)::FLOAT + (NULLIF(p.record_ct, 0) - NULLIF(SPLIT_PART(p.top_patterns, ' | ', 1), '')::BIGINT)::FLOAT/NULLIF(p.record_ct, 0)::FLOAT dq_score_risk_factor: '1' dq_dimension: Validity impact_dimension: Conformance diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml index f38b89f4..15cacaf1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. It's ideal for confirming that two datasets exactly match -- that the sum of a measure or count of a value hasn't changed or shifted between categories. Use this test to compare a raw and processed version of the same dataset, or to confirm that an aggregated table exactly matches the detail table that it's built from. An error here means that one or more value combinations fail to match. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1400' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml index 1415731d..47882a0f 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Percent.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. Use it to confirm that two datasets closely match within the tolerance you set -- that the sum of a measure or count of a value remains sufficiently consistent between categories. You could use this test compare sales per product within one month to another, when you want to be alerted if the difference for any product falls outside of the range defined as 5% below to 10% above the prior month. An error here means that one or more value combinations fail to match within the set tolerances. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1404' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml index 84f20602..c578c084 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Balance_Range.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations across two different tables. Both tables must be accessible at the same time. Use it to confirm that two datasets closely match within the tolerances you define as specific values above or below the aggregate measure for the same categories in the reference dataset -- that the sum of a measure or count of a value remains sufficiently consistent between categories. For instance, you can use this test to compare sales per product within one month to another, when you want to be alerted if the difference for any product falls outside of the range defined as 10000 dollars above or below the prior week. An error here means that one or more value combinations fail to match within the set tolerances. New categories or combinations will cause failure. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1405' diff --git a/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml b/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml index 425a72e2..4fa29cee 100644 --- a/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Aggregate_Minimum.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares sums or counts of a column rolled up to one or more category combinations, but requires a match or increase in the aggregate value, rather than an exact match, across two different tables. Both tables must be accessible at the same time. Use this to confirm that aggregate values have not dropped for any set of categories, even if some values may rise. This test is useful to compare an older and newer version of a cumulative dataset. An error here means that one or more values per category set fail to match or exceed the prior dataset. New categories or combinations are allowed (but can be restricted independently with a Combo_Match test). Both tables must be present to run this test. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1401' diff --git a/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml b/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml index 88577f60..8cd8f223 100644 --- a/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Alpha_Trunc.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Alpha Truncation tests that the longest text value in a column hasn't become shorter than the defined threshold, initially 95% of the longest value at baseline. This could indicate a problem in a cumulative dataset, where prior values should still exist unchanged. A failure here would suggest that some process changed data that you would still expect to be present and matching its value when the column was profiled. This test would not be appropriate for an incremental or windowed dataset. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7001' test_type: Alpha_Trunc diff --git a/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml index 08f801a7..30f8c7c7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Avg_Shift.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- Average Shift tests that the average of a numeric column has not significantly changed since baseline, when profiling was done. A significant shift may indicate errors in processing, differences in source data, or valid changes that may nevertheless impact assumptions in downstream data products. The test uses Cohen's D, a statistical technique to identify significant shifts in a value. Cohen's D measures the difference between the two averages, reporting results on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. Depending on your data, some difference may be expected, so it's reasonable to adjust the threshold value that triggers test failure. This test works well for measures, or even for identifiers if you expect them to increment consistently. You may want to periodically adjust the expected threshold, or even the expected average value if you expect shifting over time. Consider this test along with Variability Increase. If variability rises too, process or measurement flaws could be at work. If variability remains consistent, the issue is more likely to be with the source data itself. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's D cat_test_conditions: - id: '7002' test_type: Avg_Shift diff --git a/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml b/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml index 6005806c..099e4b05 100644 --- a/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml +++ b/testgen/template/dbsetup_test_types/test_types_CUSTOM.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This business-rule test is highly flexible, covering any error state that can be expressed by a SQL query against one or more tables in the database. In operation, the user-defined query is embedded within a parent query returning the count of error rows identified. Any row returned by the query is interpreted as a single error condition in the test. Note that this query is run independently of other tests, and that performance will be slower, depending in large part on the efficiency of the query you write. Interpretation is based on the user-defined meaning of the test. Your query might be written to return errors in individual rows identified by joining tables. Or it might return an error based on a multi-column aggregate condition returning a single row if an error is found. This query is run separately when you click `Review Source Data` from Test Results, so be sure to include enough data in your results to follow-up. Interpretation is based on the user-defined meaning of the test. active: Y + algorithm: Custom SQL cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml index f4016136..fef1ba37 100644 --- a/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Combo_Match.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test verifies that values, or combinations of values, that are present in the main table are also found in a reference table. This is a useful test for referential integrity between fact and dimension tables. You can also use it to confirm the validity of a code or category, or of combinations of values that should only be found together within each record, such as product/size/color. An error here means that one or more category combinations in the main table are not found in the reference table. Both tables must be present to run this test. active: Y + algorithm: Set / lookup cat_test_conditions: [] target_data_lookups: - id: '1402' diff --git a/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml b/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml index 9c63f169..c2b08622 100644 --- a/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Condition_Flag.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Custom Condition is a business-rule test for a user-defined error condition based on the value of one or more columns. The condition is applied to each record within the table, and the count of records failing the condition is added up. If that count exceeds a threshold of errors, the test as a whole is failed. This test is ideal for error conditions that TestGen cannot automatically infer, and any condition that involves the values of more than one column in the same record. Performance of this test is fast, since it is performed together with other aggregate tests. Interpretation is based on the user-defined meaning of the test. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '7003' test_type: Condition_Flag diff --git a/testgen/template/dbsetup_test_types/test_types_Constant.yaml b/testgen/template/dbsetup_test_types/test_types_Constant.yaml index fff389cd..0bd3351e 100644 --- a/testgen/template/dbsetup_test_types/test_types_Constant.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Constant.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Constant Match tests that a single value determined to be a constant in baseline profiling is still the only value for the column that appears in subsequent versions of the dataset. Sometimes new data or business knowledge may reveal that the value is not a constant at all, even though only one value was present at profiling. In this case, you will want to disable this test. Alternatively, you can use the Value Match test to provide a limited number of valid values for the column. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7004' test_type: Constant diff --git a/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml index df372cd6..2926c8c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Daily_Record_Ct.yaml @@ -40,6 +40,7 @@ test_types: \ of days identified without data. You can adjust the threshold to accept a number\ \ of days that you know legitimately have no records. " active: Y + algorithm: Counting cat_test_conditions: - id: '7005' test_type: Daily_Record_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml b/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml index 770d1175..bd6d5168 100644 --- a/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Dec_Trunc.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Decimal Truncation tests that the fractional (decimal) part of a numeric column has not been truncated since Baseline. This works by summing all the fractional values after the decimal point and confirming that the total is at least equal to the fractional total at baseline. This could indicate a problem in a cumulative dataset, where prior values should still exist unchanged. A failure here would suggest that some process changed data that you would still expect to be present and matching its value when the column was profiled. This test would not be appropriate for an incremental or windowed dataset. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7006' test_type: Dec_Trunc diff --git a/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml index 23967398..b77887a2 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distinct_Date_Ct.yaml @@ -29,7 +29,7 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Minimum distinct date count expected result_visualization: line_chart @@ -37,6 +37,7 @@ test_types: usage_notes: |- Date Count tests that the count of distinct dates present in the column has not dropped since baseline. The test is relevant for cumulative datasets, where old records are retained. A failure here would indicate missing records, which could be caused by a processing error or changed upstream data sources. active: Y + algorithm: Counting cat_test_conditions: - id: '7007' test_type: Distinct_Date_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml index 47e609ec..630cb901 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distinct_Value_Ct.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- Value Count tests that the count of unique values present in the column has not dropped since baseline. The test is relevant for cumulative datasets, where old records are retained, or for any dataset where you would expect a set number of distinct values should be present. A failure here would indicate missing records or a change in categories or value assignment. active: Y + algorithm: Counting cat_test_conditions: - id: '7008' test_type: Distinct_Value_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml index 666bf095..9f5efdce 100644 --- a/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Distribution_Shift.yaml @@ -38,6 +38,8 @@ test_types: usage_notes: |- This test measures the similarity of two sets of counts per categories, by using their proportional counts as probability distributions. Using Jensen-Shannon divergence, a measure of relative entropy or difference between two distributions, the test assigns a score ranging from 0, meaning that the distributions are identical, to 1, meaning that the distributions are completely unrelated. This test can be used to compare datasets that may not match exactly, but should have similar distributions. For example, it is a useful sanity check for data from different sources that you would expect to have a consistent spread, such as shipment of building materials per state and construction projects by state. Scores can be compared over time even if the distributions are not identical -- a dataset can be expected to maintain a comparable divergence score with a reference dataset over time. Both tables must be present to run this test. active: Y + algorithm: Statistical drift + statistical_technique: Jensen-Shannon Divergence cat_test_conditions: [] target_data_lookups: - id: '1403' diff --git a/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml b/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml index 1ef27125..5eeb9b81 100644 --- a/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Dupe_Rows.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test verifies that combinations of values are not repeated within the table. By default when auto-generated, the test considers all columns to protect against duplication of entire rows. If you know the minimum columns that should constitute a unique record, such as a set of ID's, you should use those to make the test as sensitive as possible. Alternatively, if you know of columns you can always exclude, such as file_date or refresh_snapshot_id, remove them to tighten the test somewhat. active: Y + algorithm: Counting cat_test_conditions: [] target_data_lookups: - id: '1409' diff --git a/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml b/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml index 6d6573b4..3a625be1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Email_Format.yaml @@ -36,6 +36,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7009' test_type: Email_Format diff --git a/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml index c8d1e3a2..15250c0a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Freshness_Trend.yaml @@ -30,7 +30,7 @@ test_types: test_scope: table dq_dimension: Recency impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected time window result_visualization: binary_chart @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test compares the current table fingerprint, calculated signature of column contents, to confirm that the table has been updated within the expectd time window. The table fingerprint is derived from a set of values and aggregates from columns most likely to change. This test allows you to track the schedule and frequency of updates and refreshes to the table. active: Y + algorithm: Freshness / time cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml b/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml index c2327843..78217c1b 100644 --- a/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Future_Date.yaml @@ -28,13 +28,14 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Conformance - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected count of future dates result_visualization: line_chart result_visualization_params: null usage_notes: null active: Y + algorithm: Boundary check cat_test_conditions: - id: '7010' test_type: Future_Date diff --git a/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml b/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml index c1ec7d6d..5e4254f9 100644 --- a/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Future_Date_1Y.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Conformance - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Expected count of future dates beyond one year result_visualization: line_chart @@ -36,6 +36,7 @@ test_types: usage_notes: |- Future Year looks for date values in the column that extend beyond one year after the test date. This would be appropriate for transactional dates where you would expect to find dates in the near future, but not beyond one year ahead. Errors could indicate invalid entries or possibly dummy dates representing blank values. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7011' test_type: Future_Date_1Y diff --git a/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml b/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml index 00b6b9f6..69b35b8a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Incr_Avg_Shift.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- This is a more sensitive test than Average Shift, because it calculates an incremental difference in the average of new values compared to the average of values at baseline. This is appropriate for a cumulative dataset only, because it calculates the average of new entries based on the assumption that the count and average of records present at baseline are still present at the time of the test. This test compares the mean of new values with the standard deviation of the baseline average to calculate a Z-score. If the new mean falls outside the Z-score threshold, a shift is detected. Potential Z-score thresholds may range from 0 to 3, depending on the sensitivity you prefer. A failed test could indicate a quality issue or a legitimate shift in new data that should be noted and assessed by business users. Consider this test along with Variability Increase. If variability rises too, process, methodology or measurement flaws could be at issue. If variability remains consistent, the problem is more likely to be with the source data itself. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's D cat_test_conditions: - id: '7012' test_type: Incr_Avg_Shift diff --git a/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml b/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml index cdd4bfda..e6b7145d 100644 --- a/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml +++ b/testgen/template/dbsetup_test_types/test_types_LOV_All.yaml @@ -34,6 +34,7 @@ test_types: usage_notes: |- This is a more restrictive form of Value Match, testing that all values in the dataset match the list provided, and also that all values present in the list appear at least once in the dataset. This would be appropriate for tables where all category values in the column are represented at least once. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7013' test_type: LOV_All diff --git a/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml b/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml index 38b8040c..463503f2 100644 --- a/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_LOV_Match.yaml @@ -140,6 +140,7 @@ test_types: usage_notes: |- This tests that all values in the column match the hard-coded list provided. This is relevant when the list of allowable values is small and not expected to change often. Even if new values might occasionally be added, this test is useful for downstream data products to provide warning that assumptions and logic may need to change. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7014' test_type: LOV_Match diff --git a/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml index 7675e9d1..2e0f56f1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Metric_Trend.yaml @@ -25,7 +25,7 @@ test_types: test_scope: table dq_dimension: Validity impact_dimension: Regularity - health_dimension: null + health_dimension: Data Drift threshold_description: |- Expected aggregate metric range. result_visualization: line_chart @@ -33,6 +33,7 @@ test_types: usage_notes: |- This test compares the aggregate metric of all or a subset of records in a table against a derived tolerance range. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '2516' test_type: Metric_Trend diff --git a/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml b/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml index 1ade805e..80dfeffa 100644 --- a/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Min_Date.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for a cumulative dataset only, because it assumes all prior values are still present. It's appropriate where new records are added with more recent dates, but old dates dates do not change. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7015' test_type: Min_Date diff --git a/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml b/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml index 90f107f4..49a3e5d4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Min_Val.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for a cumulative dataset only, assuming all prior values are still present. It is also appropriate for any measure that has an absolute, definable minimum value, or a heuristic that makes senes for valid data. active: Y + algorithm: Boundary check cat_test_conditions: - id: '7016' test_type: Min_Val diff --git a/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml index f4dd0a0a..8669d6c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Missing_Pct.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- This test uses Cohen's H, a statistical test to identify a significant difference between two ratios. Results are reported on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. An uptick in missing data may indicate a collection issue at the source. A larger change may indicate a processing failure. A drop in missing data may also be significant, if it affects assumptions built into analytic products downstream. You can refine the expected threshold value as you view legitimate results of the measure over time. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's H cat_test_conditions: - id: '7017' test_type: Missing_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml index 35580b34..29df97dc 100644 --- a/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Monthly_Rec_Ct.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Monthly Records tests that at least one record is present for every calendar month within the minimum and maximum date range for the column. The test is relevant for transactional data, where you would expect at least one transaction to be recorded each month. A failure here would suggest missing records for the number of months identified without data. You can adjust the threshold to accept a number of month that you know legitimately have no records. active: Y + algorithm: Counting cat_test_conditions: - id: '7018' test_type: Monthly_Rec_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml index 5cf1493f..2112c1b7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Above.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Accuracy impact_dimension: Regularity - health_dimension: Data Drift + health_dimension: Statistical Drift threshold_description: |- Expected maximum pct records over upper 2 SD limit result_visualization: line_chart @@ -41,6 +41,8 @@ test_types: \ you expect to see. This test uses the baseline mean rather than the mean for\ \ the latest dataset to capture systemic shift as well as individual outliers. " active: Y + algorithm: Statistical drift + statistical_technique: Outlier Detection cat_test_conditions: - id: '7019' test_type: Outlier_Pct_Above diff --git a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml index fa88ab8c..fc4d7262 100644 --- a/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Outlier_Pct_Below.yaml @@ -28,7 +28,7 @@ test_types: test_scope: column dq_dimension: Accuracy impact_dimension: Regularity - health_dimension: Data Drift + health_dimension: Statistical Drift threshold_description: |- Expected maximum pct records over lower 2 SD limit result_visualization: line_chart @@ -41,6 +41,8 @@ test_types: \ you expect to see. This test uses the baseline mean rather than the mean for\ \ the latest dataset to capture systemic shift as well as individual outliers. " active: Y + algorithm: Statistical drift + statistical_technique: Outlier Detection cat_test_conditions: - id: '7020' test_type: Outlier_Pct_Below diff --git a/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml index 6998da47..36f81737 100644 --- a/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Pattern_Match.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test is appropriate for character fields that are expected to appear in a consistent format. It uses pattern matching syntax as appropriate for your database: REGEX matching if available, otherwise LIKE expressions. The expected threshold is the number of records that fail to match the defined pattern. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7021' test_type: Pattern_Match diff --git a/testgen/template/dbsetup_test_types/test_types_Recency.yaml b/testgen/template/dbsetup_test_types/test_types_Recency.yaml index 34945e9b..b1d656c7 100644 --- a/testgen/template/dbsetup_test_types/test_types_Recency.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Recency.yaml @@ -29,7 +29,8 @@ test_types: test_scope: column dq_dimension: Timeliness impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness + statistical_technique: Predictive Model threshold_description: |- Expected maximum count of days preceding test date result_visualization: line_chart @@ -37,6 +38,7 @@ test_types: usage_notes: |- This test evaluates recency based on the latest referenced dates in the column. The test is appropriate for transactional dates and timestamps. The test can be especially valuable because timely data deliveries themselves may not assure that the most recent data is present. You can adjust the expected threshold to the maximum number of days that you expect the data to age before the dataset is refreshed. active: Y + algorithm: Freshness / time cat_test_conditions: - id: '7022' test_type: Recency diff --git a/testgen/template/dbsetup_test_types/test_types_Required.yaml b/testgen/template/dbsetup_test_types/test_types_Required.yaml index cb294860..6f785df2 100644 --- a/testgen/template/dbsetup_test_types/test_types_Required.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Required.yaml @@ -35,6 +35,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Counting cat_test_conditions: - id: '7023' test_type: Required diff --git a/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml index 06c3d62e..fee0fb9c 100644 --- a/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Row_Ct.yaml @@ -34,6 +34,7 @@ test_types: usage_notes: |- Because this tests the row count against a constant minimum threshold, it's appropriate for any dataset, as long as the number of rows doesn't radically change from refresh to refresh. But it's not responsive to change over time. You may want to adjust the threshold periodically if you are dealing with a cumulative dataset. active: Y + algorithm: Counting cat_test_conditions: - id: '7024' test_type: Row_Ct diff --git a/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml index 47cbf379..85588851 100644 --- a/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Row_Ct_Pct.yaml @@ -35,6 +35,7 @@ test_types: usage_notes: |- This test is better than Row Count for an incremental or windowed dataset where you would expect the row count to range within a percentage of baseline. active: Y + algorithm: Counting cat_test_conditions: - id: '7025' test_type: Row_Ct_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml b/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml index 4992ba2c..d492f211 100644 --- a/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Schema_Drift.yaml @@ -32,6 +32,7 @@ test_types: usage_notes: |- This test compares the current table column types with previous data, to check whether the table schema has changed. This test allows you to track any changes to the table structure. active: Y + algorithm: Schema / metadata cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml b/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml index 36b83009..ff5e2025 100644 --- a/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Street_Addr_Pattern.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- The street address pattern used in this test should match the vast majority of USA addresses. You can adjust the threshold percent of matches based on the results you are getting -- you may well want to tighten it to make the test more sensitive to invalid entries. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7026' test_type: Street_Addr_Pattern diff --git a/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml b/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml index d81e834e..11658f2a 100644 --- a/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Table_Freshness.yaml @@ -29,7 +29,7 @@ test_types: test_scope: table dq_dimension: Recency impact_dimension: Reliability - health_dimension: Recency + health_dimension: Freshness threshold_description: |- Most recent prior table fingerprint result_visualization: binary_chart @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test compares the current table fingerprint, calculated signature of column contents, to confirm that the table has been updated. The table fingerprint is derived from a set of values and aggregates from columns most likely to change. This test allows you to track the schedule and frequency of updates and refreshes to the table. active: Y + algorithm: Freshness / time cat_test_conditions: [] target_data_lookups: [] test_templates: diff --git a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml index 34329e26..55273486 100644 --- a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Gain.yaml @@ -38,6 +38,7 @@ test_types: usage_notes: |- This test checks a single transactional table to verify that categorical values or combinations that are present in the most recent time window you define include at least all those found in the prior time window of the same duration. Missing values in the latest time window will trigger the test to fail. New values are permitted. Use this test to confirm that codes or categories are not lost across successive time periods in a transactional table. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1406' diff --git a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml index 6b10231d..c378cfd3 100644 --- a/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Timeframe_Combo_Match.yaml @@ -36,6 +36,7 @@ test_types: usage_notes: |- This test checks a single transactional table (such as a fact table) to verify that categorical values or combinations that are present in the most recent time window you define match those found in the prior time window of the same duration. New or missing values in the latest time window will trigger the test to fail. Use this test to confirm the consistency in the occurrence of codes or categories across successive time periods in a transactional table. active: Y + algorithm: Aggregate reconciliation cat_test_conditions: [] target_data_lookups: - id: '1407' diff --git a/testgen/template/dbsetup_test_types/test_types_US_State.yaml b/testgen/template/dbsetup_test_types/test_types_US_State.yaml index 397611df..d936d0f8 100644 --- a/testgen/template/dbsetup_test_types/test_types_US_State.yaml +++ b/testgen/template/dbsetup_test_types/test_types_US_State.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test validates entries against a fixed list of two-character US state codes and related Armed Forces codes. active: Y + algorithm: Set / lookup cat_test_conditions: - id: '7027' test_type: US_State diff --git a/testgen/template/dbsetup_test_types/test_types_Unique.yaml b/testgen/template/dbsetup_test_types/test_types_Unique.yaml index e1b5b661..6dee77c4 100644 --- a/testgen/template/dbsetup_test_types/test_types_Unique.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Unique.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- This test is ideal when the database itself does not enforce a primary key constraint on the table. It serves as an independent check on uniqueness. If's also useful when there are a small number of exceptions to uniqueness, which can be reflected in the expected threshold count of duplicates. active: Y + algorithm: Counting cat_test_conditions: - id: '7028' test_type: Unique diff --git a/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml b/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml index 6cb5908a..d3ca67bc 100644 --- a/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Unique_Pct.yaml @@ -37,6 +37,8 @@ test_types: usage_notes: |- You can think of this as a test of similarity that measures whether the percentage of unique values is consistent with the percentage at baseline. A significant change might indicate duplication or a telling shift in cardinality between entities. The test uses Cohen's H, a statistical test to identify a significant difference between two ratios. Results are reported on a standardized scale, which can be interpreted via a rule-of-thumb from small to huge. You can refine the expected threshold value as you view legitimate results of the measure over time. active: Y + algorithm: Statistical drift + statistical_technique: Cohen's H cat_test_conditions: - id: '7029' test_type: Unique_Pct diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml index cd73a08c..d4f5d3e9 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_Characters.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: |- The acceptable number of records with invalid character values present. default_severity: Warning @@ -37,6 +38,7 @@ test_types: usage_notes: |- This test looks for the presence of non-printing ASCII characters that are considered non-standard in basic text processing. It also identifies leading spaces and values enclosed in quotes. Values that fail this test may be artifacts of data conversion, or just more difficult to process or analyze downstream. active: N + algorithm: Pattern / regex cat_test_conditions: - id: '7036' test_type: Valid_Characters diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml index fab14ff1..f1e46ae8 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_Month.yaml @@ -36,6 +36,7 @@ test_types: result_visualization_params: null usage_notes: null active: N + algorithm: Pattern / regex cat_test_conditions: - id: '7033' test_type: Valid_Month diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml index e380caef..b60bf716 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: null default_severity: Warning run_type: CAT @@ -35,6 +36,7 @@ test_types: result_visualization_params: null usage_notes: null active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7034' test_type: Valid_US_Zip diff --git a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml index 45218af9..756aefe1 100644 --- a/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Valid_US_Zip3.yaml @@ -21,7 +21,8 @@ test_types: default_parm_columns: threshold_value default_parm_values: |- 0 - default_parm_prompts: null + default_parm_prompts: |- + Threshold Invalid Value Count default_parm_help: null default_severity: Warning run_type: CAT @@ -36,6 +37,7 @@ test_types: usage_notes: |- This test looks for the presence of values that fail to match the three-digit numeric code expected for US Zip Code regional prefixes. These prefixes are often used to roll up Zip Code data to a regional level, and may be critical to anonymize detailed data and protect PID. Depending on your needs and regulatory requirements, longer zip codes could place PID at risk. active: Y + algorithm: Pattern / regex cat_test_conditions: - id: '7035' test_type: Valid_US_Zip3 diff --git a/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml b/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml index bb671fd8..4ed8791f 100644 --- a/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Variability_Decrease.yaml @@ -41,6 +41,8 @@ test_types: \ process, better precision in measurement, the elimination of outliers, or a\ \ more homogeneous cohort. " active: Y + algorithm: Statistical drift + statistical_technique: SD Shift cat_test_conditions: - id: '7032' test_type: Variability_Decrease diff --git a/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml b/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml index 54e11245..a0072b37 100644 --- a/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Variability_Increase.yaml @@ -45,6 +45,8 @@ test_types: \ that should be noted and assessed by business users. If the average does not\ \ shift, this may point to a data quality or data collection problem. " active: Y + algorithm: Statistical drift + statistical_technique: SD Shift cat_test_conditions: - id: '7031' test_type: Variability_Increase diff --git a/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml b/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml index 67c9ff29..0565a1e3 100644 --- a/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Volume_Trend.yaml @@ -27,6 +27,7 @@ test_types: dq_dimension: Completeness impact_dimension: Reliability health_dimension: Volume + statistical_technique: Predictive Model threshold_description: |- Expected row count range. result_visualization: line_chart @@ -34,6 +35,7 @@ test_types: usage_notes: |- This test compares the row count of all or a subset of records in a table against a derived tolerance range. active: Y + algorithm: Custom SQL cat_test_conditions: - id: '2515' test_type: Volume_Trend diff --git a/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml b/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml index bf0e91df..74be242d 100644 --- a/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml +++ b/testgen/template/dbsetup_test_types/test_types_Weekly_Rec_Ct.yaml @@ -37,6 +37,7 @@ test_types: usage_notes: |- Weekly Records tests that at least one record is present for every calendar week within the minimum and maximum date range for the column. The test is relevant for transactional data, where you would expect at least one transaction to be recorded each week. A failure here would suggest missing records for the number of weeks identified without data. You can adjust the threshold to accept a number of weeks that you know legitimately have no records. active: Y + algorithm: Counting cat_test_conditions: - id: '7030' test_type: Weekly_Rec_Ct diff --git a/testgen/template/dbupgrade/0193_incremental_upgrade.sql b/testgen/template/dbupgrade/0193_incremental_upgrade.sql new file mode 100644 index 00000000..29206eb2 --- /dev/null +++ b/testgen/template/dbupgrade/0193_incremental_upgrade.sql @@ -0,0 +1,7 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE data_table_chars + ADD COLUMN IF NOT EXISTS object_type VARCHAR(20); + +ALTER TABLE stg_data_chars_updates + ADD COLUMN IF NOT EXISTS object_type VARCHAR(20); diff --git a/testgen/template/dbupgrade/0194_incremental_upgrade.sql b/testgen/template/dbupgrade/0194_incremental_upgrade.sql new file mode 100644 index 00000000..683afc8c --- /dev/null +++ b/testgen/template/dbupgrade/0194_incremental_upgrade.sql @@ -0,0 +1,34 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Schema-name matching in the SQL Server DDF query is now case-sensitive, +-- consistent with TestGen's case-sensitive joins and the other flavors. +-- A case-insensitive source previously accepted a table +-- group whose configured schema name differed in case from the database; +-- such a group would now match no tables on its next profiling run. +-- +-- data_table_chars.schema_name holds the case the source database actually +-- reported (it is populated from the DDF's c.table_schema, not the entered +-- value), so realign table_group_schema to that case. Guards: +-- * only mssql connections -- the DDF change was made only for that flavor, +-- * only when the catalog reports a single unambiguous schema for the group, +-- * only a case-only difference (LOWER() match), never a different schema, +-- * '<>' is case-sensitive in PostgreSQL, so correct rows are left untouched. +-- After this, the group's next profiling run stamps profile_results with the +-- corrected case, realigning the case-sensitive downstream joins. +UPDATE table_groups tg +SET table_group_schema = actual.schema_name +FROM ( + SELECT table_groups_id, MIN(schema_name) AS schema_name + FROM data_table_chars + WHERE drop_date IS NULL + GROUP BY table_groups_id + HAVING COUNT(DISTINCT schema_name) = 1 +) actual +WHERE actual.table_groups_id = tg.id + AND tg.table_group_schema <> actual.schema_name + AND LOWER(tg.table_group_schema) = LOWER(actual.schema_name) + AND EXISTS ( + SELECT 1 FROM connections cn + WHERE cn.connection_id = tg.connection_id + AND cn.sql_flavor = 'mssql' + ); diff --git a/testgen/template/dbupgrade/0195_incremental_upgrade.sql b/testgen/template/dbupgrade/0195_incremental_upgrade.sql new file mode 100644 index 00000000..0d89f265 --- /dev/null +++ b/testgen/template/dbupgrade/0195_incremental_upgrade.sql @@ -0,0 +1,150 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Deduplicate the run tables against job_executions. +-- +-- The run primary key becomes the job execution id: job_execution_id is +-- promoted to the primary key (which is also a foreign key to +-- job_executions.id with cascade on delete), and the duplicated status, +-- end-time, log-message, and process-id columns are dropped. External +-- references are rewritten from the old run id to the job execution id. +-- Because job_execution_id already holds that value on every row, the +-- run identity is reshaped with pure DDL rather than rewriting primary +-- key values in place. + +-- Standard views join the run tables on their primary key, so they must be +-- dropped before the id column can be reshaped. _refresh_static_metadata +-- recreates them from 060_create_standard_views.sql after the upgrade. +DROP VIEW IF EXISTS v_latest_profile_results CASCADE; +DROP VIEW IF EXISTS v_inactive_anomalies CASCADE; +DROP VIEW IF EXISTS v_queued_observability_results CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_latest_by_impact_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_latest_by_impact_dimension CASCADE; +DROP VIEW IF EXISTS v_dq_profile_scoring_history_by_column CASCADE; +DROP VIEW IF EXISTS v_dq_test_scoring_history_by_column CASCADE; + +-- Remove result rows orphaned by an already-deleted run. They are +-- unreachable in the application and would block the value rewrite and the +-- new cascade constraints below. +DO $$ +DECLARE + orphan_ct BIGINT; +BEGIN + DELETE FROM test_results c + WHERE NOT EXISTS (SELECT 1 FROM test_runs r WHERE r.id = c.test_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan test_results row(s)', orphan_ct; END IF; + + DELETE FROM profile_results c + WHERE NOT EXISTS (SELECT 1 FROM profiling_runs r WHERE r.id = c.profile_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan profile_results row(s)', orphan_ct; END IF; + + DELETE FROM profile_anomaly_results c + WHERE NOT EXISTS (SELECT 1 FROM profiling_runs r WHERE r.id = c.profile_run_id); + GET DIAGNOSTICS orphan_ct = ROW_COUNT; + IF orphan_ct > 0 THEN RAISE NOTICE 'deleted % orphan profile_anomaly_results row(s)', orphan_ct; END IF; +END $$; + +-- Rewrite result-row references from the old run id to the job execution id. +UPDATE test_results c + SET test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = c.test_run_id; + +UPDATE profile_results c + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = c.profile_run_id; + +UPDATE profile_anomaly_results c + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = c.profile_run_id; + +-- Repoint the cache/reference pointers to the job execution id. These are +-- not foreign keys; rows pointing at an already-deleted run are left as-is +-- (a dangling cache that resolves to no run, exactly as before). +UPDATE table_groups tg + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = tg.last_complete_profile_run_id; + +UPDATE test_suites ts + SET last_complete_test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = ts.last_complete_test_run_id; + +UPDATE data_table_chars dtc + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = dtc.last_complete_profile_run_id; + +UPDATE data_column_chars dcc + SET last_complete_profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = dcc.last_complete_profile_run_id; + +UPDATE score_history_latest_runs sr + SET last_test_run_id = r.job_execution_id + FROM test_runs r + WHERE r.id = sr.last_test_run_id; + +UPDATE score_history_latest_runs sr + SET last_profiling_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = sr.last_profiling_run_id; + +UPDATE test_definitions td + SET profile_run_id = r.job_execution_id + FROM profiling_runs r + WHERE r.id = td.profile_run_id; + +-- Promote job_execution_id to the primary key on profiling_runs. The column +-- already holds the job execution id, so this is pure DDL. +ALTER TABLE profiling_runs DROP CONSTRAINT pk_prun_id; +ALTER TABLE profiling_runs DROP COLUMN id; +ALTER TABLE profiling_runs RENAME COLUMN job_execution_id TO id; +ALTER TABLE profiling_runs ADD CONSTRAINT pk_prun_id PRIMARY KEY (id); +ALTER TABLE profiling_runs + ADD CONSTRAINT profiling_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +-- Promote job_execution_id to the primary key on test_runs. +ALTER TABLE test_runs DROP CONSTRAINT test_runs_id_pk; +ALTER TABLE test_runs DROP COLUMN id; +ALTER TABLE test_runs RENAME COLUMN job_execution_id TO id; +ALTER TABLE test_runs ADD CONSTRAINT test_runs_id_pk PRIMARY KEY (id); +ALTER TABLE test_runs + ADD CONSTRAINT test_runs_job_executions_id_fk + FOREIGN KEY (id) REFERENCES job_executions (id) ON DELETE CASCADE; + +-- Make the result-row references real cascade foreign keys so deleting a job +-- execution removes its run and the run's results in one DB-level cascade. +ALTER TABLE test_results + ADD CONSTRAINT test_results_test_runs_id_fk + FOREIGN KEY (test_run_id) REFERENCES test_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_results + ADD CONSTRAINT profile_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +ALTER TABLE profile_anomaly_results + ADD CONSTRAINT profile_anomaly_results_profiling_runs_id_fk + FOREIGN KEY (profile_run_id) REFERENCES profiling_runs (id) ON DELETE CASCADE; + +-- Drop the columns now owned by job_executions. +ALTER TABLE test_runs + DROP COLUMN status, + DROP COLUMN test_endtime, + DROP COLUMN log_message, + DROP COLUMN process_id; + +ALTER TABLE profiling_runs + DROP COLUMN status, + DROP COLUMN profiling_endtime, + DROP COLUMN log_message, + DROP COLUMN process_id; diff --git a/testgen/template/dbupgrade/0196_incremental_upgrade.sql b/testgen/template/dbupgrade/0196_incremental_upgrade.sql new file mode 100644 index 00000000..7b712aec --- /dev/null +++ b/testgen/template/dbupgrade/0196_incremental_upgrade.sql @@ -0,0 +1,5 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE test_types + ADD COLUMN IF NOT EXISTS algorithm VARCHAR(64), + ADD COLUMN IF NOT EXISTS statistical_technique VARCHAR(64); diff --git a/testgen/template/dbupgrade/0197_incremental_upgrade.sql b/testgen/template/dbupgrade/0197_incremental_upgrade.sql new file mode 100644 index 00000000..b8a37d41 --- /dev/null +++ b/testgen/template/dbupgrade/0197_incremental_upgrade.sql @@ -0,0 +1,5 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE test_definitions + ADD COLUMN IF NOT EXISTS external_url VARCHAR, + ADD COLUMN IF NOT EXISTS custom_metadata JSONB; diff --git a/testgen/template/dbupgrade/0198_incremental_upgrade.sql b/testgen/template/dbupgrade/0198_incremental_upgrade.sql new file mode 100644 index 00000000..1677eb3e --- /dev/null +++ b/testgen/template/dbupgrade/0198_incremental_upgrade.sql @@ -0,0 +1,15 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +-- Personal access tokens are stored as oauth2_tokens rows; the name labels each +-- one in the user's token list. NULL marks a non-PAT token (auth-code / MCP flows). +ALTER TABLE oauth2_tokens ADD COLUMN name VARCHAR(255); + +-- Classify clients so personal-access-token clients are distinguishable from +-- dynamically-registered external clients (MCP apps, automation). Existing rows +-- all predate PATs and are externally registered, so they default to 'external'. +ALTER TABLE oauth2_clients ADD COLUMN client_type VARCHAR(20) NOT NULL DEFAULT 'external'; + +-- Drop the client owner: it was only ever read by the client-credentials grant +-- (removed). User identity now rides the token (oauth2_tokens.user_id + the JWT +-- username), so a client needs no owner. The FK constraint drops with the column. +ALTER TABLE oauth2_clients DROP COLUMN user_id; diff --git a/testgen/template/dbupgrade/0199_incremental_upgrade.sql b/testgen/template/dbupgrade/0199_incremental_upgrade.sql new file mode 100644 index 00000000..43aaf9c9 --- /dev/null +++ b/testgen/template/dbupgrade/0199_incremental_upgrade.sql @@ -0,0 +1,6 @@ +SET SEARCH_PATH TO {SCHEMA_NAME}; + +ALTER TABLE table_groups ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE data_table_chars ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE data_column_chars ADD COLUMN IF NOT EXISTS data_classification VARCHAR(40); +ALTER TABLE score_definition_results_breakdown ADD COLUMN IF NOT EXISTS data_classification TEXT DEFAULT NULL; diff --git a/testgen/template/execution/get_errored_autogen_monitors.sql b/testgen/template/execution/get_errored_autogen_monitors.sql index dd0fb0a5..7d54a266 100644 --- a/testgen/template/execution/get_errored_autogen_monitors.sql +++ b/testgen/template/execution/get_errored_autogen_monitors.sql @@ -1,9 +1,10 @@ WITH prev_run AS ( - SELECT id + SELECT test_runs.id FROM test_runs + INNER JOIN job_executions ON job_executions.id = test_runs.id WHERE test_suite_id = :TEST_SUITE_ID ::UUID - AND id <> :TEST_RUN_ID ::UUID - AND status = 'Complete' + AND test_runs.id <> :TEST_RUN_ID ::UUID + AND job_executions.status = 'completed' ORDER BY test_starttime DESC LIMIT 1 ) diff --git a/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql b/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql index ee2165d3..f2a6e775 100644 --- a/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/bigquery/data_chars/get_schema_ddf.sql @@ -20,8 +20,15 @@ SELECT ELSE 'X' END AS general_type, REGEXP_CONTAINS(LOWER(c.data_type), r'(decimal|numeric|bignumeric)') AS is_decimal, + CASE it.table_type + WHEN 'BASE TABLE' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, t.row_count AS approx_record_ct FROM `{DATA_SCHEMA}.INFORMATION_SCHEMA.COLUMNS` c LEFT JOIN `{DATA_SCHEMA}.__TABLES__` t ON c.table_name = t.table_id + LEFT JOIN `{DATA_SCHEMA}.INFORMATION_SCHEMA.TABLES` it ON c.table_name = it.table_name WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql b/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql index 6ae63a94..93a066ce 100644 --- a/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/databricks/data_chars/get_schema_ddf.sql @@ -22,7 +22,15 @@ SELECT WHEN c.numeric_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE t.table_type + WHEN 'MANAGED' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED_VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL' THEN 'EXTERNAL' + WHEN 'FOREIGN' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, NULL AS approx_record_ct -- table statistics unavailable FROM information_schema.columns c + LEFT JOIN information_schema.tables t ON c.table_schema = t.table_schema AND c.table_name = t.table_name WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql index 8a3ea74f..8fbda543 100644 --- a/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/mssql/data_chars/get_schema_ddf.sql @@ -49,8 +49,10 @@ SELECT 'X' END AS general_type, CASE WHEN c.numeric_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE it.table_type WHEN 'BASE TABLE' THEN 'TABLE' WHEN 'VIEW' THEN 'VIEW' ELSE 'OTHER' END AS object_type, a.approx_record_ct AS approx_record_ct FROM information_schema.columns c LEFT JOIN approx_cts a ON c.table_schema = a.schema_name AND c.table_name = a.table_name -WHERE c.table_schema = '{DATA_SCHEMA}' {TABLE_CRITERIA} + LEFT JOIN information_schema.tables it ON c.table_schema = it.table_schema AND c.table_name = it.table_name +WHERE c.table_schema = '{DATA_SCHEMA}' COLLATE Latin1_General_BIN {TABLE_CRITERIA} ORDER BY c.table_schema, c.table_name, c.ordinal_position; diff --git a/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql b/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql index d4f4f578..98acbc76 100644 --- a/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/oracle/data_chars/get_schema_ddf.sql @@ -33,6 +33,9 @@ SELECT WHEN c.data_type = 'NUMBER' AND c.data_scale > 0 THEN 1 ELSE 0 END AS is_decimal, + CASE WHEN t.table_name IS NOT NULL THEN 'TABLE' + WHEN EXISTS (SELECT 1 FROM all_views v WHERE v.owner = c.owner AND v.view_name = c.table_name) THEN 'VIEW' + ELSE 'OTHER' END AS object_type, t.num_rows AS approx_record_ct FROM all_tab_columns c LEFT JOIN all_tables t ON c.owner = t.owner AND c.table_name = t.table_name diff --git a/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql b/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql index b5fcc322..14754855 100644 --- a/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/postgresql/data_chars/get_schema_ddf.sql @@ -44,6 +44,10 @@ SELECT WHEN c.data_type = 'numeric' THEN COALESCE(numeric_scale, 1) > 0 ELSE numeric_scale > 0 END as is_decimal, + CASE p.relkind + WHEN 'r' THEN 'TABLE' WHEN 'p' THEN 'TABLE' + WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED_VIEW' WHEN 'f' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, NULLIF(p.reltuples::BIGINT, -1) AS approx_record_ct FROM information_schema.columns c LEFT JOIN pg_namespace n ON c.table_schema = n.nspname diff --git a/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql b/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql index 6bda34ec..dfd4afea 100644 --- a/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/redshift/data_chars/get_schema_ddf.sql @@ -39,6 +39,10 @@ SELECT WHEN c.data_type = 'numeric' THEN COALESCE(numeric_scale, 1) > 0 ELSE numeric_scale > 0 END AS is_decimal, + CASE p.relkind + WHEN 'r' THEN 'TABLE' WHEN 'p' THEN 'TABLE' + WHEN 'v' THEN 'VIEW' WHEN 'm' THEN 'MATERIALIZED_VIEW' + ELSE 'OTHER' END AS object_type, CASE WHEN reltuples > 0 AND reltuples < 1 THEN NULL ELSE reltuples::BIGINT diff --git a/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql b/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql index 3a6669f3..51460582 100644 --- a/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/redshift_spectrum/data_chars/get_schema_ddf.sql @@ -25,6 +25,7 @@ SELECT THEN 1 ELSE 0 END AS is_decimal, + 'EXTERNAL' AS object_type, NULL AS approx_record_ct -- Table statistics unavailable FROM svv_external_columns c WHERE c.schemaname = '{DATA_SCHEMA}' diff --git a/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql b/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql index 8a0838fd..ebb4e6f4 100644 --- a/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/sap_hana/data_chars/get_schema_ddf.sql @@ -1,3 +1,14 @@ +WITH relations AS ( + SELECT SCHEMA_NAME, TABLE_NAME, COLUMN_NAME, DATA_TYPE_NAME, LENGTH, SCALE, POSITION, + 'TABLE' AS object_type + FROM SYS.TABLE_COLUMNS + WHERE SCHEMA_NAME = '{DATA_SCHEMA}' + UNION ALL + SELECT SCHEMA_NAME, VIEW_NAME AS TABLE_NAME, COLUMN_NAME, DATA_TYPE_NAME, LENGTH, SCALE, POSITION, + 'VIEW' AS object_type + FROM SYS.VIEW_COLUMNS + WHERE SCHEMA_NAME = '{DATA_SCHEMA}' +) SELECT c.SCHEMA_NAME AS schema_name, c.TABLE_NAME AS table_name, @@ -34,8 +45,9 @@ SELECT WHEN c.DATA_TYPE_NAME IN ('DOUBLE', 'REAL', 'SMALLDECIMAL') THEN 1 ELSE 0 END AS is_decimal, + c.object_type AS object_type, t.RECORD_COUNT AS approx_record_ct -FROM SYS.TABLE_COLUMNS c +FROM relations c LEFT JOIN SYS.M_TABLES t ON c.SCHEMA_NAME = t.SCHEMA_NAME AND c.TABLE_NAME = t.TABLE_NAME -WHERE c.SCHEMA_NAME = '{DATA_SCHEMA}' {TABLE_CRITERIA} +WHERE 1=1 {TABLE_CRITERIA} ORDER BY c.SCHEMA_NAME, c.TABLE_NAME, c.POSITION diff --git a/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql b/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql index 54940da8..3062875d 100644 --- a/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql +++ b/testgen/template/flavors/snowflake/data_chars/get_schema_ddf.sql @@ -42,6 +42,12 @@ SELECT 'X' END AS general_type, numeric_scale > 0 AS is_decimal, + CASE t.table_type + WHEN 'BASE TABLE' THEN 'TABLE' + WHEN 'VIEW' THEN 'VIEW' + WHEN 'MATERIALIZED VIEW' THEN 'MATERIALIZED_VIEW' + WHEN 'EXTERNAL TABLE' THEN 'EXTERNAL' + ELSE 'OTHER' END AS object_type, t.row_count AS approx_record_ct FROM information_schema.columns c LEFT JOIN information_schema.tables t ON c.table_schema = t.table_schema AND c.table_name = t.table_name diff --git a/testgen/template/get_entities/get_profile_list.sql b/testgen/template/get_entities/get_profile_list.sql index 088bc745..8bbb83e0 100644 --- a/testgen/template/get_entities/get_profile_list.sql +++ b/testgen/template/get_entities/get_profile_list.sql @@ -5,13 +5,15 @@ Optional: table-name*/ SELECT p.id as profile_run_id, p.project_code as project_key, schema_name, p.table_groups_id, - profiling_starttime as start_time, status, + profiling_starttime as start_time, je.status, COUNT(DISTINCT table_name) as tables, COUNT(DISTINCT table_name || '.' || column_name) as columns FROM profiling_runs p INNER JOIN profile_results r ON (p.id = r.profile_run_id) +INNER JOIN job_executions je + ON (je.id = p.id) WHERE p.table_groups_id = :TABLE_GROUP_ID GROUP BY p.id, p.project_code, p.connection_id, schema_name, p.table_groups_id, - profiling_starttime, status + profiling_starttime, je.status ORDER BY profiling_starttime DESC; diff --git a/testgen/template/get_entities/get_test_run_list.sql b/testgen/template/get_entities/get_test_run_list.sql index 50f9ecc7..ea65632e 100644 --- a/testgen/template/get_entities/get_test_run_list.sql +++ b/testgen/template/get_entities/get_test_run_list.sql @@ -5,16 +5,16 @@ Optional: table-name, column-name, from-date, thru-date*/ SELECT ts.test_suite as test_suite_key, tr.test_starttime as test_time, - tr.status, + je.status, tr.id::VARCHAR as test_run_id, COUNT(DISTINCT lower(r.schema_name || '.' || table_name)) as table_ct, COUNT(*) as result_ct, SUM(CASE WHEN r.result_code = 0 THEN 1 END) as fail_ct, - SUM(CASE WHEN r.observability_status = 'Sent' THEN 1 END) as sent_to_obs, - process_id + SUM(CASE WHEN r.observability_status = 'Sent' THEN 1 END) as sent_to_obs FROM test_runs tr INNER JOIN test_results r ON tr.id = r.test_run_id INNER JOIN test_suites ts ON tr.test_suite_id = ts.id +INNER JOIN job_executions je ON je.id = tr.id WHERE ts.project_code = :PROJECT_CODE AND ts.test_suite = :TEST_SUITE AND ts.is_monitor IS NOT TRUE @@ -22,4 +22,4 @@ INNER JOIN test_suites ts ON tr.test_suite_id = ts.id ts.project_code, ts.test_suite, tr.test_starttime, - tr.status; + je.status; diff --git a/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql b/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql index ccd67104..f1d69599 100644 --- a/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql +++ b/testgen/template/rollup_scores/rollup_scores_profile_table_group.sql @@ -2,7 +2,8 @@ WITH last_profile_date AS (SELECT table_groups_id, MAX(profiling_starttime) as last_profile_run_date FROM profiling_runs - WHERE status = 'Complete' + INNER JOIN job_executions ON job_executions.id = profiling_runs.id + WHERE job_executions.status = 'completed' GROUP BY table_groups_id), score_calc AS (SELECT run.table_groups_id, run.id as profile_run_id, diff --git a/testgen/template/rollup_scores/rollup_scores_test_table_group.sql b/testgen/template/rollup_scores/rollup_scores_test_table_group.sql index 6009e5d0..dfe6d7b9 100644 --- a/testgen/template/rollup_scores/rollup_scores_test_table_group.sql +++ b/testgen/template/rollup_scores/rollup_scores_test_table_group.sql @@ -2,7 +2,8 @@ WITH last_test_date AS (SELECT r.test_suite_id, MAX(r.test_starttime) as last_test_run_date FROM test_runs r - WHERE r.status = 'Complete' + INNER JOIN job_executions je ON je.id = r.id + WHERE je.status = 'completed' GROUP BY r.test_suite_id), score_calc AS (SELECT ts.table_groups_id, diff --git a/testgen/template/score_cards/add_latest_runs.sql b/testgen/template/score_cards/add_latest_runs.sql index bf332d8c..4403cb8a 100644 --- a/testgen/template/score_cards/add_latest_runs.sql +++ b/testgen/template/score_cards/add_latest_runs.sql @@ -1,11 +1,12 @@ -- Insert latest profiling runs as of cutoff WITH ranked_profiling - AS (SELECT project_code, table_groups_id, id as profiling_run_id, + AS (SELECT r.project_code, table_groups_id, r.id as profiling_run_id, ROW_NUMBER() OVER (PARTITION BY table_groups_id ORDER BY profiling_starttime DESC) as rank FROM profiling_runs r - WHERE project_code = :project_code + INNER JOIN job_executions je ON je.id = r.id + WHERE r.project_code = :project_code AND profiling_starttime <= :score_history_cutoff_time - AND r.status = 'Complete') + AND je.status = 'completed') INSERT INTO score_history_latest_runs (definition_id, score_history_cutoff_time, table_groups_id, last_profiling_run_id) SELECT :definition_id as definition_id, :score_history_cutoff_time as score_history_cutoff_time, table_groups_id, profiling_run_id @@ -20,9 +21,11 @@ WITH ranked_test_runs FROM test_runs r INNER JOIN test_suites s ON (r.test_suite_id = s.id) + INNER JOIN job_executions je + ON (je.id = r.id) WHERE s.project_code = :project_code AND r.test_starttime <= :score_history_cutoff_time - AND r.status = 'Complete') + AND je.status = 'completed') INSERT INTO score_history_latest_runs (definition_id, score_history_cutoff_time, test_suite_id, last_test_run_id) SELECT :definition_id as definition_id, :score_history_cutoff_time as score_history_cutoff_time, test_suite_id, test_run_id diff --git a/testgen/testing/__init__.py b/testgen/testing/__init__.py new file mode 100644 index 00000000..36919a40 --- /dev/null +++ b/testgen/testing/__init__.py @@ -0,0 +1,7 @@ +"""Testing support importable by the core test suite and any plugin test suite. + +Fixtures live in ``testgen.testing.fixtures``; import the ones you need into a +``conftest.py`` to register them (pytest registers fixtures that are merely imported +into a conftest). An importable module is the only way to share fixtures across separate +test trees — sibling test directories do not share conftest fixtures. +""" diff --git a/testgen/testing/fixtures.py b/testgen/testing/fixtures.py new file mode 100644 index 00000000..0ae02560 --- /dev/null +++ b/testgen/testing/fixtures.py @@ -0,0 +1,76 @@ +"""Reusable pytest fixtures for unit-testing TestGen. + +Importable by the core test suite and any plugin test suite. To use, import the fixtures +you need into a ``conftest.py`` — pytest registers fixtures that are imported into a +conftest, preserving their ``autouse`` flag for that subtree:: + + from testgen.testing.fixtures import db_session_mock, mcp_user # noqa: F401 +""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.mcp.permissions import set_mcp_token, set_mcp_username + +# Fictional role matrix for tests. role_a has full access (but NOT view_pii — several +# tests rely on that to exercise the no-view_pii path), role_c is restricted, and +# role_d holds edit + view_pii so deny/allow pairs can be distinguished against a real +# ProjectPermissions without role_a accidentally granting view_pii. +TEST_PERM_MATRIX = { + "view": ["role_a", "role_b"], + "catalog": ["role_a", "role_b", "role_c"], + "edit": ["role_a", "role_d"], + "administer": ["role_a"], + "view_pii": ["role_d"], +} + + +def _test_roles_with_permission(permission): + return TEST_PERM_MATRIX.get(permission, []) + + +@pytest.fixture(autouse=True) +def patched_settings(): + with patch("testgen.settings.UI_BASE_URL", "http://tg-base-url"): + yield + + +@pytest.fixture +def db_session_mock(): + with patch("testgen.common.models.Session") as factory_mock: + yield factory_mock().__enter__() + + +@pytest.fixture(autouse=True) +def mcp_user(): + """Set up an authenticated MCP user for all tool tests. + + Default: user has 'role_a' on 'demo' project (full access). + The @mcp_permission decorator passes for any permission. + + Tests needing scoped access patch _compute_project_permissions directly. + """ + set_mcp_username("test_user") + set_mcp_token("test_bearer_token") + user = MagicMock() + user.id = uuid4() + + membership = MagicMock() + membership.project_code = "demo" + membership.role = "role_a" + + with ( + patch("testgen.common.auth.authorize_token", return_value=user), + patch("testgen.common.models.get_current_session", return_value=MagicMock()), + patch("testgen.mcp.permissions.ProjectMembership") as mock_membership, + patch("testgen.mcp.permissions.PluginHook") as mock_hook, + ): + mock_membership.get_memberships_for_user.return_value = [membership] + mock_hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( + _test_roles_with_permission + ) + yield user + set_mcp_username(None) + set_mcp_token(None) diff --git a/testgen/ui/app.py b/testgen/ui/app.py index 5bbd67eb..dc26e6cf 100644 --- a/testgen/ui/app.py +++ b/testgen/ui/app.py @@ -84,6 +84,7 @@ def render(log_level: int = logging.INFO): support_email=settings.SUPPORT_EMAIL, global_context=is_global_context, is_global_admin=session.auth.user_has_permission("global_admin") and bool(application.global_admin_paths), + account_path=application.account_path, ) application.router.run() diff --git a/testgen/ui/auth.py b/testgen/ui/auth.py index 8bb5b788..d86fd154 100644 --- a/testgen/ui/auth.py +++ b/testgen/ui/auth.py @@ -12,6 +12,7 @@ from testgen.ui.services.query_cache import ( get_membership_by_user_and_project, get_user, + select_projects_where, select_users_where, ) from testgen.ui.session import session @@ -55,6 +56,9 @@ def user_has_permission(self, permission: Permission, /, project_code: str | Non def user_has_project_access(self, project_code: str) -> bool: # noqa: ARG002 return True + def get_projects_with_permission(self, permission: Permission, /) -> list[str]: # noqa: ARG002 + return [p.project_code for p in select_projects_where()] + def get_jwt_hashing_key(self) -> bytes: try: return get_jwt_signing_key() diff --git a/testgen/ui/bootstrap.py b/testgen/ui/bootstrap.py index bea240d5..d5e5db54 100644 --- a/testgen/ui/bootstrap.py +++ b/testgen/ui/bootstrap.py @@ -49,13 +49,14 @@ class Application(singleton.Singleton): - def __init__(self, auth_class: Authentication, logo: plugins.Logo, router: Router, menu: Menu, logger: logging.Logger, global_admin_paths: frozenset[str]) -> None: + def __init__(self, auth_class: Authentication, logo: plugins.Logo, router: Router, menu: Menu, logger: logging.Logger, global_admin_paths: frozenset[str], account_path: str | None) -> None: self.auth_class = auth_class self.logo = logo self.router = router self.menu = menu self.logger = logger self.global_admin_paths = global_admin_paths + self.account_path = account_path def run(log_level: int = logging.INFO) -> Application: @@ -91,4 +92,5 @@ def run(log_level: int = logging.INFO) -> Application: ), logger=LOG, global_admin_paths=frozenset(page.path for page in pages if page.permission == "global_admin"), + account_path=next((page.path for page in pages if page.is_account_page), None), ) diff --git a/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js b/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js index 8689fefc..6aa6a46e 100644 --- a/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js +++ b/testgen/ui/components/frontend/js/data_profiling/data_characteristics.js @@ -36,6 +36,9 @@ const DataCharacteristicsCard = (/** @type Properties */ props, /** @type Column ); } } else { + if (item.object_type) { + attributes.push({ key: 'object_type', label: 'Object Type' }); + } attributes.push( { key: 'functional_table_type', label: `Semantic Table Type ${item.is_latest_profile ? '*' : ''}` }, ); @@ -77,6 +80,10 @@ const DataCharacteristicsCard = (/** @type Properties */ props, /** @type Column ); } else if (key === 'datatype_suggestion') { value = (value || '').toLowerCase(); + } else if (key === 'object_type') { + value = (value || '').split('_') + .map(word => word ? (word[0].toUpperCase() + word.substring(1).toLowerCase()) : '') + .join(' '); } else if (key === 'functional_table_type') { value = (value || '').split('-') .map(word => word ? (word[0].toUpperCase() + word.substring(1)) : '') diff --git a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js index 39342edc..d44255ec 100644 --- a/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js +++ b/testgen/ui/components/frontend/js/data_profiling/data_profiling_utils.js @@ -55,6 +55,7 @@ * @property {string?} transform_level * @property {string?} aggregation_level * @property {string?} data_product + * @property {string?} data_classification * * Table Tags * @property {boolean?} table_critical_data_element * @property {string?} table_data_source @@ -65,6 +66,7 @@ * @property {string?} table_transform_level * @property {string?} table_aggregation_level * @property {string?} table_data_product + * @property {string?} table_data_classification * * Table Group Tags * @property {string} table_group_data_source * @property {string} table_group_source_system @@ -73,6 +75,7 @@ * @property {string} table_group_stakeholder_group * @property {string} table_group_transform_level * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * Profile & Test Runs * @property {string?} profile_run_id * @property {number?} profile_run_date @@ -150,6 +153,7 @@ * @property {string} connection_id * @property {string} project_code * * Characteristics + * @property {string} object_type * @property {string} functional_table_type * @property {number} approx_record_ct * @property {number} record_ct @@ -168,6 +172,7 @@ * @property {string} transform_level * @property {string} aggregation_level * @property {string} data_product + * @property {string} data_classification * * Table Group Tags * @property {string} table_group_data_source * @property {string} table_group_source_system @@ -176,6 +181,7 @@ * @property {string} table_group_stakeholder_group * @property {string} table_group_transform_level * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * Profile & Test Runs * @property {string} profile_run_id * @property {number} profile_run_date diff --git a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js index fbb9ad8f..c469b0e6 100644 --- a/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js +++ b/testgen/ui/components/frontend/js/data_profiling/metadata_tags.js @@ -61,6 +61,7 @@ const TAG_KEYS = [ 'transform_level', 'aggregation_level', 'data_product', + 'data_classification', ]; const TAG_HELP = { data_source: 'Original source of the dataset', @@ -71,6 +72,7 @@ const TAG_HELP = { transform_level: 'Data warehouse processing stage, e.g., Raw, Conformed, Processed, Reporting, or Medallion level (bronze, silver, gold)', aggregation_level: 'Data granularity of the dataset, e.g. atomic, historical, snapshot, aggregated, time-rollup, rolling, summary', data_product: 'Data domain that comprises the dataset', + data_classification: 'Information sensitivity level of the dataset, e.g., Public, Internal, Confidential, Restricted', }; /** @@ -98,6 +100,7 @@ const MetadataTagsCard = (props, item) => { help: TAG_HELP[key], label: key === 'pii_flag' ? 'PII Data' : capitalize(key.replaceAll('_', ' ')), state: van.state(value), + initialValue: value, inheritTableGroup: item[`table_group_${key}`] ?? null, // Table group values inherited by table or column inheritTable: item[`table_${key}`] ?? null, // Table values inherited by column }; @@ -196,13 +199,15 @@ const MetadataTagsCard = (props, item) => { content, editingContent, onSave: () => { const items = [{ type: item.type, id: item.id }]; - const tags = attributes.reduce((object, { key, state }) => { - object[key] = state.rawVal; + const tags = attributes.reduce((object, { key, state, initialValue }) => { + if (state.rawVal !== initialValue) { + object[key] = state.rawVal; + } return object; }, {}); - warnCde.val = props.autoflagSettings.profile_flag_cdes && tags.critical_data_element !== item.critical_data_element; - warnPii.val = props.autoflagSettings.profile_flag_pii && tags.pii_flag !== item.pii_flag; + warnCde.val = props.autoflagSettings.profile_flag_cdes && 'critical_data_element' in tags; + warnPii.val = props.autoflagSettings.profile_flag_pii && 'pii_flag' in tags; if (warnCde.val || warnPii.val) { const disableFlags = []; @@ -219,8 +224,8 @@ const MetadataTagsCard = (props, item) => { } }, // Reset states to original values on cancel - onCancel: () => attributes.forEach(({ key, state }) => state.val = item[key]), - hasChanges: () => attributes.some(({ key, state }) => state.val !== item[key]), + onCancel: () => attributes.forEach(({ state, initialValue }) => state.val = initialValue), + hasChanges: () => attributes.some(({ state, initialValue }) => state.val !== initialValue), }), WarningDialog(warningDialogOpen, pendingSaveAction, warnCde, warnPii), ); diff --git a/testgen/ui/components/frontend/js/pages/data_catalog.js b/testgen/ui/components/frontend/js/pages/data_catalog.js index 3ed1c10c..92857742 100644 --- a/testgen/ui/components/frontend/js/pages/data_catalog.js +++ b/testgen/ui/components/frontend/js/pages/data_catalog.js @@ -29,6 +29,7 @@ * @property {string} transform_level * @property {string} aggregation_level * @property {string} data_product + * @property {string} data_classification * @property {string} table_data_source * @property {string} table_source_system * @property {string} table_source_process @@ -37,6 +38,15 @@ * @property {string} table_transform_level * @property {string} table_aggregation_level * @property {string} table_data_product + * @property {string} table_data_classification + * @property {string} table_group_data_source + * @property {string} table_group_source_system + * @property {string} table_group_source_process + * @property {string} table_group_business_domain + * @property {string} table_group_stakeholder_group + * @property {string} table_group_transform_level + * @property {string} table_group_data_product + * @property {string} table_group_data_classification * * @typedef Permissions * @type {object} @@ -130,7 +140,7 @@ const DataCatalog = (/** @type Properties */ props) => { criticalDataElement: !!item.table_critical_data_element, children: [], }; - TAG_KEYS.forEach(key => tables[table_id][key] = item[`table_${key}`]); + TAG_KEYS.forEach(key => tables[table_id][key] = item[`table_${key}`] ?? item[`table_group_${key}`]); } const columnNode = { id: column_id, @@ -156,7 +166,7 @@ const DataCatalog = (/** @type Properties */ props) => { excludedDataElement: !!item.excluded_data_element, piiFlag: !!item.pii_flag, }; - TAG_KEYS.forEach(key => columnNode[key] = item[key] ?? item[`table_${key}`]); + TAG_KEYS.forEach(key => columnNode[key] = item[key] ?? item[`table_${key}`] ?? item[`table_group_${key}`]); tables[table_id].children.push(columnNode); }); return Object.values(tables); diff --git a/testgen/ui/components/frontend/js/pages/monitors_dashboard.js b/testgen/ui/components/frontend/js/pages/monitors_dashboard.js index f3eb9612..2ac446fe 100644 --- a/testgen/ui/components/frontend/js/pages/monitors_dashboard.js +++ b/testgen/ui/components/frontend/js/pages/monitors_dashboard.js @@ -1,13 +1,13 @@ /** * @import { MonitorSummary } from '/app/static/js/components/monitor_anomalies_summary.js'; * @import { CronSample, FilterOption, ProjectSummary } from '../types.js'; - * + * * @typedef Schedule * @type {object} * @property {boolean} active * @property {string} cron_tz * @property {CronSample} cron_sample - * + * * @typedef Monitor * @type {object} * @property {string} table_group_id @@ -36,25 +36,25 @@ * @property {number?} column_adds * @property {number?} column_drops * @property {number?} column_mods - * + * * @typedef MonitorList * @type {object} * @property {Monitor[]} items * @property {number} current_page * @property {number} items_per_page * @property {number} total_count - * + * * @typedef MonitorListFilters * @type {object} * @property {string?} table_group_id * @property {string?} table_name_filter * @property {string?} anomaly_type_filter - * + * * @typedef MonitorListSort * @type {object} * @property {string?} sort_field * @property {('asc'|'desc')?} sort_order - * + * * @typedef Permissions * @type {object} * @property {boolean} can_edit @@ -213,17 +213,17 @@ const MonitorsDashboard = (/** @type Properties */ props) => { class: 'flex-row fx-gap-1 schema-changes', onclick: () => { const summary = getValue(props.summary); - emit('OpenSchemaChanges', { payload: { + emit('OpenSchemaChanges', { payload: { table_name: monitor.table_name, start_time: summary?.lookback_start, end_time: summary?.lookback_end, }}); }, }, - monitor.table_state === 'added' + monitor.table_state === 'added' ? Icon({size: 20, classes: 'schema-icon', filled: true}, 'add_box') : null, - monitor.table_state === 'dropped' + monitor.table_state === 'dropped' ? Icon({size: 20, classes: 'schema-icon', filled: true}, 'indeterminate_check_box') : null, monitor.column_adds ? div( @@ -245,7 +245,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { { text: div( {class: 'flex-column fx-align-flex-start'}, - monitor.table_state === 'added' + monitor.table_state === 'added' ? span({class: 'mb-1', style: 'font-size: 14px;'}, 'Table added.') : null, monitor.table_state === 'dropped' @@ -329,7 +329,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { tooltipPosition: 'bottom-left', color: 'basic', type: 'stroked', - style: 'background: var(--button-generic-background-color);', + style: 'background: var(--button-generic-background-color);', onclick: () => emit('EditNotifications', {}), }), Button({ @@ -338,7 +338,7 @@ const MonitorsDashboard = (/** @type Properties */ props) => { tooltipPosition: 'bottom-left', color: 'basic', type: 'stroked', - style: 'background: var(--button-generic-background-color);', + style: 'background: var(--button-generic-background-color);', onclick: () => emit('EditMonitorSettings', {}), }), Button({ @@ -517,14 +517,14 @@ const MonitorsDashboard = (/** @type Properties */ props) => { result: van.derive(() => getValue(props.notifications_dialog)?.result), onClose: () => emit('NotificationsDialogClosed', {}), }), - EditMonitorSettings({ emit, + EditMonitorSettings({ emit, table_group: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.table_group), schedule: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.schedule), monitor_suite: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.monitor_suite), cron_sample: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.cron_sample), dialog: van.derive(() => getValue(props.edit_monitor_settings_dialog)?.dialog), }), - TableMonitoringTrend({ emit, + TableMonitoringTrend({ emit, freshness_events: van.derive(() => getValue(props.trends_dialog)?.freshness_events ?? []), volume_events: van.derive(() => getValue(props.trends_dialog)?.volume_events ?? []), schema_events: van.derive(() => getValue(props.trends_dialog)?.schema_events ?? []), @@ -534,14 +534,14 @@ const MonitorsDashboard = (/** @type Properties */ props) => { extended_history: van.derive(() => getValue(props.trends_dialog)?.extended_history), dialog: van.derive(() => getValue(props.trends_dialog)?.dialog), }), - EditTableMonitors({ emit, + EditTableMonitors({ emit, table_name: van.derive(() => getValue(props.edit_table_monitors_dialog)?.table_name), definitions: van.derive(() => getValue(props.edit_table_monitors_dialog)?.definitions ?? []), metric_test_type: van.derive(() => getValue(props.edit_table_monitors_dialog)?.metric_test_type), result: van.derive(() => getValue(props.edit_table_monitors_dialog)?.result), dialog: van.derive(() => getValue(props.edit_table_monitors_dialog)?.dialog), }), - SchemaChangesDialog({ emit, + SchemaChangesDialog({ emit, window_start: van.derive(() => getValue(props.schema_changes_dialog)?.window_start), window_end: van.derive(() => getValue(props.schema_changes_dialog)?.window_end), data_structure_logs: van.derive(() => getValue(props.schema_changes_dialog)?.data_structure_logs), @@ -568,9 +568,6 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin const hasErrors = !!errorMessage; const content = van.derive(() => { - if (anomalies > 0) { - return span(anomalies); - } if (hasErrors) { return withTooltip( i({class: 'material-symbols-rounded'}, 'warning'), @@ -584,6 +581,9 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin }, ); } + if (anomalies > 0) { + return span(anomalies); + } if (isTraining) { return withTooltip( i({class: 'material-symbols-rounded'}, 'more_horiz'), @@ -597,7 +597,7 @@ const AnomalyTag = (anomalies, errorMessage = null, isTraining = false, isPendin { class: `anomaly-tag-wrapper flex-row p-1 ${onClick ? 'clickable' : ''}`, onclick: onClick }, div( { - class: `anomaly-tag ${anomalies > 0 ? 'has-anomalies' : ''} ${hasErrors ? 'has-errors' : ''} ${isTraining ? 'is-training' : ''}`, + class: `anomaly-tag ${hasErrors ? 'has-errors' : anomalies > 0 ? 'has-anomalies' : ''} ${isTraining ? 'is-training' : ''}`, }, content, ), @@ -646,8 +646,8 @@ const ConditionalEmptyState = (projectSummary, userCanEdit, emit) => { }, }; } - - return EmptyState({ emit, + + return EmptyState({ emit, icon: 'apps_outage', ...args, }); diff --git a/testgen/ui/components/frontend/js/pages/score_explorer.js b/testgen/ui/components/frontend/js/pages/score_explorer.js index addd28b3..7f8d8d12 100644 --- a/testgen/ui/components/frontend/js/pages/score_explorer.js +++ b/testgen/ui/components/frontend/js/pages/score_explorer.js @@ -82,6 +82,7 @@ const TRANSLATIONS = { dq_dimension: 'Quality Dimension', impact_dimension: 'Impact Dimension', data_product: 'Data Product', + data_classification: 'Data Classification', }; const ScoreExplorer = (/** @type {Properties} */ props) => { @@ -188,6 +189,7 @@ const Toolbar = ( 'dq_dimension', 'impact_dimension', 'data_product', + 'data_classification', ]; const filterableFields = categories.filter((c) => c !== 'dq_dimension' && c !== 'impact_dimension'); const filters = van.state(definition.filters.map((f, idx) => ({key: `${f.field}-${idx}-${getRandomId()}`, field: f.field, value: van.state(f.value), others: f.others ?? [] }))); diff --git a/testgen/ui/components/frontend/js/pages/table_group_list.js b/testgen/ui/components/frontend/js/pages/table_group_list.js index 3f52a24a..154010dd 100644 --- a/testgen/ui/components/frontend/js/pages/table_group_list.js +++ b/testgen/ui/components/frontend/js/pages/table_group_list.js @@ -120,6 +120,7 @@ const TableGroupList = (props) => { connections: van.derive(() => getValue(props.edit_dialog)?.connections), table_group: van.derive(() => getValue(props.edit_dialog)?.table_group), is_in_use: van.derive(() => getValue(props.edit_dialog)?.is_in_use), + can_view_pii: van.derive(() => getValue(props.permissions)?.can_view_pii), table_group_preview: van.derive(() => getValue(props.edit_dialog)?.table_group_preview), result: van.derive(() => getValue(props.edit_dialog)?.result), })); diff --git a/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js b/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js index 8e0c86cd..94180e2d 100644 --- a/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js +++ b/testgen/ui/components/frontend/js/pages/table_monitoring_trends.js @@ -123,18 +123,6 @@ const TableMonitoringTrend = (props) => { }, div( { class: '', style: 'width: 100%;' }, - () => { - const extendedHistory = getValue(props.extended_history) ?? false; - return div( - { class: 'extended-history-toggle' }, - Button({ - label: extendedHistory ? 'Show default view' : 'Show more history', - icon: extendedHistory ? 'history_toggle_off' : 'history', - width: 'auto', - onclick: () => emit('ToggleExtendedHistory', { payload: {} }), - }), - ); - }, () => { if (!getValue(props.dialog)?.open) return div(); return ChartsSection(props, { schemaChartSelection, getDataStructureLogs }); @@ -229,12 +217,22 @@ const TableMonitoringTrend = (props) => { ); const dialogTitle = van.derive(() => getValue(props.dialog)?.title ?? ''); + const historyToggle = () => { + const extendedHistory = getValue(props.extended_history) ?? false; + return Button({ + label: extendedHistory ? 'Show default view' : 'Show more history', + icon: extendedHistory ? 'history_toggle_off' : 'history', + width: 'auto', + onclick: () => emit('ToggleExtendedHistory', { payload: {} }), + }); + }; return Dialog( { title: dialogTitle, open: dialogOpen, onClose: () => { dialogOpen.val = false; emit('CloseTrendsDialog', {}); }, width: '75rem', + headerActions: historyToggle, }, content, ); @@ -825,13 +823,6 @@ stylesheet.replace(` position: relative; } - .extended-history-toggle { - position: absolute; - top: -70px; - right: 48px; - z-index: 1; - } - .table-monitoring-trend-wrapper:not(.has-sidebar) > .tg-dualpane-divider { display: none; } diff --git a/testgen/ui/components/frontend/js/pages/test_definition_summary.js b/testgen/ui/components/frontend/js/pages/test_definition_summary.js index c3cb23e5..7f25cba7 100644 --- a/testgen/ui/components/frontend/js/pages/test_definition_summary.js +++ b/testgen/ui/components/frontend/js/pages/test_definition_summary.js @@ -17,19 +17,25 @@ * @property {string} export_to_observability * @property {string?} last_manual_update * @property {string?} usage_notes + * @property {string?} external_url + * @property {object?} custom_metadata * @property {Array} attributes - * + * * @typedef Properties * @type {object} * @property {TestDefinition} test_definition */ import van from '/app/static/js/van.min.js'; -import { createEmitter, getValue, isEqual, loadStylesheet } from '/app/static/js/utils.js'; +import { createEmitter, getValue, isEqual, isHttpUrl, loadStylesheet } from '/app/static/js/utils.js'; import { Alert } from '/app/static/js/components/alert.js'; import { Attribute } from '/app/static/js/components/attribute.js'; +import { Link } from '/app/static/js/components/link.js'; const { div, strong } = van.tags; +const metadataDisplayValue = (value) => + (value !== null && typeof value === 'object') ? JSON.stringify(value) : value; + /** * @param {Properties} props * @returns @@ -112,6 +118,37 @@ const TestDefinitionSummary = (props) => { ), ), ), + testDefinition.external_url + ? Attribute({ + label: 'External URL', + value: isHttpUrl(testDefinition.external_url) + ? Link({ + href: testDefinition.external_url.trim(), + label: testDefinition.external_url.trim(), + open_new: true, + right_icon: 'open_in_new', + right_icon_size: 16, + }) + : testDefinition.external_url, + class: 'mt-4 external-url-attribute', + }) + : '', + testDefinition.custom_metadata && Object.keys(testDefinition.custom_metadata).length + ? div( + { class: 'flex-column fx-gap-3 mt-4' }, + strong({}, 'Custom Metadata'), + div( + { class: 'flex-row fx-flex-wrap fx-gap-4 test-definition-attributes' }, + Object.entries(testDefinition.custom_metadata).map(([key, value]) => + Attribute({ + label: key, + value: metadataDisplayValue(value), + class: 'fx-flex', + }) + ), + ), + ) + : '', testDefinition.usage_notes ? Alert( { type: 'info', class: 'mt-4' }, @@ -132,6 +169,20 @@ stylesheet.replace(` .test-definition-attributes > div .attribute-value { font-size: 16px; } +.external-url-attribute .attribute-value { + overflow-wrap: anywhere; +} +.external-url-attribute .tg-link { + max-width: 100%; +} +.external-url-attribute .tg-link--wrapper { + flex-wrap: wrap; +} +.external-url-attribute .tg-link--text { + overflow-wrap: anywhere; + word-break: break-all; + min-width: 0; +} `); export { TestDefinitionSummary }; diff --git a/testgen/ui/components/frontend/js/pages/test_definitions.js b/testgen/ui/components/frontend/js/pages/test_definitions.js index 47bf21d2..cbd09824 100644 --- a/testgen/ui/components/frontend/js/pages/test_definitions.js +++ b/testgen/ui/components/frontend/js/pages/test_definitions.js @@ -11,14 +11,21 @@ import { Attribute } from '/app/static/js/components/attribute.js'; import { TestDefinitionForm } from '/app/static/js/components/test_definition_form.js'; import { RunTestsDialog } from '/app/static/js/components/run_tests_dialog.js'; import { Textarea } from '/app/static/js/components/textarea.js'; +import { Tabs, Tab } from '/app/static/js/components/tabs.js'; import { Checkbox } from '/app/static/js/components/checkbox.js'; import { DropdownButton } from '/app/static/js/components/dropdown_button.js'; import { TestDefinitionNotes } from './test_definition_notes.js'; import { withTooltip } from '/app/static/js/components/tooltip.js'; import { Icon } from '/app/static/js/components/icon.js'; import { ProfilingResultsDialog } from '../shared/profiling_results_dialog.js'; +import { AXES, FACET_AXES, GROUP_BY_AXES, EMPTY, appliesToSelectedColumn } from '/app/static/js/components/test_picker_taxonomy.js'; +import { enterPage, exitPage, getPageSignal } from '/app/static/js/page_lifecycle.js'; +import { jsonObject, maxLength } from '/app/static/js/form_validators.js'; +import { capitalize } from '/app/static/js/display_utils.js'; -const { button: btn, div, i: icon, span, strong, input, label } = van.tags; +const { button: btn, div, i: icon, span, strong } = van.tags; + +const PAGE_KEY = 'testDefinitions'; const TABLE_COLUMNS = [ { name: 'table_name', label: 'Table', width: 180, sortable: true, overflow: 'hidden' }, @@ -40,8 +47,6 @@ const SEVERITY_OPTIONS = [ { label: 'Fail', value: 'Fail' }, ]; -const SCOPE_LABELS = { referential: 'Referential', table: 'Table', column: 'Column', custom: 'Custom' }; - // Blank test definition field defaults for add mode const BLANK_PARAM_FIELDS = { custom_query: null, @@ -816,6 +821,8 @@ const DetailPanel = (row) => { const paramCols = row.default_parm_columns ? row.default_parm_columns.split(',').map(c => c.trim()).filter(Boolean) : []; + const paramLabels = (row.default_parm_prompts || '').split(',').map(v => v.trim()); + const paramHelp = (row.default_parm_help || '').split('|').map(v => v.trim()); return div( { class: 'flex-column fx-gap-3 border border-radius-1 p-4 mt-2' }, @@ -826,13 +833,17 @@ const DetailPanel = (row) => { Attribute({ label: 'Schema Name', value: row.schema_name }), Attribute({ label: 'Table Name', value: row.table_name }), Attribute({ label: 'Test Focus', value: row.column_name }), - Attribute({ label: 'Test Type', value: row.test_type }), + Attribute({ label: 'Test Type', value: row.test_name_short }), Attribute({ label: 'Test Active', value: row.test_active_display }), Attribute({ label: 'Validation Status', value: row.test_definition_status }), Attribute({ label: 'Lock Refresh', value: row.lock_refresh_display }), Attribute({ label: 'Urgency', value: row.urgency }), Attribute({ label: 'Export to Observability', value: row.export_to_observability_display }), - ...paramCols.map(col => Attribute({ label: col, value: String(row[col] ?? '') })), + ...paramCols.map((col, index) => Attribute({ + label: paramLabels[index] || capitalize(col.replaceAll('_', ' ')), + help: paramHelp[index] || null, + value: String(row[col] ?? ''), + })), ), div( { class: 'flex-column fx-flex fx-gap-3' }, @@ -847,32 +858,64 @@ const DetailPanel = (row) => { ); }; -// Add dialog — mounted once, state persists across Python reruns -const AddDialogComponent = ({ open, info, validateResult: validateResultProp, onClose }, emit) => { +const TestPickerChip = (text, color) => { + const { span } = van.tags; + return span( + { + class: 'tg-test-chip', + style: `border:1px solid color-mix(in srgb, ${color} 25%, transparent);color:${color};`, + }, + text, + ); +}; + +// A keyboard-shortcut hint for the picker footer: a key chip followed by its action. +const KeyHint = (key, action) => { + const { span } = van.tags; + return span( + { class: 'tg-key-hint' }, + span({ class: 'tg-kbd' }, key), + action, + ); +}; + +// Faceted add-test picker dialog (step 1) reusing the param form (step 2) +const AddDialogComponent = ({ open, info, validateResult, onClose }, emit) => { + const { div, span, input, label, h4 } = van.tags; + const testTypes = van.derive(() => getValue(info)?.test_types ?? []); + const tableColumns = van.derive(() => getValue(info)?.table_columns ?? []); + const testSuite = van.derive(() => getValue(info)?.test_suite ?? {}); const tableGroupSchema = van.derive(() => getValue(info)?.table_group_schema ?? ''); const tableGroupsId = van.derive(() => getValue(info)?.table_groups_id ?? ''); - const testSuite = van.derive(() => getValue(info)?.test_suite ?? {}); - const tableColumns = van.derive(() => getValue(info)?.table_columns ?? []); const qualifiesTableRefsWithSchema = van.derive(() => getValue(info)?.qualifies_table_refs_with_schema ?? true); - const validateResult = van.derive(() => getValue(validateResultProp) ?? null); + const prefillColumn = van.derive(() => getValue(info)?.prefill_column ?? null); + + // selectedColumn carries the general_type needed by the type-aware filter; resolve it from + // the tableColumns entry so a prefill and a manual pick produce the identical shape. + const columnFromEntry = (c) => ({ + table_name: c.table_name, + column_name: c.column_name, + general_type: c.general_type ?? null, + }); - const scopeFilter = { - referential: van.state(true), - table: van.state(true), - column: van.state(true), - custom: van.state(true), - }; + // ---- Step + selection state ---- + const step = van.state(1); + const selectedTestType = van.state(null); + const formValues = van.state({}); + // Hoisted so the active form tab survives the form re-render on each field change. + const formActiveTab = van.state(0); - const filteredTestTypeOptions = van.derive(() => - testTypes.val - .filter(tt => tt.test_scope !== 'tablegroup' && (scopeFilter[tt.test_scope]?.val ?? true)) - .map(tt => ({ label: tt.select_name ?? tt.test_name_short, value: tt.test_type })) - ); + // ---- Picker state ---- + const searchQuery = van.state(''); + const groupBy = van.state('impact'); + const focusIndex = van.state(-1); // -1 = no row highlighted until keyboard nav + const selectedColumn = van.state(null); // { table_name, column_name, general_type } | null - const selectedTestType = van.state(null); - const formValues = van.state(null); + const facetSel = {}; + Object.keys(AXES).forEach((ax) => { facetSel[ax] = van.state([]); }); + // Build blank form values for the selected test type. const buildFormValues = (testType) => { if (!testType) return null; const tt = testTypes.rawVal.find(t => t.test_type === testType); @@ -900,80 +943,380 @@ const AddDialogComponent = ({ open, info, validateResult: validateResultProp, on }; }; - const selectTestType = (testType) => { - selectedTestType.val = testType; - formValues.val = buildFormValues(testType); + const toggleFacet = (ax, value) => { + const cur = facetSel[ax].val; + facetSel[ax].val = cur.includes(value) ? cur.filter((v) => v !== value) : [...cur, value]; + focusIndex.val = -1; + }; + const clearAll = () => { + Object.values(facetSel).forEach((s) => { s.val = []; }); + searchQuery.val = ''; + focusIndex.val = -1; }; - // Reset form state when dialog opens (transitions from closed→open) + // Reset to a fresh step-1 picker on each closed->open transition. The dialog is mounted + // once and its state otherwise persists across reopens, which would show the stale step-2 + // form. Also re-seeds the locked column per open. const wasOpen = van.state(false); van.derive(() => { const isOpen = open.val; if (isOpen && !wasOpen.val) { - selectTestType(null); + step.val = 1; + selectedTestType.val = null; + formValues.val = {}; + searchQuery.val = ''; + focusIndex.val = -1; + groupBy.val = 'impact'; + Object.values(facetSel).forEach((s) => { s.val = []; }); + const prefill = prefillColumn.rawVal; + const match = prefill + ? tableColumns.rawVal.find((c) => c.table_name === prefill.table_name && c.column_name === prefill.column_name) + : null; + selectedColumn.val = match ? columnFromEntry(match) : null; wasOpen.val = true; } else if (!isOpen) { wasOpen.val = false; } }); - return Dialog( - { title: 'Add Test', open, onClose, width: '52rem' }, - div( - { class: 'flex-column fx-gap-4 td-form-dialog' }, + const matchesSearch = (t, q) => { + if (!q) return true; + const hay = `${t.test_name_short} ${t.test_name_long} ${t.test_description}`.toLowerCase(); + return q.toLowerCase().split(/\s+/).filter(Boolean).every((tok) => hay.includes(tok)); + }; - // Test type picker — always visible - div( - { class: 'flex-column fx-gap-3' }, + // Search + column relevance only -- drives facet counts (PRD F9). A selected column filters + // the list to tests applicable to its type; clearing the column shows all tests. + const baseVisible = van.derive(() => { + const q = searchQuery.val; + const col = selectedColumn.val; + return testTypes.val.filter((t) => + matchesSearch(t, q) && (!col || appliesToSelectedColumn(t, col.general_type))); + }); + + const passesFacets = (t) => Object.entries(facetSel).every(([ax, s]) => { + const sel = s.val; + if (!sel.length) return true; + const v = AXES[ax].value(t); + return v != null && sel.includes(v); + }); + + const filtered = van.derive(() => baseVisible.val.filter(passesFacets)); + + // Group the result list by the chosen axis; null -> EMPTY bucket. + const grouped = van.derive(() => { + const axis = AXES[groupBy.val]; + const buckets = new Map(); + filtered.val.forEach((t) => { + const key = axis.value(t) || EMPTY; + if (!buckets.has(key)) buckets.set(key, []); + buckets.get(key).push(t); + }); + const order = axis.order; // sparse axes have a canonical order + const entries = [...buckets.entries()]; + entries.sort((a, b) => { + if (order) { + const ia = order.indexOf(a[0]); const ib = order.indexOf(b[0]); + return (ia < 0 ? Infinity : ia) - (ib < 0 ? Infinity : ib); // EMPTY bucket sorts last + } + return b[1].length - a[1].length; + }); + return entries; + }); + + const flatList = van.derive(() => grouped.val.flatMap(([, tests]) => tests)); + + // Per-axis value -> count, over baseVisible (not other facet selections). + const counts = (ax) => { + const m = new Map(); + baseVisible.val.forEach((t) => { + const v = AXES[ax].value(t); + if (v != null) m.set(v, (m.get(v) || 0) + 1); + }); + return m; + }; + + const selectTest = (t) => { + if (!t) return; + selectedTestType.val = t; + const fv = buildFormValues(t.test_type); + const col = selectedColumn.rawVal; + if (col) { + const scope = fv.test_scope ?? 'column'; + if (scope !== 'tablegroup') { + fv.table_name = col.table_name; + } + if (scope === 'column' || scope === 'referential' || scope === 'custom') { + fv.column_name = col.column_name; + } + } + formValues.val = fv; + formActiveTab.val = 0; + step.val = 2; + }; + + // ---- Keyboard: Cmd/Ctrl+K or "/" focus search, Up/Down move, Enter add, Esc close ---- + let searchEl = null; + const onKeyDown = (e) => { + if (step.val !== 1) return; + if ((e.key === 'k' && (e.metaKey || e.ctrlKey)) || (e.key === '/' && document.activeElement !== searchEl)) { + e.preventDefault(); searchEl?.focus(); return; + } + if (e.key === 'Escape') { onClose(); return; } + const list = flatList.val; + if (e.key === 'ArrowDown') { e.preventDefault(); focusIndex.val = Math.min(focusIndex.val + 1, Math.max(list.length - 1, 0)); } + else if (e.key === 'ArrowUp') { e.preventDefault(); focusIndex.val = Math.max(focusIndex.val - 1, 0); } + // Enter adds the focused row, or the only match when nothing is explicitly focused. + else if (e.key === 'Enter') { e.preventDefault(); selectTest(list[focusIndex.val] ?? (list.length === 1 ? list[0] : null)); } + }; + + // Listen at the document level while step 1 is open so arrow keys work without first + // clicking the picker. Attach/detach reactively; guard prevents duplicate registration. + // The page's AbortSignal removes the listener on teardown — without it, a teardown + // while the dialog is open (e.g. browser Back) would orphan onKeyDown on the document, + // since nothing toggles open/step on unmount to run the detach branch below. + let keydownAttached = false; + van.derive(() => { + const active = open.val && step.val === 1; + if (active && !keydownAttached) { + document.addEventListener('keydown', onKeyDown, { signal: getPageSignal(PAGE_KEY) ?? undefined }); + keydownAttached = true; + } else if (!active && keydownAttached) { + document.removeEventListener('keydown', onKeyDown); + keydownAttached = false; + } + }); + + // ---- Renderers ---- + const FacetGroup = (ax) => { + const axis = AXES[ax]; + // Whole group is reactive: its title is hidden when no option has a count (PRD facet review). + return () => { + const m = counts(ax); + let keys = [...m.keys()]; + if (axis.order) keys = axis.order.filter((k) => m.has(k)); + else keys.sort((a, b) => m.get(b) - m.get(a)); + if (!keys.length) return ''; + return div( + { class: 'tg-facet-group', 'data-testid': 'facet-group', 'data-axis': ax }, + h4({ class: 'tg-facet-title' }, axis.label), div( - { class: 'flex-row fx-gap-4 fx-align-flex-center fx-flex-wrap' }, - span({ class: 'text-caption' }, 'Show Types:'), - ...Object.entries(SCOPE_LABELS).map(([scope, scopeLabel]) => - Checkbox({ - label: scopeLabel, - checked: scopeFilter[scope], - onChange: (v) => { scopeFilter[scope].val = v; }, - }) - ), + { class: 'flex-column fx-gap-1' }, + ...keys.map((value) => { + const checked = van.derive(() => facetSel[ax].val.includes(value)); + // Display-only Checkbox (pointer-events disabled via .tg-facet-checkbox); the + // row's onclick is the single toggle source, so the box can't double-fire. + return div( + { + class: 'tg-facet-option', + 'data-testid': 'facet-option', + 'data-axis': ax, + 'data-value': value, + onclick: () => toggleFacet(ax, value), + }, + span({ class: 'tg-facet-checkbox' }, Checkbox({ label: '', checked })), + span({ class: 'tg-facet-dot', style: `background:${axis.color(value)}` }), + span({ class: 'tg-facet-label' }, value), + span({ class: 'tg-facet-count' }, m.get(value)), + ); + }), ), - () => Select({ - label: 'Test Type', - value: selectedTestType.val, - options: filteredTestTypeOptions.val, - allowNull: true, - filterable: true, - onChange: (value) => { selectTestType(value); }, - }), - ), + ); + }; + }; - // Form (shown after test type selected) — imperative update - // because VanJS binding replacement doesn't work inside Dialog portals - () => { - open.val; + const ResultRow = (t, isFocused) => { + const row = div( + { + class: () => `tg-test-row${isFocused.val ? ' tg-test-row--focused' : ''}`, + 'data-testid': 'test-picker-row', + 'data-test-type': t.test_type, + onclick: () => selectTest(t), + }, + div( + { class: 'flex-row fx-justify-space-between' }, + div({ class: 'tg-test-row-title' }, t.test_name_short), + ), + div({ class: 'tg-test-desc' }, t.test_description), + div( + { class: 'flex-row fx-flex-wrap fx-gap-1' }, + t.impact_dimension ? TestPickerChip(t.impact_dimension, AXES.impact.color(t.impact_dimension)) : null, + t.algorithm ? TestPickerChip(t.algorithm, AXES.algorithm.color(t.algorithm)) : null, + t.statistical_technique ? TestPickerChip(t.statistical_technique, AXES.technique.color(t.statistical_technique)) : null, + t.health_dimension ? TestPickerChip(t.health_dimension, AXES.health.color(t.health_dimension)) : null, + t.criteria ? TestPickerChip(t.criteria, AXES.criteria.color(t.criteria)) : null, + ), + ); + // Keep the keyboard-focused row visible within the scrolling list. + van.derive(() => { if (isFocused.val) row.scrollIntoView({ block: 'nearest' }); }); + return row; + }; - selectedTestType.val; - const fv = formValues.val; - const vr = validateResult.val; + // One filter row: column relevance + search on the left, group-by pushed to the right. + const ContextBar = (searchField) => div( + { class: 'tg-picker-context flex-row fx-gap-3 fx-align-flex-end' }, + () => { + // Column relevance filter. Option VALUES are the array index (string) so that + // table/column names containing spaces or dots can never break parsing. + const cols = tableColumns.val; + const options = cols.map((c, i) => ({ label: `${c.table_name}.${c.column_name}`, value: String(i) })); + const sel = selectedColumn.val; + const currentIdx = sel ? cols.findIndex((c) => c.table_name === sel.table_name && c.column_name === sel.column_name) : -1; + return Select({ + label: 'Show tests relevant to a column', + allowNull: true, + // Cap growth so a long table.column value ellipsizes instead of crowding search. + style: 'max-width: 320px;', + value: currentIdx >= 0 ? String(currentIdx) : null, + options, + onChange: (val) => { + if (val == null || val === '') { selectedColumn.val = null; return; } + selectedColumn.val = columnFromEntry(tableColumns.rawVal[Number(val)]); + }, + }); + }, + div({ class: 'tg-picker-search fx-flex', 'data-testid': 'test-picker-search' }, searchField), + Select({ + label: 'Group by', + value: groupBy.val, + options: GROUP_BY_AXES.map((ax) => ({ label: AXES[ax].label, value: ax })), + onChange: (val) => { groupBy.val = val; focusIndex.val = -1; }, + }), + ); - if (!fv) return ''; + const ActiveFilters = () => div( + { class: 'flex-row fx-flex-wrap fx-gap-1', 'data-testid': 'test-picker-active-filters' }, + () => { + const chips = []; + Object.entries(facetSel).forEach(([ax, s]) => s.val.forEach((value) => { + chips.push(span( + { class: 'tg-active-chip flex-row', onclick: () => toggleFacet(ax, value) }, + span(`${AXES[ax].label}: ${value}`), + Icon({ size: 14 }, 'close'), + )); + })); + if (!chips.length) return ''; + chips.push(span( + { 'data-testid': 'test-picker-clear' }, + Button({ type: 'basic', label: 'Clear all', onclick: clearAll }), + )); + return div({ class: 'flex-row fx-flex-wrap fx-gap-1' }, ...chips); + }, + ); - return TestDefFormContent({ - formValues: fv, - tableColumns: tableColumns.rawVal, - testSuite: testSuite.rawVal, - qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, - validateResult: vr, - mode: 'add', - onFormChange: (changes) => { - formValues.val = { ...formValues.rawVal, ...changes }; - }, - onValidate: () => emit('ValidateTest', { payload: formValues.rawVal }), - onSave: () => emit('AddTestSaved', { payload: formValues.rawVal }), - onCancel: onClose, - }); + const ResultsPane = () => div( + { class: 'tg-picker-results' }, + // Fixed header — count stays put while the list below scrolls. + div( + { class: 'tg-picker-results-header' }, + span( + { 'data-testid': 'test-picker-count', class: 'tg-picker-count' }, + () => `${filtered.val.length} matching test types${!!selectedColumn.val ? ' for selected column' : ''}` + ), + ), + div( + { class: 'tg-picker-list' }, + // One-result hint (PRD F13): when exactly one test matches, prompt Enter-to-add. + () => filtered.val.length === 1 + ? div({ class: 'tg-one-result-hint mb-2', 'data-testid': 'test-picker-one-hint' }, + `${filtered.val[0].test_name_short} — press Enter to add`) + : '', + () => { + const groups = grouped.val; + if (!groups.length) { + return div( + { class: 'tg-picker-empty', 'data-testid': 'test-picker-empty' }, + span('No tests match this filter combination.'), + Button({ type: 'stroked', label: 'Clear filters and search', onclick: clearAll }), + ); + } + let runningIndex = -1; + return div( + { class: 'flex-column fx-gap-2' }, + ...groups.map(([groupName, tests]) => div( + { class: 'tg-result-group' }, + div( + { class: 'tg-group-band' }, + span({ class: 'tg-group-band-name' }, groupName), + span({ class: 'tg-group-band-count' }, `${tests.length} test types`), + ), + ...tests.map((t) => { + runningIndex += 1; + const myIndex = runningIndex; + const isFocused = van.derive(() => focusIndex.val === myIndex); + return ResultRow(t, isFocused); + }), + )), + ); }, ), ); + + const PickerView = () => { + // Reusable Input component for search. The keyboard handler focuses the inner , + // resolved by id below (Input renders its label with this id). + const searchField = Input({ + id: 'tg-picker-search', + icon: 'search', + clearable: true, + style: 'width: 100%;', + placeholder: 'Search name or description', + value: searchQuery, + onChange: (value) => { searchQuery.val = value; focusIndex.val = -1; }, + }); + // Load-bearing on every step-1 (re)render, not just first open: refocuses search when + // returning from the step-2 form. Do not hoist out of PickerView. + requestAnimationFrame(() => { + searchEl = document.getElementById('tg-picker-search')?.querySelector('input') ?? null; + searchEl?.focus(); + }); + + return div( + { class: 'tg-test-picker', 'data-testid': 'test-picker' }, + // Fixed header: column relevance + search + group-by in one row, then active filters. + div( + { class: 'tg-picker-header' }, + ContextBar(searchField), + ActiveFilters(), + ), + div( + { class: 'tg-picker-body fx-gap-3' }, + div( + { class: 'tg-picker-rail' }, + ...FACET_AXES.map(FacetGroup), + ), + ResultsPane(), + ), + div( + { class: 'tg-picker-footer', 'data-testid': 'test-picker-footer' }, + KeyHint('↑ ↓', 'navigate'), + KeyHint('Enter', 'select'), + KeyHint('/', 'focus search'), + ), + ); + }; + + const FormView = () => TestDefFormContent({ + formValues: formValues.val, + tableColumns: tableColumns.rawVal, + testSuite: testSuite.rawVal, + qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, + validateResult: getValue(validateResult), + mode: 'add', + activeTab: formActiveTab, + onFormChange: (changes) => { formValues.val = { ...formValues.rawVal, ...changes }; }, + onValidate: () => emit('ValidateTest', { payload: formValues.rawVal }), + onSave: () => emit('AddTestSaved', { payload: formValues.rawVal }), + onBack: () => { step.val = 1; selectedTestType.val = null; }, + onCancel: onClose, + }); + + const dialogTitle = van.derive(() => `Add Test - ${step.val === 1 ? 'Pick a Test Type' : 'Configure Test'}`); + + return Dialog( + { title: dialogTitle, open, onClose, width: '68rem' }, + () => step.val === 1 ? PickerView() : FormView(), + ); }; // Edit dialog — mounted once, state persists across Python reruns @@ -985,6 +1328,8 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o const validateResult = van.derive(() => getValue(validateResultProp) ?? null); const formValues = van.state(null); + // Hoisted so the active form tab survives the form re-render on each field change. + const formActiveTab = van.state(0); const initFormFromInfo = () => { const di = dialogInfo.rawVal; @@ -997,6 +1342,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o column_name_prompt: ttRow.column_name_prompt ?? null, column_name_help: ttRow.column_name_help ?? null, }; + formActiveTab.val = 0; }; // Reset form when dialog opens (closed→open), clear when it closes @@ -1028,6 +1374,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o qualifiesTableRefsWithSchema: qualifiesTableRefsWithSchema.rawVal, validateResult: vr, mode: 'edit', + activeTab: formActiveTab, onFormChange: (changes) => { formValues.val = { ...formValues.rawVal, ...changes }; }, @@ -1041,7 +1388,7 @@ const EditDialogComponent = ({ open, info, validateResult: validateResultProp, o }; // Shared form content for add/edit dialogs -const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, onFormChange, onValidate, onSave, onCancel }) => { +const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResult, mode, qualifiesTableRefsWithSchema, activeTab, onFormChange, onValidate, onSave, onCancel, onBack }) => { const testScope = formValues.test_scope ?? 'column'; const runType = formValues.run_type ?? 'CAT'; const testType = formValues.test_type ?? ''; @@ -1054,6 +1401,13 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul onFormChange({ [key]: value }); }; + // Custom Metadata is stored as a JSON object (JSONB). Keep the raw text the user edits in a + // local state, and only commit a parsed object to the form when the JSON is valid. + const metadataText = van.state( + formValues.custom_metadata ? JSON.stringify(formValues.custom_metadata, null, 2) : '' + ); + const metadataValid = van.state(true); + const inheritedSeverity = testSuite.severity ?? formValues.default_severity ?? 'Warning'; const severityOptions = [ { label: `Inherited (${inheritedSeverity})`, value: null }, @@ -1095,8 +1449,8 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul return div( { class: 'flex-column fx-gap-3' }, - // Test type header (add mode) or read-only test type (edit mode) - mode === 'add' && formValues.test_name_short + // Header — test type identity, shown above the tabs in both add and edit modes + formValues.test_name_short ? div( { class: 'mb-1' }, div({ class: 'text-large' }, formValues.test_name_short), @@ -1106,170 +1460,203 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul ) : null, - mode === 'edit' - ? Input({ - name: 'test_type_display', - label: 'Test Type', - value: formValues.test_name_short ?? formValues.test_type ?? '', - disabled: true, - }) - : null, - formValues.usage_notes ? Alert({ type: 'info' }, strong({ class: 'mb-1' }, 'Usage Notes'), div({}, formValues.usage_notes)) : null, - // Description override - Textarea({ - name: 'test_description', - label: 'Test Description Override', - value: () => fv.val.test_description ?? '', - placeholder: `Inherited (${formValues.default_test_description ?? ''})`, - height: 72, - onChange: (value) => updateField('test_description', value || null), - }), + Tabs( + { activeTab }, + Tab( + { label: 'Parameters' }, + div( + { class: 'flex-column fx-gap-3' }, - // Checkboxes - div( - { class: 'flex-row fx-gap-4' }, - Checkbox({ - label: 'Test Active', - checked: () => fv.val.test_active ?? true, - onChange: (v) => updateField('test_active', v), - }), - Checkbox({ - label: 'Lock Refresh', - checked: () => fv.val.lock_refresh ?? false, - onChange: (v) => updateField('lock_refresh', v), - }), - ), + // Schema (read-only) + qualifiesTableRefsWithSchema + ? Input({ + name: 'schema_name', + label: 'Schema', + value: formValues.schema_name ?? '', + disabled: true, + }) + : null, - // Severity + Observability + Impact Dimension selects - div( - { class: 'flex-row fx-gap-3 fx-flex-wrap' }, - div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Urgency Override', - value: fv.val.severity ?? null, - options: severityOptions, - allowNull: false, - onChange: (value) => updateField('severity', value), - }), - ), - div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Send to Observability - Override', - value: fv.val.export_to_observability ?? null, - options: obsOptions, - allowNull: false, - onChange: (value) => updateField('export_to_observability', value), - }), + // Table name + testScope !== 'tablegroup' + ? testScope === 'custom' + ? Input({ + name: 'table_name', + label: 'Table', + value: () => fv.val.table_name ?? '', + onChange: (value) => updateField('table_name', value || null), + }) + : () => Select({ + label: 'Table', + value: fv.val.table_name ?? null, + options: tableNameOptions, + allowNull: true, + filterable: true, + disabled: mode === 'edit', + onChange: (value) => { + updateField('table_name', value); + updateField('column_name', null); + }, + }) + : null, + + // Column name (scope-dependent) + testScope === 'column' + ? () => Select({ + label: 'Column', + value: fv.val.column_name ?? null, + options: columnNameOptions.val, + allowNull: true, + filterable: true, + onChange: (value) => updateField('column_name', value), + }) + : testScope === 'referential' || testScope === 'custom' + ? Input({ + name: 'column_name', + label: columnLabel, + help: columnHelp, + value: () => fv.val.column_name ?? '', + onChange: (value) => updateField('column_name', value || null), + }) + : null, + + // Validation status (edit mode only) + mode === 'edit' && formValues.test_definition_status + ? Input({ + name: 'test_definition_status', + label: 'Validation Status', + value: formValues.test_definition_status || 'OK', + disabled: true, + }) + : null, + + // Dynamic parameter fields + div( + { class: 'td-form-params-section' }, + TestDefinitionForm({ + definition: formValues, + qualifiesTableRefsWithSchema, + hideHeader: true, + onChange: (changes) => { + if (Object.keys(changes).length === 0) return; + const updated = { ...fv.rawVal, ...changes }; + fv.val = updated; + onFormChange(changes); + }, + }), + ), + + // Skip errors (QUERY run type only) + runType === 'QUERY' + ? Input({ + name: 'skip_errors', + label: 'Threshold Error Count', + type: 'number', + value: () => fv.val.skip_errors ?? 0, + step: 1, + onChange: (value) => updateField('skip_errors', value ?? 0), + }) + : null, + ), ), - showImpactDimensionOverride ? div( - { style: 'flex: calc(50% - 8px) 0 0;' }, - () => Select({ - label: 'Impact Dimension Override', - value: fv.val.impact_dimension ?? null, - options: impactDimensionOptions, - allowNull: false, - helpText: 'Override the default impact classification for this test. Affects how the test result is categorized in score breakdowns.', - onChange: (value) => updateField('impact_dimension', value), - }), - ) : null, - ), + Tab( + { label: 'Settings' }, + div( + { class: 'flex-column fx-gap-3' }, - // Schema (read-only) - qualifiesTableRefsWithSchema - ? Input({ - name: 'schema_name', - label: 'Schema', - value: formValues.schema_name ?? '', - disabled: true, - }) - : null, + // Description override + Textarea({ + name: 'test_description', + label: 'Test Description Override', + value: () => fv.val.test_description ?? '', + placeholder: `Inherited (${formValues.default_test_description ?? ''})`, + height: 72, + onChange: (value) => updateField('test_description', value || null), + }), - // Table name - testScope !== 'tablegroup' - ? testScope === 'custom' - ? Input({ - name: 'table_name', - label: 'Table', - value: () => fv.val.table_name ?? '', - onChange: (value) => updateField('table_name', value || null), - }) - : () => Select({ - label: 'Table', - value: fv.val.table_name ?? null, - options: tableNameOptions, - allowNull: true, - filterable: true, - disabled: mode === 'edit', - onChange: (value) => { - updateField('table_name', value); - updateField('column_name', null); + // External link + custom metadata + Input({ + name: 'external_url', + label: 'External URL', + help: 'Optional link to the code, pipeline step, or system that produces this data.', + value: () => fv.val.external_url ?? '', + onChange: (value) => updateField('external_url', value || null), + }), + Textarea({ + name: 'custom_metadata', + label: 'Custom Metadata', + help: 'Optional JSON object of key-value pairs, e.g. {"pipeline": "daily_load", "task": "transform_orders"}.', + value: metadataText, + placeholder: '{\n "pipeline": "daily_load",\n "task": "transform_orders"\n}', + height: 120, + validators: [jsonObject, maxLength(10240)], + onChange: (value, state) => { + metadataText.val = value; + metadataValid.val = state.valid; + if (state.valid) { + updateField('custom_metadata', value && value.trim() ? JSON.parse(value) : null); + } }, - }) - : null, + }), - // Column name (scope-dependent) - testScope === 'column' - ? () => Select({ - label: 'Column', - value: fv.val.column_name ?? null, - options: columnNameOptions.val, - allowNull: true, - filterable: true, - onChange: (value) => updateField('column_name', value), - }) - : testScope === 'referential' || testScope === 'custom' - ? Input({ - name: 'column_name', - label: columnLabel, - help: columnHelp, - value: () => fv.val.column_name ?? '', - onChange: (value) => updateField('column_name', value || null), - }) - : null, - - // Validation status (edit mode only) - mode === 'edit' && formValues.test_definition_status - ? Input({ - name: 'test_definition_status', - label: 'Validation Status', - value: formValues.test_definition_status || 'OK', - disabled: true, - }) - : null, + // Checkboxes + div( + { class: 'flex-row fx-gap-4' }, + Checkbox({ + label: 'Test Active', + checked: () => fv.val.test_active ?? true, + onChange: (v) => updateField('test_active', v), + }), + Checkbox({ + label: 'Lock Refresh', + checked: () => fv.val.lock_refresh ?? false, + onChange: (v) => updateField('lock_refresh', v), + }), + ), - // Dynamic parameter fields - div( - { class: 'td-form-params-section' }, - TestDefinitionForm({ - definition: formValues, - qualifiesTableRefsWithSchema, - onChange: (changes) => { - if (Object.keys(changes).length === 0) return; - const updated = { ...fv.rawVal, ...changes }; - fv.val = updated; - onFormChange(changes); - }, - }), + // Severity + Observability + Impact Dimension selects + div( + { class: 'flex-row fx-gap-3 fx-flex-wrap' }, + div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Urgency Override', + value: fv.val.severity ?? null, + options: severityOptions, + allowNull: false, + onChange: (value) => updateField('severity', value), + }), + ), + div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Send to Observability - Override', + value: fv.val.export_to_observability ?? null, + options: obsOptions, + allowNull: false, + onChange: (value) => updateField('export_to_observability', value), + }), + ), + showImpactDimensionOverride ? div( + { style: 'flex: calc(50% - 8px) 0 0;' }, + () => Select({ + label: 'Impact Dimension Override', + value: fv.val.impact_dimension ?? null, + options: impactDimensionOptions, + allowNull: false, + helpText: 'Override the default impact classification for this test. Affects how the test result is categorized in score breakdowns.', + onChange: (value) => updateField('impact_dimension', value), + }), + ) : null, + ), + ), + ), ), - // Skip errors (QUERY run type only) - runType === 'QUERY' - ? Input({ - name: 'skip_errors', - label: 'Threshold Error Count', - type: 'number', - value: () => fv.val.skip_errors ?? 0, - step: 1, - onChange: (value) => updateField('skip_errors', value ?? 0), - }) - : null, - // Validate feedback validateResult ? Alert({ type: validateResult.success ? 'success' : 'error' }, validateResult.message) @@ -1278,15 +1665,28 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul // Buttons div( { class: 'flex-row fx-justify-space-between fx-gap-2' }, - isValidatable - ? Button({ - type: 'stroked', - color: 'basic', - label: 'Validate', - width: 'auto', - onclick: onValidate, - }) - : span(''), + div( + { class: 'flex-row fx-gap-2' }, + onBack + ? Button({ + type: 'stroked', + color: 'basic', + icon: 'arrow_back', + label: 'Back', + width: 'auto', + onclick: onBack, + }) + : null, + isValidatable + ? Button({ + type: 'stroked', + color: 'basic', + label: 'Validate', + width: 'auto', + onclick: onValidate, + }) + : null, + ), div( { class: 'flex-row fx-gap-2' }, Button({ @@ -1301,6 +1701,7 @@ const TestDefFormContent = ({ formValues, tableColumns, testSuite, validateResul color: 'primary', label: mode === 'edit' ? 'Save' : 'Add', width: 'auto', + disabled: () => !metadataValid.val, onclick: onSave, }), ), @@ -1313,6 +1714,7 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { const dialogInfo = van.derive(() => getValue(info) ?? null); const collision = van.derive(() => dialogInfo.val?.collision ?? null); + const targetProjectCode = van.state(null); const targetTgId = van.state(null); const targetTsId = van.state(null); const targetTableName = van.state(null); @@ -1324,6 +1726,7 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { const isOpen = open.val; if (isOpen && !wasOpen.val) { const di = dialogInfo.val; + targetProjectCode.val = di?.current_project_code ?? null; targetTgId.val = di?.current_table_group_id ?? null; targetTsId.val = null; targetTableName.val = null; @@ -1334,10 +1737,16 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { } }); - const tableGroupOptions = van.derive(() => - (dialogInfo.val?.table_groups ?? []).map(tg => ({ label: tg.table_groups_name, value: tg.id })) + const projectOptions = van.derive(() => + (dialogInfo.val?.projects ?? []).map(p => ({ label: p.project_name, value: p.project_code })) ); + const tableGroupOptions = van.derive(() => { + const project = targetProjectCode.val; + const tgs = dialogInfo.val?.table_groups_by_project?.[project] ?? []; + return tgs.map(tg => ({ label: tg.table_groups_name, value: tg.id })); + }); + const testSuiteOptions = van.derive(() => { const tg = targetTgId.val; const suites = dialogInfo.val?.test_suites_by_table_group?.[tg] ?? []; @@ -1414,6 +1823,21 @@ const CopyMoveDialogComponent = ({ open, info, onClose }, emit) => { { class: 'flex-column fx-gap-4 td-form-dialog' }, () => div({ class: 'text-caption' }, `Selected tests: ${(dialogInfo.val?.selected ?? []).length}`), + () => Select({ + label: 'Target Project', + value: targetProjectCode.val, + options: projectOptions.val, + required: true, + filterable: true, + onChange: (value) => { + targetProjectCode.val = value; + targetTgId.val = null; + targetTsId.val = null; + targetTableName.val = null; + targetColumnName.val = null; + }, + }), + () => Select({ label: 'Target Table Group', value: targetTgId.val, @@ -1530,8 +1954,44 @@ stylesheet.replace(` .td-form-params-section { border-top: 1px solid var(--border-color); padding-top: 12px; - margin-top: 4px; + margin-top: 16px; } + +.tg-test-picker { display: flex; flex-direction: column; gap: 12px; height: 70vh; min-height: 0; } +.tg-picker-header { flex: none; display: flex; flex-direction: column; gap: 12px; } +.tg-picker-body { flex: 1; min-height: 0; display: flex; flex-direction: row; align-items: stretch; } +.tg-picker-rail { flex: 0 0 260px; min-height: 0; overflow-y: auto; overflow-x: hidden; padding-right: 8px; border-right: 1px solid var(--border-color, #e0e0e0); } +.tg-picker-results { flex: 1; min-height: 0; display: flex; flex-direction: column; } +.tg-picker-results-header { flex: none; padding-bottom: 8px; margin-bottom: 8px; border-bottom: 1px solid var(--border-color, #e0e0e0); } +.tg-picker-count { font-size: 13px; color: var(--secondary-text-color, #666); } +.tg-picker-list { flex: 1; overflow-y: auto; } +.tg-facet-group { margin-bottom: 14px; } +.tg-facet-title { margin: 0 0 6px; font-size: 12px; text-transform: uppercase; letter-spacing: .06em; color: var(--secondary-text-color, #666); } +.tg-facet-option { display: flex; align-items: center; gap: 8px; cursor: pointer; font-size: 13px; color: var(--primary-text-color); } +.tg-facet-dot { width: 8px; height: 8px; border-radius: 50%; flex-shrink: 0; } +.tg-facet-label { color: var(--primary-text-color); } +.tg-facet-count { margin-left: auto; color: var(--secondary-text-color, #888); font-variant-numeric: tabular-nums; } +.tg-facet-checkbox { display: inline-flex; pointer-events: none; } +/* --empty stays distinct from the --table-hover-color row hover so the band reads as a divider. */ +.tg-group-band { position: sticky; top: 0; z-index: 1; display: flex; justify-content: space-between; align-items: center; gap: 8px; padding: 6px 10px; margin-bottom: 4px; border-radius: 6px; background: var(--empty, #e0e0e0); } +.tg-group-band-name { font-weight: 600; } +.tg-group-band-count { font-size: 12px; color: var(--secondary-text-color, #888); font-variant-numeric: tabular-nums; } +.tg-result-group { display: flex; flex-direction: column; gap: 6px; } +.tg-test-row { position: relative; padding: 10px; border: 2px solid transparent; border-radius: 8px; cursor: pointer; } +/* Divider sits in the gap below the row, not on its border, so it never overlaps the focus outline. */ +.tg-result-group .tg-test-row:not(:last-child)::after { content: ''; position: absolute; left: 8px; right: 8px; bottom: -4px; border-bottom: 1px dashed var(--border-color, #dddfe2); } +.tg-test-row:hover { background: var(--table-hover-color, #f3f4f6); } +.tg-test-row--focused { border-color: var(--primary-color, #1976d2); background: var(--table-hover-color, #f3f4f6); } +.tg-test-row-title { font-weight: 600; } +.tg-test-desc { font-size: 13px; color: var(--secondary-text-color, #555); margin: 2px 0 6px; } +.tg-test-chip { font-size: 11px; padding: 1px 7px; border-radius: 999px; white-space: nowrap; } +.tg-active-chip { font-size: 12px; line-height: 1; gap: 3px; padding: 2px 6px 2px 8px; border-radius: 999px; background: var(--table-hover-color, #eee); cursor: pointer; } +.tg-active-chip .material-symbols-rounded { cursor: pointer; } +.tg-one-result-hint { padding: 8px 10px; border-radius: 6px; background: var(--table-hover-color, #eef2ff); font-size: 13px; } +.tg-picker-empty { display: flex; flex-direction: column; gap: 12px; align-items: flex-start; padding: 40px 8px; } +.tg-picker-footer { flex: none; display: flex; gap: 16px; align-items: center; padding-top: 10px; border-top: 1px solid var(--border-color, #e0e0e0); font-size: 12px; color: var(--secondary-text-color, #666); } +.tg-key-hint { display: inline-flex; align-items: center; gap: 6px; } +.tg-kbd { display: inline-flex; align-items: center; min-width: 18px; justify-content: center; padding: 1px 6px; border: 1px solid var(--border-color, #d0d0d0); border-radius: 4px; background: var(--empty, #e0e0e0); font-family: monospace; font-size: 11px; color: var(--primary-text-color); } `); export { TestDefinitions, EditDialogComponent }; @@ -1547,6 +2007,7 @@ export default (component) => { } parentElement.state = componentState; componentState.emit = createEmitter(setTriggerValue); + componentState.signal = enterPage(PAGE_KEY); van.add(parentElement, TestDefinitions(componentState)); } else { for (const [key, value] of Object.entries(data)) { @@ -1556,5 +2017,8 @@ export default (component) => { } } - return () => { parentElement.state = null; }; + return () => { + exitPage(PAGE_KEY); + parentElement.state = null; + }; }; diff --git a/testgen/ui/components/frontend/js/pages/test_results.js b/testgen/ui/components/frontend/js/pages/test_results.js index 8fdc3bab..5c986d0e 100644 --- a/testgen/ui/components/frontend/js/pages/test_results.js +++ b/testgen/ui/components/frontend/js/pages/test_results.js @@ -20,6 +20,7 @@ * @property {string?} table_groups_id * @property {string?} severity * @property {string} test_type + * @property {string?} external_url * * @typedef Properties * @type {object} @@ -41,7 +42,7 @@ * @property {object} filter_options */ import van from '/app/static/js/van.min.js'; -import { createEmitter, getValue, isEqual, loadStylesheet, parseDate } from '/app/static/js/utils.js'; +import { createEmitter, getValue, isEqual, isHttpUrl, loadStylesheet, parseDate } from '/app/static/js/utils.js'; import { Table } from '/app/static/js/components/table.js'; import { Select } from '/app/static/js/components/select.js'; import { Tabs, Tab } from '/app/static/js/components/tabs.js'; @@ -49,6 +50,7 @@ import { Button } from '/app/static/js/components/button.js'; import { Checkbox } from '/app/static/js/components/checkbox.js'; import { DropdownButton } from '/app/static/js/components/dropdown_button.js'; import { Icon } from '/app/static/js/components/icon.js'; +import { Link } from '/app/static/js/components/link.js'; import { SummaryBar } from '/app/static/js/components/summary_bar.js'; import { Dialog } from '/app/static/js/components/dialog.js'; import { Toggle } from '/app/static/js/components/toggle.js'; @@ -106,6 +108,7 @@ const DATA_COLUMNS = [ { name: 'flagged_display', label: 'Flagged', width: 80, align: 'center' }, { name: 'notes_count', label: 'Notes', width: 70, align: 'center' }, { name: 'result_message', label: 'Details', width: 200, overflow: 'hidden' }, + { name: 'external_link', label: 'Link', width: 60, align: 'center' }, ]; const HISTORY_COLUMNS = [ @@ -149,6 +152,24 @@ const formatNumber = (v) => { return n.toLocaleString(undefined, { maximumFractionDigits: 5 }); }; +const buildExternalLinkCell = (url) => { + if (!isHttpUrl(url)) { + return ''; + } + // Stop the click from also selecting the row; the anchor still opens in a new tab. + return span( + { class: 'flex-row fx-justify-center', onclick: (event) => event.stopPropagation() }, + Link({ + href: url.trim(), + label: '', + open_new: true, + left_icon: 'open_in_new', + left_icon_size: 18, + tooltip: 'Open external link', + }), + ); +}; + const buildTableRow = (item) => ({ id: item.test_result_id, table_name: item.table_name ?? '', @@ -171,6 +192,7 @@ const buildTableRow = (item) => ({ span(item.notes_count), ) : '', result_message: item.result_message ?? '', + external_link: buildExternalLinkCell(item.external_url), }); const ExportMenu = (statusFilter, tableFilter, columnFilter, testTypeFilter, actionFilter, flaggedFilter, hasSelection, getSelectedIds, emit) => { diff --git a/testgen/ui/components/widgets/sidebar.py b/testgen/ui/components/widgets/sidebar.py index 6806d22c..2b68a3b6 100644 --- a/testgen/ui/components/widgets/sidebar.py +++ b/testgen/ui/components/widgets/sidebar.py @@ -25,6 +25,7 @@ def sidebar( support_email: str | None = None, global_context: bool = False, is_global_admin: bool = False, + account_path: str | None = None, ) -> None: """ Testgen custom component to display a styled menu over streamlit's @@ -53,6 +54,7 @@ def sidebar( "support_email": support_email, "global_context": global_context, "is_global_admin": is_global_admin, + "account_path": account_path, }, on_Navigate_change=_on_navigate, ) diff --git a/testgen/ui/navigation/page.py b/testgen/ui/navigation/page.py index 11e93f06..9f4de9a8 100644 --- a/testgen/ui/navigation/page.py +++ b/testgen/ui/navigation/page.py @@ -21,6 +21,7 @@ class Page(abc.ABC): menu_item: MenuItem | None = None permission: Permission | None = "view" can_activate: typing.ClassVar[list[CanActivateGuard] | None] = None + is_account_page: bool = False def __init__(self, router: testgen.ui.navigation.router.Router) -> None: self.router = router diff --git a/testgen/ui/pdf/hygiene_issue_report.py b/testgen/ui/pdf/hygiene_issue_report.py index a6d24fe3..b754ab19 100644 --- a/testgen/ui/pdf/hygiene_issue_report.py +++ b/testgen/ui/pdf/hygiene_issue_report.py @@ -138,6 +138,7 @@ def build_summary_table(document, hi_data): "transform_level", "aggregation_level", "data_product", + "data_classification", ] if hi_data[tag] ], diff --git a/testgen/ui/pdf/test_result_report.py b/testgen/ui/pdf/test_result_report.py index 9b57ff73..455b2bdf 100644 --- a/testgen/ui/pdf/test_result_report.py +++ b/testgen/ui/pdf/test_result_report.py @@ -154,6 +154,7 @@ def build_summary_table(document, tr_data): "transform_level", "aggregation_level", "data_product", + "data_classification", ] if tr_data[tag] ], diff --git a/testgen/ui/queries/profiling_queries.py b/testgen/ui/queries/profiling_queries.py index 70e5f681..49087d13 100644 --- a/testgen/ui/queries/profiling_queries.py +++ b/testgen/ui/queries/profiling_queries.py @@ -3,19 +3,10 @@ import pandas as pd import streamlit as st +from testgen.common.data_catalog_service import TAG_FIELDS from testgen.ui.services.database_service import fetch_all_from_db, fetch_df_from_db, fetch_one_from_db from testgen.utils import is_uuid4 -TAG_FIELDS = [ - "data_source", - "source_system", - "source_process", - "business_domain", - "stakeholder_group", - "transform_level", - "aggregation_level", - "data_product", -] COLUMN_PROFILING_FIELDS = """ -- Value Counts profile_results.record_ct, @@ -255,6 +246,7 @@ def get_tables_by_condition( table_chars.schema_name, table_chars.table_groups_id::VARCHAR AS table_group_id, -- Characteristics + table_chars.object_type, functional_table_type, approx_record_ct, table_chars.record_ct, @@ -621,7 +613,8 @@ def get_profiling_anomalies( COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM profile_anomaly_results r INNER JOIN profile_anomaly_types t ON r.anomaly_id = t.id @@ -677,7 +670,7 @@ def get_profiling_anomalies_by_ids(anomaly_ids: list[str]) -> pd.DataFrame: t.anomaly_description, r.detail, t.detail_redactable, t.suggested_action, t.dq_dimension, r.impact_dimension, r.anomaly_id, r.table_groups_id::VARCHAR, r.id::VARCHAR, p.profiling_starttime, r.profile_run_id::VARCHAR, - p.job_execution_id::VARCHAR as job_execution_id, + p.id::VARCHAR as job_execution_id, tg.table_groups_name, tg.project_code, dcc.functional_data_type, dcc.description as column_description, @@ -690,7 +683,8 @@ def get_profiling_anomalies_by_ids(anomaly_ids: list[str]) -> pd.DataFrame: COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM profile_anomaly_results r INNER JOIN profile_anomaly_types t ON r.anomaly_id = t.id diff --git a/testgen/ui/queries/scoring_queries.py b/testgen/ui/queries/scoring_queries.py index 26ff23a5..23f9e00b 100644 --- a/testgen/ui/queries/scoring_queries.py +++ b/testgen/ui/queries/scoring_queries.py @@ -42,7 +42,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list groups.table_groups_name, results.disposition, results.profile_run_id::VARCHAR, - runs.job_execution_id::VARCHAR, + runs.id::VARCHAR AS job_execution_id, types.suggested_action, results.table_groups_id::VARCHAR, results.project_code, @@ -59,6 +59,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list COALESCE(column_chars.transform_level, table_chars.transform_level, groups.transform_level) as transform_level, COALESCE(column_chars.aggregation_level, table_chars.aggregation_level) as aggregation_level, COALESCE(column_chars.data_product, table_chars.data_product, groups.data_product) as data_product, + COALESCE(column_chars.data_classification, table_chars.data_classification, groups.data_classification) as data_classification, types.impact_dimension, types.dq_dimension FROM profile_anomaly_results results @@ -108,7 +109,7 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list ELSE 'Passed' END as disposition, results.test_run_id::VARCHAR, - test_runs.job_execution_id::VARCHAR, + test_runs.id::VARCHAR AS job_execution_id, types.usage_notes, types.test_type, results.auto_gen, @@ -128,7 +129,9 @@ def get_score_card_issue_reports(selected_issues: list["SelectedIssue"]) -> list COALESCE(column_chars.transform_level, table_chars.transform_level, groups.transform_level) as transform_level, COALESCE(column_chars.aggregation_level, table_chars.aggregation_level) as aggregation_level, COALESCE(column_chars.data_product, table_chars.data_product, groups.data_product) as data_product, - COALESCE(results.impact_dimension, types.impact_dimension) as impact_dimension FROM test_results results + COALESCE(column_chars.data_classification, table_chars.data_classification, groups.data_classification) as data_classification, + COALESCE(results.impact_dimension, types.impact_dimension) as impact_dimension + FROM test_results results INNER JOIN test_types types ON (results.test_type = types.test_type) INNER JOIN test_suites suites @@ -179,6 +182,7 @@ def get_score_category_values(project_code: str) -> dict[ScoreCategory, list[str "stakeholder_group", "transform_level", "data_product", + "data_classification", ] quote = lambda v: f"'{v}'" diff --git a/testgen/ui/queries/table_group_queries.py b/testgen/ui/queries/table_group_queries.py index 0db27e12..1d4cdf6f 100644 --- a/testgen/ui/queries/table_group_queries.py +++ b/testgen/ui/queries/table_group_queries.py @@ -1,98 +1,60 @@ +"""UI cache adapter around the table-group preview service. + +The implementation lives in ``testgen.common.database.table_group_service`` so +MCP tools and the CLI can reuse identical logic without a Streamlit runtime. +This module wraps the service in ``@st.cache_data`` for the Streamlit pages. +""" + from collections.abc import Callable -from datetime import UTC, datetime -from typing import TypedDict from uuid import UUID import streamlit as st -from testgen.commands.queries.refresh_data_chars_query import RefreshDataCharsSQL -from testgen.commands.run_refresh_data_chars import write_data_chars -from testgen.common.database.column_chars import ColumnChars -from testgen.common.database.flavor.flavor_service import resolve_connection_params +from testgen.common.database.table_group_service import ( + TableGroupPreview, + make_save_data_chars, + preview_table_group, +) from testgen.common.models.connection import Connection from testgen.common.models.table_group import TableGroup -from testgen.ui.services.database_service import fetch_from_target_db from testgen.ui.services.query_cache import get_connection -class StatsPreview(TypedDict): - id: UUID - table_groups_name: str - table_group_schema: str - table_ct: int | None - column_ct: int | None - approx_record_ct: int | None - approx_data_point_ct: int | None - -class TablePreview(TypedDict): - column_ct: int - approx_record_ct: int | None - approx_data_point_ct: int | None - can_access: bool | None - - -class TableGroupPreview(TypedDict): - stats: StatsPreview - tables: dict[str, TablePreview] - success: bool - message: str | None - - def get_table_group_preview( table_group: TableGroup, connection: Connection | None = None, verify_table_access: bool = False, -) -> tuple[TableGroupPreview, Callable[[UUID], None]]: - table_group_preview: TableGroupPreview = { - "stats": { - "id": table_group.id, - "table_groups_name": table_group.table_groups_name, - "table_group_schema": table_group.table_group_schema, - }, - "tables": {}, - "success": True, - "message": None, - } - save_data_chars = None - - if connection or table_group.connection_id: - try: - connection = connection or get_connection(table_group.connection_id) - table_group_preview, data_chars, sql_generator = _get_preview(table_group, connection) - - def save_data_chars(table_group_id: UUID) -> None: - # Unsaved table groups will not have an ID, so we have to update it after saving - sql_generator.table_group.id = table_group_id - write_data_chars(data_chars, sql_generator, datetime.now(UTC)) - - if verify_table_access: - tables_preview = table_group_preview["tables"] - for table_name in tables_preview.keys(): - try: - results = fetch_from_target_db(connection, *sql_generator.verify_access(table_name)) - except Exception as error: - tables_preview[table_name]["can_access"] = False - else: - tables_preview[table_name]["can_access"] = results is not None and len(results) > 0 - - if not all(table["can_access"] for table in tables_preview.values()): - table_group_preview["message"] = ( - "Some tables were not accessible. Please the check the database permissions." - ) - except Exception as error: - table_group_preview["success"] = False - table_group_preview["message"] = error.args[0] - else: - table_group_preview["success"] = False - table_group_preview["message"] = ( - "No connection selected. Please select a connection to preview the Table Group." +) -> tuple[TableGroupPreview, Callable[[UUID], None] | None]: + """Streamlit-cached wrapper around ``preview_table_group``. + + The service returns ``(preview, data_chars, sql_generator)`` — all picklable — + so the cache can store the result. The ``save_data_chars`` closure is built + here, outside the cached function, because local closures can't be pickled. + + When the caller does not supply a ``Connection`` and the table group has a + ``connection_id``, the connection is resolved via the Streamlit cache so + repeated previews on the same page don't re-fetch it. + """ + if connection is None and table_group.connection_id: + connection = get_connection(table_group.connection_id) + + if verify_table_access: + preview, data_chars, sql_generator = preview_table_group( + table_group, connection=connection, verify_access=True, ) + else: + preview, data_chars, sql_generator = _cached_preview(table_group, connection) - return table_group_preview, save_data_chars + save = ( + make_save_data_chars(data_chars, sql_generator) + if data_chars is not None and sql_generator is not None + else None + ) + return preview, save def reset_table_group_preview() -> None: - _get_preview.clear() + _cached_preview.clear() @st.cache_data( @@ -107,61 +69,5 @@ def reset_table_group_preview() -> None: Connection: lambda x: x.to_dict(), }, ) -def _get_preview( - table_group: TableGroup, - connection: Connection, -) -> tuple[TableGroupPreview, list[ColumnChars], RefreshDataCharsSQL]: - sql_generator = RefreshDataCharsSQL(connection, table_group) - if sql_generator.flavor_service.metadata_via_api: - params = resolve_connection_params(connection.__dict__) - api_columns = sql_generator.flavor_service.get_schema_columns(params, table_group.table_group_schema) or [] - data_chars = sql_generator.filter_schema_columns(api_columns) - else: - rows = fetch_from_target_db(connection, *sql_generator.get_schema_ddf()) - data_chars = [ColumnChars(**column) for column in rows] - - preview: TableGroupPreview = { - "stats": { - "id": table_group.id, - "table_groups_name": table_group.table_groups_name, - "table_group_schema": table_group.table_group_schema, - "table_ct": 0, - "column_ct": 0, - "approx_record_ct": None, - "approx_data_point_ct": None, - }, - "tables": {}, - "success": True, - "message": None, - } - stats = preview["stats"] - tables = preview["tables"] - - for column in data_chars: - if not tables.get(column.table_name): - tables[column.table_name] = { - "column_ct": 0, - "approx_record_ct": column.approx_record_ct, - "approx_data_point_ct": None, - "can_access": None, - } - stats["table_ct"] += 1 - if column.approx_record_ct is not None: - stats["approx_record_ct"] = (stats["approx_record_ct"] or 0) + column.approx_record_ct - - stats["column_ct"] += 1 - tables[column.table_name]["column_ct"] += 1 - if column.approx_record_ct is not None: - stats["approx_data_point_ct"] = (stats["approx_data_point_ct"] or 0) + column.approx_record_ct - tables[column.table_name]["approx_data_point_ct"] = ( - tables[column.table_name]["approx_data_point_ct"] or 0 - ) + column.approx_record_ct - - if len(data_chars) <= 0: - preview["success"] = False - preview["message"] = ( - "No tables found matching the criteria. Please check the Table Group configuration" - " or the database permissions." - ) - - return preview, data_chars, sql_generator +def _cached_preview(table_group, connection): + return preview_table_group(table_group, connection=connection, verify_access=False) diff --git a/testgen/ui/queries/test_result_queries.py b/testgen/ui/queries/test_result_queries.py index 82ebc23a..ed3c520c 100644 --- a/testgen/ui/queries/test_result_queries.py +++ b/testgen/ui/queries/test_result_queries.py @@ -115,6 +115,7 @@ def get_test_results( r.test_definition_id::VARCHAR, r.auto_gen, td.flagged, + td.external_url, (SELECT COUNT(*) FROM test_definition_notes tdn WHERE tdn.test_definition_id = td.id) as notes_count, -- These are used in the PDF report @@ -130,7 +131,8 @@ def get_test_results( COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM run_results r INNER JOIN test_types tt ON (r.test_type = tt.test_type) @@ -199,7 +201,7 @@ def get_test_results_by_ids(test_result_ids: list[str]) -> pd.DataFrame: END as execution_error_ct, p.project_code, r.table_groups_id::VARCHAR, r.id::VARCHAR as test_result_id, r.test_run_id::VARCHAR, - tr.job_execution_id::VARCHAR as job_execution_id, + tr.id::VARCHAR as job_execution_id, c.id::VARCHAR as connection_id, r.test_suite_id::VARCHAR, r.test_definition_id::VARCHAR, r.auto_gen, @@ -217,7 +219,8 @@ def get_test_results_by_ids(test_result_ids: list[str]) -> pd.DataFrame: COALESCE(dcc.stakeholder_group, dtc.stakeholder_group, tg.stakeholder_group) as stakeholder_group, COALESCE(dcc.transform_level, dtc.transform_level, tg.transform_level) as transform_level, COALESCE(dcc.aggregation_level, dtc.aggregation_level) as aggregation_level, - COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product + COALESCE(dcc.data_product, dtc.data_product, tg.data_product) as data_product, + COALESCE(dcc.data_classification, dtc.data_classification, tg.data_classification) as data_classification FROM test_results r INNER JOIN test_runs tr ON (r.test_run_id = tr.id) diff --git a/testgen/ui/services/query_cache.py b/testgen/ui/services/query_cache.py index 55be1e67..2e7367a7 100644 --- a/testgen/ui/services/query_cache.py +++ b/testgen/ui/services/query_cache.py @@ -17,7 +17,7 @@ from testgen.common.models.profiling_run import ProfilingRun, ProfilingRunMinimal, ProfilingRunSummary from testgen.common.models.project import Project, ProjectSummary from testgen.common.models.project_membership import ProjectMembership -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import ( TableGroup, TableGroupMinimal, @@ -146,10 +146,7 @@ def get_profiling_run_summaries( @st.cache_data(show_spinner=False) def get_monitor_schedule(monitor_suite_id: str | UUID) -> JobSchedule | None: - return JobSchedule.get( - JobSchedule.key == RUN_MONITORS_JOB_KEY, - JobSchedule.kwargs["test_suite_id"].astext == str(monitor_suite_id), - ) + return JobSchedule.get_for_monitor_suite(monitor_suite_id) # -- Connection --------------------------------------------------------------- diff --git a/testgen/ui/static/css/shared.css b/testgen/ui/static/css/shared.css index 4d5b1f0b..4a0ff545 100644 --- a/testgen/ui/static/css/shared.css +++ b/testgen/ui/static/css/shared.css @@ -26,6 +26,20 @@ body { --empty-dark: #BDBDBD; --empty-teal: #E7F1F0; + /* Faceted test-picker taxonomy palette (test_picker_taxonomy.js). Deep, saturated hues tuned + for light backgrounds; overridden with lighter, higher-contrast variants in dark mode below. */ + --facet-crimson: #9F1239; + --facet-blue: #1E40AF; + --facet-amber: #B45309; + --facet-purple: #6B21A8; + --facet-rust: #7C2D12; + --facet-teal: #0E5F4F; + --facet-bronze: #92400E; + --facet-indigo: #312E81; + --facet-orange: #C75D2C; + --facet-emerald: #1B7C68; + --facet-grey: #6A6F7E; + --primary-text-color: #000000de; --secondary-text-color: #0000008a; --disabled-text-color: #00000042; @@ -97,6 +111,19 @@ body { --empty-dark: #757575; --empty-teal: #242E2D; + /* Lighter, higher-contrast facet hues for dark backgrounds (chip text + facet dots). */ + --facet-crimson: #FB7185; + --facet-blue: #60A5FA; + --facet-amber: #FBBF24; + --facet-purple: #C084FC; + --facet-rust: #F4845F; + --facet-teal: #2DD4BF; + --facet-bronze: #D9A05B; + --facet-indigo: #A5B4FC; + --facet-orange: #FB923C; + --facet-emerald: #34D399; + --facet-grey: #9CA3AF; + --primary-text-color: rgba(255, 255, 255); --secondary-text-color: rgba(255, 255, 255, .7); --disabled-text-color: rgba(255, 255, 255, .5); diff --git a/testgen/ui/static/js/components/connection_form.js b/testgen/ui/static/js/components/connection_form.js index 1fadaacf..c0d3b1d7 100644 --- a/testgen/ui/static/js/components/connection_form.js +++ b/testgen/ui/static/js/components/connection_form.js @@ -784,7 +784,7 @@ const AzureMSSQLForm = ( ); }; -const SynapseMSSQLForm = RedshiftForm; +const SynapseMSSQLForm = AzureMSSQLForm; const MSSQLForm = RedshiftForm; diff --git a/testgen/ui/static/js/components/dialog.js b/testgen/ui/static/js/components/dialog.js index 465c3faf..0f285fd3 100644 --- a/testgen/ui/static/js/components/dialog.js +++ b/testgen/ui/static/js/components/dialog.js @@ -5,6 +5,7 @@ * @property {import('../van.min.js').State} open - Reactive open state * @property {Function} onClose - Called when the dialog is closed (backdrop click or X button) * @property {string} [width] - CSS width value, default '30rem' + * @property {(Element | Function)} [headerActions] - Optional node rendered right-aligned in header row */ import van from '../van.min.js'; import { getValue, loadStylesheet } from '../utils.js'; @@ -28,7 +29,7 @@ const { button, div, i, span } = van.tags; * @param {DialogProps} props * @param {...(Element | string)} children - Content rendered in the dialog body */ -const Dialog = ({ title, open, onClose, width = '30rem' }, ...children) => { +const Dialog = ({ title, open, onClose, width = '30rem', headerActions }, ...children) => { loadStylesheet('dialog', stylesheet); const overlay = div( @@ -50,6 +51,7 @@ const Dialog = ({ title, open, onClose, width = '30rem' }, ...children) => { div( { class: 'tg-dialog-header' }, span({ 'data-testid': 'dialog-title', class: 'tg-dialog-title' }, title), + headerActions ? div({ class: 'tg-dialog-header-actions' }, headerActions) : null, ), div({ 'data-testid': 'dialog-content', class: 'tg-dialog-content' }, ...children), button( @@ -117,6 +119,12 @@ stylesheet.replace(` flex-shrink: 0; } +.tg-dialog-header-actions { + margin-left: auto; + font-size: initial; + font-weight: initial; +} + .tg-dialog-content { padding: 0.75rem 1.5rem 1.5rem; overflow-y: auto; diff --git a/testgen/ui/static/js/components/link.js b/testgen/ui/static/js/components/link.js index c78e821c..e709f601 100644 --- a/testgen/ui/static/js/components/link.js +++ b/testgen/ui/static/js/components/link.js @@ -20,6 +20,7 @@ * @property {((event: any) => void)?} onClick */ import { getValue, loadStylesheet } from '../utils.js'; +import { Tooltip } from './tooltip.js'; import van from '../van.min.js'; const { a, div, i, span } = van.tags; diff --git a/testgen/ui/static/js/components/score_breakdown.js b/testgen/ui/static/js/components/score_breakdown.js index 0b843405..cb069ee5 100644 --- a/testgen/ui/static/js/components/score_breakdown.js +++ b/testgen/ui/static/js/components/score_breakdown.js @@ -194,6 +194,7 @@ const CATEGORIES = { stakeholder_group: 'Stakeholder Group', transform_level: 'Transform Level', data_product: 'Data Product', + data_classification: 'Data Classification', }; const BREAKDOWN_COLUMN_LABEL = { diff --git a/testgen/ui/static/js/components/table_group_edit_dialog.js b/testgen/ui/static/js/components/table_group_edit_dialog.js index a5af0c5d..ac00f375 100644 --- a/testgen/ui/static/js/components/table_group_edit_dialog.js +++ b/testgen/ui/static/js/components/table_group_edit_dialog.js @@ -14,6 +14,7 @@ * @property {Connection[]} connections * @property {TableGroup} table_group * @property {boolean} is_in_use + * @property {boolean?} can_view_pii * @property {TableGroupPreview?} table_group_preview * @property {EditResult?} result */ @@ -61,6 +62,7 @@ const TableGroupEditDialog = (props) => { showConnectionSelector: connections.length > 1, disableConnectionSelector: false, disableSchemaField: getValue(props.is_in_use) ?? false, + disablePiiFlag: !(getValue(props.can_view_pii) ?? false), onChange: (updatedTableGroup, state) => { tableGroupState.val = updatedTableGroup; formValid.val = state.valid; diff --git a/testgen/ui/static/js/components/table_group_form.js b/testgen/ui/static/js/components/table_group_form.js index 93becaa9..739f9210 100644 --- a/testgen/ui/static/js/components/table_group_form.js +++ b/testgen/ui/static/js/components/table_group_form.js @@ -29,6 +29,7 @@ * @property {string?} stakeholder_group * @property {string?} transform_level * @property {string?} data_product + * @property {string?} data_classification * * @typedef FormState * @type {object} @@ -55,8 +56,12 @@ import { required } from '../form_validators.js'; import { Select } from './select.js'; import { Caption } from './caption.js'; import { Textarea } from './textarea.js'; +import { Icon } from './icon.js'; -const { div } = van.tags; +// Flavors whose profiling sample clause cannot sample views — views are profiled in full. +const SAMPLE_SKIPS_VIEWS_FLAVORS = ['mssql', 'postgresql']; + +const { div, span } = van.tags; const normalizeTableSet = (value) => { return value?.split(/[,\n]/) @@ -100,6 +105,7 @@ const TableGroupForm = (props) => { const stakeholderGroup = van.state(tableGroup.stakeholder_group); const transformLevel = van.state(tableGroup.transform_level); const dataProduct = van.state(tableGroup.data_product); + const dataClassification = van.state(tableGroup.data_classification); const connectionOptions = van.derive(() => { const connections = getValue(props.connections) ?? []; @@ -119,6 +125,13 @@ const TableGroupForm = (props) => { return flavor === 'salesforce_data360'; }); + const sampleSkipsViews = van.derive(() => { + const connections = getValue(props.connections) ?? []; + const selected = connections.find(c => c.connection_id === tableGroupConnectionId.val); + const flavor = selected?.sql_flavor ?? getValue(props.sqlFlavor); + return SAMPLE_SKIPS_VIEWS_FLAVORS.includes(flavor); + }); + const updatedTableGroup = van.derive(() => { return { id: tableGroup.id, @@ -148,6 +161,7 @@ const TableGroupForm = (props) => { stakeholder_group: stakeholderGroup.val, transform_level: transformLevel.val, data_product: dataProduct.val, + data_classification: dataClassification.val, }; }); const dirty = van.derive(() => !isEqual(updatedTableGroup.val, tableGroup)); @@ -206,7 +220,7 @@ const TableGroupForm = (props) => { addScorecardDefinition, ), SamplingForm( - { setValidity: setFieldValidity }, + { setValidity: setFieldValidity, sampleSkipsViews }, profileUseSampling, profileSamplePercent, profileSampleMinCount, @@ -222,6 +236,7 @@ const TableGroupForm = (props) => { stakeholderGroup, transformLevel, dataProduct, + dataClassification, ), ); }; @@ -447,6 +462,13 @@ const SamplingForm = ( }, }), ), + () => (getValue(options.sampleSkipsViews) && profileUseSampling.val) + ? div( + { class: 'flex-row fx-align-center fx-gap-1' }, + Icon({ style: 'color: var(--orange); font-size: 18px;' }, 'warning'), + span('Views are profiled in full on this database — sampling applies only to tables and materialized views.'), + ) + : '', ), ); }; @@ -462,6 +484,7 @@ const TaggingForm = ( stakeholderGroup, transformLevel, dataProduct, + dataClassification, ) => { return ExpansionPanel( { title: 'Table Group Tags' }, @@ -557,6 +580,16 @@ const TaggingForm = ( options.setValidity?.('data_product', state.valid); }, }), + Input({ + name: 'data_classification', + label: 'Data Classification', + value: dataClassification, + help: 'Information sensitivity level of the dataset, e.g., Public, Internal, Confidential, Restricted', + onChange: (value, state) => { + dataClassification.val = value; + options.setValidity?.('data_classification', state.valid); + }, + }), ), ); }; diff --git a/testgen/ui/static/js/components/tabs.js b/testgen/ui/static/js/components/tabs.js index d6cbc51c..ae23f7b1 100644 --- a/testgen/ui/static/js/components/tabs.js +++ b/testgen/ui/static/js/components/tabs.js @@ -20,6 +20,8 @@ const Tab = ({ label }, ...children) => ({ /** * @typedef {Object} TabsProps * @property {string?} class + * @property {van.State?} activeTab Optional external state for the active tab index. Pass one + * to control or persist the selection across re-renders; omit to use internal state. * * @param {TabsProps} props * @param {...Tab} tabs @@ -27,9 +29,9 @@ const Tab = ({ label }, ...children) => ({ const Tabs = (props, ...tabs) => { loadStylesheet('tabs', stylesheet); - const { ...restProps } = props; + const { activeTab: activeTabProp, ...restProps } = props; - const activeTab = van.state(0); + const activeTab = activeTabProp ?? van.state(0); let labelsContainerEl; const highlightEl = span({ class: "tg-tabs--highlight" }); diff --git a/testgen/ui/static/js/components/test_definition_form.js b/testgen/ui/static/js/components/test_definition_form.js index e7fa14d6..88dae2ba 100644 --- a/testgen/ui/static/js/components/test_definition_form.js +++ b/testgen/ui/static/js/components/test_definition_form.js @@ -61,6 +61,7 @@ * @property {TestDefinition} definition * @property {string?} class * @property {boolean} qualifiesTableRefsWithSchema + * @property {boolean?} hideHeader Hide the test-type name/description header (when the host already shows it). * @property {(changes: object, valid: boolean) => void} onChange */ @@ -72,6 +73,7 @@ import { Textarea } from './textarea.js'; import { RadioGroup } from './radio_group.js'; import { Caption } from './caption.js'; import { numberBetween, required } from '../form_validators.js'; +import { capitalize } from '../display_utils.js'; const { div, span } = van.tags; @@ -107,7 +109,7 @@ const TestDefinitionForm = (/** @type Properties */ props) => { .map((column, index) => ({ ...(PARAMETER_CONFIG[column] || { type: 'text' }), column, - label: paramLabels[index] || column.replaceAll('_', ' '), + label: paramLabels[index] || capitalize(column.replaceAll('_', ' ')), help: paramHelp[index] || null, validators: paramRequired[index] ? [required] : undefined, })) @@ -143,13 +145,15 @@ const TestDefinitionForm = (/** @type Properties */ props) => { return div( { class: props.class }, - div( - { class: 'mb-2' }, - div({ class: 'text-large' }, definition.test_name_short), - definition.test_description || definition.default_test_description - ? span({ class: 'text-caption mt-2' }, definition.test_description ?? definition.default_test_description) - : null, - ), + getValue(props.hideHeader) + ? null + : div( + { class: 'mb-2' }, + div({ class: 'text-large' }, definition.test_name_short), + definition.test_description || definition.default_test_description + ? span({ class: 'text-caption mt-2' }, definition.test_description ?? definition.default_test_description) + : null, + ), () => div( { class: 'flex-row fx-flex-wrap fx-gap-3' }, dynamicParamColumns.map(config => { diff --git a/testgen/ui/static/js/components/test_picker_taxonomy.js b/testgen/ui/static/js/components/test_picker_taxonomy.js new file mode 100644 index 00000000..64673b65 --- /dev/null +++ b/testgen/ui/static/js/components/test_picker_taxonomy.js @@ -0,0 +1,108 @@ +// Faceted Add-Test picker taxonomy: colors, axis definitions, and the column-aware pre-filter. +// All seven facets read a column on test_types (no derivation here): `impact_dimension`, +// `dq_dimension`, `health_dimension`, `algorithm`, `statistical_technique`, `test_scope`, and +// `criteria`. `criteria` is derived once in Python (derive_test_criteria) and shared by UI + MCP, +// so this file holds presentation only. Colors reference the --facet-* CSS variables defined in +// shared.css, which carry light values plus higher-contrast dark-mode overrides. + +const IMPACT_COLOR = { + 'Conformance': 'var(--facet-crimson)', 'Reliability': 'var(--facet-blue)', 'Regularity': 'var(--facet-amber)', 'Usability': 'var(--facet-purple)', +}; + +const QD_COLOR = { + 'Validity': 'var(--facet-crimson)', 'Completeness': 'var(--facet-blue)', 'Consistency': 'var(--facet-rust)', 'Accuracy': 'var(--facet-teal)', + 'Timeliness': 'var(--facet-amber)', 'Uniqueness': 'var(--facet-purple)', 'Recency': 'var(--facet-bronze)', +}; + +const HD_COLOR = { + 'Schema Drift': 'var(--facet-indigo)', 'Data Drift': 'var(--facet-orange)', 'Volume': 'var(--facet-teal)', 'Freshness': 'var(--facet-amber)', +}; + +const ALGO_COLOR = { + 'Boundary check': 'var(--facet-rust)', + 'Counting': 'var(--facet-blue)', + 'Pattern / regex': 'var(--facet-crimson)', + 'Set / lookup': 'var(--facet-purple)', + 'Statistical drift': 'var(--facet-orange)', + 'Aggregate reconciliation': 'var(--facet-bronze)', + 'Freshness / time': 'var(--facet-amber)', + 'Schema / metadata': 'var(--facet-indigo)', + 'Custom SQL': 'var(--facet-teal)', +}; + +const TECH_COLOR = { + "Cohen's D": 'var(--facet-blue)', + "Cohen's H": 'var(--facet-purple)', + 'Outlier Detection': 'var(--facet-orange)', + 'SD Shift': 'var(--facet-rust)', + 'Jensen-Shannon Divergence': 'var(--facet-teal)', +}; + +// Canonical order for the sparse Statistical Technique facet. +const TECHNIQUE_ORDER = ["Cohen's D", "Cohen's H", 'Outlier Detection', 'SD Shift', 'Jensen-Shannon Divergence']; + +const CRITERIA_COLOR = { + 'Defined Rule': 'var(--facet-teal)', 'Defined Threshold': 'var(--facet-blue)', 'Defined Value': 'var(--facet-emerald)', + 'List of Values': 'var(--facet-purple)', 'Reference Dataset': 'var(--facet-bronze)', 'Custom Criteria': 'var(--facet-crimson)', +}; +// Canonical order for the Criteria facet (least to most setup effort). +const CRITERIA_ORDER = [ + 'Defined Rule', 'Defined Threshold', 'Defined Value', + 'List of Values', 'Reference Dataset', 'Custom Criteria', +]; + +const SCOPE_LABEL = { + column: 'Column', table: 'Table', referential: 'Referential', custom: 'Custom', tablegroup: 'Table Group', +}; +const SCOPE_ORDER = ['Column', 'Table', 'Referential', 'Custom', 'Table Group']; + +const FALLBACK = 'var(--facet-grey)'; +const EMPTY = '—'; + +// ── Column-aware pre-filter (PRD §9) ────────────────────────────────────────── +// Not a facet: the separate "hide tests that don't apply to this column" toggle. +// `generalType` is data_column_chars.general_type: 'N' Numeric, 'D' Datetime, 'T' Time, +// 'A' Alpha, 'B' Boolean, 'X' Other. Lists enumerate type affinity by test code; gating on +// scope first excludes table/referential-scoped codes regardless. +const NUMERIC_ONLY = ['Min_Val', 'Avg_Shift', 'Incr_Avg_Shift', 'Variability_Increase', 'Variability_Decrease', + 'Outlier_Pct_Above', 'Outlier_Pct_Below', 'Dec_Trunc', 'Distribution_Shift', 'Aggregate_Minimum']; +const DATE_ONLY = ['Min_Date', 'Future_Date', 'Future_Date_1Y', 'Recency', 'Distinct_Date_Ct', + 'Daily_Record_Ct', 'Weekly_Rec_Ct', 'Monthly_Rec_Ct', 'Valid_Month', 'Freshness_Trend', 'Table_Freshness']; + +function appliesToSelectedColumn(t, generalType) { + if (['table', 'tablegroup', 'referential'].includes(t.test_scope)) return false; + // Unknown column type (null/undefined/empty): we can't assert a test is inapplicable, so + // don't hide it. Without this, an absent general_type fails every type check below and + // silently hides all numeric/date tests. Known non-numeric/non-date codes ('A','B','X') + // are truthy and still filter correctly. + if (!generalType) return true; + if (NUMERIC_ONLY.includes(t.test_type) && generalType !== 'N') return false; + const isDate = generalType === 'D' || generalType === 'T'; + if (DATE_ONLY.includes(t.test_type) && !isDate) return false; + return true; +} + +// ── Axis registry ───────────────────────────────────────────────────────────── +// `value(t)` returns the facet value for a test (or null when absent). `color(key)` returns a +// hex color. `sparse` axes filter out null-valued tests when active. `order`, when present, +// fixes the facet's display order. + +const AXES = { + impact: { label: 'Impact Dimension', value: (t) => t.impact_dimension || null, color: (k) => IMPACT_COLOR[k] || FALLBACK }, + dq: { label: 'Quality Dimension', value: (t) => t.dq_dimension || null, color: (k) => QD_COLOR[k] || FALLBACK }, + health: { label: 'Health Dimension', value: (t) => t.health_dimension || null, color: (k) => HD_COLOR[k] || FALLBACK }, + algorithm: { label: 'Algorithm', value: (t) => t.algorithm || null, color: (k) => ALGO_COLOR[k] || FALLBACK }, + technique: { label: 'Statistical Technique', value: (t) => t.statistical_technique || null, color: (k) => TECH_COLOR[k] || FALLBACK, sparse: true, order: TECHNIQUE_ORDER }, + scope: { label: 'Test Scope', value: (t) => SCOPE_LABEL[t.test_scope] || 'Other', color: () => FALLBACK, order: SCOPE_ORDER }, + criteria: { label: 'Criteria', value: (t) => t.criteria || null, color: (k) => CRITERIA_COLOR[k] || FALLBACK, order: CRITERIA_ORDER }, +}; + +// Facet order in the rail (PRD facet review). All facets are always visible. +const FACET_AXES = ['impact', 'dq', 'health', 'algorithm', 'technique', 'scope', 'criteria']; +// Group-by options offered in the result-list dropdown. +const GROUP_BY_AXES = ['impact', 'dq', 'health', 'algorithm', 'technique', 'scope', 'criteria']; + +export { + AXES, FACET_AXES, GROUP_BY_AXES, TECHNIQUE_ORDER, CRITERIA_ORDER, SCOPE_LABEL, EMPTY, + appliesToSelectedColumn, +}; diff --git a/testgen/ui/static/js/form_validators.js b/testgen/ui/static/js/form_validators.js index 58c085bb..799292aa 100644 --- a/testgen/ui/static/js/form_validators.js +++ b/testgen/ui/static/js/form_validators.js @@ -140,6 +140,28 @@ function notIn(values, options) { }; } +/** + * Validate that the value is empty (optional) or a well-formed JSON object of key-value pairs. + * + * @param {any} value + * @returns {string} + */ +function jsonObject(value) { + if (!value || !String(value).trim()) { + return null; + } + let parsed; + try { + parsed = JSON.parse(value); + } catch { + return 'Must be valid JSON.'; + } + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + return 'Must be a JSON object of key-value pairs.'; + } + return null; +} + export { maxLength, minLength, @@ -149,4 +171,5 @@ export { requiredIf, sizeLimit, notIn, + jsonObject, }; diff --git a/testgen/ui/static/js/sidebar.js b/testgen/ui/static/js/sidebar.js index 9ebd4b2f..89006e5d 100644 --- a/testgen/ui/static/js/sidebar.js +++ b/testgen/ui/static/js/sidebar.js @@ -36,6 +36,7 @@ * @property {string} support_email * @property {boolean} global_context * @property {boolean} is_global_admin + * @property {(string|null)} account_path */ const van = window.top.van; const { a, button, div, i, img, label, option, select, span } = van.tags; @@ -93,11 +94,7 @@ const Sidebar = (/** @type {Properties} */ props) => { }, ), div( - div( - { class: 'menu--user' }, - span({class: 'menu--username', title: props.username}, props.username), - span({class: 'menu--role'}, () => props.role.val?.replace('_', ' ')), - ), + () => UserBlock(props, currentProject), div( { class: 'menu--buttons' }, button( @@ -167,6 +164,36 @@ const ProjectSelect = (/** @type Project[] */ projects, /** @type string */ curr ); }; +const UserBlock = (/** @type Properties */ props, currentProject) => { + const roleAndProject = () => { + const role = props.role.val?.replace('_', ' '); + const projectName = currentProject.val?.name; + return [role, projectName].filter(Boolean).join(' · '); + }; + + const userInfo = div( + { class: 'menu--user--info' }, + span({ class: 'menu--username', title: props.username }, props.username), + span({ class: 'menu--role' }, roleAndProject), + ); + + const accountPath = props.account_path?.val; + if (!accountPath) { + return div({ class: 'menu--user' }, userInfo); + } + + return a( + { + class: () => isCurrentPage(accountPath, props.current_page?.val) ? 'menu--user--link active' : 'menu--user--link', + href: `/${accountPath}?${PROJECT_CODE_QUERY_PARAM}=${currentProject.val?.code ?? ''}`, + onclick: (event) => navigate(event, accountPath, { [PROJECT_CODE_QUERY_PARAM]: currentProject.val?.code }), + }, + i({ class: 'menu--user--icon material-symbols-rounded' }, 'manage_accounts'), + userInfo, + i({ class: 'menu--user--chevron material-symbols-rounded' }, 'chevron_right'), + ); +}; + const MenuSection = ( /** @type {MenuItem} */ item, /** @type {string} */ currentPage, @@ -336,6 +363,49 @@ stylesheet.replace(` padding: 16px; } +.menu .menu--user--link { + display: flex; + flex-direction: row; + align-items: center; + gap: 12px; + margin: 4px 8px; + padding: 8px 12px; + border-radius: 8px; + color: var(--primary-text-color); + text-decoration: none; + cursor: pointer; +} + +.menu .menu--user--link:hover { + background: var(--sidebar-item-hover-color); +} + +.menu .menu--user--link.active { + color: var(--primary-color); + background: var(--sidebar-active-item-color); +} + +.menu .menu--user--info { + display: flex; + flex-direction: column; + flex: 1; + min-width: 0; +} + +.menu .menu--user--icon { + flex: none; + font-size: 24px; + line-height: 24px; +} + +.menu .menu--user--chevron { + flex: none; + margin-left: auto; + font-size: 18px; + line-height: 18px; + color: var(--secondary-text-color); +} + .menu .menu--username { overflow-x: hidden; text-overflow: ellipsis; @@ -346,6 +416,9 @@ stylesheet.replace(` text-transform: uppercase; font-size: 12px; color: var(--secondary-text-color); + overflow-x: hidden; + text-overflow: ellipsis; + text-wrap: nowrap; } .menu .content > .menu--section > .menu--section--label { diff --git a/testgen/ui/static/js/utils.js b/testgen/ui/static/js/utils.js index 71eb7403..bcd45d3b 100644 --- a/testgen/ui/static/js/utils.js +++ b/testgen/ui/static/js/utils.js @@ -189,6 +189,12 @@ function isDataURL(/** @type string */ url) { return url.startsWith('data:'); } +// Guards a user-supplied value before it is rendered as a link href, blocking +// javascript:/data:/vbscript: URIs. Only http(s) URLs may become clickable links. +function isHttpUrl(/** @type any */ value) { + return typeof value === 'string' && /^https?:\/\//i.test(value.trim()); +} + function checkIsRequired(validators) { let isRequired = validators.some(v => v.name === 'required'); if (!isRequired) { @@ -228,4 +234,4 @@ function createEmitter(setTriggerValue) { }; } -export { afterMount, createEmitter, debounce, fillViewportHeight, getRandomId, getValue, getParents, isEqual, isState, loadStylesheet, friendlyPercent, slugify, isDataURL, checkIsRequired, onFrameResized, parseDate }; +export { afterMount, createEmitter, debounce, fillViewportHeight, getRandomId, getValue, getParents, isEqual, isState, loadStylesheet, friendlyPercent, slugify, isDataURL, isHttpUrl, checkIsRequired, onFrameResized, parseDate }; diff --git a/testgen/ui/views/connections.py b/testgen/ui/views/connections.py index 619e10ed..53fa5f7d 100644 --- a/testgen/ui/views/connections.py +++ b/testgen/ui/views/connections.py @@ -1,25 +1,17 @@ import base64 import logging -import random import typing -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass import streamlit as st -from testgen.commands.test_generation import run_monitor_generation -from testgen.ui.queries import table_group_queries - -try: - from pyodbc import Error as PyODBCError -except ImportError: - PyODBCError = None -from sqlalchemy.exc import DatabaseError, DBAPIError - -import testgen.ui.services.database_service as db from testgen import settings -from testgen.common.database.database_service import empty_cache, get_flavor_service +from testgen.commands.test_generation import run_monitor_generation +from testgen.common.database.connection_service import ConnectionStatus, test_connection_status +from testgen.common.database.database_service import get_flavor_service from testgen.common.database.flavor.flavor_service import resolve_connection_params from testgen.common.enums import JobSource +from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL from testgen.common.models import get_current_session, with_database_session from testgen.common.models.connection import Connection, ConnectionMinimal from testgen.common.models.job_execution import JobExecution @@ -30,6 +22,7 @@ from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page +from testgen.ui.queries import table_group_queries from testgen.ui.services.query_cache import ( get_connection, select_connections_where, @@ -256,43 +249,8 @@ def _format_connection(self, connection: Connection, should_test: bool = False) formatted_connection["status"] = asdict(self.test_connection(connection)) return formatted_connection - def test_connection(self, connection: Connection) -> "ConnectionStatus": - empty_cache() - try: - flavor_service = get_flavor_service(connection.sql_flavor) - results = db.fetch_from_target_db(connection, flavor_service.test_query) - connection_successful = len(results) == 1 and next(iter(results[0].values())) == 1 - - if not connection_successful: - return ConnectionStatus(message="Error completing a query to the database server.", successful=False) - return ConnectionStatus(message="The connection was successful.", successful=True) - except KeyError: - return ConnectionStatus( - message="Error attempting the connection. ", - details="Complete all the required fields.", - successful=False, - ) - except DatabaseError as error: - LOG.exception("Error testing database connection") - return ConnectionStatus(message="Error attempting the connection.", details=str(error.orig), successful=False) - except DBAPIError as error: - LOG.exception("Error testing database connection") - details = str(error.orig) - if PyODBCError and isinstance(error.orig, PyODBCError) and error.orig.args: - details = error.orig.args[1] - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) - except (TypeError, ValueError) as error: - LOG.exception("Error testing database connection") - details = str(error) - if is_open_ssl_error(error): - details = error.args[0] - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) - except Exception as error: - details = "Something went wrong while testing the connection." - if connection.connect_by_key and not connection.private_key: - details = "The private key is missing." - LOG.exception("Error testing database connection") - return ConnectionStatus(message="Error attempting the connection.", details=details, successful=False) + def test_connection(self, connection: Connection) -> ConnectionStatus: + return test_connection_status(connection) @with_database_session def setup_data_configuration(self, project_code: str, connection_id: str) -> None: @@ -539,24 +497,6 @@ def on_close_clicked(_params: dict) -> None: return wizard_data, wizard_handlers -@dataclass(frozen=True, slots=True) -class ConnectionStatus: - message: str - successful: bool - details: str | None = field(default=None) - _: float = field(default_factory=random.random) - - -def is_open_ssl_error(error: Exception): - return ( - error.args - and len(error.args) > 1 - and isinstance(error.args[1], list) - and len(error.args[1]) > 0 - and type(error.args[1][0]).__name__ == "OpenSSLError" - ) - - def format_connection(connection: Connection | ConnectionMinimal) -> dict: formatted_connection = connection.to_dict(json_safe=True) @@ -582,79 +522,31 @@ class ConnectionFlavor: flavor: str +# Labels and families come from the shared source in ``common/flavors.py`` (the +# same labels MCP uses); only the icon + display order are UI-specific. +_FLAVOR_ICONS: dict[str, str] = { + "redshift": "flavors/redshift.svg", + "redshift_spectrum": "flavors/redshift.svg", + "azure_mssql": "flavors/azure_sql.svg", + "synapse_mssql": "flavors/azure_synapse_table.svg", + "databricks": "flavors/databricks.svg", + "bigquery": "flavors/bigquery.svg", + "mssql": "flavors/mssql.svg", + "oracle": "flavors/oracle.svg", + "postgresql": "flavors/postgresql.svg", + "sap_hana": "flavors/sap_hana.svg", + "salesforce_data360": "flavors/salesforce_data360.svg", + "snowflake": "flavors/snowflake.svg", +} + FLAVOR_OPTIONS = [ ConnectionFlavor( - label="Amazon Redshift", - value="redshift", - flavor="redshift", - icon=get_asset_data_url("flavors/redshift.svg"), - ), - ConnectionFlavor( - label="Amazon Redshift Spectrum", - value="redshift_spectrum", - flavor="redshift_spectrum", - icon=get_asset_data_url("flavors/redshift.svg"), - ), - ConnectionFlavor( - label="Azure SQL Database", - value="azure_mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/azure_sql.svg"), - ), - ConnectionFlavor( - label="Azure Synapse Analytics", - value="synapse_mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/azure_synapse_table.svg"), - ), - ConnectionFlavor( - label="Databricks", - value="databricks", - flavor="databricks", - icon=get_asset_data_url("flavors/databricks.svg"), - ), - ConnectionFlavor( - label="Google BigQuery", - value="bigquery", - flavor="bigquery", - icon=get_asset_data_url("flavors/bigquery.svg"), - ), - ConnectionFlavor( - label="Microsoft SQL Server", - value="mssql", - flavor="mssql", - icon=get_asset_data_url("flavors/mssql.svg"), - ), - ConnectionFlavor( - label="Oracle", - value="oracle", - flavor="oracle", - icon=get_asset_data_url("flavors/oracle.svg"), - ), - ConnectionFlavor( - label="PostgreSQL", - value="postgresql", - flavor="postgresql", - icon=get_asset_data_url("flavors/postgresql.svg"), - ), - ConnectionFlavor( - label="SAP HANA", - value="sap_hana", - flavor="sap_hana", - icon=get_asset_data_url("flavors/sap_hana.svg"), - ), - ConnectionFlavor( - label="Salesforce Data 360", - value="salesforce_data360", - flavor="salesforce_data360", - icon=get_asset_data_url("flavors/salesforce_data360.svg"), - ), - ConnectionFlavor( - label="Snowflake", - value="snowflake", - flavor="snowflake", - icon=get_asset_data_url("flavors/snowflake.svg"), - ), + value=code, + label=str(FLAVOR_CODE_TO_LABEL[code]), + flavor=FLAVOR_CODE_TO_FAMILY[code], + icon=get_asset_data_url(icon), + ) + for code, icon in _FLAVOR_ICONS.items() ] # SAP HANA is hidden in the Docker image because pyhdbcli is glibc-only and fails to load on Alpine/musl. diff --git a/testgen/ui/views/data_catalog.py b/testgen/ui/views/data_catalog.py index d316f237..ca068ebd 100644 --- a/testgen/ui/views/data_catalog.py +++ b/testgen/ui/views/data_catalog.py @@ -4,16 +4,26 @@ from collections import defaultdict from datetime import datetime from functools import partial +from uuid import UUID import pandas as pd import streamlit as st from sqlalchemy.sql.expression import func as sa_func from streamlit.delta_generator import DeltaGenerator -from testgen.common.database.database_service import get_flavor_service +from testgen.common import date_service +from testgen.common.data_catalog_service import ( + apply_column_metadata, + apply_table_metadata, + build_create_table_script, + disable_autoflags, + fetch_table_sample, +) from testgen.common.enums import JobSource from testgen.common.models import database_session, with_database_session from testgen.common.models.connection import Connection +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable from testgen.common.models.job_execution import JobExecution from testgen.common.models.profiling_run import ProfilingRun from testgen.common.models.table_group import TableGroup, TableGroupMinimal @@ -22,7 +32,6 @@ get_pii_columns, mask_hygiene_detail, mask_profiling_pii, - mask_source_data_pii, ) from testgen.common.profile_top_values import parse_top_freq_values, parse_top_patterns from testgen.ui.components import widgets as testgen @@ -46,7 +55,7 @@ get_tables_by_id, get_tables_by_table_group, ) -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db, fetch_from_target_db +from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db from testgen.ui.services.query_cache import ( get_profiling_run_summaries, get_project_summary, @@ -61,7 +70,6 @@ build_import_preview_props, parse_import_csv, ) -from testgen.ui.views.dialogs.table_create_script_dialog import generate_create_script from testgen.utils import friendly_score, is_uuid4, make_json_safe, score LOG = logging.getLogger("testgen") @@ -255,12 +263,6 @@ def on_data_preview_clicked(item) -> None: item["table_name"], item.get("column_name"), ) - if preview_data.get("rows") and not session.auth.user_has_permission("view_pii"): - pii_columns = get_pii_columns(item["table_group_id"], item["schema_name"], item["table_name"]) - if pii_columns: - df = pd.DataFrame(preview_data["rows"], columns=preview_data["columns"]) - mask_source_data_pii(df, pii_columns) - preview_data["rows"] = make_json_safe(df.values.tolist()) st.session_state[DC_DATA_PREVIEW_DIALOG_KEY] = preview_data def on_data_preview_dialog_closed(*_) -> None: @@ -302,7 +304,7 @@ def on_history_dialog_closed(*_) -> None: create_script_dialog_data = None if create_script_item := st.session_state.get(DC_CREATE_SCRIPT_DIALOG_KEY): - script = generate_create_script(create_script_item["table_name"], columns) + script = build_create_table_script(table_group_id, create_script_item["table_name"], annotate_changes=True) create_script_dialog_data = { "title": f"Table CREATE Script: {create_script_item['table_name']}", "table_name": create_script_item["table_name"], @@ -491,10 +493,10 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta for key in ["min_date", "max_date", "add_date", "last_mod_date", "drop_date"]: data[key] = data[key].apply( - lambda val: val.strftime("%b %-d %Y, %-I:%M %p") if not pd.isna(val) and not isinstance(val, str) else val + lambda val: date_service.format_friendly_datetime(val, "%b %-d %Y, %-I:%M %p") if not pd.isna(val) and not isinstance(val, str) else val ) - for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product"]: + for key in ["data_source", "source_system", "source_process", "business_domain", "stakeholder_group", "transform_level", "aggregation_level", "data_product", "data_classification"]: data[key] = data.apply( lambda row: row[key] or row[f"table_{key}"] or row.get(f"table_group_{key}"), axis=1, @@ -585,6 +587,7 @@ def get_excel_report_data(update_progress: PROGRESS_UPDATE_TYPE, table_group: Ta "transform_level": {}, "aggregation_level": {}, "data_product": {}, + "data_classification": {}, } return get_excel_file_data( data, @@ -611,60 +614,33 @@ def remove_table_dialog(item: dict) -> None: @with_database_session -def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> FILE_DATA_TYPE: - attributes = ["description"] - attributes.extend(TAG_FIELDS) - +def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> None: tags = payload["tags"] - set_attributes = [ f"{key} = NULLIF(:{key}, '')" for key in attributes if key in tags ] - params = { key: tags.get(key) or "" for key in attributes if key in tags } + + # Empty string clears the field (matches the prior NULLIF semantics). + shared_fields: dict = {key: (tags.get(key) or None) for key in ["description", *TAG_FIELDS] if key in tags} if "critical_data_element" in tags: - set_attributes.append("critical_data_element = :critical_data_element") - params["critical_data_element"] = tags.get("critical_data_element") + shared_fields["critical_data_element"] = tags.get("critical_data_element") # pii_flag and excluded_data_element are column-only fields (not in data_table_chars) - column_set_attributes = list(set_attributes) + column_fields = dict(shared_fields) if "pii_flag" in tags: - column_set_attributes.append("pii_flag = :pii_flag") - params["pii_flag"] = tags.get("pii_flag") + column_fields["pii_flag"] = tags.get("pii_flag") if "excluded_data_element" in tags: - column_set_attributes.append("excluded_data_element = :excluded_data_element") - params["excluded_data_element"] = tags.get("excluded_data_element") + column_fields["excluded_data_element"] = tags.get("excluded_data_element") - params["table_ids"] = [ item["id"] for item in payload["items"] if item["type"] == "table" ] - params["column_ids"] = [ item["id"] for item in payload["items"] if item["type"] == "column" ] + table_ids = [ UUID(item["id"]) for item in payload["items"] if item["type"] == "table" ] + column_ids = [ UUID(item["id"]) for item in payload["items"] if item["type"] == "column" ] with spinner_container: with st.spinner("Saving tags"): - if params["table_ids"] and set_attributes: - execute_db_query( - f""" - WITH selected as ( - SELECT UNNEST(ARRAY [:table_ids]) AS table_id - ) - UPDATE data_table_chars - SET {', '.join(set_attributes)} - FROM data_table_chars dtc - INNER JOIN selected ON (dtc.table_id = selected.table_id::UUID) - WHERE dtc.table_id = data_table_chars.table_id; - """, - params, - ) + if table_ids and shared_fields: + for table in DataTable.select_where(DataTable.id.in_(table_ids)): + apply_table_metadata(table, shared_fields) - if params["column_ids"] and column_set_attributes: - execute_db_query( - f""" - WITH selected as ( - SELECT UNNEST(ARRAY [:column_ids]) AS column_id - ) - UPDATE data_column_chars - SET {', '.join(column_set_attributes)} - FROM data_column_chars dcc - INNER JOIN selected ON (dcc.column_id = selected.column_id::UUID) - WHERE dcc.column_id = data_column_chars.column_id; - """, - params, - ) + if column_ids and column_fields: + for column in DataColumnChars.select_where(DataColumnChars.id.in_(column_ids)): + apply_column_metadata(column, column_fields) # Disable autodetection flags on table group if requested disable_flags = payload.get("disable_flags", []) @@ -672,14 +648,12 @@ def on_tags_changed(spinner_container: DeltaGenerator, payload: dict) -> FILE_DA table_group_id = st.query_params.get("table_group_id") if table_group_id: table_group = get_table_group(table_group_id) - changed = False - if "profile_flag_cdes" in disable_flags and table_group.profile_flag_cdes: - table_group.profile_flag_cdes = False - changed = True - if "profile_flag_pii" in disable_flags and table_group.profile_flag_pii: - table_group.profile_flag_pii = False - changed = True - if changed: + disabled = disable_autoflags( + table_group, + wrote_cde="profile_flag_cdes" in disable_flags, + wrote_pii="profile_flag_pii" in disable_flags, + ) + if disabled: table_group.save() for func in [ get_table_group_columns, get_table_by_id, get_column_by_id, get_tag_values, select_table_groups_minimal_where ]: @@ -713,7 +687,8 @@ def get_table_group_columns(table_group_id: str) -> list[dict]: column_chars.pii_flag, column_chars.excluded_data_element, {", ".join([ f"column_chars.{tag}" for tag in TAG_FIELDS ])}, - {", ".join([ f"table_chars.{tag} AS table_{tag}" for tag in TAG_FIELDS ])} + {", ".join([ f"table_chars.{tag} AS table_{tag}" for tag in TAG_FIELDS ])}, + {", ".join([ f"table_groups.{tag} AS table_group_{tag}" for tag in TAG_FIELDS if tag != "aggregation_level" ])} FROM data_column_chars column_chars LEFT JOIN data_table_chars table_chars ON ( column_chars.table_id = table_chars.table_id @@ -723,6 +698,9 @@ def get_table_group_columns(table_group_id: str) -> list[dict]: AND column_chars.table_name = profile_results.table_name AND column_chars.column_name = profile_results.column_name ) + LEFT JOIN table_groups ON ( + column_chars.table_groups_id = table_groups.id + ) WHERE column_chars.table_groups_id = :table_group_id ORDER BY LOWER(table_chars.table_name), ordinal_position; """ @@ -906,35 +884,47 @@ def get_preview_data( if not connection: return {"title": title, "status": "ERR", "message": "Connection not found."} - flavor_service = get_flavor_service(connection.sql_flavor) - prefix, suffix = flavor_service.row_limit_clauses(100) - quote = flavor_service.quote_character - table_ref = flavor_service.get_table_ref(schema_name, table_name) - query = f""" - SELECT DISTINCT - {prefix} - {f"{quote}{column_name}{quote}" if column_name else "*"} - FROM {table_ref} - {suffix} - """ + result = fetch_table_sample( + connection, + table_group_id, + schema_name, + table_name, + limit=100, + mask_pii=not session.auth.user_has_permission("view_pii"), + column_name=column_name, + ) - try: - results = fetch_from_target_db(connection, query) - except Exception: + if result.status == "ERR": return {"title": title, "status": "ERR", "message": "The preview data could not be loaded."} - - if not results: + if result.status == "ND" or result.df is None: return {"title": title, "status": "ND", "message": "No data found."} - columns_list = list(results[0].keys()) - rows = [list(row.values()) for row in results] return { "title": title, - "columns": columns_list, - "rows": make_json_safe(rows), + "columns": list(result.df.columns), + "rows": make_json_safe(result.df.values.tolist()), } +TAG_FIELD_DEFAULTS: dict[str, list[str]] = { + "data_classification": ["Public", "Internal", "Confidential", "Restricted"], +} + + +def merge_tag_defaults( + values: dict[str, list[str]], + defaults: dict[str, list[str]], +) -> dict[str, list[str]]: + result: dict[str, list[str]] = defaultdict(list, {k: list(v) for k, v in values.items()}) + for tag, tag_defaults in defaults.items(): + existing_lower = {v.lower() for v in result[tag]} + for default in tag_defaults: + if default.lower() not in existing_lower: + result[tag].append(default) + result[tag] = sorted(result[tag], key=str.lower) + return result + + @st.cache_data(show_spinner=False) def get_tag_values() -> dict[str, list[str]]: quote = lambda v: f"'{v}'" @@ -960,8 +950,8 @@ def get_tag_values() -> dict[str, list[str]]: """ results = fetch_all_from_db(query) - values = defaultdict(list) + values: dict[str, list[str]] = defaultdict(list) for row in results: if row.tag and row.value: values[row.tag].append(row.value) - return values + return merge_tag_defaults(values, TAG_FIELD_DEFAULTS) diff --git a/testgen/ui/views/dialogs/import_metadata_dialog.py b/testgen/ui/views/dialogs/import_metadata_dialog.py index b6b1f4c8..b7147240 100644 --- a/testgen/ui/views/dialogs/import_metadata_dialog.py +++ b/testgen/ui/views/dialogs/import_metadata_dialog.py @@ -4,8 +4,11 @@ import pandas as pd +from testgen.common.data_catalog_service import apply_column_metadata, apply_table_metadata, disable_autoflags +from testgen.common.models.data_column import DataColumnChars +from testgen.common.models.data_table import DataTable from testgen.ui.queries.profiling_queries import TAG_FIELDS -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db +from testgen.ui.services.database_service import fetch_all_from_db from testgen.ui.services.query_cache import get_table_group from testgen.ui.session import session @@ -29,7 +32,8 @@ "stakeholder group": "stakeholder_group", "transform level": "transform_level", "aggregation level": "aggregation_level", - "data product": "data_product", + "data product": "data_product", + "data classification": "data_classification", } METADATA_COLUMNS = ["description", "critical_data_element", "excluded_data_element", "pii_flag", *TAG_FIELDS] @@ -296,53 +300,50 @@ def _set_row_status(preview_row: dict, bad_cde: bool, bad_xde: bool, bad_pii: bo def apply_metadata_import(preview: dict, table_group_id: str | None = None) -> dict: + metadata_columns = preview["metadata_columns"] table_count = 0 column_count = 0 for row in preview.get("table_rows", []): - set_clauses, params = _build_update_params(row, preview["metadata_columns"], is_column=False) - if not set_clauses: + fields = _build_metadata_fields(row, metadata_columns, is_column=False) + if not fields: continue - params["table_id"] = row["table_id"] - execute_db_query( - f"UPDATE data_table_chars SET {', '.join(set_clauses)} WHERE table_id = CAST(:table_id AS UUID)", - params, - ) + table = DataTable.get(row["table_id"]) + if table is None: + continue + apply_table_metadata(table, fields) table_count += 1 for row in preview.get("column_rows", []): - set_clauses, params = _build_update_params(row, preview["metadata_columns"], is_column=True) - if not set_clauses: + fields = _build_metadata_fields(row, metadata_columns, is_column=True) + if not fields: + continue + column = DataColumnChars.get(row["column_id"]) + if column is None: continue - params["column_id"] = row["column_id"] - execute_db_query( - f"UPDATE data_column_chars SET {', '.join(set_clauses)} WHERE column_id = CAST(:column_id AS UUID)", - params, - ) + apply_column_metadata(column, fields) column_count += 1 if table_group_id: - _disable_autoflags(table_group_id, preview.get("metadata_columns", [])) + _disable_autoflags(table_group_id, metadata_columns) return {"table_count": table_count, "column_count": column_count} def _disable_autoflags(table_group_id: str, metadata_columns: list[str]) -> None: table_group = get_table_group(table_group_id) - changed = False - if "critical_data_element" in metadata_columns and table_group.profile_flag_cdes: - table_group.profile_flag_cdes = False - changed = True - if "pii_flag" in metadata_columns and table_group.profile_flag_pii: - table_group.profile_flag_pii = False - changed = True - if changed: + disabled = disable_autoflags( + table_group, + wrote_cde="critical_data_element" in metadata_columns, + wrote_pii="pii_flag" in metadata_columns, + ) + if disabled: table_group.save() -def _build_update_params(row: dict, metadata_columns: list[str], is_column: bool = False) -> tuple[list[str], dict]: - set_clauses = [] - params = {} +def _build_metadata_fields(row: dict, metadata_columns: list[str], is_column: bool = False) -> dict: + """Map a preview row to a {column_name: value} update dict (empty string clears the field).""" + fields: dict = {} for col in metadata_columns: if col not in row: @@ -350,22 +351,18 @@ def _build_update_params(row: dict, metadata_columns: list[str], is_column: bool value = row[col] if col == "critical_data_element": - set_clauses.append("critical_data_element = :critical_data_element") - params["critical_data_element"] = value + fields[col] = value elif col == "excluded_data_element": if is_column: - set_clauses.append("excluded_data_element = :excluded_data_element") - params["excluded_data_element"] = value + fields[col] = value elif col == "pii_flag": # Prevent user from editing PII flag if they cannot view PII if is_column and session.auth.user_has_permission("view_pii"): - set_clauses.append("pii_flag = :pii_flag") - params["pii_flag"] = value + fields[col] = value else: - set_clauses.append(f"{col} = NULLIF(:{col}, '')") - params[col] = value if value is not None else "" + fields[col] = value or None - return set_clauses, params + return fields def build_import_preview_props(preview: dict) -> dict: diff --git a/testgen/ui/views/dialogs/manage_schedules.py b/testgen/ui/views/dialogs/manage_schedules.py index 4d917f7c..e22d3dd8 100644 --- a/testgen/ui/views/dialogs/manage_schedules.py +++ b/testgen/ui/views/dialogs/manage_schedules.py @@ -6,6 +6,7 @@ from sqlalchemy import select from sqlalchemy.exc import IntegrityError +from testgen.common import date_service from testgen.common.models import Session, with_database_session from testgen.common.models.scheduler import JobSchedule from testgen.ui.session import session @@ -55,7 +56,7 @@ def build_data(self) -> dict: "readableExpr": cron_descriptor.get_description(job.cron_expr), "cronTz": job.cron_tz_str, "sample": [ - sample.strftime("%a %b %-d, %-I:%M %p") + date_service.format_friendly_datetime(sample, "%a %b %-d, %-I:%M %p") for sample in job.get_sample_triggering_timestamps(CRON_SAMPLE_COUNT + 1) ], "active": job.active, diff --git a/testgen/ui/views/dialogs/table_create_script_dialog.py b/testgen/ui/views/dialogs/table_create_script_dialog.py deleted file mode 100644 index b121e4e4..00000000 --- a/testgen/ui/views/dialogs/table_create_script_dialog.py +++ /dev/null @@ -1,26 +0,0 @@ -def generate_create_script(table_name: str, data: list[dict]) -> str | None: - table_data = [col for col in data if col["table_name"] == table_name] - if not table_data: - return None - - max_name = max(len(col["column_name"]) for col in table_data) + 3 - max_type = max(len(col["datatype_suggestion"] or "") for col in table_data) + 3 - - col_defs = [] - for index, col in enumerate(table_data): - comment = ( - f"-- WAS {col['db_data_type']}" - if isinstance(col["db_data_type"], str) - and isinstance(col["datatype_suggestion"], str) - and col["db_data_type"].lower() != col["datatype_suggestion"].lower() - else "" - ) - col_type = col["datatype_suggestion"] or col["db_data_type"] or "" - separator = " " if index == len(table_data) - 1 else "," - col_defs.append(f"{col['column_name']:<{max_name}} {(col_type):<{max_type}}{separator} {comment}") - - col_defs_joined = "\n ".join(col_defs) - return f""" -CREATE TABLE {table_data[0]['schema_name']}.{table_data[0]['table_name']} ( - {col_defs_joined} -);""" diff --git a/testgen/ui/views/monitors_dashboard.py b/testgen/ui/views/monitors_dashboard.py index e4fd3719..08838ba6 100644 --- a/testgen/ui/views/monitors_dashboard.py +++ b/testgen/ui/views/monitors_dashboard.py @@ -1,30 +1,34 @@ +import dataclasses import logging -from datetime import UTC, date, datetime +from datetime import UTC, datetime from math import ceil -from typing import Any, ClassVar, Literal +from typing import Any, ClassVar, Literal, cast -import pandas as pd import streamlit as st -from testgen.commands.test_generation import run_monitor_generation from testgen.common.cron_service import get_cron_sample -from testgen.common.freshness_service import add_business_minutes, get_schedule_params, resolve_holiday_dates -from testgen.common.models import get_current_session, with_database_session +from testgen.common.models import with_database_session +from testgen.common.models.data_structure_log import DataStructureLog from testgen.common.models.notification_settings import ( MonitorNotificationSettings, MonitorNotificationTrigger, NotificationEvent, ) -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import JobSchedule from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_definition import TestDefinition, TestDefinitionSummary from testgen.common.models.test_suite import PredictSensitivity, TestSuite +from testgen.common.monitor_forecast import ( + gated_forecast_prediction, + next_update_window, +) +from testgen.common.monitor_service import disable_monitoring, enable_monitoring, update_monitoring from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page from testgen.ui.navigation.router import Router from testgen.ui.queries.profiling_queries import get_tables_by_table_group -from testgen.ui.services.database_service import execute_db_query, fetch_all_from_db, fetch_one_from_db +from testgen.ui.services.database_service import fetch_all_from_db from testgen.ui.services.query_cache import ( get_monitor_schedule, get_project_summary, @@ -41,6 +45,15 @@ from testgen.ui.views.dialogs.manage_notifications import NotificationSettingsDialogBase from testgen.utils import make_json_safe +# Maps the user-facing ``anomaly_type_filter`` values supplied by the dashboard +# frontend to the internal ``test_type`` codes the model method expects. +_DASHBOARD_ANOMALY_TYPE_TO_DB: dict[str, str] = { + "freshness": "Freshness_Trend", + "volume": "Volume_Trend", + "schema": "Schema_Drift", + "metrics": "Metric_Trend", +} + PAGE_ICON = "apps_outage" PAGE_TITLE = "Monitors" LOG = logging.getLogger("testgen") @@ -106,7 +119,7 @@ def render( all_monitored_tables_count = 0 monitor_changes_summary = None auto_open_table = None - + current_page = int(current_page) items_per_page = int(items_per_page) page_start = current_page * items_per_page @@ -123,7 +136,7 @@ def render( if sort_field and sort_field not in ALLOWED_SORT_FIELDS: sort_field = None - monitored_tables_page = get_monitor_changes_by_tables( + monitored_tables_page, all_monitored_tables_count = get_monitor_changes_by_tables( table_group_id, table_name_filter=table_name_filter, anomaly_type_filter=anomaly_type_filter, @@ -132,11 +145,6 @@ def render( limit=int(items_per_page), offset=page_start, ) - all_monitored_tables_count = count_monitor_changes_by_tables( - table_group_id, - table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, - ) monitor_changes_summary = summarize_monitor_changes(table_group_id) monitored_table_names = {table["table_name"] for table in monitored_tables_page} @@ -348,249 +356,54 @@ def get_monitor_changes_by_tables( sort_order: Literal["asc"] | Literal["desc"] | None = None, limit: int | None = None, offset: int | None = None, -) -> list[dict]: - query, params = _monitor_changes_by_tables_query( - table_group_id, - table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, - sort_field=sort_field, - sort_order=sort_order, - limit=limit, - offset=offset, - ) - - results = fetch_all_from_db(query, params) - return [ dict(row) for row in results ] - - -@st.cache_data(show_spinner=False) -def count_monitor_changes_by_tables( - table_group_id: str, - table_name_filter: str | None = None, - anomaly_type_filter: list[str] | None = None, -) -> int: - query, params = _monitor_changes_by_tables_query( +) -> tuple[list[dict], int]: + """Per-monitored-table summaries shaped for the dashboard's JSON payload. + + Returns ``(rows, total)`` so the dashboard can fill its pager without an extra + round-trip to the model (which would re-run the heavy CTE twice — once for the + rows, once for the count — just to throw the rows away). Rows are dicts (rather + than ``MonitorTableSummary`` dataclasses) because the monitor-dashboard widget + consumes the payload via ``make_json_safe``. Each row is augmented with + ``table_group_id`` to match the historical payload shape. + """ + page = 1 + (offset // limit) if limit and offset else 1 + summaries, total = TableGroup.list_monitor_table_summaries( table_group_id, + anomaly_types=_dashboard_anomaly_types(anomaly_type_filter), + sort_by=_dashboard_sort_to_model(sort_field, sort_order), table_name_filter=table_name_filter, - anomaly_type_filter=anomaly_type_filter, + page=page, + limit=limit or 1000, ) - count_query = f"SELECT COUNT(*) AS count FROM ({query}) AS subquery" - result = execute_db_query(count_query, params) - return result or 0 + rows = [{**dataclasses.asdict(s), "table_group_id": table_group_id} for s in summaries] + return rows, total @st.cache_data(show_spinner=False) def summarize_monitor_changes(table_group_id: str) -> dict: - query, params = _monitor_changes_by_tables_query(table_group_id) - count_query = f""" - SELECT - lookback, - MIN(lookback_start) AS lookback_start, - MAX(lookback_end) AS lookback_end, - SUM(freshness_anomalies)::INTEGER AS freshness_anomalies, - SUM(volume_anomalies)::INTEGER AS volume_anomalies, - SUM(schema_anomalies)::INTEGER AS schema_anomalies, - SUM(metric_anomalies)::INTEGER AS metric_anomalies, - BOOL_OR(freshness_error_message IS NOT NULL) AS freshness_has_errors, - BOOL_OR(volume_error_message IS NOT NULL) AS volume_has_errors, - BOOL_OR(schema_error_message IS NOT NULL) AS schema_has_errors, - BOOL_OR(metric_error_message IS NOT NULL) AS metric_has_errors, - BOOL_OR(freshness_is_training) AND BOOL_AND(freshness_is_training OR freshness_is_pending) AS freshness_is_training, - BOOL_OR(volume_is_training) AND BOOL_AND(volume_is_training OR volume_is_pending) AS volume_is_training, - BOOL_OR(metric_is_training) AND BOOL_AND(metric_is_training OR metric_is_pending) AS metric_is_training, - BOOL_AND(freshness_is_pending) AS freshness_is_pending, - BOOL_AND(volume_is_pending) AS volume_is_pending, - BOOL_AND(schema_is_pending) AS schema_is_pending, - BOOL_AND(metric_is_pending) AS metric_is_pending - FROM ({query}) AS subquery - GROUP BY lookback - """ - - result = fetch_one_from_db(count_query, params) - return {**result} if result else { - "lookback": 0, - "freshness_anomalies": 0, - "volume_anomalies": 0, - "schema_anomalies": 0, - "metric_anomalies": 0, - "freshness_is_training": False, - "volume_is_training": False, - "metric_is_training": False, - "freshness_is_pending": False, - "volume_is_pending": False, - "schema_is_pending": False, - "metric_is_pending": False, - "freshness_has_errors": False, - "volume_has_errors": False, - "schema_has_errors": False, - "metric_has_errors": False, - } - - -def _monitor_changes_by_tables_query( - table_group_id: str, - table_name_filter: str | None = None, - anomaly_type_filter: list[str] | None = None, - sort_field: str | None = None, - sort_order: Literal["asc"] | Literal["desc"] | None = None, - limit: int | None = None, - offset: int | None = None, -) -> tuple[str, dict]: - query = f""" - WITH ranked_test_runs AS ( - SELECT - test_runs.id, - test_runs.test_starttime, - COALESCE(test_suites.monitor_lookback, 1) AS lookback, - ROW_NUMBER() OVER (PARTITION BY test_runs.test_suite_id ORDER BY test_runs.test_starttime DESC) AS position - FROM table_groups - INNER JOIN test_runs - ON (test_runs.test_suite_id = table_groups.monitor_test_suite_id) - INNER JOIN test_suites - ON (table_groups.monitor_test_suite_id = test_suites.id) - WHERE table_groups.id = :table_group_id - ), - lookback_window AS ( - SELECT MIN(test_starttime) AS lookback_start - FROM ranked_test_runs - WHERE position <= lookback - ), - latest_tables AS ( - SELECT DISTINCT - table_chars.schema_name, - table_chars.table_name - FROM data_table_chars table_chars - CROSS JOIN lookback_window - WHERE table_chars.table_groups_id = :table_group_id - -- Include current tables and tables dropped within lookback window - AND (table_chars.drop_date IS NULL OR table_chars.drop_date >= lookback_window.lookback_start) - {"AND table_chars.table_name ILIKE :table_name_filter" if table_name_filter else ''} - ), - monitor_results AS ( - SELECT - latest_tables.table_name, - results.test_time, - results.test_type, - results.result_code, - ranked_test_runs.lookback, - ranked_test_runs.position, - ranked_test_runs.test_starttime, - -- result_code = -1 indicates training mode - CASE WHEN results.result_code = -1 THEN 1 ELSE 0 END AS is_training, - CASE WHEN results.test_type = 'Freshness_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS freshness_anomaly, - CASE WHEN results.test_type = 'Volume_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS volume_anomaly, - CASE WHEN results.test_type = 'Schema_Drift' AND results.result_code = 0 THEN 1 ELSE 0 END AS schema_anomaly, - CASE WHEN results.test_type = 'Metric_Trend' AND results.result_code = 0 THEN 1 ELSE 0 END AS metric_anomaly, - CASE WHEN results.test_type = 'Freshness_Trend' THEN results.result_signal ELSE NULL END AS freshness_interval, - CASE WHEN results.test_type = 'Volume_Trend' THEN results.result_signal::BIGINT ELSE NULL END AS row_count, - CASE WHEN results.test_type = 'Schema_Drift' THEN SPLIT_PART(results.result_signal, '|', 1) ELSE NULL END AS table_change, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 2), '')::INT ELSE 0 END AS col_adds, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 3), '')::INT ELSE 0 END AS col_drops, - CASE WHEN results.test_type = 'Schema_Drift' THEN NULLIF(SPLIT_PART(results.result_signal, '|', 4), '')::INT ELSE 0 END AS col_mods, - CASE WHEN results.result_status = 'Error' THEN results.result_message ELSE NULL END AS error_message - FROM latest_tables - LEFT JOIN ranked_test_runs ON TRUE - LEFT JOIN test_results AS results - ON results.test_run_id = ranked_test_runs.id - AND results.table_name = latest_tables.table_name - WHERE ranked_test_runs.position IS NULL - -- Also capture 1 run before the lookback to get baseline results - OR ranked_test_runs.position <= ranked_test_runs.lookback + 1 - ), - monitor_tables AS ( - SELECT - :table_group_id AS table_group_id, - table_name, - MAX(lookback) AS lookback, - SUM(freshness_anomaly) AS freshness_anomalies, - SUM(volume_anomaly) AS volume_anomalies, - SUM(schema_anomaly) AS schema_anomalies, - SUM(metric_anomaly) AS metric_anomalies, - MAX(test_time - (COALESCE(NULLIF(freshness_interval, 'Unknown')::INTEGER, 0) * INTERVAL '1 minute')) - FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS latest_update, - MAX(row_count) FILTER (WHERE position = 1) AS row_count, - SUM(col_adds) AS column_adds, - SUM(col_drops) AS column_drops, - SUM(col_mods) AS column_mods, - MAX(error_message) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Schema_Drift' AND position = 1) AS schema_error_message, - MAX(error_message) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_error_message, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Freshness_Trend' AND position = 1) AS freshness_is_training, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Volume_Trend' AND position = 1) AS volume_is_training, - BOOL_OR(is_training = 1) FILTER (WHERE test_type = 'Metric_Trend' AND position = 1) AS metric_is_training, - BOOL_OR(test_type = 'Freshness_Trend') IS NOT TRUE AS freshness_is_pending, - BOOL_OR(test_type = 'Volume_Trend') IS NOT TRUE AS volume_is_pending, - -- Schema monitor only creates results on schema changes (Failed) - -- Mark it as pending only if there are no results of any test type - BOOL_OR(test_time IS NOT NULL) IS NOT TRUE AS schema_is_pending, - BOOL_OR(test_type = 'Metric_Trend') IS NOT TRUE AS metric_is_pending, - CASE - -- Mark as Dropped if latest Schema Drift result for the table indicates it was dropped - WHEN (ARRAY_AGG(table_change ORDER BY test_time DESC) FILTER (WHERE table_change IS NOT NULL))[1] = 'D' - THEN 'dropped' - -- Only mark as Added if latest change does not indicate a drop - WHEN MAX(CASE WHEN table_change = 'A' THEN 1 ELSE 0 END) = 1 - THEN 'added' - WHEN SUM(schema_anomaly) > 0 - THEN 'modified' - ELSE NULL - END AS table_state - FROM monitor_results - -- Only aggregate within lookback runs - WHERE position IS NULL OR position <= COALESCE(lookback, 1) - GROUP BY table_name - ), - table_bounds AS ( - SELECT - table_name, - MIN(position) AS min_position, - MAX(position) AS max_position - FROM monitor_results - WHERE position IS NOT NULL - GROUP BY table_name - ), - baseline_tables AS ( - SELECT - monitor_results.table_name, - MIN(monitor_results.test_starttime) FILTER ( - WHERE monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) - ) AS lookback_start, - MAX(monitor_results.test_starttime) FILTER ( - WHERE monitor_results.position = GREATEST(1, table_bounds.min_position) - ) AS lookback_end, - MAX(monitor_results.row_count) FILTER ( - WHERE monitor_results.test_type = 'Volume_Trend' - AND monitor_results.position = LEAST(monitor_results.lookback + 1, table_bounds.max_position) - ) AS previous_row_count - FROM monitor_results - JOIN table_bounds ON monitor_results.table_name = table_bounds.table_name - GROUP BY monitor_results.table_name - ) - SELECT - monitor_tables.*, - baseline_tables.lookback_start, - baseline_tables.lookback_end, - baseline_tables.previous_row_count - FROM monitor_tables - LEFT JOIN baseline_tables ON monitor_tables.table_name = baseline_tables.table_name - {f"WHERE ({' OR '.join(f'{ANOMALY_TYPE_FILTERS[t]} > 0' for t in anomaly_type_filter)})" if anomaly_type_filter else ""} - ORDER BY {"LOWER(monitor_tables.table_name)" if not sort_field or sort_field == "table_name" else f"monitor_tables.{sort_field}"} - {"DESC" if sort_order == "desc" else "ASC"} NULLS LAST - {"LIMIT :limit" if limit else ""} - {"OFFSET :offset" if offset else ""} - """ + return dataclasses.asdict(TableGroup.get_monitor_group_summary(table_group_id)) - escaped_table_name_filter = table_name_filter.replace("_", "\\_") if table_name_filter else None - params = { - "table_group_id": table_group_id, - "table_name_filter": f"%{escaped_table_name_filter}%" if escaped_table_name_filter else None, - "sort_field": sort_field, - "limit": limit, - "offset": offset, - } - return query, params +def _dashboard_anomaly_types(anomaly_type_filter: list[str] | None) -> list[str] | None: + """Map dashboard-form anomaly type labels to internal ``test_type`` codes.""" + if not anomaly_type_filter: + return None + return [ + _DASHBOARD_ANOMALY_TYPE_TO_DB[t] + for t in anomaly_type_filter + if t in _DASHBOARD_ANOMALY_TYPE_TO_DB + ] or None + + +def _dashboard_sort_to_model( + sort_field: str | None, + sort_order: Literal["asc"] | Literal["desc"] | None, +) -> str | None: + """Translate the dashboard's (sort_field, sort_order) pair into the model's + ``sort_by`` form (the field name, optionally suffixed with ``_desc``).""" + if not sort_field: + return None + return f"{sort_field}_desc" if sort_order == "desc" else sort_field def set_param_values(payload: dict) -> None: @@ -627,39 +440,25 @@ def on_save_settings_clicked(payload: dict) -> None: get_monitor_suite, set_monitor_suite = temp_value(f"monitors:updated_suite:{monitor_suite_id}", default={}) if should_save(): - for key, value in get_monitor_suite().items(): - setattr(monitor_suite, key, value) - - is_new = not monitor_suite.id - monitor_suite.save() - - new_schedule_config = get_schedule() - if ( # Check if schedule has to be created/recreated - not schedule - or schedule.cron_tz != new_schedule_config["cron_tz"] - or schedule.cron_expr != new_schedule_config["cron_expr"] - ): - if schedule: - JobSchedule.delete(schedule.id) - - new_schedule = JobSchedule( - project_code=table_group.project_code, - key=RUN_MONITORS_JOB_KEY, - kwargs={"test_suite_id": str(monitor_suite.id)}, - **new_schedule_config, + schedule_config = get_schedule() + if monitor_suite_id: + # An existing monitor suite always has a run-monitors schedule. + update_monitoring( + monitor_suite, + cast(JobSchedule, schedule), + suite_attrs=get_monitor_suite(), + cron_expr=schedule_config["cron_expr"], + cron_tz=schedule_config["cron_tz"], + active=schedule_config["active"], + ) + else: + enable_monitoring( + get_table_group(table_group.id), + schedule_config["cron_expr"], + schedule_config["cron_tz"], + suite_attrs=get_monitor_suite(), + active=schedule_config["active"], ) - new_schedule.save() - - elif schedule.active != new_schedule_config["active"]: # Only active status changed - JobSchedule.update_active(schedule.id, new_schedule_config["active"]) - - if is_new: - updated_table_group = get_table_group(table_group.id) - updated_table_group.monitor_test_suite_id = monitor_suite.id - updated_table_group.save() - # Commit needed to make test suite visible to run_monitor_generation's separate DB connection - get_current_session().commit() - run_monitor_generation(monitor_suite.id, ["Volume_Trend", "Schema_Drift"]) st.session_state.pop(EDIT_MONITOR_SETTINGS_DIALOG_KEY, None) safe_rerun() @@ -687,7 +486,7 @@ def on_save_settings_clicked(payload: dict) -> None: def delete_monitor_suite(table_group: TableGroupMinimal) -> None: try: monitor_suite = get_test_suite(table_group.monitor_test_suite_id) - TestSuite.cascade_delete([monitor_suite.id]) + disable_monitoring(monitor_suite) st.cache_data.clear() except Exception: LOG.exception("Failed to delete monitor suite") @@ -712,12 +511,27 @@ def build_schema_changes_data(table_group: TableGroupMinimal, payload: dict) -> return data, handlers -def _resolve_holiday_dates(test_suite: TestSuite) -> set[date] | None: - if not test_suite.holiday_codes_list: - return None - now = pd.Timestamp.now("UTC") - idx = pd.DatetimeIndex([now - pd.Timedelta(days=7), now + pd.Timedelta(days=30)]) - return resolve_holiday_dates(test_suite.holiday_codes_list, idx) +def _freshness_next_update_window( + freshness_definition: TestDefinition | None, + events: dict, + test_suite: TestSuite, + monitor_schedule, +) -> dict | None: + """Predicted next freshness-update window for the table — drives the + Freshness_Trend display window and couples the Volume/Metric forecast to the + expected next refresh. Extracts the last detected update from the chart's + event payload and delegates the business-time math to the shared service.""" + last_update_events = [ + e for e in events["freshness_events"] + if e["changed"] and not e["is_training"] and not e["is_pending"] + ] + last_detection_time = max((e["time"] for e in last_update_events), default=None) + return next_update_window( + freshness_definition, + last_detection_time, + test_suite=test_suite, + cron_tz=monitor_schedule.cron_tz if monitor_schedule else None, + ) @with_database_session @@ -775,9 +589,20 @@ def on_close_trends(_payload=None): metric_definition_id = metric_group["test_definition_id"] last_run_time_per_test_key[f"metric:{metric_definition_id}"] = max(e["time"] for e in metric_group["events"]) + # Predicted next freshness-update window for the table — shared by the Freshness_Trend + # display window and the freshness-gated Volume/Metric forecast (which expects no change + # until a refresh lands in this window). + freshness_definition = next((d for d in definitions if d.test_type == "Freshness_Trend"), None) + freshness_window = _freshness_next_update_window(freshness_definition, events, test_suite, monitor_schedule) + for definition in definitions: test_key = f"metric:{definition.id}" if definition.test_type == "Metric_Trend" else definition.test_type.lower() - if definition.history_calculation == "PREDICT" and definition.prediction and (base_mean_predictions := definition.prediction.get("mean")): + if ( + definition.history_calculation == "PREDICT" + and definition.prediction + and not definition.prediction.get("freshness_gated") + and (base_mean_predictions := definition.prediction.get("mean")) + ): predicted_times = sorted([datetime.fromtimestamp(int(timestamp) / 1000.0, UTC) for timestamp in base_mean_predictions.keys()]) # Limit predictions to 1/3 of the lookback, with minimum 3 points predicted_times = [str(int(t.timestamp() * 1000)) for idx, t in enumerate(predicted_times) if idx < 3 or idx < monitor_lookback / 3] @@ -796,6 +621,43 @@ def on_close_trends(_payload=None): "lower_tolerance": lower_tolerance_predictions, "upper_tolerance": upper_tolerance_predictions, } + elif ( + definition.history_calculation == "PREDICT" + and definition.prediction + and definition.prediction.get("freshness_gated") + and (definition.lower_tolerance is not None or definition.upper_tolerance is not None) + ): + # A freshness-gated monitor holds at its baseline between refreshes (the stale-period + # check is value == baseline), so it must never render the rising forecast cone. + gated_prediction = gated_forecast_prediction( + definition, freshness_window, last_run_time_per_test_key.get(test_key), + ) + if gated_prediction is not None: + predictions[test_key] = gated_prediction + else: + # No freshness window available — fall back to a flat band at the next-refresh + # tolerance sampled across upcoming scheduled runs. + cron_sample = get_cron_sample( + monitor_schedule.cron_expr, + monitor_schedule.cron_tz, + sample_count=ceil(min(max(3, monitor_lookback / 3), 10)), + reference_time=last_run_time_per_test_key.get(test_key), + ) + mean_predictions: dict = {} + lower_tolerance_predictions: dict = {} + upper_tolerance_predictions: dict = {} + sample_next_runs = [timestamp * 1000 for timestamp in (cron_sample.get("samples") or [])] + for timestamp in sample_next_runs: + mean_predictions[timestamp] = None + lower_tolerance_predictions[timestamp] = definition.lower_tolerance + upper_tolerance_predictions[timestamp] = definition.upper_tolerance + + predictions[test_key] = { + "method": "static", + "mean": mean_predictions, + "lower_tolerance": lower_tolerance_predictions, + "upper_tolerance": upper_tolerance_predictions, + } elif definition.history_calculation is None and (definition.lower_tolerance is not None or definition.upper_tolerance is not None): cron_sample = get_cron_sample( monitor_schedule.cron_expr, @@ -818,46 +680,11 @@ def on_close_trends(_payload=None): "lower_tolerance": lower_tolerance_predictions, "upper_tolerance": upper_tolerance_predictions, } - elif ( - definition.test_type == "Freshness_Trend" - and definition.history_calculation == "PREDICT" - and (not definition.prediction or definition.prediction.get("schedule_stage")) - and definition.upper_tolerance is not None - ): - last_update_events = [ - e for e in events["freshness_events"] - if e["changed"] and not e["is_training"] and not e["is_pending"] - ] - if last_update_events: - last_detection_time = max(e["time"] for e in last_update_events) - holiday_dates = _resolve_holiday_dates(test_suite) - tz = monitor_schedule.cron_tz or "UTC" if monitor_schedule else None - sched = get_schedule_params(definition.prediction) - - window_end = add_business_minutes( - pd.Timestamp(last_detection_time), - float(definition.upper_tolerance), - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, - ) - window_start = None - if lower_minutes := float(definition.lower_tolerance) if definition.lower_tolerance else None: - window_start = add_business_minutes( - pd.Timestamp(last_detection_time), - lower_minutes, - test_suite.predict_exclude_weekends, - holiday_dates, tz, - excluded_days=sched.excluded_days, - ) - - predictions["freshness_trend"] = { - "method": "freshness_window", - "window": { - "start": int(window_start.timestamp() * 1000) if window_start else None, - "end": int(window_end.timestamp() * 1000), - }, - } + elif definition.test_type == "Freshness_Trend" and freshness_window is not None: + predictions["freshness_trend"] = { + "method": "freshness_window", + "window": freshness_window, + } data = { **make_json_safe(events), @@ -913,7 +740,7 @@ def get_monitor_events_for_table(test_suite_id: str, table_name: str, lookback_m CROSS JOIN target_tests tt LEFT JOIN test_results AS results ON ( - results.test_run_id = active_runs.id + results.test_run_id = active_runs.id AND results.test_type = tt.test_type AND results.table_name = :table_name ) @@ -995,29 +822,35 @@ def get_monitor_events_for_table(test_suite_id: str, table_name: str, lookback_m @st.cache_data(show_spinner=False) -def get_data_structure_logs(table_group_id: str, table_name: str, start_time: str, end_time: str): - query = """ - SELECT - change_date, - change, - old_data_type, - new_data_type, - column_name - FROM data_structure_log - WHERE table_groups_id = :table_group_id - AND table_name = :table_name - AND change_date > :start_time ::TIMESTAMP - AND change_date <= :end_time ::TIMESTAMP; - """ - params = { - "table_group_id": str(table_group_id), - "table_name": table_name, - "start_time": datetime.fromtimestamp(start_time, UTC), - "end_time": datetime.fromtimestamp(end_time, UTC), - } +def get_data_structure_logs(table_group_id: str, table_name: str, start_time: float, end_time: float): + """Schema-change rows for a (table-group, table) inside an epoch-seconds window. - results = fetch_all_from_db(query, params) - return [ dict(row) for row in results ] + The dashboard's chart layer sends timestamps as epoch seconds (the JS-side + convention); we convert to UTC datetimes for the model method, which expects + real date/datetime bounds. + """ + since = datetime.fromtimestamp(start_time, UTC) + until = datetime.fromtimestamp(end_time, UTC) + # Half-open window (exclusive lower, inclusive upper) so a change at the + # window-start boundary (the previous run's timestamp) is attributed to the + # earlier interval, matching the chart's data-point selection. + entries, _total = DataStructureLog.list_for_table_group( + table_group_id, + DataStructureLog.table_name == table_name, + DataStructureLog.change_date > since, + DataStructureLog.change_date <= until, + limit=None, + ) + return [ + { + "change_date": entry.change_date, + "change": entry.change, + "old_data_type": entry.old_data_type, + "new_data_type": entry.new_data_type, + "column_name": entry.column_name, + } + for entry in entries + ] @with_database_session diff --git a/testgen/ui/views/profiling_runs.py b/testgen/ui/views/profiling_runs.py index 5707838a..78d24df2 100644 --- a/testgen/ui/views/profiling_runs.py +++ b/testgen/ui/views/profiling_runs.py @@ -289,9 +289,6 @@ def on_cancel_run(payload: dict) -> None: job_exec = JobExecution.get(job_execution_id) if job_exec and job_exec.request_cancel(): - # Stopgap: also update the run status so the UI reflects cancellation immediately. - if profiling_run_id := payload.get("profiling_run_id"): - ProfilingRun.cancel_run(profiling_run_id) get_profiling_run_summaries.clear() fm.reset_post_updates(str_message=":green[Cancellation requested.]", as_toast=True) else: @@ -307,10 +304,11 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: continue if job_exec.status in (JobStatus.PENDING, JobStatus.CLAIMED, JobStatus.RUNNING, JobStatus.CANCEL_REQUESTED): job_exec.request_cancel() - profiling_run = next(iter(select_profiling_runs_where(ProfilingRun.job_execution_id == je_id)), None) + profiling_run = next(iter(select_profiling_runs_where(ProfilingRun.id == je_id)), None) if profiling_run: ProfilingRun.cascade_delete([str(profiling_run.id)]) - get_current_session().delete(job_exec) + else: + get_current_session().delete(job_exec) get_profiling_run_summaries.clear() Router().set_query_params({"page": 1}) except Exception: diff --git a/testgen/ui/views/table_groups.py b/testgen/ui/views/table_groups.py index bb9e186c..6d568198 100644 --- a/testgen/ui/views/table_groups.py +++ b/testgen/ui/views/table_groups.py @@ -6,17 +6,17 @@ import streamlit as st from sqlalchemy.exc import IntegrityError -from testgen.commands.test_generation import run_monitor_generation from testgen.common.enums import JobSource from testgen.common.models import get_current_session, with_database_session from testgen.common.models.connection import Connection from testgen.common.models.job_execution import JobExecution from testgen.common.models.notification_settings import ProfilingRunNotificationSettings from testgen.common.models.profiling_run import ProfilingRun -from testgen.common.models.scheduler import RUN_MONITORS_JOB_KEY, RUN_TESTS_JOB_KEY, JobSchedule +from testgen.common.models.scheduler import RUN_TESTS_JOB_KEY, JobSchedule from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_run import TestRun from testgen.common.models.test_suite import TestSuite +from testgen.common.monitor_service import enable_monitoring from testgen.ui.components import widgets as testgen from testgen.ui.navigation.menu import MenuItem from testgen.ui.navigation.page import Page @@ -419,33 +419,19 @@ def on_close_clicked(_params: dict) -> None: monitor_test_suite_data = get_monitor_test_suite_data() or {} if monitor_test_suite_data.get("generate"): generate_monitor_suite = True - monitor_test_suite = TestSuite( - project_code=project_code, - test_suite=f"{table_group.table_groups_name} Monitors", - connection_id=table_group.connection_id, - table_groups_id=table_group.id, - export_to_observability=False, - dq_score_exclude=True, - is_monitor=True, - monitor_lookback=monitor_test_suite_data.get("monitor_lookback") or 14, - monitor_regenerate_freshness=monitor_test_suite_data.get("monitor_regenerate_freshness") or True, - predict_min_lookback=monitor_test_suite_data.get("predict_min_lookback") or 30, - predict_sensitivity=monitor_test_suite_data.get("predict_sensitivity") or "medium", - predict_exclude_weekends=monitor_test_suite_data.get("predict_exclude_weekends") or False, - predict_holiday_codes=monitor_test_suite_data.get("predict_holiday_codes") or None, + monitor_test_suite, _ = enable_monitoring( + table_group, + monitor_test_suite_data.get("schedule"), + monitor_test_suite_data.get("timezone") or "UTC", + suite_attrs={ + "monitor_lookback": monitor_test_suite_data.get("monitor_lookback"), + "monitor_regenerate_freshness": monitor_test_suite_data.get("monitor_regenerate_freshness"), + "predict_min_lookback": monitor_test_suite_data.get("predict_min_lookback"), + "predict_sensitivity": monitor_test_suite_data.get("predict_sensitivity"), + "predict_exclude_weekends": monitor_test_suite_data.get("predict_exclude_weekends"), + "predict_holiday_codes": monitor_test_suite_data.get("predict_holiday_codes"), + }, ) - monitor_test_suite.save() - # Commit needed to make test suite visible to run_monitor_generation's separate DB connection - get_current_session().commit() - run_monitor_generation(monitor_test_suite.id, ["Volume_Trend", "Schema_Drift"]) - - JobSchedule( - project_code=project_code, - key=RUN_MONITORS_JOB_KEY, - cron_expr=monitor_test_suite_data.get("schedule"), - cron_tz=monitor_test_suite_data.get("timezone"), - kwargs={"test_suite_id": str(monitor_test_suite.id)}, - ).save() if standard_test_suite or monitor_test_suite: table_group.default_test_suite_id = standard_test_suite.id if standard_test_suite else None @@ -539,10 +525,14 @@ def on_close_edit(_params: dict) -> None: is_in_use = TableGroup.is_in_use([table_group_id]) edit_tg_data = get_edit_tg() + can_view_pii = session.auth.user_has_permission("view_pii") add_scorecard_definition = False for key, value in edit_tg_data.items(): if key == "add_scorecard_definition": add_scorecard_definition = value + elif key == "profile_flag_pii" and not can_view_pii: + # Users without view_pii cannot change the PII flag — keep the stored value. + continue else: setattr(table_group, key, value) diff --git a/testgen/ui/views/test_definitions.py b/testgen/ui/views/test_definitions.py index f460d873..ef28ff6f 100644 --- a/testgen/ui/views/test_definitions.py +++ b/testgen/ui/views/test_definitions.py @@ -12,6 +12,7 @@ from testgen.common.enums import JobSource from testgen.common.models import with_database_session from testgen.common.models.job_execution import JobExecution +from testgen.common.models.project import Project from testgen.common.models.table_group import TableGroup, TableGroupMinimal from testgen.common.models.test_definition import ( TestDefinition, @@ -19,6 +20,7 @@ TestDefinitionNote, TestDefinitionSummary, TestType, + derive_test_criteria, ) from testgen.common.models.test_suite import TestSuite from testgen.common.pii_masking import get_pii_columns, mask_profiling_pii @@ -37,6 +39,7 @@ get_connection, get_table_group_minimal, get_test_suite, + select_projects_where, select_table_groups_minimal_where, select_test_definitions_minimal_where, select_test_definitions_page, @@ -157,11 +160,6 @@ def render( test_types = run_test_type_lookup_query().to_dict("records") table_columns = get_columns(str(table_group.id)) filter_columns_df = get_test_suite_columns(test_suite_id) - table_groups = select_table_groups_minimal_where(TableGroup.project_code == project_code) - all_test_suites = select_test_suites_minimal_where( - TestSuite.table_groups_id.in_([str(tg.id) for tg in table_groups]), - TestSuite.is_monitor.isnot(True), - ) # Build filter options table_options = sorted(filter_columns_df["table_name"].dropna().unique().tolist(), key=str.lower) @@ -198,6 +196,15 @@ def render( add_dialog = None if st.session_state.get(TD_ADD_DIALOG_KEY): + # Pre-fill the picker's column filter from the page's table+column filter when a + # specific column is filtered. The frontend resolves general_type from table_columns, + # so only the column identity is passed here. The prefilled column stays editable -- + # the user can change or clear it inside the picker. + prefill_column = ( + {"table_name": table_name, "column_name": column_name} + if table_name and column_name + else None + ) add_dialog = { "open": True, "test_types": test_types, @@ -206,6 +213,7 @@ def render( "table_group_schema": table_group.table_group_schema, "test_suite": test_suite_info, "qualifies_table_refs_with_schema": qualifies_table_refs_with_schema, + "prefill_column": prefill_column, } edit_dialog = None @@ -230,15 +238,38 @@ def render( copy_move_dialog = None if selected := st.session_state.get(TD_COPY_MOVE_DIALOG_KEY): + editable_project_codes = session.auth.get_projects_with_permission("edit") + editable_projects = ( + select_projects_where(Project.project_code.in_(editable_project_codes)) + if editable_project_codes + else [] + ) + editable_table_groups = select_table_groups_minimal_where( + TableGroup.project_code.in_(editable_project_codes) + ) + editable_test_suites = select_test_suites_minimal_where( + TestSuite.table_groups_id.in_([str(tg.id) for tg in editable_table_groups]), + TestSuite.is_monitor.isnot(True), + ) + tgs_by_project: dict[str, list] = {} + for tg in editable_table_groups: + tgs_by_project.setdefault(tg.project_code, []).append( + {"id": str(tg.id), "table_groups_name": tg.table_groups_name} + ) suites_by_tg: dict[str, list] = {} - for ts in all_test_suites: + for ts in editable_test_suites: suites_by_tg.setdefault(str(ts.table_groups_id), []).append( {"id": str(ts.id), "test_suite": ts.test_suite} ) copy_move_dialog = { "open": True, "selected": selected, - "table_groups": [{"id": str(tg.id), "table_groups_name": tg.table_groups_name} for tg in table_groups], + "projects": [ + {"project_code": p.project_code, "project_name": p.project_name} + for p in editable_projects + ], + "current_project_code": project_code, + "table_groups_by_project": tgs_by_project, "current_table_group_id": str(table_group.id), "current_test_suite_id": str(test_suite.id), "test_suites_by_table_group": suites_by_tg, @@ -402,6 +433,10 @@ def on_update_attribute_all(payload: dict) -> None: TestDefinition.set_status_attribute(attribute, all_ids, value) st.cache_data.clear() + def _resolve_target_project(target_tg_id: str) -> str | None: + target_tg = TableGroup.get_minimal(target_tg_id) + return target_tg.project_code if target_tg else None + @with_database_session def on_copy_confirmed(payload: dict) -> None: ids = payload["ids"] @@ -409,6 +444,11 @@ def on_copy_confirmed(payload: dict) -> None: target_ts_id = payload["target_test_suite_id"] target_table = payload.get("target_table_name") target_col = payload.get("target_column_name") + target_project = _resolve_target_project(target_tg_id) + if not target_project or not session.auth.user_has_permission("edit", target_project): + LOG.warning("Refusing copy to table group %s — user lacks edit permission", target_tg_id) + st.toast("You don't have edit permission for the target project.", icon=":material/error:") + return overwrite_ids = st.session_state.pop(TD_COPY_MOVE_OVERWRITE_KEY, []) if overwrite_ids: TestDefinition.delete_where(TestDefinition.id.in_(overwrite_ids)) @@ -426,6 +466,11 @@ def on_move_confirmed(payload: dict) -> None: target_ts_id = payload["target_test_suite_id"] target_table = payload.get("target_table_name") target_col = payload.get("target_column_name") + target_project = _resolve_target_project(target_tg_id) + if not target_project or not session.auth.user_has_permission("edit", target_project): + LOG.warning("Refusing move to table group %s — user lacks edit permission", target_tg_id) + st.toast("You don't have edit permission for the target project.", icon=":material/error:") + return overwrite_ids = st.session_state.pop(TD_COPY_MOVE_OVERWRITE_KEY, []) if overwrite_ids: TestDefinition.delete_where(TestDefinition.id.in_(overwrite_ids)) @@ -734,7 +779,7 @@ def get_excel_report_data( for key in ["profiling_as_of_date", "last_manual_update"]: data[key] = data[key].apply( - lambda val: datetime.strptime(val, "%Y-%m-%d %H:%M:%S").strftime("%b %-d %Y, %-I:%M %p") + lambda val: date_service.format_friendly_datetime(datetime.strptime(val, "%Y-%m-%d %H:%M:%S"), "%b %-d %Y, %-I:%M %p") if (val and not pd.isna(val) and val != "NaT") else None ) @@ -771,6 +816,7 @@ def run_test_type_lookup_query(test_type: str | None = None) -> pd.DataFrame: tt.measure_uom, COALESCE(tt.measure_uom_description, '') as measure_uom_description, tt.default_parm_columns, tt.default_severity, tt.run_type, tt.test_scope, tt.dq_dimension, tt.impact_dimension, tt.threshold_description, + tt.health_dimension, tt.algorithm, tt.statistical_technique, tt.column_name_prompt, tt.column_name_help, tt.default_parm_prompts, tt.default_parm_help, tt.usage_notes, CASE tt.test_scope @@ -802,7 +848,14 @@ def run_test_type_lookup_query(test_type: str | None = None) -> pd.DataFrame: END, tt.test_name_short; """ - return fetch_df_from_db(query, {"test_type": test_type}) + df = fetch_df_from_db(query, {"test_type": test_type}) + if not df.empty: + # Criteria facet is derived (not stored) via the shared classifier so UI and MCP agree. + df["criteria"] = df.apply( + lambda row: str(derive_test_criteria(row["test_type"], row["test_scope"], row["algorithm"])), + axis=1, + ) + return df @st.cache_data(show_spinner=False) @@ -960,7 +1013,7 @@ def get_test_definitions_collision( def get_columns(table_groups_id: str) -> list[dict]: results = fetch_all_from_db( """ - SELECT table_name, column_name + SELECT table_name, column_name, general_type FROM data_column_chars WHERE table_groups_id = :table_groups_id AND drop_date IS NULL diff --git a/testgen/ui/views/test_results.py b/testgen/ui/views/test_results.py index fd4d7b2c..1814b08c 100644 --- a/testgen/ui/views/test_results.py +++ b/testgen/ui/views/test_results.py @@ -13,6 +13,10 @@ from testgen.common.models.test_definition import TestDefinition, TestDefinitionNote, TestDefinitionSummary from testgen.common.models.test_suite import TestSuiteMinimal from testgen.common.pii_masking import get_pii_columns, mask_profiling_pii +from testgen.common.test_result_disposition_service import ( + coerce_ui_disposition, + set_test_results_disposition, +) from testgen.ui.components import widgets as testgen from testgen.ui.components.widgets.download_dialog import ( FILE_DATA_TYPE, @@ -31,7 +35,7 @@ get_test_issue_source_query, get_test_issue_source_query_custom, ) -from testgen.ui.services.database_service import execute_db_query, fetch_df_from_db, fetch_one_from_db +from testgen.ui.services.database_service import fetch_df_from_db, fetch_one_from_db from testgen.ui.services.query_cache import ( get_table_group_minimal, get_test_definition, @@ -716,6 +720,8 @@ def readable_boolean(v: bool) -> str: "locked": readable_boolean(test_definition.lock_refresh), "active": readable_boolean(test_definition.test_active), "usage_notes": test_definition.usage_notes, + "external_url": test_definition.external_url, + "custom_metadata": test_definition.custom_metadata, "last_manual_update": ( test_definition.last_manual_update.isoformat() if test_definition.last_manual_update else None ), @@ -911,43 +917,4 @@ def update_result_disposition( test_result_ids: list[str], disposition: str, ) -> None: - execute_db_query( - """ - WITH selects - AS (SELECT UNNEST(ARRAY [:test_result_ids]) AS selected_id) - UPDATE test_results - SET disposition = NULLIF(:disposition, 'No Decision') - FROM test_results r - INNER JOIN selects s - ON (r.id = s.selected_id::UUID) - WHERE r.id = test_results.id - AND r.result_status != 'Passed'; - """, - { - "test_result_ids": test_result_ids, - "disposition": disposition, - }, - ) - - execute_db_query( - """ - WITH selects - AS (SELECT UNNEST(ARRAY [:test_result_ids]) AS selected_id) - UPDATE test_definitions - SET test_active = :test_active, - last_manual_update = CURRENT_TIMESTAMP AT TIME ZONE 'UTC', - lock_refresh = :lock_refresh - FROM test_definitions d - INNER JOIN test_results r - ON (d.id = r.test_definition_id) - INNER JOIN selects s - ON (r.id = s.selected_id::UUID) - WHERE d.id = test_definitions.id - AND r.result_status != 'Passed'; - """, - { - "test_result_ids": test_result_ids, - "test_active": "N" if disposition == "Inactive" else "Y", - "lock_refresh": "Y" if disposition == "Inactive" else "N", - }, - ) + set_test_results_disposition(test_result_ids, coerce_ui_disposition(disposition)) diff --git a/testgen/ui/views/test_runs.py b/testgen/ui/views/test_runs.py index f6563b94..fe52d051 100644 --- a/testgen/ui/views/test_runs.py +++ b/testgen/ui/views/test_runs.py @@ -301,9 +301,6 @@ def on_cancel_run(payload: dict) -> None: job_exec = JobExecution.get(job_execution_id) if job_exec and job_exec.request_cancel(): - # Stopgap: also update the run status so the UI reflects cancellation immediately. - if test_run_id := payload.get("test_run_id"): - TestRun.cancel_run(test_run_id) get_test_run_summaries.clear() fm.reset_post_updates(str_message=":green[Cancellation requested.]", as_toast=True) else: @@ -319,10 +316,11 @@ def on_delete_runs(job_execution_ids: list[str]) -> None: continue if job_exec.status in (JobStatus.PENDING, JobStatus.CLAIMED, JobStatus.RUNNING, JobStatus.CANCEL_REQUESTED): job_exec.request_cancel() - test_run = next(iter(select_test_runs_where(TestRun.job_execution_id == je_id)), None) + test_run = next(iter(select_test_runs_where(TestRun.id == je_id)), None) if test_run: TestRun.cascade_delete([str(test_run.id)]) - get_current_session().delete(job_exec) + else: + get_current_session().delete(job_exec) get_test_run_summaries.clear() Router().set_query_params({"page": 1}) except Exception: diff --git a/testgen/utils/__init__.py b/testgen/utils/__init__.py index 5218d3d0..75760994 100644 --- a/testgen/utils/__init__.py +++ b/testgen/utils/__init__.py @@ -181,6 +181,7 @@ def format_score_card(score_card: ScoreCard | None) -> ScoreCard: "dq_dimension": "Quality Dimension", "impact_dimension": "Impact Dimension", "data_product": "Data Product", + "data_classification": "Data Classification", } if not score_card: return { diff --git a/testgen/utils/plugins.py b/testgen/utils/plugins.py index f4bfbe7f..07a87c8b 100644 --- a/testgen/utils/plugins.py +++ b/testgen/utils/plugins.py @@ -3,7 +3,7 @@ import importlib import importlib.metadata import inspect -from collections.abc import Generator +from collections.abc import Callable, Generator from typing import ClassVar, get_args from testgen.ui.assets import get_asset_path @@ -61,6 +61,16 @@ def configure_ui(cls) -> None: is actually running. Called by ``Plugin.load_streamlit()``, never by ``Plugin.load()``. """ + @classmethod + def get_mcp_tools(cls) -> list[Callable]: + """Return MCP tool callables to register into the MCP server. + + Override in plugins to register additional MCP tools. Implementations import their + tool modules inside the method (like ``configure_ui``) so plugin import stays light + for processes that never build the MCP server. + """ + return [] + class PluginHook: """Singleton holding resolved plugin values, pre-loaded with defaults.""" diff --git a/tests/unit/api/oauth/test_server.py b/tests/unit/api/oauth/test_server.py index c9810e15..83e5bf20 100644 --- a/tests/unit/api/oauth/test_server.py +++ b/tests/unit/api/oauth/test_server.py @@ -5,12 +5,10 @@ from uuid import uuid4 import pytest -from authlib.oauth2.rfc6749 import grants from testgen import settings from testgen.api.oauth.server import ( AuthorizationCodeGrant, - ClientCredentialsGrant, RefreshTokenGrant, TestGenAuthorizationServer, TestGenRevocationEndpoint, @@ -335,63 +333,6 @@ def test_generate_bearer_token_no_scope_omits_field(mock_jwt): assert "scope" not in token -# --- ClientCredentialsGrant --- - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_resolves_owner(mock_get_session): - mock_session = MagicMock() - mock_get_session.return_value = mock_session - - mock_owner = MagicMock() - mock_owner.username = "owner_user" - mock_session.scalars.return_value.first.return_value = mock_owner - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = uuid4() - grant.request.client.check_grant_type.return_value = True - - with patch.object(grants.ClientCredentialsGrant, "validate_token_request"): - grant.validate_token_request() - - assert grant.request.user is mock_owner - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_rejects_client_without_owner(mock_get_session): - from authlib.oauth2.rfc6749.errors import InvalidGrantError - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = None - - with ( - patch.object(grants.ClientCredentialsGrant, "validate_token_request"), - pytest.raises(InvalidGrantError), - ): - grant.validate_token_request() - - -@patch("testgen.api.oauth.server.get_current_session") -def test_client_credentials_rejects_deleted_owner(mock_get_session): - from authlib.oauth2.rfc6749.errors import InvalidGrantError - - mock_session = MagicMock() - mock_get_session.return_value = mock_session - mock_session.scalars.return_value.first.return_value = None - - grant = ClientCredentialsGrant.__new__(ClientCredentialsGrant) - grant.request = MagicMock() - grant.request.client.user_id = uuid4() - - with ( - patch.object(grants.ClientCredentialsGrant, "validate_token_request"), - pytest.raises(InvalidGrantError), - ): - grant.validate_token_request() - - # --- RevocationEndpoint --- @@ -460,6 +401,6 @@ def test_create_authorization_server_returns_configured_server(): assert isinstance(server, TestGenAuthorizationServer) token_grant_classes = [entry[0] for entry in server._token_grants] assert RefreshTokenGrant in token_grant_classes - assert ClientCredentialsGrant in token_grant_classes + assert AuthorizationCodeGrant in token_grant_classes assert "revocation" in server._endpoints diff --git a/tests/unit/api/test_jobs.py b/tests/unit/api/test_jobs.py index 6af04eb4..0659bb04 100644 --- a/tests/unit/api/test_jobs.py +++ b/tests/unit/api/test_jobs.py @@ -7,6 +7,7 @@ import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from sqlalchemy.dialects import postgresql from testgen.api.deps import db_session, get_authorized_user from testgen.api.jobs import ( @@ -18,6 +19,8 @@ submit_test_generation, submit_test_run, ) +from testgen.common.enums import JobKey, PublicJobKey +from testgen.common.models.job_execution import JobExecution pytestmark = pytest.mark.unit @@ -236,3 +239,44 @@ def test_list_jobs_accepts_valid_status(mock_je_cls, _mock_perm): assert resp.status_code == 200 # Verify the status string was forwarded to the model layer. assert mock_je_cls.list_for_project.call_args.kwargs["status"] == "completed" + + +# --- internal job-kind exclusion --- +# JobResponse.job_key is typed PublicJobKey, which is only safe because every path +# that builds a JobResponse excludes internal job kinds. These pin that invariant. + + +def test_public_job_key_excludes_internal_kinds(): + public = {member.value for member in PublicJobKey} + assert public == {"run-profile", "run-tests", "run-test-generation"} + internal = {"run-monitors", "run-score-update", "recalculate-project-scores", "run-data-cleanup"} + assert public.isdisjoint(internal) + # Every public value must be a real JobKey. + assert public <= {member.value for member in JobKey} + + +@patch.object(JobExecution, "list_for_project", return_value=([], 0)) +def test_list_jobs_filters_to_public_job_kinds(_mock_list): + list_jobs(project_code="DEFAULT", job_key=None, status=None, page=1, limit=20) + + clause = _mock_list.call_args.args[1] + sql = str(clause.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})) + assert "run-profile" in sql + assert "run-tests" in sql + assert "run-test-generation" in sql + for internal in ("run-monitors", "run-score-update", "recalculate-project-scores", "run-data-cleanup"): + assert internal not in sql + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch(f"{MODULE}.JobExecution") +def test_list_jobs_rejects_internal_job_key(mock_je_cls, _mock_perm): + mock_je_cls.list_for_project.return_value = ([], 0) + client = TestClient(_client_with_overrides()) + + resp = client.get("/api/v1/projects/DEFAULT/jobs?job_key=run-data-cleanup") + + assert resp.status_code == 422 + body = resp.json() + assert body["detail"][0]["loc"] == ["query", "job_key"] + assert body["detail"][0]["type"] == "enum" diff --git a/tests/unit/api/test_runs.py b/tests/unit/api/test_runs.py index 0fe8143c..1dc1a825 100644 --- a/tests/unit/api/test_runs.py +++ b/tests/unit/api/test_runs.py @@ -5,10 +5,15 @@ from uuid import uuid4 import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from sqlalchemy.dialects import postgresql -from testgen.api.runs import get_profiling_run, get_test_run -from testgen.common.models.hygiene_issue import IssueLikelihoodCounts -from testgen.common.models.test_result import ResultStatusCounts +from testgen.api.deps import db_session, get_authorized_user +from testgen.api.enums import Disposition, ResultStatus +from testgen.api.runs import get_profiling_run, get_test_run, list_test_run_results, router +from testgen.common.models.hygiene_issue import HygieneIssueCounts, IssueCounts, PotentialPiiCounts +from testgen.common.models.test_result import ResultStatusCounts, TestResult, TestResultStatus, TestRunResultRow pytestmark = pytest.mark.unit @@ -18,6 +23,46 @@ TABLE_GROUP_ID = uuid4() +def _mock_result_row(**overrides): + defaults = { + "test_definition_id": uuid4(), + "test_type": "Unique", + "schema_name": "demo", + "table_name": "orders", + "column_names": "amount", + "status": TestResultStatus.Failed, + "result_measure": "3", + "threshold_value": "0", + "message": "Duplicate values: 3", + "test_time": datetime.now(UTC), + "disposition": None, + } + defaults.update(overrides) + return TestRunResultRow(**defaults) + + +def _result_clauses_sql(mock_list) -> str: + """Compile the WHERE clauses passed to list_for_run (args after test_run_id) to SQL.""" + clauses = mock_list.call_args.args[1:] + return " ".join( + str(c.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True})) for c in clauses + ) + + +def _no_filters(**overrides): + kwargs = { + "status": None, + "table_name": None, + "column_name": None, + "test_type": None, + "disposition": None, + "page": 1, + "limit": 20, + } + kwargs.update(overrides) + return kwargs + + def _mock_job(**overrides): defaults = { "id": uuid4(), @@ -70,7 +115,7 @@ def _mock_profiling_run(**overrides): @patch(f"{MODULE}.TestRun") def test_get_test_run_completed(mock_tr_cls, mock_result_cls, mock_session): job = _mock_job() - mock_tr_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_tr_cls.get.return_value = _mock_test_run() mock_result_cls.count_by_status.return_value = ResultStatusCounts( passed=90, failed=5, warning=3, error=2, log=0, dismissed=12, ) @@ -84,15 +129,15 @@ def test_get_test_run_completed(mock_tr_cls, mock_result_cls, mock_session): assert result.table_group_id == TABLE_GROUP_ID assert result.result is not None assert result.result.score == 0.95 - assert result.result.tests.passed == 90 - assert result.result.tests.failed == 5 - assert result.result.tests.dismissed == 12 + assert result.result.result_counts.passed == 90 + assert result.result.result_counts.failed == 5 + assert result.result.result_counts.dismissed == 12 @patch(f"{MODULE}.TestRun") def test_get_test_run_pending_no_run(mock_tr_cls): job = _mock_job(status="pending", started_at=None, completed_at=None) - mock_tr_cls.get_by_id_or_job.return_value = None + mock_tr_cls.get.return_value = None result = get_test_run(job) @@ -110,9 +155,11 @@ def test_get_test_run_pending_no_run(mock_tr_cls): @patch(f"{MODULE}.ProfilingRun") def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): job = _mock_job() - mock_pr_cls.get_by_id_or_job.return_value = _mock_profiling_run() - mock_issue_cls.count_by_likelihood.return_value = IssueLikelihoodCounts( - definite=5, likely=3, possible=8, dismissed=2, + mock_pr_cls.get.return_value = _mock_profiling_run() + mock_issue_cls.count_for_run.return_value = IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=5, likely=3, possible=8), + potential_pii=PotentialPiiCounts(high=4, moderate=6), + dismissed=2, ) result = get_profiling_run(job) @@ -123,15 +170,18 @@ def test_get_profiling_run_completed(mock_pr_cls, mock_issue_cls): assert result.result is not None assert result.result.score == 0.88 assert result.result.table_ct == 10 - assert result.result.issues.definite == 5 - assert result.result.issues.likely == 3 - assert result.result.issues.dismissed == 2 + assert result.result.issue_counts.hygiene_issues.definite == 5 + assert result.result.issue_counts.hygiene_issues.likely == 3 + assert result.result.issue_counts.hygiene_issues.possible == 8 + assert result.result.issue_counts.potential_pii.high == 4 + assert result.result.issue_counts.potential_pii.moderate == 6 + assert result.result.issue_counts.dismissed == 2 @patch(f"{MODULE}.ProfilingRun") def test_get_profiling_run_pending_no_run(mock_pr_cls): job = _mock_job(status="pending", started_at=None, completed_at=None) - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None result = get_profiling_run(job) @@ -139,3 +189,176 @@ def test_get_profiling_run_pending_no_run(mock_pr_cls): assert result.status == "pending" assert result.table_group_id is None assert result.result is None + + +# --- list_test_run_results: envelope + field/enum mapping --- + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_envelope_and_mapping(mock_list): + job = _mock_job() + row = _mock_result_row(test_type="Unique", status=TestResultStatus.Failed, message="dup", disposition=None) + mock_list.return_value = ([row], 7) + + result = list_test_run_results(job, **_no_filters(page=2, limit=5)) + + assert result.total == 7 + assert result.page == 2 + assert result.limit == 5 + assert len(result.items) == 1 + item = result.items[0] + assert item.test_type == "Unique" # raw code, not a display name + assert item.result_status == ResultStatus.failed # title-case DB value -> lowercase API enum + assert item.result_message == "dup" # ORM attr `message` -> field `result_message` + assert item.disposition == Disposition.no_decision # NULL -> no_decision (not confirmed) + # test_run_id (job.id) is the scope; page/limit are forwarded as kwargs. + assert mock_list.call_args.args[0] == job.id + assert mock_list.call_args.kwargs == {"page": 2, "limit": 5} + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_empty_envelope(mock_list): + result = list_test_run_results(_mock_job(), **_no_filters()) + assert result.items == [] + assert result.total == 0 + assert result.page == 1 + + +@pytest.mark.parametrize( + "db_status,expected", + [ + (TestResultStatus.Passed, ResultStatus.passed), + (TestResultStatus.Failed, ResultStatus.failed), + (TestResultStatus.Warning, ResultStatus.warning), + (TestResultStatus.Error, ResultStatus.error), + (TestResultStatus.Log, ResultStatus.log), + (None, None), + ], +) +@patch.object(TestResult, "list_for_run") +def test_list_results_status_render(mock_list, db_status, expected): + mock_list.return_value = ([_mock_result_row(status=db_status)], 1) + item = list_test_run_results(_mock_job(), **_no_filters()).items[0] + assert item.result_status == expected + + +@pytest.mark.parametrize( + "db_disposition,expected", + [ + (None, Disposition.no_decision), + ("", Disposition.no_decision), + ("Confirmed", Disposition.confirmed), + ("Dismissed", Disposition.dismissed), + ("Inactive", Disposition.muted), + ("Bogus", Disposition.no_decision), + ], +) +@patch.object(TestResult, "list_for_run") +def test_list_results_disposition_render(mock_list, db_disposition, expected): + mock_list.return_value = ([_mock_result_row(disposition=db_disposition)], 1) + item = list_test_run_results(_mock_job(), **_no_filters()).items[0] + assert item.disposition == expected + + +# --- list_test_run_results: filter clause building --- + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_list_results_builds_filter_clauses(mock_list): + list_test_run_results( + _mock_job(), + **_no_filters(status=ResultStatus.failed, table_name="orders", column_name="amount", test_type="Unique"), + ) + sql = _result_clauses_sql(mock_list) + assert "'Failed'" in sql # status mapped to DB value + assert "'orders'" in sql + assert "'amount'" in sql # column_name -> column_names column + assert "'Unique'" in sql + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_omitted_returns_active(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=None)) + sql = _result_clauses_sql(mock_list) + # Active = confirmed + no_decision (NULL); dismissed/muted excluded. + assert "IS NULL" in sql + assert "'Confirmed'" in sql + assert "'Dismissed'" not in sql + assert "'Inactive'" not in sql + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_confirmed_excludes_nulls(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=Disposition.confirmed)) + sql = _result_clauses_sql(mock_list) + assert "'Confirmed'" in sql + assert "IS NULL" not in sql # explicit confirmed no longer sweeps in undecided rows + + +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_no_decision_is_null(mock_list): + list_test_run_results(_mock_job(), **_no_filters(disposition=Disposition.no_decision)) + sql = _result_clauses_sql(mock_list) + assert "IS NULL" in sql + assert "'Confirmed'" not in sql + + +@pytest.mark.parametrize( + "disposition,db_value", + [(Disposition.dismissed, "'Dismissed'"), (Disposition.muted, "'Inactive'")], +) +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_disposition_dismissed_muted_map_to_db(mock_list, disposition, db_value): + list_test_run_results(_mock_job(), **_no_filters(disposition=disposition)) + assert db_value in _result_clauses_sql(mock_list) + + +# --- list_test_run_results: HTTP-level query validation --- + + +def _client() -> TestClient: + app = FastAPI() + app.include_router(router, prefix="/api/v1") + app.dependency_overrides[db_session] = lambda: iter([None]) + app.dependency_overrides[get_authorized_user] = lambda: MagicMock(id=uuid4()) + return TestClient(app) + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_unknown_status(_mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?status=BOGUS") + assert resp.status_code == 422 + assert resp.json()["detail"][0]["loc"] == ["query", "status"] + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_unknown_disposition(_mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?disposition=foo") + assert resp.status_code == 422 + assert resp.json()["detail"][0]["loc"] == ["query", "disposition"] + + +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_accepts_no_decision(mock_list, mock_sess, _mock_perm): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?disposition=no_decision") + assert resp.status_code == 200 + mock_list.assert_called_once() + + +@pytest.mark.parametrize("query", ["page=0", "limit=0", "limit=101"]) +@patch("testgen.api.deps.has_project_permission", return_value=True) +@patch("testgen.api.deps.get_current_session") +@patch.object(TestResult, "list_for_run", return_value=([], 0)) +def test_http_rejects_out_of_range_pagination(_mock_list, mock_sess, _mock_perm, query): + mock_sess.return_value.scalars.return_value.first.return_value = _mock_job() + resp = _client().get(f"/api/v1/test-runs/{uuid4()}/results?{query}") + assert resp.status_code == 422 diff --git a/tests/unit/api/test_td_export_import.py b/tests/unit/api/test_td_export_import.py index ca3a1b0b..ba448956 100644 --- a/tests/unit/api/test_td_export_import.py +++ b/tests/unit/api/test_td_export_import.py @@ -6,15 +6,22 @@ import pytest from fastapi import HTTPException +from pydantic import ValidationError -from testgen.api.schemas import ( +from testgen.api.schemas import ImportRequest +from testgen.common.test_definition_export_import_service import ( ExportDocument, + ExportSource, ImportAction, ImportConfig, + ImportErrorCode, ImportMode, ImportPayload, ImportReason, ImportResponse, + ImportStrictViolation, + ImportSummary, + InvalidImportPayload, OnAbsence, OnMatch, OnNew, @@ -24,7 +31,7 @@ pytestmark = pytest.mark.unit -SERVICE_MODULE = "testgen.api.test_definition_service" +SERVICE_MODULE = "testgen.common.test_definition_export_import_service" ENDPOINT_MODULE = "testgen.api.test_definitions" @@ -162,7 +169,7 @@ def test_export_builds_document(self, mock_session_fn, mock_tg_cls, mock_setting session.scalars.return_value.all.return_value = [td_obj] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions result = export_definitions(ts, Origin.both, None, None) @@ -186,7 +193,7 @@ def test_export_assigns_external_id_to_manual_tds(self, mock_session_fn, mock_tg ts = _make_test_suite() session.scalars.return_value.all.return_value = [] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions export_definitions(ts, Origin.manual, None, None) @@ -207,7 +214,7 @@ def test_export_skips_external_id_assignment_for_auto_only(self, mock_session_fn ts = _make_test_suite() session.scalars.return_value.all.return_value = [] - from testgen.api.test_definition_service import export_definitions + from testgen.common.test_definition_export_import_service import export_definitions export_definitions(ts, Origin.auto, None, None) @@ -237,7 +244,7 @@ def test_auto_td_matches_by_natural_key(self, mock_session_fn, mock_tg_cls, mock td = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -265,7 +272,7 @@ def test_manual_td_matches_by_external_id(self, mock_session_fn, mock_tg_cls, mo td = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -287,7 +294,7 @@ def test_no_match_creates(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -317,7 +324,7 @@ def test_on_match_skip(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td() config = _make_config(on_match=OnMatch.skip) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -341,7 +348,7 @@ def test_on_match_overwrite_unlocked_skips_locked(self, mock_session_fn, mock_tg td = _make_import_td() config = _make_config(on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -365,7 +372,7 @@ def test_on_match_overwrite_all_ignores_lock(self, mock_session_fn, mock_tg_cls, td = _make_import_td() config = _make_config(on_match=OnMatch.overwrite_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -386,7 +393,7 @@ def test_on_new_skip(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td() config = _make_config(on_new=OnNew.skip) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -412,7 +419,7 @@ def test_on_absence_delete_all(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config(on_new=OnNew.create, on_absence=OnAbsence.delete_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -436,7 +443,7 @@ def test_on_absence_delete_unlocked_spares_locked(self, mock_session_fn, mock_tg td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(on_new=OnNew.create, on_absence=OnAbsence.delete_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -465,7 +472,7 @@ def test_invalid_test_type_skipped(self, mock_session_fn, mock_tg_cls, mock_dt_c td = _make_import_td(test_type="NonExistent", table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -487,7 +494,7 @@ def test_invalid_table_skipped(self, mock_session_fn, mock_tg_cls, mock_dt_cls): td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -510,7 +517,7 @@ def test_manual_td_without_external_id_skipped(self, mock_session_fn, mock_tg_cl td = _make_import_td(last_auto_gen_date=None, external_id=None, table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -533,7 +540,7 @@ def test_td_with_null_table_name_passes_validation(self, mock_session_fn, mock_t td = _make_import_td(test_type="Alpha", table_name=None) config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -555,13 +562,13 @@ def test_duplicate_auto_keys_rejected(self, mock_session_fn, mock_tg_cls, mock_d td2 = _make_import_td(test_type="Alpha", table_name="t1", column_name="c1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(InvalidImportPayload) as exc_info: import_definitions(ts, config, _make_payload(td1, td2)) - assert exc_info.value.status_code == 400 - assert "duplicate_natural_key" in str(exc_info.value.detail) + assert exc_info.value.code == "duplicate_natural_key" + assert "Duplicate auto-gen key" in str(exc_info.value) @patch(f"{SERVICE_MODULE}.DataTable") @patch(f"{SERVICE_MODULE}.TableGroup") @@ -579,12 +586,12 @@ def test_duplicate_external_ids_rejected(self, mock_session_fn, mock_tg_cls, moc td2 = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config() - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(InvalidImportPayload) as exc_info: import_definitions(ts, config, _make_payload(td1, td2)) - assert exc_info.value.status_code == 400 + assert exc_info.value.code == "duplicate_natural_key" # --- Import: apply_strict mode --- @@ -607,14 +614,17 @@ def test_apply_strict_does_not_apply_when_skips_exist(self, mock_session_fn, moc bad_td = _make_import_td(test_type="NonExistent", table_name="t1") config = _make_config(mode=ImportMode.apply_strict) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): - result = import_definitions(ts, config, _make_payload(good_td, bad_td)) + with pytest.raises(ImportStrictViolation) as exc_info: + import_definitions(ts, config, _make_payload(good_td, bad_td)) - # Result reports what would happen, but nothing was applied + # The exception carries the projected result, but nothing was applied + result = exc_info.value.result assert result.summary.created == 1 assert result.summary.skipped == 1 + assert "1 test definitions would be skipped" in str(exc_info.value) # No session.add calls (create not applied) session.add.assert_not_called() @@ -632,7 +642,7 @@ def test_apply_strict_applies_when_no_skips(self, mock_session_fn, mock_tg_cls, td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.apply_strict) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -663,7 +673,7 @@ def test_create_auto_td_sets_last_auto_gen_date_to_now(self, mock_session_fn, mo td = _make_import_td(test_type="Alpha", table_name="t1", last_auto_gen_date=datetime(2024, 1, 1, tzinfo=UTC)) config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions before = datetime.now(UTC) with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): @@ -690,7 +700,7 @@ def test_create_manual_td_sets_null_last_auto_gen_date(self, mock_session_fn, mo td = _make_import_td(last_auto_gen_date=None, external_id=ext_id, table_name="t1") config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -714,7 +724,7 @@ def test_create_and_lock_forces_lock_for_auto_td(self, mock_session_fn, mock_tg_ td = _make_import_td(test_type="Alpha", table_name="t1", lock_refresh=False) config = _make_config(mode=ImportMode.apply, on_new=OnNew.create_and_lock) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: @@ -738,7 +748,7 @@ def test_create_sets_last_manual_update(self, mock_session_fn, mock_tg_cls, mock td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions before = datetime.now(UTC) with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): @@ -775,7 +785,7 @@ def test_update_excludes_identity_fields(self, mock_session_fn, mock_tg_cls, moc ) config = _make_config(mode=ImportMode.apply, on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): import_definitions(ts, config, _make_payload(td)) @@ -810,7 +820,7 @@ def test_update_inherits_external_id_if_target_has_none(self, mock_session_fn, m ) config = _make_config(mode=ImportMode.apply, on_match=OnMatch.overwrite_unlocked) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): import_definitions(ts, config, _make_payload(td)) @@ -848,7 +858,7 @@ def test_invalid_td_still_protects_match_from_absence_delete(self, mock_session_ td = _make_import_td(test_type="InvalidType", table_name="t1", column_name="c1") config = _make_config(on_absence=OnAbsence.delete_all) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -877,7 +887,7 @@ def test_preview_does_not_apply(self, mock_session_fn, mock_tg_cls, mock_dt_cls) td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.preview) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -900,7 +910,7 @@ def test_preview_target_id_is_none_for_creates(self, mock_session_fn, mock_tg_cl td = _make_import_td(test_type="Alpha", table_name="t1") config = _make_config(mode=ImportMode.preview) - from testgen.api.test_definition_service import import_definitions + from testgen.common.test_definition_export_import_service import import_definitions with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): result = import_definitions(ts, config, _make_payload(td)) @@ -915,7 +925,6 @@ def test_preview_target_id_is_none_for_creates(self, mock_session_fn, mock_tg_cl class Test_import_endpoint_strict: def test_strict_raises_400_on_skips(self): - from testgen.api.schemas import ImportRequest from testgen.api.test_definitions import import_test_definitions ts = _make_test_suite() @@ -923,14 +932,14 @@ def test_strict_raises_400_on_skips(self): payload = _make_payload(_make_import_td(test_type="Alpha", table_name="t1")) request = ImportRequest(config=config, payload=payload) - # Mock the service to return a response with skips + # The service signals strict violations with an exception carrying the projected result mock_result = ImportResponse( summary=ImportSummary(created=1, skipped=1), items=[], ) - with patch(f"{ENDPOINT_MODULE}.test_definition_service") as mock_service: - mock_service.import_definitions.return_value = mock_result + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.side_effect = ImportStrictViolation(mock_result) with pytest.raises(HTTPException) as exc_info: import_test_definitions(body=request, test_suite=ts) @@ -939,7 +948,6 @@ def test_strict_raises_400_on_skips(self): assert "import_result" in exc_info.value.detail def test_strict_returns_200_when_no_skips(self): - from testgen.api.schemas import ImportRequest from testgen.api.test_definitions import import_test_definitions ts = _make_test_suite() @@ -952,13 +960,99 @@ def test_strict_returns_200_when_no_skips(self): items=[], ) - with patch(f"{ENDPOINT_MODULE}.test_definition_service") as mock_service: - mock_service.import_definitions.return_value = mock_result + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.return_value = mock_result result = import_test_definitions(body=request, test_suite=ts) assert result.summary.created == 1 +def test_invalid_payload_adapted_to_400(): + from testgen.api.test_definitions import import_test_definitions + + ts = _make_test_suite() + request = ImportRequest( + config=_make_config(), + payload=_make_payload(_make_import_td(test_type="Alpha", table_name="t1")), + ) + + with patch(f"{ENDPOINT_MODULE}.import_definitions") as mock_import: + mock_import.side_effect = InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, "Duplicate external_id at index 1" + ) + with pytest.raises(HTTPException) as exc_info: + import_test_definitions(body=request, test_suite=ts) + + assert exc_info.value.status_code == 400 + assert "duplicate_natural_key" in str(exc_info.value.detail) + assert "Duplicate external_id at index 1" in str(exc_info.value.detail) + + +@patch(f"{SERVICE_MODULE}.DataTable") +@patch(f"{SERVICE_MODULE}.TableGroup") +@patch(f"{SERVICE_MODULE}.get_current_session") +def test_unsupported_payload_version_rejected(mock_session_fn, mock_tg_cls, mock_dt_cls): + mock_session_fn.return_value = MagicMock() + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + + ts = _make_test_suite() + payload = _make_payload(_make_import_td(test_type="Alpha", table_name="t1")) + payload.version = 2 + + from testgen.common.test_definition_export_import_service import import_definitions + + with pytest.raises(InvalidImportPayload) as exc_info: + import_definitions(ts, _make_config(), payload) + + assert exc_info.value.code == ImportErrorCode.unsupported_version + assert "version 2" in str(exc_info.value) + + +@patch(f"{SERVICE_MODULE}.DataTable") +@patch(f"{SERVICE_MODULE}.TableGroup") +@patch(f"{SERVICE_MODULE}.get_current_session") +def test_apply_with_skips_still_applies_valid_actions(mock_session_fn, mock_tg_cls, mock_dt_cls): + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + session.execute.return_value.all.return_value = [] + + ts = _make_test_suite() + good_td = _make_import_td(test_type="Alpha", table_name="t1") + bad_td = _make_import_td(test_type="NonExistent", table_name="t1") + config = _make_config(mode=ImportMode.apply) + + from testgen.common.test_definition_export_import_service import import_definitions + + with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): + with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: + mock_td_cls.return_value = MagicMock(id=uuid4()) + result = import_definitions(ts, config, _make_payload(good_td, bad_td)) + + # Plain apply tolerates skips: the valid create is applied, no strict violation raised + assert result.summary.created == 1 + assert result.summary.skipped == 1 + session.add.assert_called_once() + + +def test_export_document_version_survives_default_omitting_serialization(): + doc = ExportDocument( + version=1, + source=ExportSource( + project_code="DEFAULT", + test_suite="suite", + table_group="tg", + table_group_schema="public", + exported_at=datetime(2026, 6, 11, tzinfo=UTC), + ), + definitions=[], + ) + + assert '"version":1' in doc.model_dump_json(exclude_defaults=True) + + # --- Schema: TestDefinitionExport --- @@ -988,6 +1082,96 @@ def test_defaults_match_expected_values(self): assert td.window_days == 0 assert td.history_lookback == 0 + def test_click_through_fields_default_none(self): + td = TestDefinitionExport(test_type="Alpha") + assert td.external_url is None + assert td.custom_metadata is None + assert "external_url" not in td.model_fields_set + assert "custom_metadata" not in td.model_fields_set + + def test_custom_metadata_accepts_object(self): + td = TestDefinitionExport(test_type="Alpha", custom_metadata={"pipeline": "daily_load", "task": "t1"}) + assert td.custom_metadata == {"pipeline": "daily_load", "task": "t1"} + + @pytest.mark.parametrize("bad_value", ["not-json", ["a", "b"], 42]) + def test_custom_metadata_rejects_non_object(self, bad_value): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata=bad_value) + + def test_custom_metadata_rejects_too_many_keys(self): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata={f"k{i}": "v" for i in range(51)}) + + def test_custom_metadata_rejects_oversized(self): + with pytest.raises(ValidationError): + TestDefinitionExport(test_type="Alpha", custom_metadata={"blob": "x" * 10_241}) + + +# --- Click-through fields: export + import --- + + +class Test_click_through_fields: + + @patch(f"{SERVICE_MODULE}.settings") + @patch(f"{SERVICE_MODULE}.TableGroup") + @patch(f"{SERVICE_MODULE}.get_current_session") + def test_export_includes_external_url_and_custom_metadata(self, mock_session_fn, mock_tg_cls, mock_settings): + mock_settings.VERSION = "5.12.0" + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + + ts = _make_test_suite() + + td_obj = MagicMock(spec=[]) + for field in TestDefinitionExport.model_fields: + setattr(td_obj, field, None) + td_obj.test_type = "Alpha" + td_obj.test_active = True + td_obj.lock_refresh = False + td_obj.skip_errors = 0 + td_obj.window_days = 0 + td_obj.history_lookback = 0 + td_obj.external_url = "https://example.com/notebook" + td_obj.custom_metadata = {"pipeline": "daily_load", "task": "transform_orders"} + + session.scalars.return_value.all.return_value = [td_obj] + + from testgen.common.test_definition_export_import_service import export_definitions + + result = export_definitions(ts, Origin.both, None, None) + + exported = result.definitions[0] + assert exported.external_url == "https://example.com/notebook" + assert exported.custom_metadata == {"pipeline": "daily_load", "task": "transform_orders"} + + @patch(f"{SERVICE_MODULE}.DataTable") + @patch(f"{SERVICE_MODULE}.TableGroup") + @patch(f"{SERVICE_MODULE}.get_current_session") + def test_create_sets_click_through_fields(self, mock_session_fn, mock_tg_cls, mock_dt_cls): + session = MagicMock() + mock_session_fn.return_value = session + mock_tg_cls.get.return_value = _make_table_group() + mock_dt_cls.select_table_names.return_value = ["t1"] + session.execute.return_value.all.return_value = [] + + ts = _make_test_suite() + td = _make_import_td( + test_type="Alpha", + table_name="t1", + external_url="https://example.com/task", + custom_metadata={"node": "n1"}, + ) + config = _make_config(mode=ImportMode.apply, on_new=OnNew.create) + + from testgen.common.test_definition_export_import_service import import_definitions + + with patch(f"{SERVICE_MODULE}._load_valid_test_types", return_value={"Alpha"}): + with patch(f"{SERVICE_MODULE}.TestDefinition") as mock_td_cls: + created_td = MagicMock(id=uuid4()) + mock_td_cls.return_value = created_td + import_definitions(ts, config, _make_payload(td)) + + assert created_td.external_url == "https://example.com/task" + assert created_td.custom_metadata == {"node": "n1"} -# Need to import this here for the endpoint test -from testgen.api.schemas import ImportSummary diff --git a/tests/unit/commands/queries/test_execute_tests_query.py b/tests/unit/commands/queries/test_execute_tests_query.py index 4839f99b..0c1276a9 100644 --- a/tests/unit/commands/queries/test_execute_tests_query.py +++ b/tests/unit/commands/queries/test_execute_tests_query.py @@ -1,5 +1,5 @@ from datetime import UTC, datetime -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pytest @@ -11,6 +11,8 @@ group_cat_tests, parse_cat_results, ) +from testgen.common.database.database_service import get_flavor_service +from testgen.common.models.connection import Connection pytestmark = pytest.mark.unit @@ -478,3 +480,100 @@ def test_resolve_cat_no_freshness_result_uses_band_check(_mock_changed): assert operator == "NOT BETWEEN" +def test_aggregate_cat_tests_handles_null_max_query_chars(): + """A connection with NULL max_query_chars must not crash CAT batching — the + `- 400` headroom subtraction falls back to DEFAULT_MAX_QUERY_CHARS.""" + instance = _make_execution_sql() + instance.connection = Connection(sql_flavor="postgresql", max_query_chars=None) + instance.flavor = "postgresql" + instance.flavor_service = get_flavor_service("postgresql") + + td = _make_td(measure_expression="m_expr", condition_expression="c_expr") + queries, grouped_defs = instance.aggregate_cat_tests([td], single=True) + + assert len(queries) == 1 + assert grouped_defs == [[td]] + + +# --- TestExecutionSQL._get_params baseline guards --- + + +def _make_params_execution_sql() -> TestExecutionSQL: + """Build a minimal TestExecutionSQL for exercising _get_params without a database.""" + instance = TestExecutionSQL.__new__(TestExecutionSQL) + flavor_service = MagicMock() + flavor_service.quote_character = '"' + flavor_service.varchar_type = "VARCHAR" + instance.flavor_service = flavor_service + instance.flavor = "postgresql" + instance.table_group = MagicMock(id=uuid4()) + instance.test_run = MagicMock(test_suite_id=uuid4(), id=uuid4()) + instance.run_date = datetime(2026, 1, 1, tzinfo=UTC) + return instance + + +def test_get_params_empty_baseline_counts_become_null(): + """Empty baseline counts must render as NULL, not "", to avoid CAST( AS FLOAT) syntax errors.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct="", baseline_value_ct="")) + assert params["BASELINE_CT"] == "NULL" + assert params["BASELINE_VALUE_CT"] == "NULL" + + +def test_get_params_none_baseline_counts_become_null(): + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct=None, baseline_value_ct=None)) + assert params["BASELINE_CT"] == "NULL" + assert params["BASELINE_VALUE_CT"] == "NULL" + + +def test_get_params_populated_baseline_counts_pass_through(): + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Missing_Pct", baseline_ct="1000", baseline_value_ct="950")) + assert params["BASELINE_CT"] == "1000" + assert params["BASELINE_VALUE_CT"] == "950" + + +def test_get_params_zero_baseline_count_is_not_nulled(): + """A real 0 is a meaningful value and must not be coerced to NULL.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Row_Ct_Pct", baseline_ct=0)) + assert params["BASELINE_CT"] == 0 + + +def test_get_params_empty_numeric_baselines_become_null(): + """All numeric baseline params render NULL when empty.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td( + test_type="Avg_Shift", + baseline_unique_ct="", baseline_avg="", baseline_sd="", baseline_sum="", + )) + assert params["BASELINE_UNIQUE_CT"] == "NULL" + assert params["BASELINE_AVG"] == "NULL" + assert params["BASELINE_SD"] == "NULL" + # Non-Freshness test types null-guard BASELINE_SUM (numeric use in Incr_Avg_Shift) + assert params["BASELINE_SUM"] == "NULL" + + +def test_get_params_freshness_baseline_sum_kept_raw_when_empty(): + """Freshness_Trend quotes BASELINE_SUM (NULLIF('', '') in template) — must stay empty, not 'NULL'.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Freshness_Trend", baseline_sum="")) + assert params["BASELINE_SUM"] == "" + + +def test_get_params_baseline_value_left_unguarded(): + """BASELINE_VALUE has non-uniform usage (quoted/number/IN-list) — not coerced to NULL.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Constant", baseline_value="")) + assert params["BASELINE_VALUE"] == "" + + +def test_get_params_empty_tolerances_become_null(): + """Tolerances use the same NULL guard.""" + instance = _make_params_execution_sql() + params = instance._get_params(_make_td(test_type="Volume_Trend", lower_tolerance="", upper_tolerance="")) + assert params["LOWER_TOLERANCE"] == "NULL" + assert params["UPPER_TOLERANCE"] == "NULL" + + diff --git a/tests/unit/commands/queries/test_refresh_data_chars_query.py b/tests/unit/commands/queries/test_refresh_data_chars_query.py index 5dc3ff16..427587e8 100644 --- a/tests/unit/commands/queries/test_refresh_data_chars_query.py +++ b/tests/unit/commands/queries/test_refresh_data_chars_query.py @@ -217,3 +217,16 @@ def test_filter_schema_columns_no_filters_returns_all(): filtered = sql_generator.filter_schema_columns(columns) assert {c.table_name for c in filtered} == {"users", "orders"} + + +def test_get_row_counts_handles_null_max_query_chars(): + """A connection with NULL max_query_chars must not crash chunking — it falls + back to DEFAULT_MAX_QUERY_CHARS so the UNION ALL count queries still build.""" + connection = Connection(sql_flavor="postgresql", max_query_chars=None) + table_group = TableGroup(table_group_schema="test_schema") + sql_generator = RefreshDataCharsSQL(connection, table_group) + + result = sql_generator.get_row_counts(["orders", "customers"]) + + assert result + assert all(isinstance(query, str) and query for query, _ in result) diff --git a/tests/unit/commands/test_run_data_cleanup.py b/tests/unit/commands/test_run_data_cleanup.py index 6c6e4304..07b463b2 100644 --- a/tests/unit/commands/test_run_data_cleanup.py +++ b/tests/unit/commands/test_run_data_cleanup.py @@ -51,17 +51,16 @@ def _patch_orchestrator( } started = {name: p.start() for name, p in patches.items()} - started["ProfilingRun"].find_latest_per_table_group.return_value = protected_profiling or set() - # get_job_execution_ids returns dict[run_id, je_id]; orchestrator filters nulls. - started["ProfilingRun"].get_job_execution_ids.return_value = { - uuid4(): je_id for je_id in (protected_profiling_jes or set()) - } + # The run id IS the job execution id, so find_latest_* returns the protected + # JE ids directly. The *_jes params let a test name the same set as JE ids. + started["ProfilingRun"].find_latest_per_table_group.return_value = ( + protected_profiling_jes or protected_profiling or set() + ) started["ProfilingRun"].delete_older_than.return_value = deleted_profiling - started["TestRun"].find_latest_per_test_suite.return_value = protected_tests or set() - started["TestRun"].get_job_execution_ids.return_value = { - uuid4(): je_id for je_id in (protected_test_jes or set()) - } + started["TestRun"].find_latest_per_test_suite.return_value = ( + protected_test_jes or protected_tests or set() + ) started["TestRun"].delete_older_than.return_value = deleted_tests started["JobExecution"].delete_older_than.return_value = deleted_job_executions diff --git a/tests/unit/commands/test_thresholds_prediction.py b/tests/unit/commands/test_thresholds_prediction.py index 37891a8c..d9f1f2ae 100644 --- a/tests/unit/commands/test_thresholds_prediction.py +++ b/tests/unit/commands/test_thresholds_prediction.py @@ -369,6 +369,36 @@ def test_freshness_gating_fits_on_filtered_series(mock_forecast): assert len(fitted_history) == len(freshness_updates) +@patch(MOCK_TARGET) +def test_freshness_gating_filtered_fit_uses_event_space(mock_forecast): + """The freshness-filtered fit must run in event-space (event_space=True): the filtered + series has one point per refresh and is irregularly spaced, so calendar resampling would + interpolate the refresh jumps into uniform increments and bias the forecast low.""" + mock_forecast.return_value = _make_forecast([220.0], [1.0]) + timestamps = [f"2026-01-{day:02d}" for day in range(1, 21)] + run_ids = [f"run_{i:02d}" for i in range(len(timestamps))] + history = _history_with_run_ids(timestamps, run_ids, value=220.0) + + compute_volume_or_metric_threshold(history, run_ids[:8], PredictSensitivity.medium) + + assert mock_forecast.call_args.kwargs["event_space"] is True + + +@patch(MOCK_TARGET) +def test_freshness_gating_raw_fallback_uses_calendar(mock_forecast): + """The raw-history fallback (when the filtered fit fails) keeps calendar regularization.""" + mock_forecast.side_effect = [NotEnoughData("not enough"), _make_forecast([220.0], [1.0])] + timestamps = [f"2026-01-{day:02d}" for day in range(1, 21)] + run_ids = [f"run_{i:02d}" for i in range(len(timestamps))] + history = _history_with_run_ids(timestamps, run_ids, value=220.0) + + compute_volume_or_metric_threshold(history, run_ids[:5], PredictSensitivity.medium) + + assert mock_forecast.call_count == 2 # filtered (event-space) failed, raw retried + assert mock_forecast.call_args_list[0].kwargs["event_space"] is True + assert mock_forecast.call_args_list[1].kwargs["event_space"] is False + + @patch(MOCK_TARGET) def test_freshness_gating_baseline_from_filtered_when_events_extend_past_history(mock_forecast): """When freshness_updates includes runs beyond the (retention-trimmed) history window, diff --git a/tests/unit/common/conftest.py b/tests/unit/common/conftest.py index 875d147b..9b895ce9 100644 --- a/tests/unit/common/conftest.py +++ b/tests/unit/common/conftest.py @@ -294,6 +294,24 @@ def _gen_subdaily_gap_schedule_phase() -> list[tuple[pd.Timestamp, float]]: return _to_csv_rows(_make_observations(start, end, 2, updates)) +def _gen_twice_daily_outage() -> list[tuple[pd.Timestamp, float]]: + """Updates every 12h (00:00/12:00 UTC), 7 days a week, 5 weeks, then updates stop. + + The 12h cadence classifies as "daily" frequency with a full-week active schedule, + so no excluded days or sub-daily window apply — staleness must come from the + first-pass gap thresholds. + """ + start = datetime(2025, 10, 6, 0, 0) + end = datetime(2025, 11, 11, 12, 0) + outage_start = datetime(2025, 11, 10, 0, 0) + updates: set[datetime] = set() + d = start + while d < outage_start: + updates.add(d) + d += timedelta(hours=12) + return _to_csv_rows(_make_observations(start, end, 12, updates)) + + def _gen_weekly_early() -> list[tuple[pd.Timestamp, float]]: start = datetime(2025, 8, 7, 10, 0) end = datetime(2025, 11, 6, 22, 0) diff --git a/tests/unit/api/oauth/test_models.py b/tests/unit/common/models/test_oauth.py similarity index 53% rename from tests/unit/api/oauth/test_models.py rename to tests/unit/common/models/test_oauth.py index d546feb6..bf9a5f7f 100644 --- a/tests/unit/api/oauth/test_models.py +++ b/tests/unit/common/models/test_oauth.py @@ -1,8 +1,8 @@ -"""Tests for testgen.api.oauth.models — OAuth2 ORM model business logic.""" +"""Tests for testgen.common.models.oauth — OAuth2 ORM model business logic.""" import time -from testgen.api.oauth.models import OAuth2Token +from testgen.common.models.oauth import OAuth2Token, PersonalAccessTokenStatus def _make_token(**overrides): @@ -37,3 +37,31 @@ def test_is_refresh_token_active_ignores_access_revocation(): def test_is_refresh_token_active_returns_false_when_expired(): token = _make_token(issued_at=int(time.time()) - (31 * 86400)) assert token.is_refresh_token_active() is False + + +# --- status (hybrid_property) --- + + +def test_status_active_when_unrevoked_and_unexpired(): + token = _make_token(issued_at=int(time.time()) - 10, expires_in=3600) + assert token.status == PersonalAccessTokenStatus.ACTIVE + + +def test_status_expired_when_past_expiry(): + token = _make_token(issued_at=int(time.time()) - 3600, expires_in=60) + assert token.status == PersonalAccessTokenStatus.EXPIRED + + +def test_status_revoked_when_access_revoked(): + token = _make_token(access_token_revoked_at=int(time.time())) + assert token.status == PersonalAccessTokenStatus.REVOKED + + +def test_status_revoked_takes_precedence_over_expired(): + # Past expiry AND revoked: revoked wins. + token = _make_token( + issued_at=int(time.time()) - 3600, + expires_in=60, + access_token_revoked_at=int(time.time()), + ) + assert token.status == PersonalAccessTokenStatus.REVOKED diff --git a/tests/unit/common/models/test_test_definition.py b/tests/unit/common/models/test_test_definition.py index f733d1b6..8a689b00 100644 --- a/tests/unit/common/models/test_test_definition.py +++ b/tests/unit/common/models/test_test_definition.py @@ -7,6 +7,8 @@ import pytest from testgen.common.models.test_definition import ( + CUSTOM_METADATA_MAX_BYTES, + CUSTOM_METADATA_MAX_KEYS, InvalidTestDefinitionFields, Severity, TestDefinition, @@ -99,6 +101,13 @@ def test_editable_fields_includes_param_columns(): assert {"threshold_value", "baseline_value"} <= accepted +def test_editable_fields_includes_click_through_fields(): + tt = make_test_type(param_columns=set(), default_parm_columns=None) + td = make_td() + accepted = td.editable_fields(tt) + assert {"external_url", "custom_metadata"} <= accepted + + def test_editable_fields_includes_impact_dimension_only_for_custom_or_referential_scope(): """impact_dimension is overridable only for user-defined-semantic scopes.""" td = make_td() @@ -116,8 +125,9 @@ def test_editable_fields_includes_impact_dimension_only_for_custom_or_referentia assert "impact_dimension" not in td.editable_fields(table_tt) -def test_editable_fields_includes_column_name_only_for_column_or_custom_scope(): - """column_name is meaningful for column-scope (column under test) and custom-scope (label).""" +def test_editable_fields_includes_column_name_except_for_table_scope(): + """column_name is meaningful for column (column under test), custom (label), and referential + (aggregate expression / categorical column list) scopes — but not table scope.""" td = make_td() column_tt = make_test_type(scope="column", param_columns={"threshold_value"}) @@ -126,12 +136,12 @@ def test_editable_fields_includes_column_name_only_for_column_or_custom_scope(): custom_tt = make_test_type(scope="custom", param_columns={"custom_query"}) assert "column_name" in td.editable_fields(custom_tt) + referential_tt = make_test_type(scope="referential", param_columns={"match_column_names"}) + assert "column_name" in td.editable_fields(referential_tt) + table_tt = make_test_type(scope="table", param_columns=set()) assert "column_name" not in td.editable_fields(table_tt) - referential_tt = make_test_type(scope="referential", param_columns={"match_column_names"}) - assert "column_name" not in td.editable_fields(referential_tt) - def test_editable_fields_does_not_leak_identity_or_internal_columns(): tt = make_test_type(param_columns={"threshold_value"}) @@ -170,6 +180,13 @@ def test_validate_wrong_scope_column_name_rejected(): assert "column_name" in exc_info.value.errors +def test_validate_referential_scope_accepts_column_name(): + # Referential tests use column_name as the aggregate expression / categorical column list. + tt = make_test_type(code="Aggregate_Balance", scope="referential", param_columns={"match_column_names"}) + td = make_td(column_name="SUM(total_amount)", match_column_names="SUM(total_amount)") + td.validate(tt) # no raise + + def test_validate_custom_scope_accepts_column_name_as_label(): # CUSTOM uses column_name as a "Test Focus" label — must be accepted. tt = make_test_type( @@ -220,6 +237,40 @@ def test_validate_severity_empty_string_treated_as_unset(): td.validate(tt) # empty severity is OK — falls back to test type default +def test_validate_custom_metadata_accepts_object_and_none(): + tt = make_test_type() + make_td(column_name="email", threshold_value="10", custom_metadata={"pipeline": "p1"}).validate(tt) + make_td(column_name="email", threshold_value="10", custom_metadata=None).validate(tt) + make_td(column_name="email", threshold_value="10", custom_metadata={}).validate(tt) + + +@pytest.mark.parametrize("bad_value", ["a string", ["a", "b"], 42]) +def test_validate_custom_metadata_rejects_non_object(bad_value): + tt = make_test_type() + td = make_td(column_name="email", threshold_value="10", custom_metadata=bad_value) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + +def test_validate_custom_metadata_rejects_too_many_keys(): + tt = make_test_type() + too_many = {f"k{i}": "v" for i in range(CUSTOM_METADATA_MAX_KEYS + 1)} + td = make_td(column_name="email", threshold_value="10", custom_metadata=too_many) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + +def test_validate_custom_metadata_rejects_oversized(): + tt = make_test_type() + oversized = {"blob": "x" * (CUSTOM_METADATA_MAX_BYTES + 1)} + td = make_td(column_name="email", threshold_value="10", custom_metadata=oversized) + with pytest.raises(InvalidTestDefinitionFields) as exc_info: + td.validate(tt) + assert "custom_metadata" in exc_info.value.errors + + def test_validate_aggregates_errors(): tt = make_test_type(scope="column") td = make_td(severity="critical", custom_query="SELECT 1") # no column_name @@ -294,6 +345,8 @@ def _make_summary_row(table_name: str = "my_table") -> dict: "prediction": None, "flagged": False, "impact_dimension": None, + "external_url": None, + "custom_metadata": None, "test_name_short": "Custom", "default_test_description": "A test", "measure_uom": "", diff --git a/tests/unit/common/models/test_test_type_columns.py b/tests/unit/common/models/test_test_type_columns.py new file mode 100644 index 00000000..f7b97095 --- /dev/null +++ b/tests/unit/common/models/test_test_type_columns.py @@ -0,0 +1,49 @@ +import pytest +from sqlalchemy.dialects.postgresql import dialect as pg_dialect + +from testgen.common.models.test_definition import StatisticalTechnique, TestAlgorithm, TestType + +pytestmark = pytest.mark.unit + +DIALECT = pg_dialect() + + +@pytest.mark.parametrize( + ("column_name", "enum_cls", "member"), + [ + ("algorithm", TestAlgorithm, TestAlgorithm.BOUNDARY_CHECK), + ("statistical_technique", StatisticalTechnique, StatisticalTechnique.COHENS_D), + ], +) +def test_enum_column_reads_db_value_as_enum_member(column_name, enum_cls, member): + column_type = TestType.__table__.c[column_name].type + process = column_type.result_processor(DIALECT, None) + + result = process(member.value) + + assert result is member + assert isinstance(result, enum_cls) + + +@pytest.mark.parametrize( + ("column_name", "member"), + [ + ("algorithm", TestAlgorithm.STATISTICAL_DRIFT), + ("statistical_technique", StatisticalTechnique.JENSEN_SHANNON_DIVERGENCE), + ], +) +def test_enum_column_writes_enum_as_db_value(column_name, member): + column_type = TestType.__table__.c[column_name].type + process = column_type.bind_processor(DIALECT) + + assert process(member) == member.value + + +@pytest.mark.parametrize("column_name", ["algorithm", "statistical_technique"]) +def test_enum_column_round_trips_none(column_name): + column_type = TestType.__table__.c[column_name].type + read = column_type.result_processor(DIALECT, None) + write = column_type.bind_processor(DIALECT) + + assert read(None) is None + assert write(None) is None diff --git a/tests/unit/common/notifications/test_monitor_run_notifications.py b/tests/unit/common/notifications/test_monitor_run_notifications.py index 248fcd41..9e3c9266 100644 --- a/tests/unit/common/notifications/test_monitor_run_notifications.py +++ b/tests/unit/common/notifications/test_monitor_run_notifications.py @@ -29,6 +29,12 @@ def create_test_result(table_name, test_type, message, result_code=0): return mock +def make_monitor_run(completed_at="2024-01-15T10:30:00Z"): + run = TestRun(id="monitor-run-id", test_suite_id="monitor-suite-id") + run.job_execution = Mock(completed_at=completed_at) + return run + + @pytest.fixture def ns_select_result(): return [ @@ -110,11 +116,7 @@ def test_send_monitor_notifications( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-id" @@ -193,11 +195,7 @@ def test_send_monitor_notifications_early_exit( test_result_select_where_mock, send_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() if not has_notifications: ns_select_patched.return_value = [] @@ -219,11 +217,7 @@ def test_send_monitor_notifications_anomaly_counts( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-id" @@ -270,11 +264,7 @@ def test_send_monitor_notifications_url_construction( send_mock, persisted_setting_mock, ): - test_run = TestRun( - id="monitor-run-id", - test_suite_id="monitor-suite-id", - test_endtime="2024-01-15T10:30:00Z", - ) + test_run = make_monitor_run() table_group = Mock(spec=TableGroup) table_group.id = "tg-123" diff --git a/tests/unit/common/notifications/test_profiling_run_notifications.py b/tests/unit/common/notifications/test_profiling_run_notifications.py index 15cd4e8b..16351212 100644 --- a/tests/unit/common/notifications/test_profiling_run_notifications.py +++ b/tests/unit/common/notifications/test_profiling_run_notifications.py @@ -4,6 +4,7 @@ import pytest +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus from testgen.common.models.hygiene_issue import IssueCount from testgen.common.models.notification_settings import ( ProfilingRunNotificationSettings, @@ -15,6 +16,10 @@ pytestmark = pytest.mark.unit +def make_job_execution(status, started_at=None, completed_at=None, error_message=None): + return Mock(status=status, started_at=started_at, completed_at=completed_at, error_message=error_message) + + def create_ns(**kwargs): with patch("testgen.common.notifications.profiling_run.ProfilingRunNotificationSettings.save"): return ProfilingRunNotificationSettings.create("proj", None, **kwargs) @@ -79,12 +84,12 @@ def get_prev_mock(): @pytest.mark.parametrize( ("profiling_run_status", "has_prev_run", "issue_count", "new_issue_count", "expected_triggers"), ( - ("Error", True, 25, 0, ("always", "on_changes")), - ("Error", True, 0, 0, ("always", "on_changes")), - ("Cancelled", True, 50, 10, ("always", "on_changes")), - ("Complete", True, 50, 10, ("always", "on_changes")), - ("Complete", True, 15, 0, ("always",)), - ("Complete", False, 15, 15, ("always", "on_changes")), + (JobStatus.ERROR, True, 25, 0, ("always", "on_changes")), + (JobStatus.ERROR, True, 0, 0, ("always", "on_changes")), + (JobStatus.CANCELED, True, 50, 10, ("always", "on_changes")), + (JobStatus.COMPLETED, True, 50, 10, ("always", "on_changes")), + (JobStatus.COMPLETED, True, 15, 0, ("always",)), + (JobStatus.COMPLETED, False, 15, 15, ("always", "on_changes")), ), ) def test_send_profiling_run_notification( @@ -102,11 +107,10 @@ def test_send_profiling_run_notification( ): profiling_run = ProfilingRun( id="pr-id", - job_execution_id="pr-id", table_groups_id="tg-id", - status=profiling_run_status, project_code="proj", ) + profiling_run.job_execution = make_job_execution(profiling_run_status) get_prev_mock.return_value = ProfilingRun(id="pr-prev-id") if has_prev_run else None new_count = iter(count()) priorities = ("Definite", "Likely", "Possible", "High", "Moderate") @@ -144,6 +148,7 @@ def test_send_profiling_run_notification( "start_time": None, "end_time": None, "status": profiling_run_status, + "status_label": JOB_STATUS_LABEL[profiling_run_status], "log_message": None, "table_ct": None, "column_ct": None, diff --git a/tests/unit/common/notifications/test_test_run_notifications.py b/tests/unit/common/notifications/test_test_run_notifications.py index beb7a2e8..d791037e 100644 --- a/tests/unit/common/notifications/test_test_run_notifications.py +++ b/tests/unit/common/notifications/test_test_run_notifications.py @@ -3,6 +3,7 @@ import pytest +from testgen.common.enums import JobStatus from testgen.common.models.notification_settings import TestRunNotificationSettings, TestRunNotificationTrigger from testgen.common.models.test_result import TestResultStatus from testgen.common.models.test_run import TestRun @@ -11,6 +12,10 @@ pytestmark = pytest.mark.unit +def make_job_execution(status): + return Mock(status=status) + + def create_ns(**kwargs): with patch("testgen.common.notifications.test_run.TestRunNotificationSettings.save"): return TestRunNotificationSettings.create("proj", None, **kwargs) @@ -98,19 +103,19 @@ def select_summary_mock(): "failed_expected", "warning_expected", "error_expected", "expected_triggers" ), [ - ("Complete", 0, 0, 0, {}, 0, 0, 0, ["always"]), - ("Complete", 0, 5, 0, {}, 0, 5, 0, ["always", "on_warnings"]), - ("Complete", 1, 1, 1, {}, 1, 1, 1, ["always", "on_failures", "on_warnings"]), - ("Complete", 50, 50, 50, {"failed": 2, "warning": 3}, 10, 5, 5, [ + (JobStatus.COMPLETED, 0, 0, 0, {}, 0, 0, 0, ["always"]), + (JobStatus.COMPLETED, 0, 5, 0, {}, 0, 5, 0, ["always", "on_warnings"]), + (JobStatus.COMPLETED, 1, 1, 1, {}, 1, 1, 1, ["always", "on_failures", "on_warnings"]), + (JobStatus.COMPLETED, 50, 50, 50, {"failed": 2, "warning": 3}, 10, 5, 5, [ "always", "on_failures", "on_warnings", "on_changes", ]), - ("Complete", 0, 0, 50, {"error": 50}, 0, 0, 20, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Complete", 50, 0, 0, None, 20, 0, 0, ["always", "on_failures", "on_warnings"]), - ("Complete", 50, 0, 10, {"failed": 5}, 15, 0, 5, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Error", 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Error", 20, 10, 0, None, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Cancelled", 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), - ("Cancelled", 30, 20, 0, {}, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.COMPLETED, 0, 0, 50, {"error": 50}, 0, 0, 20, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.COMPLETED, 50, 0, 0, None, 20, 0, 0, ["always", "on_failures", "on_warnings"]), + (JobStatus.COMPLETED, 50, 0, 10, {"failed": 5}, 15, 0, 5, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.ERROR, 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.ERROR, 20, 10, 0, None, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.CANCELED, 0, 0, 0, {}, 0, 0, 0, ["always", "on_failures", "on_warnings", "on_changes"]), + (JobStatus.CANCELED, 30, 20, 0, {}, 15, 5, 0, ["always", "on_failures", "on_warnings", "on_changes"]), ] ) def test_send_test_run_notification( @@ -134,13 +139,12 @@ def test_send_test_run_notification( test_run = TestRun( id="tr-id", - job_execution_id="tr-id", - status=test_run_status, test_suite_id="ts-id", failed_ct=failed_ct, warning_ct=warning_ct, error_ct=error_ct, ) + test_run.job_execution = make_job_execution(test_run_status) # SA 2.0 Row objects expose ._mapping; mock them accordingly def _make_row(): diff --git a/tests/unit/common/test_auth.py b/tests/unit/common/test_auth.py index 3d6bab34..6e926558 100644 --- a/tests/unit/common/test_auth.py +++ b/tests/unit/common/test_auth.py @@ -7,6 +7,7 @@ import pytest from testgen.common.auth import ( + AuthError, authorize_token, check_permission, create_jwt_token, @@ -50,14 +51,14 @@ def test_decode_jwt_token_decodes_valid_token(mock_settings): def test_decode_jwt_token_raises_for_expired_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY token = _make_token(exp_seconds=-3600) - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): decode_jwt_token(token) @patch("testgen.common.auth.settings") def test_decode_jwt_token_raises_for_invalid_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): decode_jwt_token("not-a-valid-token") @@ -99,7 +100,7 @@ def test_authorize_token_rejects_revoked(): mock_token_record.access_token_revoked_at = 1700000000 _set_scalars_results(mock_session, mock_user, mock_token_record) - with pytest.raises(ValueError, match="Token has been revoked"): + with pytest.raises(AuthError, match="Token has been revoked"): authorize_token("revoked_token", "testuser", mock_session) @@ -118,7 +119,7 @@ def test_authorize_token_raises_when_user_not_found(): # User lookup returns None — token check is never reached _set_scalars_results(mock_session, None) - with pytest.raises(ValueError, match="User not found"): + with pytest.raises(AuthError, match="User not found"): authorize_token("some_token", "ghost", mock_session) diff --git a/tests/unit/common/test_connection_service.py b/tests/unit/common/test_connection_service.py new file mode 100644 index 00000000..824fce6a --- /dev/null +++ b/tests/unit/common/test_connection_service.py @@ -0,0 +1,280 @@ +"""Tests for the connection_service common-layer module. + +Domain only: ``ConnectionStatus`` (with its load-bearing random field), the +``test_connection_status`` runner, and auth-path normalization. The label-bearing +connection-parameter schema / validation lives in ``mcp/tools/common.py`` and is +tested in ``tests/unit/mcp/test_connection_schema.py``. +""" + +from unittest.mock import patch + +import pytest +from sqlalchemy.exc import DatabaseError, DBAPIError + +from testgen.common.database.connection_service import ( + ConnectionStatus, + apply_connection_defaults, + normalize_auth_fields, +) + +# Aliased on import so pytest doesn't try to collect the ``test_*`` function as a test. +from testgen.common.database.connection_service import test_connection_status as run_connection_test +from testgen.common.models.connection import Connection + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.database.connection_service" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _conn(**overrides) -> Connection: + """Build a Connection without touching the DB. Defaults to a complete PG config.""" + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_name": "My DB", + "project_host": "localhost", + "project_port": "5432", + "project_db": "testgen_local", + "project_user": "testgen", + "project_pw_encrypted": "pw", + "connect_by_url": False, + "connect_by_key": False, + "connect_with_identity": False, + "max_threads": 4, + "max_query_chars": 20000, + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# ConnectionStatus dataclass +# --------------------------------------------------------------------------- + + +def test_connection_status_random_field_breaks_equality(): + """Two instances with identical (message, successful, details) MUST compare unequal. + + The random ``_`` field is required so Streamlit's reactive system re-renders + on repeated failed-test clicks producing the same error. Removing the field + would silently swallow the second click. + """ + a = ConnectionStatus(message="Error attempting the connection.", successful=False, details="boom") + b = ConnectionStatus(message="Error attempting the connection.", successful=False, details="boom") + assert a != b + + +# --------------------------------------------------------------------------- +# test_connection_status — happy path and exception branches +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[{"col": 1}]) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_success(mock_flavor, mock_fetch, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is True + assert status.message == "The connection was successful." + assert status.details is None + mock_empty_cache.assert_called_once() # service owns the cache reset + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[{"col": 0}]) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_query_returns_wrong_result(mock_flavor, mock_fetch, mock_empty_cache): + """``SELECT 1`` returns something other than 1 → 'Error completing a query'.""" + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error completing a query to the database server." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.fetch_from_target_db", side_effect=KeyError("host")) +@patch(f"{MODULE}.get_flavor_service") +def test_connection_status_key_error(mock_flavor, mock_fetch, mock_empty_cache): + """Missing required field → 'Complete all the required fields.'""" + mock_flavor.return_value.test_query = "SELECT 1" + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection. " + assert status.details == "Complete all the required fields." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_database_error(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + orig = Exception("FATAL: password authentication failed") + mock_fetch.side_effect = DatabaseError("stmt", {}, orig) + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert "password authentication failed" in str(status.details) + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_dbapi_error(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + orig = Exception("driver-level failure") + mock_fetch.side_effect = DBAPIError("stmt", {}, orig) + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert "driver-level failure" in str(status.details) + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_open_ssl_error(mock_fetch, mock_flavor, mock_empty_cache): + """A TypeError whose args[1][0] is an OpenSSLError uses args[0] as details.""" + mock_flavor.return_value.test_query = "SELECT 1" + + class OpenSSLError: # name matches what is_open_ssl_error sniffs for + pass + + err = TypeError("bad key", [OpenSSLError()]) + mock_fetch.side_effect = err + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "bad key" + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_missing_private_key(mock_fetch, mock_flavor, mock_empty_cache): + """connect_by_key=True with no private_key → 'The private key is missing.'""" + mock_flavor.return_value.test_query = "SELECT 1" + mock_fetch.side_effect = RuntimeError("something") + + conn = _conn(sql_flavor="snowflake", sql_flavor_code="snowflake", connect_by_key=True, private_key=None) + status = run_connection_test(conn) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "The private key is missing." + + +@patch(f"{MODULE}.empty_cache") +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.fetch_from_target_db") +def test_connection_status_generic_exception(mock_fetch, mock_flavor, mock_empty_cache): + mock_flavor.return_value.test_query = "SELECT 1" + mock_fetch.side_effect = RuntimeError("unexpected") + + status = run_connection_test(_conn()) + + assert status.successful is False + assert status.message == "Error attempting the connection." + assert status.details == "Something went wrong while testing the connection." + + +# --------------------------------------------------------------------------- +# normalize_auth_fields +# --------------------------------------------------------------------------- + + +def test_normalize_clears_password_when_connect_by_key_non_databricks(): + """Snowflake connect_by_key=True → project_pw_encrypted cleared.""" + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + connect_by_key=True, + project_pw_encrypted="old_pw", + private_key="key", + ) + normalize_auth_fields(conn) + assert conn.project_pw_encrypted in (None, "") + assert conn.private_key == "key" # untouched + + +def test_normalize_keeps_password_for_databricks_connect_by_key(): + """Databricks OAuth uses connect_by_key but stores the secret in project_pw_encrypted.""" + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + connect_by_key=True, + project_pw_encrypted="client_secret_xyz", + ) + normalize_auth_fields(conn) + assert conn.project_pw_encrypted == "client_secret_xyz" + assert not conn.private_key + assert not conn.private_key_passphrase + + +def test_normalize_clears_private_key_fields_when_password_auth(): + """connect_by_key=False → private_key + passphrase cleared.""" + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + connect_by_key=False, + private_key="old_key", + private_key_passphrase="old_phrase", # noqa: S106 + ) + normalize_auth_fields(conn) + assert not conn.private_key + assert not conn.private_key_passphrase + + +def test_normalize_clears_user_password_when_identity(): + """connect_with_identity=True → project_user + project_pw_encrypted cleared.""" + conn = _conn( + sql_flavor_code="azure_mssql", + sql_flavor="mssql", + connect_with_identity=True, + project_user="leftover_user", + project_pw_encrypted="leftover_pw", + ) + normalize_auth_fields(conn) + assert not conn.project_user + assert not conn.project_pw_encrypted + + +# --------------------------------------------------------------------------- +# apply_connection_defaults +# --------------------------------------------------------------------------- + + +def test_defaults_fill_max_query_chars(): + conn = _conn(max_query_chars=None) + apply_connection_defaults(conn) + assert conn.max_query_chars == 20000 + + +def test_defaults_salesforce_lower_max_query_chars(): + """Salesforce Data 360's Hyper engine gets the lower 15000 default.""" + conn = _conn(sql_flavor="salesforce_data360", sql_flavor_code="salesforce_data360", max_query_chars=None) + apply_connection_defaults(conn) + assert conn.max_query_chars == 15000 + + +def test_defaults_keep_explicit_max_query_chars(): + conn = _conn(max_query_chars=30000) + apply_connection_defaults(conn) + assert conn.max_query_chars == 30000 diff --git a/tests/unit/common/test_data_catalog_service.py b/tests/unit/common/test_data_catalog_service.py new file mode 100644 index 00000000..2d0c7f9a --- /dev/null +++ b/tests/unit/common/test_data_catalog_service.py @@ -0,0 +1,305 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from testgen.common.data_catalog_service import ( + DESCRIPTION_MAX_LENGTH, + TAG_FIELDS, + TAG_MAX_LENGTH, + TableSampleResult, + apply_column_metadata, + apply_table_metadata, + build_create_table_script, + disable_autoflags, + fetch_table_sample, + render_create_table_script, + validate_metadata_fields, +) +from testgen.common.models.data_column import CreateScriptColumn + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.data_catalog_service" + + +def _flavor_service(table_ref='"demo"."orders"', limit_clauses=("", "LIMIT 100")): + fs = MagicMock() + fs.quote_character = '"' + fs.get_table_ref.return_value = table_ref + fs.row_limit_clauses.return_value = limit_clauses + return fs + + +def _column(name, db_type="varchar(50)", suggestion="VARCHAR(20)"): + return CreateScriptColumn(column_name=name, db_data_type=db_type, datatype_suggestion=suggestion) + + +# --------------------------------------------------------------------------- +# render_create_table_script +# --------------------------------------------------------------------------- + +def test_render_uses_suggested_type_and_quotes_identifiers(): + fs = _flavor_service() + columns = [_column("id", "int", "INTEGER"), _column("name", "text", "VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert script.startswith('CREATE TABLE "demo"."orders" (') + assert '"id"' in script + assert '"name"' in script + assert "INTEGER" in script + assert "VARCHAR(20)" in script + + +def test_render_omits_was_annotations_by_default(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert "-- WAS" not in script + + +def test_render_includes_was_annotations_when_enabled(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs, annotate_changes=True) + + assert "-- WAS text" in script + + +def test_render_no_was_annotation_when_type_unchanged(): + fs = _flavor_service() + columns = [_column("name", db_type="VARCHAR(20)", suggestion="VARCHAR(20)")] + + script = render_create_table_script("demo", "orders", columns, fs, annotate_changes=True) + + assert "-- WAS" not in script + + +def test_render_falls_back_to_db_type_when_no_suggestion(): + fs = _flavor_service() + columns = [_column("name", db_type="text", suggestion=None)] + + script = render_create_table_script("demo", "orders", columns, fs) + + assert "text" in script + + +def test_render_only_last_column_has_no_trailing_comma(): + fs = _flavor_service() + columns = [_column("a", "int", "INTEGER"), _column("b", "int", "INTEGER")] + + script = render_create_table_script("demo", "orders", columns, fs) + + body_lines = [line for line in script.splitlines() if '"a"' in line or '"b"' in line] + assert body_lines[0].rstrip().endswith(",") + assert not body_lines[1].rstrip().endswith(",") + + +# --------------------------------------------------------------------------- +# build_create_table_script +# --------------------------------------------------------------------------- + +@patch(f"{MODULE}.DataColumnChars") +def test_build_returns_none_when_no_columns(mock_columns): + mock_columns.list_for_create_script.return_value = (None, []) + + assert build_create_table_script("tg-id", "orders") is None + + +@patch(f"{MODULE}.get_flavor_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.DataColumnChars") +def test_build_renders_script_for_existing_table(mock_columns, mock_conn, mock_flavor): + mock_columns.list_for_create_script.return_value = ("demo", [_column("id", "int", "INTEGER")]) + mock_conn.get_by_table_group.return_value = MagicMock(sql_flavor="postgresql") + mock_flavor.return_value = _flavor_service() + + script = build_create_table_script("tg-id", "orders") + + assert script is not None + assert "CREATE TABLE" in script + assert "INTEGER" in script + + +# --------------------------------------------------------------------------- +# fetch_table_sample +# --------------------------------------------------------------------------- + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_ok(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "OK" + assert len(result.df) == 2 + assert result.pii_redacted is False + + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_builds_distinct_limited_query(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service(limit_clauses=("", "LIMIT 50")) + mock_fetch.return_value = [{"id": 1}] + connection = MagicMock(sql_flavor="postgresql") + + fetch_table_sample(connection, "tg-id", "demo", "orders", limit=50, mask_pii=False) + + query = mock_fetch.call_args.args[1] + assert "SELECT DISTINCT" in query + assert "LIMIT 50" in query + assert '"demo"."orders"' in query + + +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_empty_returns_nd(mock_flavor, mock_fetch): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "ND" + + +@patch(f"{MODULE}.fetch_from_target_db", side_effect=Exception("connection refused")) +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_error_returns_err(mock_flavor, _mock_fetch): + mock_flavor.return_value = _flavor_service() + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=False) + + assert result.status == "ERR" + + +@patch(f"{MODULE}.get_pii_columns", return_value={"name"}) +@patch(f"{MODULE}.fetch_from_target_db") +@patch(f"{MODULE}.get_flavor_service") +def test_fetch_sample_masks_pii_columns(mock_flavor, mock_fetch, _mock_pii): + mock_flavor.return_value = _flavor_service() + mock_fetch.return_value = [{"id": 1, "name": "secret"}] + connection = MagicMock(sql_flavor="postgresql") + + result = fetch_table_sample(connection, "tg-id", "demo", "orders", limit=100, mask_pii=True) + + assert result.status == "OK" + assert result.pii_redacted is True + assert "secret" not in result.df["name"].tolist() + + +def test_table_sample_result_defaults(): + result = TableSampleResult("ND") + assert result.df is None + assert result.pii_redacted is False + + +# --------------------------------------------------------------------------- +# validate_metadata_fields +# --------------------------------------------------------------------------- + +def test_validate_accepts_values_within_limits(): + fields = {"description": "x" * DESCRIPTION_MAX_LENGTH, "business_domain": "y" * TAG_MAX_LENGTH} + assert validate_metadata_fields(fields) == [] + + +def test_validate_flags_description_too_long(): + errors = validate_metadata_fields({"description": "x" * (DESCRIPTION_MAX_LENGTH + 1)}) + assert len(errors) == 1 + assert "description" in errors[0] + + +def test_validate_flags_each_tag_too_long(): + fields = {"data_source": "a" * (TAG_MAX_LENGTH + 1), "data_product": "b" * (TAG_MAX_LENGTH + 1)} + errors = validate_metadata_fields(fields) + assert len(errors) == 2 + assert any("data_source" in e for e in errors) + assert any("data_product" in e for e in errors) + + +def test_validate_ignores_none_and_absent(): + # None clears the field — never a length violation; absent keys are skipped. + assert validate_metadata_fields({"description": None, "source_system": None}) == [] + assert validate_metadata_fields({}) == [] + + +# --------------------------------------------------------------------------- +# apply_table_metadata / apply_column_metadata +# --------------------------------------------------------------------------- + +def _table(): + return SimpleNamespace(description="old", business_domain="old", critical_data_element=False) + + +def _metadata_column(): + return SimpleNamespace( + description="old", data_product="old", critical_data_element=False, + excluded_data_element=False, pii_flag=None, + ) + + +def test_apply_table_sets_present_keys_only(): + table = _table() + apply_table_metadata(table, {"description": "new", "critical_data_element": True}) + assert table.description == "new" + assert table.critical_data_element is True + assert table.business_domain == "old" # absent → unchanged + + +def test_apply_table_none_clears(): + table = _table() + apply_table_metadata(table, {"business_domain": None}) + assert table.business_domain is None + + +def test_apply_column_sets_column_only_fields(): + column = _metadata_column() + apply_column_metadata(column, {"pii_flag": "MANUAL", "excluded_data_element": True, "data_product": "CRM"}) + assert column.pii_flag == "MANUAL" + assert column.excluded_data_element is True + assert column.data_product == "CRM" + + +def test_apply_column_none_clears_pii(): + column = _metadata_column() + column.pii_flag = "MANUAL" + apply_column_metadata(column, {"pii_flag": None}) + assert column.pii_flag is None + + +# --------------------------------------------------------------------------- +# disable_autoflags +# --------------------------------------------------------------------------- + +def test_disable_autoflags_disables_only_written_and_enabled(): + tg = SimpleNamespace(profile_flag_cdes=True, profile_flag_pii=True) + disabled = disable_autoflags(tg, wrote_cde=True, wrote_pii=False) + assert disabled == ["profile_flag_cdes"] + assert tg.profile_flag_cdes is False + assert tg.profile_flag_pii is True + + +def test_disable_autoflags_noop_when_already_disabled(): + tg = SimpleNamespace(profile_flag_cdes=False, profile_flag_pii=False) + assert disable_autoflags(tg, wrote_cde=True, wrote_pii=True) == [] + + +def test_disable_autoflags_disables_both(): + tg = SimpleNamespace(profile_flag_cdes=True, profile_flag_pii=True) + disabled = disable_autoflags(tg, wrote_cde=True, wrote_pii=True) + assert disabled == ["profile_flag_cdes", "profile_flag_pii"] + + +def test_tag_fields_has_nine_entries(): + assert len(TAG_FIELDS) == 9 + assert "data_source" in TAG_FIELDS + assert "aggregation_level" in TAG_FIELDS + assert "data_classification" in TAG_FIELDS diff --git a/tests/unit/common/test_enums.py b/tests/unit/common/test_enums.py new file mode 100644 index 00000000..2525839d --- /dev/null +++ b/tests/unit/common/test_enums.py @@ -0,0 +1,18 @@ +import pytest + +from testgen.common.enums import JOB_STATUS_LABEL, JobStatus +from testgen.common.models.profiling_run import ProfilingRunSummary +from testgen.common.models.test_run import TestRunSummary + +pytestmark = pytest.mark.unit + + +def test_job_status_label_covers_every_status(): + # A missing entry would surface a raw lowercase status code (e.g. "cancel_requested") + # in run lists, emails, and MCP output instead of a display label. + assert set(JOB_STATUS_LABEL) == set(JobStatus) + + +def test_run_summaries_share_the_job_status_label_map(): + assert TestRunSummary.STATUS_LABEL is JOB_STATUS_LABEL + assert ProfilingRunSummary.STATUS_LABEL is JOB_STATUS_LABEL diff --git a/tests/unit/common/test_flavor_sampling.py b/tests/unit/common/test_flavor_sampling.py new file mode 100644 index 00000000..06177d4f --- /dev/null +++ b/tests/unit/common/test_flavor_sampling.py @@ -0,0 +1,47 @@ +"""Per-flavor sampleable object types. + +Profiling skips sampling for object types a flavor's sample clause cannot handle +(see run_profiling._compute_sampling_params). These tests lock in the per-flavor +sampleable_object_types and the object-type signal the skip depends on. +""" +import pytest + +from testgen.common.database.column_chars import ColumnChars, ObjectType +from testgen.common.database.database_service import get_flavor_service + +pytestmark = pytest.mark.unit + + +# Verified live: TABLESAMPLE errors on a view for these flavors, so only physical relations +# are sampleable; views/external/other are profiled in full. +def test_mssql_samples_only_base_tables(): + # SQL Server has no materialized views, so the sampleable set is just base tables. + assert get_flavor_service("mssql").sampleable_object_types == {ObjectType.TABLE} + + +def test_postgresql_samples_tables_and_materialized_views(): + assert get_flavor_service("postgresql").sampleable_object_types == { + ObjectType.TABLE, + ObjectType.MATERIALIZED_VIEW, + } + + +@pytest.mark.parametrize("flavor", ["redshift", "snowflake", "databricks", "oracle", "bigquery"]) +def test_sampleable_unrestricted_where_sampling_works_on_views(flavor): + # None = no restriction. Row-based samplers, or engines that accept the sample clause on + # views (verified live for snowflake/databricks/oracle). + assert get_flavor_service(flavor).sampleable_object_types is None + + +def test_object_type_enum_values_match_ddf_strings(): + # The DDF emits these literal strings; the skip compares column.object_type against them. + assert ObjectType.TABLE == "TABLE" + assert ObjectType.VIEW == "VIEW" + assert ObjectType.MATERIALIZED_VIEW == "MATERIALIZED_VIEW" + assert ObjectType.EXTERNAL == "EXTERNAL" + assert ObjectType.OTHER == "OTHER" + + +def test_column_chars_object_type_defaults_none(): + # Flavors whose DDF does not supply object_type leave it None. + assert ColumnChars(schema_name="s", table_name="t", column_name="c").object_type is None diff --git a/tests/unit/common/test_flavors.py b/tests/unit/common/test_flavors.py new file mode 100644 index 00000000..4f334817 --- /dev/null +++ b/tests/unit/common/test_flavors.py @@ -0,0 +1,35 @@ +"""Tests for the shared flavor-identity source of truth.""" + +import pytest + +from testgen.common.flavors import ( + FLAVOR_CODE_TO_FAMILY, + FLAVOR_CODE_TO_LABEL, + SqlFlavorLabel, +) + +pytestmark = pytest.mark.unit + + +def test_label_and_family_maps_cover_the_same_codes(): + assert set(FLAVOR_CODE_TO_LABEL) == set(FLAVOR_CODE_TO_FAMILY) + + +def test_every_label_is_a_known_enum_member(): + assert set(FLAVOR_CODE_TO_LABEL.values()) == set(SqlFlavorLabel) + + +def test_labels_are_unique_per_code(): + labels = list(FLAVOR_CODE_TO_LABEL.values()) + assert len(labels) == len(set(labels)) + + +def test_azure_variants_share_the_mssql_family(): + assert FLAVOR_CODE_TO_FAMILY["azure_mssql"] == "mssql" + assert FLAVOR_CODE_TO_FAMILY["synapse_mssql"] == "mssql" + assert FLAVOR_CODE_TO_FAMILY["mssql"] == "mssql" + + +def test_label_renders_as_plain_string(): + assert f"{FLAVOR_CODE_TO_LABEL['postgresql']}" == "PostgreSQL" + assert str(FLAVOR_CODE_TO_LABEL["azure_mssql"]) == "Azure SQL Database" diff --git a/tests/unit/common/test_freshness_scenarios.py b/tests/unit/common/test_freshness_scenarios.py index a20d2a52..a9e1f289 100644 --- a/tests/unit/common/test_freshness_scenarios.py +++ b/tests/unit/common/test_freshness_scenarios.py @@ -25,6 +25,7 @@ _gen_subdaily_gap_schedule_phase, _gen_subdaily_regular, _gen_training_only, + _gen_twice_daily_outage, _gen_weekly_early, _run_scenario, ) @@ -472,3 +473,48 @@ def test_recovery_passes(self, results: list[ScenarioPoint]) -> None: assert post_recovery[0].result_code == 0 # Second update after recovery should pass assert post_recovery[1].result_code == 1 + + +# ─── Scenario 9: Twice-Daily Outage (full-week schedule) ──────────── + + +class Test_TwiceDailyOutage: + """Updates every 12h around the clock for 5 weeks, then updates stop. + + Full-week "daily"-frequency schedule: no excluded days and no sub-daily + window, so the staleness threshold must come from the first-pass gap + computation. With a 720-min median gap, staleness = 720 * 0.85 = 612 min, + so the first missed check (gap 720 min) flags Late. + """ + + @pytest.fixture(scope="class") + def results(self) -> list[ScenarioPoint]: + rows = _gen_twice_daily_outage() + return _run_scenario(rows, PredictSensitivity.medium, exclude_weekends=False, tz="UTC") + + def test_schedule_goes_active_all_days(self, results: list[ScenarioPoint]) -> None: + sched = _schedule(_updates(results)[-1]) + assert sched is not None + assert sched.get("schedule_stage") == "active" + assert sched.get("active_days") == [0, 1, 2, 3, 4, 5, 6] + + def test_staleness_stored_once_active(self, results: list[ScenarioPoint]) -> None: + last_update = _updates(results)[-1] + assert last_update.staleness == pytest.approx(720 * 0.85) + + def test_no_anomalies_before_outage(self, results: list[ScenarioPoint]) -> None: + outage_start = pd.Timestamp("2025-11-10") + assert all(p.timestamp >= outage_start for p in _anomalies(results)) + + def test_first_missed_check_flags_late(self, results: list[ScenarioPoint]) -> None: + anomalies = _anomalies(results) + assert len(anomalies) > 0 + first = anomalies[0] + assert first.timestamp == pd.Timestamp("2025-11-10 00:00") + assert first.value == 720 + + def test_all_missed_checks_flag(self, results: list[ScenarioPoint]) -> None: + outage_start = pd.Timestamp("2025-11-10") + missed_checks = [p for p in results if p.timestamp >= outage_start] + assert len(missed_checks) == 4 + assert all(p.result_code == 0 for p in missed_checks) diff --git a/tests/unit/common/test_holiday_service.py b/tests/unit/common/test_holiday_service.py new file mode 100644 index 00000000..a2b7ddf8 --- /dev/null +++ b/tests/unit/common/test_holiday_service.py @@ -0,0 +1,28 @@ +"""Tests for holiday-code validation in ``common/holiday_service.py``.""" + +import pytest + +from testgen.common.holiday_service import is_supported_holiday_code + +pytestmark = pytest.mark.unit + + +@pytest.mark.parametrize("code", ["US", "GB", "CA", "USA", "NYSE", "ECB", " us "]) +def test_supported_codes(code): + """ISO country codes, uppercase aliases, and financial-market codes resolve.""" + assert is_supported_holiday_code(code) is True + + +@pytest.mark.parametrize( + "code", + [ + "US_FEDERAL", # not a holidays-package calendar + "CA_STAT", + "Canada", # the upper-cased "CANADA" is not a valid key — documented quirk + "NOTAREALCODE", + "", + " ", + ], +) +def test_unsupported_codes(code): + assert is_supported_holiday_code(code) is False diff --git a/tests/unit/common/test_mixpanel_service.py b/tests/unit/common/test_mixpanel_service.py new file mode 100644 index 00000000..58bc598a --- /dev/null +++ b/tests/unit/common/test_mixpanel_service.py @@ -0,0 +1,146 @@ +import queue +from unittest.mock import MagicMock + +import pytest + +from testgen import settings +from testgen.common import mixpanel_service as mp_module +from testgen.common.mixpanel_service import MixpanelService +from testgen.utils.singleton import SingletonType + + +@pytest.fixture +def service(monkeypatch): + """Fresh, isolated MixpanelService with analytics on and no real network/DB.""" + SingletonType._instances.pop(MixpanelService, None) + monkeypatch.setattr(settings, "ANALYTICS_ENABLED", True) + monkeypatch.setattr(settings, "MIXPANEL_TOKEN", "tok") + # Avoid Streamlit + DB: session.auth falsy, instance_id pre-seeded. + monkeypatch.setattr(mp_module, "session", MagicMock(auth=None)) + svc = MixpanelService() + svc.__dict__["instance_id"] = "iid" + yield svc + svc.drain() + SingletonType._instances.pop(MixpanelService, None) + + +def test_send_event_enqueues_and_worker_flushes_one_event(service, monkeypatch): + captured = [] + monkeypatch.setattr(service, "send_mp_request", lambda endpoint, payload: captured.append((endpoint, payload))) + + service.send_event("nav-home") + service.drain() + + assert len(captured) == 1 + endpoint, payload = captured[0] + assert endpoint == "track?ip=1" + assert isinstance(payload, list) + assert len(payload) == 1 + assert payload[0]["event"] == "nav-home" + assert payload[0]["properties"]["token"] == "tok" # noqa: S105 (test token, not a secret) + assert payload[0]["properties"]["instance_id"] == "iid" + assert "time" in payload[0]["properties"] + + +def test_next_batch_caps_at_batch_size(service, monkeypatch): + monkeypatch.setattr(mp_module, "_BATCH_SIZE", 2) + service._queue = queue.Queue() + for i in range(5): + service._queue.put({"event": f"e{i}", "properties": {}}) + + batch, stop = service._next_batch() + + assert [e["event"] for e in batch] == ["e0", "e1"] + assert stop is False + + +def test_next_batch_stops_on_sentinel_and_returns_pending(service): + service._queue = queue.Queue() + service._queue.put({"event": "a", "properties": {}}) + service._queue.put({"event": "b", "properties": {}}) + service._queue.put(mp_module._SHUTDOWN) + + batch, stop = service._next_batch() + + assert [e["event"] for e in batch] == ["a", "b"] + assert stop is True + + +def test_flush_chunks_over_batch_size(service, monkeypatch): + monkeypatch.setattr(mp_module, "_BATCH_SIZE", 2) + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(len(payload))) + + service._flush([{"event": f"e{i}", "properties": {}} for i in range(3)]) + + assert posts == [2, 1] + + +def test_queue_overflow_drops_and_warns(service, monkeypatch, caplog): + # Pretend the worker is running so _enqueue does not start a real one, + # and give it a size-1 queue that nothing drains. + service._started = True + service._queue = queue.Queue(maxsize=1) + + service._enqueue({"event": "a", "properties": {}}) + with caplog.at_level("WARNING"): + service._enqueue({"event": "b", "properties": {}}) + + assert service._queue.qsize() == 1 + assert "analytics queue full" in caplog.text + + +def test_drain_stops_worker_and_is_idempotent(service, monkeypatch): + monkeypatch.setattr(service, "send_mp_request", lambda *_: None) + service.send_event("x") # starts worker + + service.drain() + assert service._worker is not None + assert not service._worker.is_alive() + + service.drain() # second call is a no-op, must not raise + + +def test_analytics_disabled_never_starts_worker(service, monkeypatch): + monkeypatch.setattr(settings, "ANALYTICS_ENABLED", False) + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(payload)) + + service.send_event("nav-home") + service.drain() + + assert service._started is False + assert service._worker is None + assert posts == [] + + +def test_send_event_after_drain_drops_and_warns(service, monkeypatch, caplog): + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda _endpoint, payload: posts.append(payload)) + + service.send_event("pre-drain") + service.drain() + + posts.clear() + with caplog.at_level("WARNING"): + service.send_event("late") + + assert "stopped" in caplog.text + assert posts == [] + + +def test_send_feedback_posts_synchronously_without_worker(service, monkeypatch): + posts = [] + monkeypatch.setattr(service, "send_mp_request", lambda endpoint, payload: posts.append((endpoint, payload))) + + service.send_feedback(comment="great tool") + + # Posted immediately (no drain needed), as a single dict (not a list), worker never started. + assert len(posts) == 1 + endpoint, payload = posts[0] + assert endpoint == "track?ip=1" + assert isinstance(payload, dict) + assert payload["event"] == "feedback" + assert payload["properties"]["comment"] == "great tool" + assert service._started is False + assert service._worker is None diff --git a/tests/unit/common/test_monitor_service.py b/tests/unit/common/test_monitor_service.py new file mode 100644 index 00000000..11c17c50 --- /dev/null +++ b/tests/unit/common/test_monitor_service.py @@ -0,0 +1,118 @@ +"""Tests for ``common/monitor_service.py`` — enable / update / disable orchestration.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.monitor_service" + + +def _table_group(**overrides) -> MagicMock: + tg = MagicMock() + tg.id = overrides.get("id", uuid4()) + tg.project_code = overrides.get("project_code", "demo") + tg.table_groups_name = overrides.get("table_groups_name", "Sales") + tg.connection_id = overrides.get("connection_id", 1) + tg.monitor_test_suite_id = overrides.get("monitor_test_suite_id", None) + return tg + + +# --------------------------------------------------------------------------- +# enable_monitoring +# --------------------------------------------------------------------------- + + +def test_enable_monitoring_rejects_when_already_enabled(): + from testgen.common.monitor_service import enable_monitoring + + with pytest.raises(ValueError, match="already enabled"): + enable_monitoring(_table_group(monitor_test_suite_id=uuid4()), "0 6 * * *") + + +@patch(f"{MODULE}.run_monitor_generation") +@patch(f"{MODULE}.get_current_session") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestSuite") +def test_enable_monitoring_merges_defaults_skipping_none(mock_ts, mock_js, mock_session, mock_gen): + tg = _table_group() + suite = mock_ts.return_value + suite.id = uuid4() + mock_session.return_value.scalar.return_value = 4 + + from testgen.common.monitor_service import enable_monitoring + + returned_suite, count = enable_monitoring( + tg, + "0 6 * * *", + "UTC", + # monitor_lookback=None must fall back to the default; False must be honored (not None). + suite_attrs={"monitor_lookback": None, "predict_sensitivity": "high", "predict_exclude_weekends": False}, + ) + + kwargs = mock_ts.call_args.kwargs + assert kwargs["is_monitor"] is True + assert kwargs["monitor_lookback"] == 14 # None override skipped → default + assert kwargs["predict_sensitivity"] == "high" # override applied + assert kwargs["predict_exclude_weekends"] is False # False is not None → applied + assert kwargs["predict_min_lookback"] == 30 # untouched default + assert returned_suite is suite + assert count == 4 + mock_gen.assert_called_once_with(suite.id, ["Volume_Trend", "Schema_Drift"]) + tg.save.assert_called() + assert tg.monitor_test_suite_id == suite.id + + +# --------------------------------------------------------------------------- +# update_monitoring +# --------------------------------------------------------------------------- + + +def test_update_monitoring_whitelists_suite_attrs_and_edits_schedule_in_place(): + suite = SimpleNamespace(save=MagicMock()) + schedule = SimpleNamespace(cron_expr="0 6 * * *", cron_tz="UTC", active=True, save=MagicMock()) + + from testgen.common.monitor_service import update_monitoring + + update_monitoring( + suite, + schedule, + suite_attrs={"predict_sensitivity": "high", "predict_holiday_codes": None, "is_monitor": False}, + cron_expr="0 */12 * * *", + active=False, + ) + + # Whitelisted columns applied (a present None clears); non-whitelisted key ignored. + assert suite.predict_sensitivity == "high" + assert suite.predict_holiday_codes is None + assert not hasattr(suite, "is_monitor") + suite.save.assert_called_once() + + # Schedule edited in place — same object, no recreate. + assert schedule.cron_expr == "0 */12 * * *" + assert schedule.cron_tz == "UTC" # not supplied → unchanged + assert schedule.active is False + schedule.save.assert_called_once() + + +# --------------------------------------------------------------------------- +# disable_monitoring +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.get_current_session") +def test_disable_monitoring_counts_before_cascade(mock_session, mock_ts): + suite = MagicMock() + suite.id = uuid4() + mock_session.return_value.scalar.side_effect = [4, 2, 1] # monitors, events, runs + + from testgen.common.monitor_service import disable_monitoring + + counts = disable_monitoring(suite) + + assert counts == {"monitors": 4, "events": 2, "runs": 1} + mock_ts.cascade_delete.assert_called_once_with([suite.id]) diff --git a/tests/unit/common/test_table_group_service.py b/tests/unit/common/test_table_group_service.py new file mode 100644 index 00000000..b8d051ba --- /dev/null +++ b/tests/unit/common/test_table_group_service.py @@ -0,0 +1,330 @@ +"""Tests for the table_group_service common-layer module. + +Covers ``validate_table_group_fields`` and ``preview_table_group`` (happy path, +verify-access flag, partial-inaccessible footer, no-connection branch, +connection failure, empty result). +""" + +from unittest.mock import patch + +import pytest + +from testgen.common.database.table_group_service import ( + preview_table_group, + validate_table_group_fields, +) +from testgen.common.models.connection import Connection +from testgen.common.models.table_group import TableGroup + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.database.table_group_service" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _tg(**overrides) -> TableGroup: + """Build a TableGroup without touching the DB. Defaults to a valid baseline.""" + defaults = { + "project_code": "demo", + "connection_id": 7, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "profile_use_sampling": False, + "profile_sample_percent": "30", + "profile_sample_min_count": 100000, + "profiling_delay_days": "0", + } + defaults.update(overrides) + return TableGroup(**defaults) + + +def _conn(**overrides) -> Connection: + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_id": 7, + "project_code": "demo", + "connection_name": "Local PG", + "project_host": "localhost", + "project_port": "5432", + "project_db": "demo", + "project_user": "demo", + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# validate_table_group_fields — happy paths and per-field errors +# --------------------------------------------------------------------------- + + +def test_validate_passes_baseline(): + assert validate_table_group_fields(_tg()) == [] + + +def test_validate_missing_name(): + errors = validate_table_group_fields(_tg(table_groups_name=None)) + assert "`table_group_name` is required." in errors + + +def test_validate_blank_name(): + errors = validate_table_group_fields(_tg(table_groups_name=" ")) + assert "`table_group_name` is required." in errors + + +def test_validate_name_too_short(): + errors = validate_table_group_fields(_tg(table_groups_name="ab")) + assert "`table_group_name` must be between 3 and 40 characters." in errors + + +def test_validate_name_too_long(): + errors = validate_table_group_fields(_tg(table_groups_name="a" * 41)) + assert "`table_group_name` must be between 3 and 40 characters." in errors + + +def test_validate_name_at_lower_bound_passes(): + assert validate_table_group_fields(_tg(table_groups_name="abc")) == [] + + +def test_validate_name_at_upper_bound_passes(): + assert validate_table_group_fields(_tg(table_groups_name="a" * 40)) == [] + + +def test_validate_missing_schema(): + errors = validate_table_group_fields(_tg(table_group_schema=None)) + assert "`schema` is required." in errors + + +def test_validate_blank_schema(): + errors = validate_table_group_fields(_tg(table_group_schema=" ")) + assert "`schema` is required." in errors + + +@pytest.mark.parametrize("bad_delay", ["-1", "-100"]) +def test_validate_delay_days_negative(bad_delay): + errors = validate_table_group_fields(_tg(profiling_delay_days=bad_delay)) + assert "`profiling_delay_days` must be a non-negative integer." in errors + + +def test_validate_delay_days_non_integer(): + errors = validate_table_group_fields(_tg(profiling_delay_days="abc")) + assert "`profiling_delay_days` must be a non-negative integer." in errors + + +def test_validate_delay_days_int_input_accepted(): + """Validator accepts int input as well — apply-args casts but validators should be tolerant.""" + assert validate_table_group_fields(_tg(profiling_delay_days=0)) == [] + + +@pytest.mark.parametrize("bad_pct", ["0", "101", "200", "-1"]) +def test_validate_sample_percent_out_of_range(bad_pct): + errors = validate_table_group_fields(_tg(profile_sample_percent=bad_pct)) + assert "`profile_sample_percent` must be between 1 and 100." in errors + + +def test_validate_sample_percent_non_integer(): + errors = validate_table_group_fields(_tg(profile_sample_percent="abc")) + assert "`profile_sample_percent` must be between 1 and 100." in errors + + +def test_validate_sample_percent_int_input_accepted(): + assert validate_table_group_fields(_tg(profile_sample_percent=50)) == [] + + +def test_validate_sample_min_count_negative(): + errors = validate_table_group_fields(_tg(profile_sample_min_count=-5)) + assert "`profile_sample_min_count` must be a non-negative integer." in errors + + +def test_validate_aggregates_all_errors(): + """Multiple failures are collected and surfaced together — never one-at-a-time.""" + errors = validate_table_group_fields(_tg( + table_groups_name="", + table_group_schema="", + profiling_delay_days="-1", + profile_sample_percent="200", + profile_sample_min_count=-5, + )) + joined = "\n".join(errors) + assert "`table_group_name` is required." in joined + assert "`schema` is required." in joined + assert "`profiling_delay_days` must be a non-negative integer." in joined + assert "`profile_sample_percent` must be between 1 and 100." in joined + assert "`profile_sample_min_count` must be a non-negative integer." in joined + + +# --------------------------------------------------------------------------- +# preview_table_group — lifted UI logic +# --------------------------------------------------------------------------- + + +def _column_row(table_name: str, column_name: str, approx_record_ct: int | None = None) -> dict: + return { + "schema_name": "public", + "table_name": table_name, + "column_name": column_name, + "ordinal_position": 1, + "general_type": "A", + "column_type": "varchar", + "db_data_type": "VARCHAR", + "is_decimal": False, + "approx_record_ct": approx_record_ct, + "record_ct": None, + } + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_returns_picklable_shape(mock_fetch, mock_sql_cls): + """Contract: return ``(preview, data_chars, sql_generator)`` — three picklable values. + + A local ``save_data_chars`` closure as the second element would fail + pickling (``Can't get local object 'preview_table_group..save_data_chars'``) + in any caller that caches the result. Closures are built by callers via + ``make_save_data_chars``. + """ + from inspect import isfunction + + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + mock_fetch.return_value = [_column_row("customer", "id", approx_record_ct=100)] + + result = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert isinstance(result, tuple) and len(result) == 3, ( + "preview_table_group must return (preview, data_chars, sql_generator)" + ) + preview, data_chars, sql_gen = result + assert isinstance(preview, dict) + assert not isfunction(data_chars), "data_chars must not be a local closure" + assert not isfunction(sql_gen), "sql_generator must not be a local closure" + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_happy_path_aggregates_stats(mock_fetch, mock_sql_cls): + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + mock_fetch.return_value = [ + _column_row("customer", "id", approx_record_ct=100), + _column_row("customer", "email", approx_record_ct=100), + _column_row("rental", "id", approx_record_ct=50), + ] + + preview, data_chars, sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is True + assert preview["message"] is None + assert preview["stats"]["table_ct"] == 2 + assert preview["stats"]["column_ct"] == 3 + assert preview["stats"]["approx_record_ct"] == 150 # 100 + 50, counted once per table + assert "customer" in preview["tables"] + assert "rental" in preview["tables"] + assert preview["tables"]["customer"]["column_ct"] == 2 + assert data_chars is not None and len(data_chars) == 3 + assert sql_gen is sql_generator + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_verify_access_marks_all_accessible(mock_fetch, mock_sql_cls): + """When verify_access=True and every table is reachable, no footer message is set.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + sql_generator.verify_access.return_value = ("SELECT 1", None) + mock_fetch.side_effect = [ + [_column_row("customer", "id", approx_record_ct=100)], # initial DDF + [{"col": 1}], # verify customer + ] + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=True) + + assert preview["success"] is True + assert preview["tables"]["customer"]["can_access"] is True + assert preview["message"] is None + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db") +def test_preview_verify_access_partial_failure_sets_footer(mock_fetch, mock_sql_cls): + """One inaccessible table → can_access=False AND the UI-verbatim footer message.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + sql_generator.verify_access.side_effect = lambda name: (f"SELECT 1 FROM {name}", None) + + def fetch_side_effect(_conn_arg, query, *_args, **_kwargs): + if "SELECT ..." in query: + return [ + _column_row("customer", "id", approx_record_ct=100), + _column_row("rental", "id", approx_record_ct=50), + ] + if "SELECT 1 FROM customer" in query: + return [{"col": 1}] + if "SELECT 1 FROM rental" in query: + raise RuntimeError("permission denied") + raise AssertionError(f"unexpected query: {query}") + + mock_fetch.side_effect = fetch_side_effect + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=True) + + assert preview["success"] is True + assert preview["tables"]["customer"]["can_access"] is True + assert preview["tables"]["rental"]["can_access"] is False + assert preview["message"] == ( + "Some tables were not accessible. Please the check the database permissions." + ) + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db", side_effect=RuntimeError("connection refused")) +def test_preview_connection_failure_marks_unsuccessful(mock_fetch, mock_sql_cls): + """When the DDF query blows up, preview comes back ``success=False`` carrying the driver text.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + + preview, data_chars, sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is False + assert "connection refused" in (preview["message"] or "") + assert data_chars is None + assert sql_gen is None + + +@patch(f"{MODULE}.RefreshDataCharsSQL") +@patch(f"{MODULE}.fetch_from_target_db", return_value=[]) +def test_preview_empty_result_uses_ui_verbatim_message(mock_fetch, mock_sql_cls): + """Zero matching tables → UI-verbatim "No tables found matching the criteria." message.""" + sql_generator = mock_sql_cls.return_value + sql_generator.flavor_service.metadata_via_api = False + sql_generator.get_schema_ddf.return_value = ("SELECT ...", {}) + + preview, _data_chars, _sql_gen = preview_table_group(_tg(), connection=_conn(), verify_access=False) + + assert preview["success"] is False + assert preview["message"] == ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." + ) + + +def test_preview_no_connection_returns_failure_no_io(): + """No connection passed AND no ``table_group.connection_id`` → fail fast, no DB calls attempted.""" + tg = _tg(connection_id=None) + preview, data_chars, sql_gen = preview_table_group(tg, connection=None, verify_access=False) + + assert preview["success"] is False + assert preview["message"] is not None + assert "connection" in preview["message"].lower() + assert data_chars is None + assert sql_gen is None diff --git a/tests/unit/common/test_test_result_disposition_service.py b/tests/unit/common/test_test_result_disposition_service.py new file mode 100644 index 00000000..83d3de62 --- /dev/null +++ b/tests/unit/common/test_test_result_disposition_service.py @@ -0,0 +1,83 @@ +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.common.enums import Disposition +from testgen.common.test_result_disposition_service import ( + DispositionUpdate, + coerce_ui_disposition, + coupled_test_definition_state, + set_test_results_disposition, +) + +pytestmark = pytest.mark.unit + +MODULE = "testgen.common.test_result_disposition_service" + + +class Test_coupled_test_definition_state: + def test_muted_deactivates_and_locks(self): + # (test_active, lock_refresh) + assert coupled_test_definition_state(Disposition.INACTIVE) == (False, True) + + @pytest.mark.parametrize("disposition", [Disposition.CONFIRMED, Disposition.DISMISSED, None]) + def test_other_values_reactivate_and_unlock(self, disposition): + assert coupled_test_definition_state(disposition) == (True, False) + + +class Test_coerce_ui_disposition: + @pytest.mark.parametrize( + "value,expected", + [ + ("Confirmed", Disposition.CONFIRMED), + ("Dismissed", Disposition.DISMISSED), + ("Inactive", Disposition.INACTIVE), + ("No Decision", None), + (None, None), + ("", None), + ], + ) + def test_maps_ui_string_to_stored_value(self, value, expected): + assert coerce_ui_disposition(value) is expected + + +class Test_set_test_results_disposition: + def test_empty_ids_is_noop(self): + with patch(f"{MODULE}.get_current_session") as get_session: + result = set_test_results_disposition([], Disposition.CONFIRMED) + assert result == DispositionUpdate(matched=0, passed_skipped=0) + get_session.assert_not_called() + + def test_returns_matched_and_passed_counts(self): + session = MagicMock() + # First call: COUNT of passed rows -> 2. Then the TR update returns rowcount 3. + session.scalar.return_value = 2 + tr_update_result = MagicMock(rowcount=3) + session.execute.return_value = tr_update_result + with patch(f"{MODULE}.get_current_session", return_value=session): + result = set_test_results_disposition([uuid4(), uuid4()], Disposition.DISMISSED) + assert result == DispositionUpdate(matched=3, passed_skipped=2) + # One TR update + one TD update. + assert session.execute.call_count == 2 + + def test_muted_couples_td_to_inactive(self): + session = MagicMock() + session.scalar.return_value = 0 + session.execute.return_value = MagicMock(rowcount=1) + with patch(f"{MODULE}.get_current_session", return_value=session): + set_test_results_disposition([uuid4()], Disposition.INACTIVE) + td_stmt = session.execute.call_args_list[1].args[0] + params = td_stmt.compile().params + # YNString stores 'N'/'Y'; bound values are the Python bools before type processing. + assert params["test_active"] is False + assert params["lock_refresh"] is True + + def test_clear_passes_null_disposition(self): + session = MagicMock() + session.scalar.return_value = 0 + session.execute.return_value = MagicMock(rowcount=1) + with patch(f"{MODULE}.get_current_session", return_value=session): + set_test_results_disposition([uuid4()], None) + tr_stmt = session.execute.call_args_list[0].args[0] + assert tr_stmt.compile().params["disposition"] is None diff --git a/tests/unit/common/test_time_series_service.py b/tests/unit/common/test_time_series_service.py index c8a5d66a..b34bea76 100644 --- a/tests/unit/common/test_time_series_service.py +++ b/tests/unit/common/test_time_series_service.py @@ -617,3 +617,37 @@ def test_without_exclusions_timezone_has_no_effect(self): forecast_with_tz = get_sarimax_forecast(history, num_forecast=3, exclude_weekends=False, tz="America/New_York") pd.testing.assert_frame_equal(forecast_no_tz, forecast_with_tz) + + +class Test_GetSarimaxForecast_EventSpace: + """Event-space (event_space=True) vs default calendar-regularized fitting for irregular series.""" + + @staticmethod + def _irregular_refresh_history() -> pd.DataFrame: + """A refresh-driven series: irregular gaps (1-3 days), jumps that scale with the gap. + This is the shape that calendar interpolation smooths into uniform daily increments.""" + times = pd.to_datetime([ + "2026-01-01", "2026-01-02", "2026-01-03", "2026-01-05", "2026-01-06", + "2026-01-09", "2026-01-10", "2026-01-12", "2026-01-15", "2026-01-16", + "2026-01-18", "2026-01-21", + ]) + values = [100.0, 110, 121, 150, 162, 200, 212, 240, 290, 303, 330, 375] + return pd.DataFrame({"result_signal": values}, index=times) + + def test_event_space_centers_on_full_refresh_jump(self): + history = self._irregular_refresh_history() + + regularized = get_sarimax_forecast(history, num_forecast=1, event_space=False) + event_space = get_sarimax_forecast(history, num_forecast=1, event_space=True) + + last = history["result_signal"].iloc[-1] + # Calendar interpolation forecasts a single per-day increment; event-space forecasts a + # full refresh-to-refresh jump, so its next-point prediction sits meaningfully higher. + assert (event_space["mean"].iloc[0] - last) > (regularized["mean"].iloc[0] - last) + + def test_event_space_returns_requested_forecast_length(self): + history = self._irregular_refresh_history() + forecast = get_sarimax_forecast(history, num_forecast=5, event_space=True) + assert len(forecast) == 5 + assert forecast["mean"].notna().all() + assert forecast["se"].notna().all() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 57794a74..fc7a48d3 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1,15 +1,2 @@ -from unittest.mock import patch - -import pytest - - -@pytest.fixture(autouse=True) -def patched_settings(): - with patch("testgen.settings.UI_BASE_URL", "http://tg-base-url"): - yield - - -@pytest.fixture -def db_session_mock(): - with patch("testgen.common.models.Session") as factory_mock: - yield factory_mock().__enter__() +# Importing the fixtures registers them for this test subtree (autouse flags preserved). +from testgen.testing.fixtures import db_session_mock, patched_settings # noqa: F401 diff --git a/tests/unit/mcp/conftest.py b/tests/unit/mcp/conftest.py index cc932043..a2906ff5 100644 --- a/tests/unit/mcp/conftest.py +++ b/tests/unit/mcp/conftest.py @@ -1,50 +1,2 @@ -from unittest.mock import MagicMock, patch -from uuid import uuid4 - -import pytest - -from testgen.mcp.permissions import set_mcp_token, set_mcp_username - -# Fictional role matrix for tests. role_a has full access, role_c is restricted. -TEST_PERM_MATRIX = { - "view": ["role_a", "role_b"], - "catalog": ["role_a", "role_b", "role_c"], - "edit": ["role_a"], -} - - -def _test_roles_with_permission(permission): - return TEST_PERM_MATRIX.get(permission, []) - - -@pytest.fixture(autouse=True) -def mcp_user(): - """Set up an authenticated MCP user for all tool tests. - - Default: user has 'role_a' on 'demo' project (full access). - The @mcp_permission decorator passes for any permission. - - Tests needing scoped access patch _compute_project_permissions directly. - """ - set_mcp_username("test_user") - set_mcp_token("test_bearer_token") - user = MagicMock() - user.id = uuid4() - - membership = MagicMock() - membership.project_code = "demo" - membership.role = "role_a" - - with ( - patch("testgen.common.auth.authorize_token", return_value=user), - patch("testgen.common.models.get_current_session", return_value=MagicMock()), - patch("testgen.mcp.permissions.ProjectMembership") as mock_membership, - patch("testgen.mcp.permissions.PluginHook") as mock_hook, - ): - mock_membership.get_memberships_for_user.return_value = [membership] - mock_hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( - _test_roles_with_permission - ) - yield user - set_mcp_username(None) - set_mcp_token(None) +# Importing the fixture registers the autouse authenticated-MCP-user setup for these tests. +from testgen.testing.fixtures import mcp_user # noqa: F401 diff --git a/tests/unit/mcp/test_auth.py b/tests/unit/mcp/test_auth.py index 29cafbdb..e8125948 100644 --- a/tests/unit/mcp/test_auth.py +++ b/tests/unit/mcp/test_auth.py @@ -7,6 +7,7 @@ import jwt import pytest +from testgen.common.auth import AuthError from testgen.mcp.auth import authenticate_user, validate_token from testgen.mcp.server import JWTTokenVerifier @@ -83,7 +84,7 @@ def test_validate_token_returns_user(mock_user_cls, mock_settings): def test_validate_token_raises_for_expired_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): validate_token(_make_token(exp_seconds=-3600)) @@ -91,7 +92,7 @@ def test_validate_token_raises_for_expired_token(mock_settings): def test_validate_token_raises_for_invalid_token(mock_settings): mock_settings.JWT_HASHING_KEY_B64 = JWT_KEY - with pytest.raises(ValueError, match="Invalid token"): + with pytest.raises(AuthError, match="Invalid token"): validate_token("not-a-valid-token") diff --git a/tests/unit/mcp/test_connection_schema.py b/tests/unit/mcp/test_connection_schema.py new file mode 100644 index 00000000..75919cfc --- /dev/null +++ b/tests/unit/mcp/test_connection_schema.py @@ -0,0 +1,649 @@ +"""Tests for the connection-parameter contract in ``mcp/tools/common.py``: +the per-flavor schema (``schema_for`` / ``resolve_mode``), label-keyed param +mapping (``apply_connection_params``), schema-driven validation +(``validate_connection_fields``), and mode inference (``infer_mode``). + +The schema mirrors the per-flavor JS forms in +``ui/static/js/components/connection_form.js`` — these tests encode that contract. +""" + +import pytest + +from testgen.common.models.connection import Connection +from testgen.mcp.exceptions import MCPUserError +from testgen.mcp.tools.common import ( + ConnectionMode, + Req, + apply_connection_params, + infer_mode, + resolve_mode, + schema_for, + validate_connection_fields, +) + +pytestmark = pytest.mark.unit + + +def _conn(**overrides) -> Connection: + """Build a Connection without touching the DB. Defaults to a complete PG config.""" + defaults = { + "sql_flavor": "postgresql", + "sql_flavor_code": "postgresql", + "connection_name": "My DB", + "project_host": "localhost", + "project_port": "5432", + "project_db": "testgen_local", + "project_user": "testgen", + "project_pw_encrypted": "pw", + "connect_by_url": False, + "connect_by_key": False, + "connect_with_identity": False, + "max_threads": 4, + "max_query_chars": 20000, + } + defaults.update(overrides) + return Connection(**defaults) + + +# --------------------------------------------------------------------------- +# schema_for +# --------------------------------------------------------------------------- + + +def test_schema_for_unknown_code_raises(): + with pytest.raises(KeyError): + schema_for("not_a_flavor") + + +def test_schema_for_postgresql_single_mode_host_fields(): + schema = schema_for("postgresql") + assert schema.label == "PostgreSQL" + assert len(schema.modes) == 1 + mode = schema.modes[0] + assert mode.mode is None + labels = {f.label: f for f in mode.fields} + assert set(labels) == {"Host", "Port", "Database", "Username", "Password"} + assert labels["Host"].requirement is Req.REQUIRED_UNLESS_URL + assert labels["Username"].requirement is Req.REQUIRED + assert labels["Password"].requirement is Req.OPTIONAL + assert labels["Password"].secret is True + assert mode.supports_url is True + assert schema.url_field is not None and schema.url_field.label == "URL" + assert labels["Host"].column == "project_host" + assert labels["Username"].column == "project_user" + assert labels["Password"].column == "project_pw_encrypted" + + +def test_schema_for_oracle_uses_service_name_label(): + labels = {f.label: f for f in schema_for("oracle").modes[0].fields} + assert "Service Name" in labels + assert "Database" not in labels + assert labels["Service Name"].column == "project_db" + + +def test_schema_for_sap_hana_uses_database_label(): + labels = {f.label for f in schema_for("sap_hana").modes[0].fields} + assert "Database" in labels + + +def test_schema_for_snowflake_two_modes_with_warehouse(): + schema = schema_for("snowflake") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.KEY_PAIR, ConnectionMode.PASSWORD} + key_labels = {f.label: f for f in modes[ConnectionMode.KEY_PAIR].fields} + pw_labels = {f.label: f for f in modes[ConnectionMode.PASSWORD].fields} + assert "Warehouse" in key_labels and key_labels["Warehouse"].requirement is Req.OPTIONAL + assert key_labels["Username"].requirement is Req.REQUIRED + assert key_labels["Private Key"].requirement is Req.REQUIRED + assert key_labels["Private Key"].secret is True + assert key_labels["Private Key Passphrase"].requirement is Req.OPTIONAL + assert "Password" not in key_labels + assert pw_labels["Password"].requirement is Req.REQUIRED + assert "Private Key" not in pw_labels + assert all(m.supports_url for m in schema.modes) + + +def test_schema_for_databricks_modes_and_url_support(): + schema = schema_for("databricks") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.ACCESS_TOKEN, ConnectionMode.SERVICE_PRINCIPAL} + pat = modes[ConnectionMode.ACCESS_TOKEN] + oauth = modes[ConnectionMode.SERVICE_PRINCIPAL] + pat_labels = {f.label: f for f in pat.fields} + oauth_labels = {f.label: f for f in oauth.fields} + assert pat_labels["Catalog"].column == "project_db" + assert pat_labels["Catalog"].requirement is Req.REQUIRED_UNLESS_URL + assert pat_labels["HTTP Path"].requirement is Req.REQUIRED_UNLESS_URL + assert pat_labels["Access Token"].column == "project_pw_encrypted" + assert pat_labels["Access Token"].requirement is Req.REQUIRED + assert "Username" not in pat_labels # auto-set to 'token' + assert pat.supports_url is True + assert oauth_labels["Client ID"].column == "project_user" + assert oauth_labels["Client Secret"].column == "project_pw_encrypted" + assert oauth_labels["Host"].requirement is Req.REQUIRED + assert oauth.supports_url is False + + +def test_schema_for_bigquery_single_field_no_url(): + schema = schema_for("bigquery") + assert len(schema.modes) == 1 + labels = {f.label: f for f in schema.modes[0].fields} + assert set(labels) == {"Service Account Key"} + assert labels["Service Account Key"].column == "service_account_key" + assert labels["Service Account Key"].secret is True + assert schema.modes[0].supports_url is False + assert schema.url_field is None + + +def test_schema_for_salesforce_two_modes_field_mapping(): + schema = schema_for("salesforce_data360") + modes = {m.mode: m for m in schema.modes} + assert set(modes) == {ConnectionMode.JWT_BEARER, ConnectionMode.CLIENT_CREDENTIALS} + jwt = {f.label: f for f in modes[ConnectionMode.JWT_BEARER].fields} + cc = {f.label: f for f in modes[ConnectionMode.CLIENT_CREDENTIALS].fields} + assert jwt["Login URL"].column == "project_host" + assert jwt["Consumer Key"].column == "project_user" + assert jwt["Username"].column == "project_db" + assert jwt["Private Key"].column == "private_key" + assert "Consumer Secret" not in jwt + assert cc["Consumer Secret"].column == "project_pw_encrypted" + assert "Username" not in cc and "Private Key" not in cc + assert all(not m.supports_url for m in schema.modes) + + +# --------------------------------------------------------------------------- +# resolve_mode +# --------------------------------------------------------------------------- + + +def test_resolve_mode_single_mode_no_label(): + assert resolve_mode("postgresql", None).mode is None + + +def test_resolve_mode_single_mode_rejects_label(): + with pytest.raises(MCPUserError): + resolve_mode("postgresql", "Password") + + +def test_resolve_mode_multi_mode_requires_label(): + with pytest.raises(MCPUserError) as exc: + resolve_mode("snowflake", None) + assert "Key-Pair" in str(exc.value) and "Password" in str(exc.value) + + +def test_resolve_mode_multi_mode_invalid_label(): + with pytest.raises(MCPUserError) as exc: + resolve_mode("snowflake", "Bogus") + assert "Key-Pair" in str(exc.value) + + +def test_resolve_mode_multi_mode_valid_label(): + assert resolve_mode("snowflake", "Key-Pair").mode is ConnectionMode.KEY_PAIR + + +# --------------------------------------------------------------------------- +# apply_connection_params +# --------------------------------------------------------------------------- + + +def test_apply_params_postgresql_maps_labels_to_columns(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + apply_connection_params( + conn, + "postgresql", + None, + {"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + assert conn.project_host == "h" + assert conn.project_port == "5432" # cast to str + assert conn.project_db == "d" + assert conn.project_user == "u" + assert conn.project_pw_encrypted == "p" + assert conn.connect_by_url is False + + +def test_apply_params_url_sets_connect_by_url(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + apply_connection_params(conn, "postgresql", None, {"URL": "host:5432/db", "Username": "u"}) + assert conn.connect_by_url is True + assert conn.url == "host:5432/db" + assert conn.project_user == "u" + + +def test_apply_params_url_conflicts_with_host_group(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "postgresql", None, {"URL": "x", "Host": "h"}) + + +def test_apply_params_url_on_unsupported_flavor_rejected(): + conn = Connection(sql_flavor="bigquery", sql_flavor_code="bigquery") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "bigquery", None, {"URL": "x"}) + + +def test_apply_params_unknown_key_rejected(): + conn = Connection(sql_flavor="postgresql", sql_flavor_code="postgresql") + with pytest.raises(MCPUserError) as exc: + apply_connection_params(conn, "postgresql", None, {"Hostname": "h"}) + assert "Hostname" in str(exc.value) + + +def test_apply_params_snowflake_key_pair_sets_flag(): + conn = Connection(sql_flavor="snowflake", sql_flavor_code="snowflake") + apply_connection_params( + conn, + "snowflake", + "Key-Pair", + {"Host": "h", "Port": 443, "Database": "d", "Username": "u", "Private Key": "KEY"}, + ) + assert conn.connect_by_key is True + assert conn.private_key == "KEY" + + +def test_apply_params_databricks_pat_auto_token_username(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + apply_connection_params( + conn, + "databricks", + "Access Token", + {"Host": "h", "Port": 443, "Catalog": "main", "HTTP Path": "/sql/1.0/abc", "Access Token": "tok"}, + ) + assert conn.project_user == "token" + assert conn.connect_by_key is False + assert conn.project_db == "main" + assert conn.http_path == "/sql/1.0/abc" + assert conn.project_pw_encrypted == "tok" + + +def test_apply_params_databricks_oauth_client_id_secret(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + apply_connection_params( + conn, + "databricks", + "Service Principal (OAuth)", + {"Host": "h", "Port": 443, "Catalog": "main", "HTTP Path": "/p", "Client ID": "cid", "Client Secret": "csec"}, + ) + assert conn.connect_by_key is True + assert conn.project_user == "cid" + assert conn.project_pw_encrypted == "csec" + + +def test_apply_params_databricks_oauth_rejects_url(): + conn = Connection(sql_flavor="databricks", sql_flavor_code="databricks") + with pytest.raises(MCPUserError): + apply_connection_params(conn, "databricks", "Service Principal (OAuth)", {"URL": "x"}) + + +def test_apply_params_azure_managed_identity_sets_flag(): + conn = Connection(sql_flavor="mssql", sql_flavor_code="azure_mssql") + apply_connection_params(conn, "azure_mssql", "Managed Identity", {"Host": "h", "Port": 1433, "Database": "d"}) + assert conn.connect_with_identity is True + + +def test_apply_params_salesforce_jwt_field_mapping(): + conn = Connection(sql_flavor="salesforce_data360", sql_flavor_code="salesforce_data360") + apply_connection_params( + conn, + "salesforce_data360", + "JWT Bearer Flow", + {"Login URL": "https://my.salesforce.com", "Consumer Key": "ck", "Username": "u@x.com", "Private Key": "PK"}, + ) + assert conn.project_host == "https://my.salesforce.com" + assert conn.project_user == "ck" + assert conn.project_db == "u@x.com" + assert conn.private_key == "PK" + assert conn.connect_by_key is True + + +# --------------------------------------------------------------------------- +# validate_connection_fields — happy paths +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flavor_code", ["postgresql", "redshift", "redshift_spectrum", "mssql", "oracle", "sap_hana"]) +def test_validate_passes_basic_host_auth_flavors(flavor_code): + conn = _conn(sql_flavor_code=flavor_code, sql_flavor=flavor_code) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_azure_mssql_with_user_password(): + assert validate_connection_fields(_conn(sql_flavor_code="azure_mssql", sql_flavor="mssql")) == [] + + +def test_validate_passes_azure_mssql_with_identity(): + conn = _conn( + sql_flavor_code="azure_mssql", + sql_flavor="mssql", + connect_with_identity=True, + project_user=None, + project_pw_encrypted=None, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_snowflake_password_auth(): + conn = _conn(sql_flavor_code="snowflake", sql_flavor="snowflake", project_port="443", connect_by_key=False) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_snowflake_key_pair_auth(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + project_pw_encrypted=None, + connect_by_key=True, + private_key="-----BEGIN PRIVATE KEY-----\n...", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_databricks_pat(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_db="main", + project_user="token", + http_path="/sql/1.0/warehouses/abc", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_bigquery(): + conn = _conn( + sql_flavor_code="bigquery", + sql_flavor="bigquery", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + project_pw_encrypted=None, + service_account_key={"type": "service_account", "project_id": "demo"}, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_salesforce_jwt(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db="user@x.com", + project_user="consumer_key", + project_pw_encrypted=None, + connect_by_key=True, + private_key="-----BEGIN PRIVATE KEY-----\n...", + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_salesforce_client_credentials(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db=None, + project_user="consumer_key", + project_pw_encrypted="consumer_secret", + connect_by_key=False, + private_key=None, + ) + assert validate_connection_fields(conn) == [] + + +def test_validate_passes_connect_by_url_keeps_username(): + """URL mode: host/port/db not required, but Username STILL required (matches UI).""" + conn = _conn(connect_by_url=True, url="localhost:5432/mydb", project_host=None, project_port=None, project_db=None) + assert validate_connection_fields(conn) == [] + + +# --------------------------------------------------------------------------- +# validate_connection_fields — divergence fix + per-flavor errors +# --------------------------------------------------------------------------- + + +def test_validate_url_mode_missing_username_fails(): + """The PR-flagged divergence: URL mode must still require Username.""" + conn = _conn( + connect_by_url=True, + url="localhost:5432/mydb", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + ) + assert "`Username` is required for PostgreSQL." in validate_connection_fields(conn) + + +@pytest.mark.parametrize( + "label,model_attr", + [("Host", "project_host"), ("Port", "project_port"), ("Database", "project_db"), ("Username", "project_user")], +) +def test_validate_postgresql_missing_field(label, model_attr): + conn = _conn(**{model_attr: None}) + assert f"`{label}` is required for PostgreSQL." in validate_connection_fields(conn) + + +def test_validate_oracle_missing_service_name(): + conn = _conn(sql_flavor_code="oracle", sql_flavor="oracle", project_db=None) + assert "`Service Name` is required for Oracle." in validate_connection_fields(conn) + + +def test_validate_postgresql_missing_url_when_connect_by_url(): + conn = _conn(connect_by_url=True, url=None, project_host=None, project_port=None) + assert "`URL` is required for PostgreSQL." in validate_connection_fields(conn) + + +def test_validate_snowflake_missing_password_when_not_connect_by_key(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + connect_by_key=False, + project_pw_encrypted=None, + ) + assert "`Password` is required for Snowflake." in validate_connection_fields(conn) + + +def test_validate_snowflake_missing_private_key_when_connect_by_key(): + conn = _conn( + sql_flavor_code="snowflake", + sql_flavor="snowflake", + project_port="443", + project_pw_encrypted=None, + connect_by_key=True, + private_key=None, + ) + assert "`Private Key` is required for Snowflake." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_catalog(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db=None, + http_path="/sql/1.0/warehouses/abc", + ) + assert "`Catalog` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_http_path(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db="main", + http_path=None, + ) + assert "`HTTP Path` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_pat_missing_access_token(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_user="token", + project_db="main", + http_path="/p", + project_pw_encrypted=None, + ) + assert "`Access Token` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_databricks_oauth_missing_client_id(): + conn = _conn( + sql_flavor_code="databricks", + sql_flavor="databricks", + project_port="443", + project_db="main", + http_path="/p", + connect_by_key=True, + project_user=None, + project_pw_encrypted="secret", + ) + assert "`Client ID` is required for Databricks." in validate_connection_fields(conn) + + +def test_validate_bigquery_missing_service_account_key(): + conn = _conn( + sql_flavor_code="bigquery", + sql_flavor="bigquery", + project_host=None, + project_port=None, + project_db=None, + project_user=None, + project_pw_encrypted=None, + service_account_key=None, + ) + assert "`Service Account Key` is required for Google BigQuery." in validate_connection_fields(conn) + + +def test_validate_salesforce_jwt_missing_private_key(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db="user@x.com", + project_user="ck", + project_pw_encrypted=None, + connect_by_key=True, + private_key=None, + ) + assert "`Private Key` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_salesforce_jwt_missing_login_url(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host=None, + project_port=None, + project_db="user@x.com", + project_user="ck", + project_pw_encrypted=None, + connect_by_key=True, + private_key="key", + ) + assert "`Login URL` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_salesforce_client_credentials_missing_secret(): + conn = _conn( + sql_flavor_code="salesforce_data360", + sql_flavor="salesforce_data360", + project_host="https://my.salesforce.com", + project_port=None, + project_db=None, + project_user="ck", + project_pw_encrypted=None, + connect_by_key=False, + private_key=None, + ) + assert "`Consumer Secret` is required for Salesforce Data 360." in validate_connection_fields(conn) + + +def test_validate_azure_mssql_password_auth_missing_user(): + conn = _conn(sql_flavor_code="azure_mssql", sql_flavor="mssql", connect_with_identity=False, project_user=None) + assert "`Username` is required for Azure SQL Database." in validate_connection_fields(conn) + + +# --------------------------------------------------------------------------- +# validate_connection_fields — name / threads / query-chars +# --------------------------------------------------------------------------- + + +def test_validate_missing_connection_name(): + assert "`connection_name` is required." in validate_connection_fields(_conn(connection_name=None)) + + +def test_validate_connection_name_too_short(): + assert "`connection_name` must be between 3 and 40 characters." in validate_connection_fields(_conn(connection_name="ab")) + + +def test_validate_connection_name_too_long(): + assert "`connection_name` must be between 3 and 40 characters." in validate_connection_fields(_conn(connection_name="a" * 41)) + + +@pytest.mark.parametrize("bad", [0, -1, 9, 100]) +def test_validate_max_threads_out_of_range(bad): + assert "`max_threads` must be between 1 and 8." in validate_connection_fields(_conn(max_threads=bad)) + + +@pytest.mark.parametrize("bad", [499, 0, 50001, 100000]) +def test_validate_max_query_chars_out_of_range(bad): + assert "`max_query_chars` must be between 500 and 50000." in validate_connection_fields(_conn(max_query_chars=bad)) + + +def test_validate_aggregates_all_errors(): + conn = _conn(connection_name="", project_host=None, project_port=None, max_threads=99) + joined = "\n".join(validate_connection_fields(conn)) + assert "`connection_name` is required." in joined + assert "`Host` is required for PostgreSQL." in joined + assert "`Port` is required for PostgreSQL." in joined + assert "`max_threads` must be between 1 and 8." in joined + + +# --------------------------------------------------------------------------- +# infer_mode +# --------------------------------------------------------------------------- + + +def test_infer_mode_single_mode_flavor_is_none(): + assert infer_mode(_conn(sql_flavor_code="postgresql")) is None + + +def test_infer_mode_snowflake_key_pair(): + assert infer_mode(_conn(sql_flavor_code="snowflake", sql_flavor="snowflake", connect_by_key=True)) is ConnectionMode.KEY_PAIR + + +def test_infer_mode_snowflake_password(): + assert infer_mode(_conn(sql_flavor_code="snowflake", sql_flavor="snowflake", connect_by_key=False)) is ConnectionMode.PASSWORD + + +def test_infer_mode_azure_identity(): + conn = _conn(sql_flavor_code="azure_mssql", sql_flavor="mssql", connect_with_identity=True) + assert infer_mode(conn) is ConnectionMode.MANAGED_IDENTITY + + +def test_infer_mode_databricks_oauth(): + conn = _conn(sql_flavor_code="databricks", sql_flavor="databricks", connect_by_key=True) + assert infer_mode(conn) is ConnectionMode.SERVICE_PRINCIPAL + + +def test_infer_mode_databricks_pat(): + conn = _conn(sql_flavor_code="databricks", sql_flavor="databricks", connect_by_key=False) + assert infer_mode(conn) is ConnectionMode.ACCESS_TOKEN + + +def test_infer_mode_salesforce_jwt(): + conn = _conn(sql_flavor_code="salesforce_data360", sql_flavor="salesforce_data360", connect_by_key=True) + assert infer_mode(conn) is ConnectionMode.JWT_BEARER diff --git a/tests/unit/mcp/test_model_data_column.py b/tests/unit/mcp/test_model_data_column.py index aa8b0181..abdcae6f 100644 --- a/tests/unit/mcp/test_model_data_column.py +++ b/tests/unit/mcp/test_model_data_column.py @@ -4,6 +4,32 @@ from testgen.common.models.data_column import ColumnProfileDetail, DataColumnChars +_CATALOG_METADATA_COLUMNS = { + "description", + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", + "data_classification", +} + + +def test_catalog_metadata_columns_are_mapped(): + mapped = set(DataColumnChars.__table__.columns.keys()) + assert _CATALOG_METADATA_COLUMNS <= mapped + + +def test_catalog_metadata_attributes_settable(): + column = DataColumnChars() + column.description = "Primary email" + column.data_product = "CRM" + assert column.description == "Primary email" + assert column.data_product == "CRM" + def _detail_row(**overrides) -> dict: """Build a dict matching every ColumnProfileDetail field.""" diff --git a/tests/unit/mcp/test_model_data_table.py b/tests/unit/mcp/test_model_data_table.py index 0f9f10e2..2f04386f 100644 --- a/tests/unit/mcp/test_model_data_table.py +++ b/tests/unit/mcp/test_model_data_table.py @@ -1,9 +1,48 @@ from unittest.mock import patch from uuid import uuid4 +from sqlalchemy import select +from sqlalchemy.dialects import postgresql + from testgen.common.models.data_table import DataTable +def test_default_order_by_compiles_for_entity_select(): + # Regression: the inherited Entity default ordered by the textual label "id", which + # fails to compile here because the PK column is "table_id". A full-entity select + # (e.g. DataTable.select_where) must compile and order by a real column. + stmt = select(DataTable).order_by(*DataTable._default_order_by) + compiled = str(stmt.compile(dialect=postgresql.dialect())) + assert "ORDER BY" in compiled + assert "table_name" in compiled + +_CATALOG_METADATA_COLUMNS = { + "description", + "data_source", + "source_system", + "source_process", + "business_domain", + "stakeholder_group", + "transform_level", + "aggregation_level", + "data_product", + "data_classification", +} + + +def test_catalog_metadata_columns_are_mapped(): + mapped = set(DataTable.__table__.columns.keys()) + assert _CATALOG_METADATA_COLUMNS <= mapped + + +def test_catalog_metadata_attributes_settable(): + table = DataTable() + table.description = "Customer master table" + table.business_domain = "Sales" + assert table.description == "Customer master table" + assert table.business_domain == "Sales" + + @patch("testgen.common.models.data_table.get_current_session") def test_select_table_names_returns_list(session_mock): session_mock.return_value.scalars.return_value.all.return_value = ["customers", "orders", "products"] diff --git a/tests/unit/mcp/test_model_hygiene_issue.py b/tests/unit/mcp/test_model_hygiene_issue.py index bf90e2b6..74c7de55 100644 --- a/tests/unit/mcp/test_model_hygiene_issue.py +++ b/tests/unit/mcp/test_model_hygiene_issue.py @@ -95,14 +95,14 @@ def _stub_paginate(session_mock, *, total=0, rows=()): # --------------------------------------------------------------------------- -def test_list_for_run_filters_by_je_id_not_run_pk(session_mock): +def test_list_for_run_filters_by_run_id(session_mock): _stub_paginate(session_mock) HygieneIssue.list_for_run(uuid4()) sql = _all_compiled_sql(session_mock) - assert "profiling_runs.job_execution_id =" in sql - # The legacy run PK must NOT be the filter: + # The run id is the job execution id; filter targets the profiling run id directly. + assert "profiling_runs.id =" in sql assert "profile_anomaly_results.profile_run_id =" not in sql diff --git a/tests/unit/mcp/test_model_profiling_run.py b/tests/unit/mcp/test_model_profiling_run.py index 94d88b1b..888c50fa 100644 --- a/tests/unit/mcp/test_model_profiling_run.py +++ b/tests/unit/mcp/test_model_profiling_run.py @@ -53,13 +53,13 @@ def test_get_latest_complete_je_id_orders_desc_limit_1(session_mock): assert "LIMIT" in sql -def test_get_latest_complete_je_id_selects_je_id_not_run_pk(session_mock): +def test_get_latest_complete_je_id_selects_run_id(session_mock): """Pin the docstring contract — does NOT read ``table_groups.last_complete_profile_run_id``, - selects the JE id directly from ``profiling_runs``.""" + selects the run id (which is the JE id) directly from ``profiling_runs``.""" session_mock.scalar.return_value = None ProfilingRun.get_latest_complete_je_id_for_table_group(uuid4()) sql = _compiled_sql(session_mock.scalar.call_args[0][0]) - assert "SELECT profiling_runs.job_execution_id" in sql + assert "SELECT profiling_runs.id" in sql assert "table_groups.last_complete_profile_run_id" not in sql diff --git a/tests/unit/mcp/test_model_test_definition.py b/tests/unit/mcp/test_model_test_definition.py index bc77630a..1cb30f84 100644 --- a/tests/unit/mcp/test_model_test_definition.py +++ b/tests/unit/mcp/test_model_test_definition.py @@ -1,10 +1,15 @@ +from datetime import UTC, datetime +from types import SimpleNamespace from unittest.mock import patch from uuid import uuid4 import pytest from sqlalchemy.dialects import postgresql -from testgen.common.models.test_definition import TestDefinition +from testgen.common.models.test_definition import ( + TestDefinition, + ThresholdMode, +) @pytest.fixture @@ -54,3 +59,151 @@ def test_list_for_suite_excludes_monitor_suites(session_mock): sql_joined = "\n".join(_compiled_sql(q) for q in queries) assert "test_suites.is_monitor IS NOT true" in sql_joined assert "JOIN test_suites" in sql_joined + + +# --------------------------------------------------------------------------- +# _derive_threshold_mode — must match the UI form's detection at +# test_definition_form.js:274 (history_calculation == "PREDICT" → Prediction; +# any other non-empty value → Historical; empty → Static; schema → N/A). +# ``threshold_value`` must never leak into the returned bounds. +# --------------------------------------------------------------------------- + + +def _td(**overrides): + defaults = { + "test_type": "Volume_Trend", + "history_calculation": None, + "history_calculation_upper": None, + "lower_tolerance": None, + "upper_tolerance": None, + "threshold_value": None, + "prediction": None, + } + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def test_derive_threshold_mode_schema_returns_none(): + td = _td(test_type="Schema_Drift", history_calculation="PREDICT", lower_tolerance="5") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (ThresholdMode.NONE, None, None) + + +def test_derive_threshold_mode_prediction_keys_on_history_calculation_not_prediction_jsonb(): + """``history_calculation == "PREDICT"`` is the canonical signal — the + ``prediction`` JSONB is empty during training or after a failed run, so + keying off it would misreport prediction monitors as Static.""" + td = _td(test_type="Volume_Trend", history_calculation="PREDICT", prediction=None) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == ThresholdMode.PREDICTION + assert lower is None and upper is None + + +def test_derive_threshold_mode_historical_for_non_freshness(): + td = _td( + test_type="Volume_Trend", + history_calculation="Minimum", + history_calculation_upper="Maximum", + ) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (ThresholdMode.HISTORICAL, "Minimum", "Maximum") + + +def test_derive_threshold_mode_historical_blocked_for_freshness(): + """Freshness does not support Historical mode — a stray + ``history_calculation`` value falls through to Static.""" + td = _td( + test_type="Freshness_Trend", + history_calculation="Minimum", + upper_tolerance="720", + ) + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == ThresholdMode.STATIC + assert lower is None # Freshness has no lower bound + assert upper == "720" + + +def test_derive_threshold_mode_static_freshness_uses_upper_only(): + td = _td(test_type="Freshness_Trend", lower_tolerance="60", upper_tolerance="720") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == ThresholdMode.STATIC + assert lower is None # lower is silently dropped for Freshness even if populated + assert upper == "720" + + +def test_derive_threshold_mode_static_volume_uses_both_tolerances(): + td = _td(test_type="Volume_Trend", lower_tolerance="900", upper_tolerance="1100") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert (mode, lower, upper) == (ThresholdMode.STATIC, "900", "1100") + + +def test_derive_threshold_mode_threshold_value_never_returned(): + """``threshold_value`` is not a monitor configuration — even when no + tolerances are set it must not leak into the bounds.""" + td = _td(test_type="Volume_Trend", threshold_value="42") + mode, lower, upper = TestDefinition._derive_threshold_mode(td) + assert mode == ThresholdMode.STATIC + assert lower is None + assert upper is None + + +# --------------------------------------------------------------------------- +# forecast_points_from_prediction — extracts forecast rows from a monitor's +# stored prediction JSONB. Keys are epoch-millisecond integer strings; the +# helper must agree with the dashboard's reader at ``monitors_dashboard.py``. +# Standalone (not bound to the ORM class) so it works against either the +# ORM row or the ``TestDefinitionSummary`` dataclass. +# --------------------------------------------------------------------------- + + +def _epoch_ms(ts: datetime) -> str: + return str(int(ts.timestamp() * 1000)) + + +def test_forecast_points_returns_sorted_per_sensitivity_band(): + from testgen.common.models.test_definition import forecast_points_from_prediction + + t1 = datetime(2026, 7, 1, 12, 0, tzinfo=UTC) + t2 = datetime(2026, 7, 8, 12, 0, tzinfo=UTC) + t3 = datetime(2026, 7, 15, 12, 0, tzinfo=UTC) + prediction = { + "lower_tolerance|medium": {_epoch_ms(t2): 90.0, _epoch_ms(t1): 100.0, _epoch_ms(t3): 80.0}, + "upper_tolerance|medium": {_epoch_ms(t1): 110.0, _epoch_ms(t2): 120.0, _epoch_ms(t3): 130.0}, + # Other sensitivities present but not selected: + "lower_tolerance|low": {_epoch_ms(t1): 200.0}, + } + + points = forecast_points_from_prediction(prediction, "medium") + + assert [p.test_time for p in points] == [t1, t2, t3] + assert [p.lower_bound for p in points] == [100.0, 90.0, 80.0] + assert [p.upper_bound for p in points] == [110.0, 120.0, 130.0] + + +def test_forecast_points_returns_empty_when_no_prediction(): + from testgen.common.models.test_definition import forecast_points_from_prediction + assert forecast_points_from_prediction(None, "medium") == [] + assert forecast_points_from_prediction({}, "medium") == [] + + +def test_forecast_points_returns_empty_for_unknown_sensitivity(): + from testgen.common.models.test_definition import forecast_points_from_prediction + prediction = {"lower_tolerance|medium": {_epoch_ms(datetime(2026, 7, 1, tzinfo=UTC)): 1.0}} + assert forecast_points_from_prediction(prediction, "high") == [] + + +def test_forecast_points_handles_one_sided_bounds(): + """When a sensitivity has only one of lower/upper stored, the other + bound on each point should come through as ``None`` rather than skipping + the point entirely.""" + from testgen.common.models.test_definition import forecast_points_from_prediction + + t1 = datetime(2026, 7, 1, 12, 0, tzinfo=UTC) + prediction = {"upper_tolerance|medium": {_epoch_ms(t1): 130.0}} + + points = forecast_points_from_prediction(prediction, "medium") + + assert len(points) == 1 + assert points[0].test_time == t1 + assert points[0].lower_bound is None + assert points[0].upper_bound == 130.0 diff --git a/tests/unit/mcp/test_model_test_result.py b/tests/unit/mcp/test_model_test_result.py index ec037820..8abd9018 100644 --- a/tests/unit/mcp/test_model_test_result.py +++ b/tests/unit/mcp/test_model_test_result.py @@ -4,7 +4,12 @@ import pytest from sqlalchemy.dialects import postgresql -from testgen.common.models.test_result import TestResult, TestResultStatus +from testgen.common.models.test_result import ( + TestResult, + TestResultStatus, + TestRunResultRow, + _parse_kv_pairs, +) @pytest.fixture @@ -128,6 +133,26 @@ def test_select_results_excludes_monitor_suites_with_project_codes(session_mock) assert "test_suites.project_code IN" in sql +def test_list_for_run_excludes_monitor_suites(): + with patch.object(TestResult, "_paginate", return_value=([], 0)) as mock_paginate: + TestResult.list_for_run(uuid4()) + + sql = _compiled_sql(mock_paginate.call_args.args[0]) + assert "test_suites.is_monitor IS NOT true" in sql + assert "JOIN test_suites" in sql + + +def test_list_for_run_applies_caller_clauses_and_pagination(): + with patch.object(TestResult, "_paginate", return_value=([], 0)) as mock_paginate: + TestResult.list_for_run(uuid4(), TestResult.table_name == "orders", page=3, limit=10) + + sql = _compiled_sql(mock_paginate.call_args.args[0]) + assert "test_results.table_name = " in sql + assert mock_paginate.call_args.kwargs["page"] == 3 + assert mock_paginate.call_args.kwargs["limit"] == 10 + assert mock_paginate.call_args.kwargs["data_class"] is TestRunResultRow + + def test_select_failures_excludes_monitor_suites(session_mock): session_mock.execute.return_value.all.return_value = [] @@ -146,3 +171,98 @@ def test_select_history_excludes_monitor_suites(session_mock): sql = _compiled_sql(session_mock.scalars.call_args[0][0]) assert "test_suites.is_monitor IS NOT true" in sql assert "JOIN test_suites" in sql + + +# --------------------------------------------------------------------------- +# _parse_kv_pairs — contract with the producer of ``input_parameters`` +# --------------------------------------------------------------------------- + + +def test_parse_kv_pairs_splits_on_semicolon_for_multi_field_input(): + """``input_parameters`` is built with ``"; ".join(...)`` on the writer side + and consumed by ``dict_from_kv`` with ``;`` as the default separator. The + parser must agree, otherwise a row carrying both ``lower_tolerance`` and + ``upper_tolerance`` returns one entry with the second pair embedded in the + first value, and the upper bound never surfaces.""" + raw = "lower_tolerance=5; upper_tolerance=10; baseline_value=42" + assert _parse_kv_pairs(raw) == { + "lower_tolerance": "5", + "upper_tolerance": "10", + "baseline_value": "42", + } + + +def test_parse_kv_pairs_returns_empty_for_empty_input(): + assert _parse_kv_pairs(None) == {} + assert _parse_kv_pairs("") == {} + + +def test_parse_kv_pairs_tolerates_whitespace_and_skips_unkeyed_entries(): + raw = " k1 = v1 ;k2=v2; just_text ; k3=v3" + assert _parse_kv_pairs(raw) == {"k1": "v1", "k2": "v2", "k3": "v3"} + + +# --------------------------------------------------------------------------- +# list_monitor_events_for_table — ORDER BY needs a stable tiebreaker because +# the result is then paginated in Python; without it a row sharing ``test_time`` +# with another can duplicate or skip across pages. +# --------------------------------------------------------------------------- + + +def test_list_monitor_events_for_table_orders_by_stable_tiebreaker(session_mock): + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_monitor_events_for_table(uuid4(), "orders") + + raw_sql = str(session_mock.execute.call_args[0][0]) + order_clause = raw_sql.split("ORDER BY", 1)[1] + assert "results.id" in order_clause + assert "active_runs.id" in order_clause + + +# --------------------------------------------------------------------------- +# list_metric_monitor_events — scoped to one test_definition_id, separate +# query path from the run-by-type CTE (no synthesized pending rows). +# --------------------------------------------------------------------------- + + +def test_list_metric_monitor_events_scopes_to_test_definition_id(session_mock): + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + suite_id = uuid4() + monitor_id = uuid4() + TestResult.list_metric_monitor_events(suite_id, monitor_id) + + raw_sql = str(session_mock.execute.call_args[0][0]) + params = session_mock.execute.call_args[0][1] + assert "results.test_definition_id = :test_definition_id" in raw_sql + assert "results.test_suite_id = :test_suite_id" in raw_sql + assert params["test_definition_id"] == str(monitor_id) + assert params["test_suite_id"] == str(suite_id) + + +def test_list_metric_monitor_events_has_stable_order_with_tiebreaker(session_mock): + """Like the multi-type CTE, this path paginates in Python — the ORDER BY + must include a tiebreaker so rows sharing ``test_time`` don't shuffle + between calls.""" + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_metric_monitor_events(uuid4(), uuid4()) + + raw_sql = str(session_mock.execute.call_args[0][0]) + order_clause = raw_sql.split("ORDER BY", 1)[1] + assert "results.test_time DESC" in order_clause + assert "results.id" in order_clause + + +def test_list_metric_monitor_events_bounds_by_suite_lookback(session_mock): + """The suite's ``monitor_lookback`` (defaulting to 1) caps how many rows + come back. The query baking the LIMIT directly via a CTE means the model + doesn't need a second round-trip to read the lookback value.""" + session_mock.execute.return_value.mappings.return_value.all.return_value = [] + + TestResult.list_metric_monitor_events(uuid4(), uuid4()) + + raw_sql = str(session_mock.execute.call_args[0][0]) + assert "monitor_lookback" in raw_sql + assert "lookback_multiplier" in raw_sql diff --git a/tests/unit/mcp/test_permissions.py b/tests/unit/mcp/test_permissions.py index 0f97cbe4..f0815d4a 100644 --- a/tests/unit/mcp/test_permissions.py +++ b/tests/unit/mcp/test_permissions.py @@ -3,7 +3,8 @@ import pytest -from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible +from testgen.common.auth import AuthError +from testgen.mcp.exceptions import MCPAuthenticationError, MCPPermissionDenied, MCPResourceNotAccessible from testgen.mcp.permissions import ( _NOT_SET, ProjectPermissions, @@ -38,12 +39,14 @@ def test_get_authorized_mcp_user_raises_when_no_username(): @patch("testgen.common.models.get_current_session") @patch("testgen.common.auth.authorize_token") -def test_get_authorized_mcp_user_raises_when_user_not_found(mock_authorize, mock_get_session): - mock_authorize.side_effect = ValueError("User not found") +def test_get_authorized_mcp_user_translates_unknown_user_to_auth_error(mock_authorize, mock_get_session): + """A valid-signature token whose user no longer exists surfaces a clean MCP auth error, + not the generic 'unexpected error' the boundary returns for raw exceptions.""" + mock_authorize.side_effect = AuthError("User not found") set_mcp_username("ghost") set_mcp_token("some_token") - with pytest.raises(ValueError, match="User not found"): + with pytest.raises(MCPAuthenticationError, match="no longer valid"): get_authorized_mcp_user() @@ -65,12 +68,12 @@ def test_get_authorized_mcp_user_returns_user(mock_authorize, mock_get_session): @patch("testgen.common.models.get_current_session") @patch("testgen.common.auth.authorize_token") -def test_get_authorized_mcp_user_rejects_revoked_token(mock_authorize, mock_get_session): - mock_authorize.side_effect = ValueError("Token has been revoked") +def test_get_authorized_mcp_user_translates_revoked_token_to_auth_error(mock_authorize, mock_get_session): + mock_authorize.side_effect = AuthError("Token has been revoked") set_mcp_username("admin") set_mcp_token("revoked_token") - with pytest.raises(ValueError, match="Token has been revoked"): + with pytest.raises(MCPAuthenticationError, match="no longer valid"): get_authorized_mcp_user() @@ -329,3 +332,51 @@ def my_tool(x: int, y: str = "default") -> str: assert my_tool.__name__ == "my_tool" assert my_tool.__doc__ == "Tool docstring." + + +# --- mcp_permission("global_admin") --- + + +def test_mcp_permission_global_admin_allows_when_flag_set(mcp_user): + """global_admin gates on User.is_global_admin, not on project memberships.""" + set_mcp_username("test") + mcp_user.is_global_admin = True + + captured = {} + + @mcp_permission("global_admin") + def tool_fn(): + captured["perms"] = get_project_permissions() + return "ok" + + assert tool_fn() == "ok" + assert captured["perms"].permission == "global_admin" + assert captured["perms"].memberships == {} + + +def test_mcp_permission_global_admin_denies_when_flag_false(mcp_user): + set_mcp_username("test") + mcp_user.is_global_admin = False + + @mcp_permission("global_admin") + def tool_fn(): + raise AssertionError("Should not be called") + + with pytest.raises(MCPPermissionDenied, match="necessary permission"): + tool_fn() + + +def test_mcp_permission_global_admin_allows_with_zero_memberships(mcp_user): + """A global admin with no project memberships still passes the gate — the standard + 'no allowed_codes' bailout would otherwise refuse a fresh super-user.""" + set_mcp_username("test") + mcp_user.is_global_admin = True + + # global_admin path must NOT call _compute_project_permissions — the tool still runs + # even though the conftest's membership mock would yield zero allowed_codes for an + # unknown permission. + @mcp_permission("global_admin") + def tool_fn(): + return "ok" + + assert tool_fn() == "ok" diff --git a/tests/unit/mcp/test_server_instrumentation.py b/tests/unit/mcp/test_server_instrumentation.py new file mode 100644 index 00000000..38a968bb --- /dev/null +++ b/tests/unit/mcp/test_server_instrumentation.py @@ -0,0 +1,125 @@ +from unittest.mock import patch + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import set_mcp_username +from testgen.mcp.server import HandlerKind, MCPCallStatus, _instrument + + +def _emit_capture(): + calls = [] + + def fake_send_event(self, event_name, include_usage=False, **properties): + calls.append((event_name, properties)) + + return calls, fake_send_event + + +def test_instrument_emits_success_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_tool(): + return "ok" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_tool, HandlerKind.TOOL) + assert wrapped() == "ok" + + assert len(calls) == 1 + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.TOOL + assert props["handler_name"] == "my_tool" + assert props["username"] == "alice" + assert props["status"] == MCPCallStatus.SUCCESS + assert isinstance(props["latency_ms"], int) + + +def test_instrument_resource_emits_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_resource(): + return "doc" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_resource, HandlerKind.RESOURCE) + assert wrapped() == "doc" + + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.RESOURCE + assert props["handler_name"] == "my_resource" + assert props["status"] == MCPCallStatus.SUCCESS + + +def test_instrument_prompt_emits_event(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def my_prompt(): + return "template" + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(my_prompt, HandlerKind.PROMPT) + assert wrapped() == "template" + + event_name, props = calls[0] + assert event_name == "mcp-call" + assert props["kind"] == HandlerKind.PROMPT + assert props["handler_name"] == "my_prompt" + assert props["status"] == MCPCallStatus.SUCCESS + + +def test_instrument_permission_denied_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def denied_tool(): + raise MCPResourceNotAccessible("Project", "demo") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(denied_tool, HandlerKind.TOOL) + with pytest.raises(MCPPermissionDenied): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.PERMISSION_DENIED + + +def test_instrument_user_error_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def bad_input(): + raise MCPUserError("invalid uuid") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(bad_input, HandlerKind.TOOL) + with pytest.raises(MCPUserError): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.USER_ERROR + + +def test_instrument_error_status(): + set_mcp_username("alice") + calls, fake = _emit_capture() + + def boom(): + raise ValueError("nope") + + with patch("testgen.common.mixpanel_service.MixpanelService.send_event", fake): + wrapped = _instrument(boom, HandlerKind.TOOL) + with pytest.raises(ValueError): + wrapped() + + assert calls[0][1]["status"] == MCPCallStatus.ERROR + + +def test_instrument_preserves_name(): + def my_tool(): + return "ok" + + assert _instrument(my_tool, HandlerKind.TOOL).__name__ == "my_tool" diff --git a/tests/unit/mcp/test_tools_common.py b/tests/unit/mcp/test_tools_common.py index 3b5e7f32..142946f8 100644 --- a/tests/unit/mcp/test_tools_common.py +++ b/tests/unit/mcp/test_tools_common.py @@ -5,6 +5,7 @@ from testgen.common.enums import Disposition, ImpactDimension, IssueLikelihood, PiiRisk, QualityDimension from testgen.common.models.scores import ScoreCategory +from testgen.common.models.test_definition import Severity from testgen.common.models.test_result import TestResultStatus from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.tools.common import ( @@ -28,7 +29,9 @@ parse_score_filter_field, parse_score_group_by, parse_score_type, + parse_severity, parse_uuid, + resolve_hygiene_issue, resolve_issue_type, resolve_profiling_run, resolve_test_note, @@ -87,6 +90,32 @@ def test_parse_result_status_invalid_lists_valid_values(): assert status.value in str(exc_info.value) +# --- parse_severity --- + + +def test_parse_severity_valid(): + assert parse_severity("Fail") == Severity.FAIL + assert parse_severity("Warning") == Severity.WARNING + + +def test_parse_severity_invalid_names_value(): + with pytest.raises(MCPUserError, match="Invalid severity `Critical`"): + parse_severity("Critical") + + +def test_parse_severity_invalid_lists_valid_values(): + with pytest.raises(MCPUserError, match="Valid values:") as exc_info: + parse_severity("nope") + for severity in Severity: + assert severity.value in str(exc_info.value) + + +def test_parse_severity_case_sensitive(): + """Severity stores 'Fail' / 'Warning' verbatim; lowercase rejected by design.""" + with pytest.raises(MCPUserError, match="Invalid severity `fail`"): + parse_severity("fail") + + # --- validate_page --- @@ -310,7 +339,7 @@ def _mock_perms(allowed_projects=("demo",)): def test_resolve_profiling_run_happy_path(mock_pr_cls, mock_get_perms, db_session_mock): run = MagicMock() run.project_code = "demo" - mock_pr_cls.get_by_id_or_job.return_value = run + mock_pr_cls.get.return_value = run mock_get_perms.return_value = _mock_perms(allowed_projects=("demo",)) result = resolve_profiling_run(str(uuid4())) @@ -321,7 +350,7 @@ def test_resolve_profiling_run_happy_path(mock_pr_cls, mock_get_perms, db_sessio @patch("testgen.mcp.tools.common.get_project_permissions") @patch("testgen.mcp.tools.common.ProfilingRun") def test_resolve_profiling_run_unknown_run_id(mock_pr_cls, mock_get_perms, db_session_mock): - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None mock_get_perms.return_value = _mock_perms() with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -334,7 +363,7 @@ def test_resolve_profiling_run_inaccessible_project(mock_pr_cls, mock_get_perms, """Run exists but caller can't access its project — same unified error as unknown run.""" run = MagicMock() run.project_code = "forbidden" - mock_pr_cls.get_by_id_or_job.return_value = run + mock_pr_cls.get.return_value = run mock_get_perms.return_value = _mock_perms(allowed_projects=("demo",)) with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -379,6 +408,35 @@ def test_resolve_test_note_invalid_uuid(): resolve_test_note("not-a-uuid") +# --- resolve_hygiene_issue --- + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.HygieneIssue") +def test_resolve_hygiene_issue_happy_path(mock_hi_cls, mock_get_perms, db_session_mock): + issue = MagicMock() + mock_hi_cls.get.return_value = issue + mock_get_perms.return_value = _mock_perms() + + assert resolve_hygiene_issue(str(uuid4())) is issue + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.HygieneIssue") +def test_resolve_hygiene_issue_missing_or_inaccessible(mock_hi_cls, mock_get_perms, db_session_mock): + """Missing issue and forbidden-project issue both collapse to one error (project scoped in the query).""" + mock_hi_cls.get.return_value = None + mock_get_perms.return_value = _mock_perms() + + with pytest.raises(MCPResourceNotAccessible, match=r"Hygiene issue .* not found or not accessible"): + resolve_hygiene_issue(str(uuid4())) + + +def test_resolve_hygiene_issue_invalid_uuid(): + with pytest.raises(MCPUserError, match="Invalid issue_id"): + resolve_hygiene_issue("not-a-uuid") + + # --- parse_pii_category --- @@ -652,6 +710,7 @@ def test_parse_score_type_invalid_lists_valid_values(): ("Stakeholder Group", ScoreCategory.stakeholder_group), ("Transform Level", ScoreCategory.transform_level), ("Data Product", ScoreCategory.data_product), + ("Data Classification", ScoreCategory.data_classification), ], ) def test_parse_category_display_form_returns_column_form_enum(display_value, expected): @@ -680,6 +739,7 @@ def test_parse_category_translation_dict_covers_all_args(): "stakeholder_group", "transform_level", "data_product", + "data_classification", ], ) def test_parse_category_rejects_column_form_input(internal): @@ -711,3 +771,430 @@ def test_score_chain_leaf_field_values(): def test_score_chain_leaf_to_column_mapping(): assert SCORE_CHAIN_LEAF_TO_COLUMN[ScoreChainLeafField.TABLE] == "table_name" assert SCORE_CHAIN_LEAF_TO_COLUMN[ScoreChainLeafField.COLUMN] == "column_name" + + +# --- SqlFlavorLabel --- + + +def test_sql_flavor_label_set_matches_common_layer(): + """Codes covered by the MCP enum and the common-layer maps must stay in sync.""" + from testgen.common.flavors import FLAVOR_CODE_TO_FAMILY, FLAVOR_CODE_TO_LABEL + from testgen.mcp.tools.common import SQL_FLAVOR_CODE_TO_LABEL, SQL_FLAVOR_LABEL_TO_CODE + + assert set(SQL_FLAVOR_CODE_TO_LABEL) == set(FLAVOR_CODE_TO_LABEL) + assert set(SQL_FLAVOR_LABEL_TO_CODE.values()) == set(FLAVOR_CODE_TO_FAMILY.keys()) + + +def test_parse_sql_flavor_returns_label_code_family(): + from testgen.mcp.tools.common import SqlFlavorLabel, parse_sql_flavor + + label, code, family = parse_sql_flavor("Azure SQL Database") + assert label == SqlFlavorLabel.AZURE_MSSQL + assert code == "azure_mssql" + assert family == "mssql" + + +def test_parse_sql_flavor_invalid_lists_display_values(): + from testgen.mcp.tools.common import SqlFlavorLabel, parse_sql_flavor + + with pytest.raises(MCPUserError, match="Invalid sql_flavor `bogus`") as exc: + parse_sql_flavor("bogus") + msg = str(exc.value) + for member in SqlFlavorLabel: + assert member.value in msg + + +# --- parse_test_result_disposition --- + + +@pytest.mark.parametrize( + "user_label,expected", + [ + ("Confirmed", Disposition.CONFIRMED), + ("Dismissed", Disposition.DISMISSED), + ("Muted", Disposition.INACTIVE), + ("No Decision", None), + ], +) +def test_parse_test_result_disposition_user_labels(user_label, expected): + from testgen.mcp.tools.common import parse_test_result_disposition + + assert parse_test_result_disposition(user_label) is expected + + +def test_parse_test_result_disposition_rejects_unknown_and_lists_accepted(): + from testgen.mcp.tools.common import parse_test_result_disposition + + with pytest.raises(MCPUserError) as exc: + parse_test_result_disposition("Inactive") # DB value, not user-facing + msg = str(exc.value) + for label in ("Confirmed", "Dismissed", "Muted", "No Decision"): + assert label in msg + + +# --- resolve_test_result --- + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.get_current_session") +def test_resolve_test_result_happy_path(mock_session, mock_perms, db_session_mock): + from testgen.mcp.tools.common import resolve_test_result + + result = MagicMock() + mock_session.return_value.scalars.return_value.first.return_value = result + mock_perms.return_value = _mock_perms(allowed_projects=("demo",)) + + assert resolve_test_result(str(uuid4())) is result + + +@patch("testgen.mcp.tools.common.get_project_permissions") +@patch("testgen.mcp.tools.common.get_current_session") +def test_resolve_test_result_missing_or_inaccessible(mock_session, mock_perms, db_session_mock): + from testgen.mcp.tools.common import resolve_test_result + + mock_session.return_value.scalars.return_value.first.return_value = None + mock_perms.return_value = _mock_perms(allowed_projects=("demo",)) + + with pytest.raises(MCPResourceNotAccessible, match=r"Test result .* not found or not accessible"): + resolve_test_result(str(uuid4())) + + +def test_resolve_test_result_invalid_uuid(): + from testgen.mcp.tools.common import resolve_test_result + + with pytest.raises(MCPUserError, match="Invalid test_result_id"): + resolve_test_result("not-a-uuid") + + +# --- Monitor helpers --- + + +@pytest.mark.parametrize( + "label,expected_value", + [ + ("freshness", "Freshness_Trend"), + ("volume", "Volume_Trend"), + ("schema", "Schema_Drift"), + ("metric", "Metric_Trend"), + ], +) +def test_parse_monitor_type_user_labels(label, expected_value): + from testgen.mcp.tools.common import parse_monitor_type + + assert parse_monitor_type(label).value == expected_value + + +def test_parse_monitor_type_rejects_db_codes(): + """Internal codes like ``Freshness_Trend`` are not accepted on the input boundary — + only the lowercase user-facing short labels are.""" + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match="Invalid monitor_type"): + parse_monitor_type("Freshness_Trend") + + +def test_parse_monitor_type_lists_valid_values_on_error(): + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match="Valid values:") as exc: + parse_monitor_type("metrics") + msg = str(exc.value) + for label in ("freshness", "volume", "schema", "metric"): + assert label in msg + + +def test_parse_monitor_type_label_override(): + """``label`` argument lets callers tailor the error to their public arg name + (e.g. ``list_monitored_tables`` exposes it as ``anomaly_type``).""" + from testgen.mcp.tools.common import parse_monitor_type + + with pytest.raises(MCPUserError, match=r"Invalid anomaly_type `bogus`"): + parse_monitor_type("bogus", "anomaly_type") + + +@pytest.mark.parametrize( + "value", + ["table_name", "anomaly_count_desc", "latest_update_desc", "row_count_change_desc"], +) +def test_parse_monitor_table_sort_accepts_documented_values(value): + from testgen.mcp.tools.common import parse_monitor_table_sort + + assert parse_monitor_table_sort(value).value == value + + +def test_parse_monitor_table_sort_rejects_unknown(): + from testgen.mcp.tools.common import parse_monitor_table_sort + + with pytest.raises(MCPUserError, match="Invalid sort_by") as exc: + parse_monitor_table_sort("alphabetical") + msg = str(exc.value) + for valid in ("table_name", "anomaly_count_desc", "latest_update_desc", "row_count_change_desc"): + assert valid in msg + + +def test_parse_monitor_table_sort_rejects_legacy_row_count_desc(): + """Guard against the pre-review-feedback ``row_count_desc`` name accidentally + coming back: the rename to ``row_count_change_desc`` is the canonical signal + that the column shows a delta, not the raw current count. Drop this only when + introducing a deliberate replacement.""" + from testgen.mcp.tools.common import parse_monitor_table_sort + + with pytest.raises(MCPUserError, match="Invalid sort_by") as exc: + parse_monitor_table_sort("row_count_desc") + assert "row_count_change_desc" in str(exc.value) + + +def test_resolve_monitored_table_group_returns_suite(): + from testgen.common.models.table_group import TableGroup + from testgen.common.models.test_suite import TestSuite + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = uuid4() + suite = MagicMock(spec=TestSuite) + suite.is_monitor = True + + with ( + patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg), + patch("testgen.mcp.tools.common.TestSuite.get", return_value=suite) as mock_get, + ): + out_tg, out_suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert out_suite is suite + assert mock_get.call_args.args[0] == tg.monitor_test_suite_id + + +def test_resolve_monitored_table_group_returns_none_when_unlinked(): + """Table group exists but has no monitor suite pointer.""" + from testgen.common.models.table_group import TableGroup + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = None + + with patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg): + out_tg, suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert suite is None + + +def test_resolve_monitored_table_group_returns_none_when_pointer_stale(): + """Pointer set, but suite no longer exists or no longer ``is_monitor=True``.""" + from testgen.common.models.table_group import TableGroup + from testgen.mcp.tools.common import resolve_monitored_table_group + + tg = MagicMock(spec=TableGroup) + tg.id = uuid4() + tg.monitor_test_suite_id = uuid4() + + with ( + patch("testgen.mcp.tools.common.resolve_table_group", return_value=tg), + patch("testgen.mcp.tools.common.TestSuite.get", return_value=None), + ): + out_tg, suite = resolve_monitored_table_group(str(tg.id)) + + assert out_tg is tg + assert suite is None + + +def test_resolve_monitored_table_group_raises_when_tg_inaccessible(): + """Inaccessible TG propagates ``MCPResourceNotAccessible`` from ``resolve_table_group`` + — the "not monitored" path must not mask an access failure.""" + from testgen.mcp.tools.common import resolve_monitored_table_group + + bad_id = str(uuid4()) + with ( + patch( + "testgen.mcp.tools.common.resolve_table_group", + side_effect=MCPResourceNotAccessible("Table group", bad_id), + ), + pytest.raises(MCPResourceNotAccessible), + ): + resolve_monitored_table_group(bad_id) + + + +# --- resolve_project --- + + +def test_resolve_project_returns_project_when_in_scope(db_session_mock): + """Happy path: project_code in allowed_codes, Project.get returns the row.""" + from testgen.mcp.permissions import ProjectPermissions, _mcp_project_permissions + from testgen.mcp.tools.common import resolve_project + + project = MagicMock() + project.project_code = "demo" + + with patch("testgen.mcp.tools.common.Project") as mock_project_cls: + mock_project_cls.get.return_value = project + mock_project_cls.project_code = MagicMock() # for the .in_() filter clause + perms = ProjectPermissions(memberships={"demo": "admin"}, permission="administer", username="t") + token = _mcp_project_permissions.set(perms) + try: + with patch( + "testgen.mcp.permissions.PluginHook" + ) as mock_hook: + mock_hook.instance.return_value.rbac.get_roles_with_permission.return_value = ["admin"] + assert resolve_project("demo") is project + finally: + _mcp_project_permissions.reset(token) + + +def test_resolve_project_raises_unified_when_get_returns_none(db_session_mock): + """Project not in scope (or absent) → unified ``MCPResourceNotAccessible``.""" + from testgen.mcp.permissions import ProjectPermissions, _mcp_project_permissions + from testgen.mcp.tools.common import resolve_project + + with patch("testgen.mcp.tools.common.Project") as mock_project_cls: + mock_project_cls.get.return_value = None + mock_project_cls.project_code = MagicMock() + perms = ProjectPermissions(memberships={"demo": "admin"}, permission="administer", username="t") + token = _mcp_project_permissions.set(perms) + try: + with patch( + "testgen.mcp.permissions.PluginHook" + ) as mock_hook: + mock_hook.instance.return_value.rbac.get_roles_with_permission.return_value = ["admin"] + with pytest.raises(MCPResourceNotAccessible) as exc: + resolve_project("secret") + assert "Project `secret` not found or not accessible" in str(exc.value) + finally: + _mcp_project_permissions.reset(token) + + +# --- render_diff_table / _default_render_diff_value --- + + +def test_render_diff_table_emits_rows_for_changed_attrs(db_session_mock): + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"name": "Old", "active": True, "label": "x"} + after = {"name": "New", "active": True, "label": "y"} + + rendered = render_diff_table( + doc, before, after, + attrs=("name", "active", "label"), + labels={"name": "Name", "active": "Active", "label": "Label"}, + ) + + assert rendered is True + out = doc.render() + assert "| Field | Before | After |" in out + assert "Old" in out and "New" in out + assert "Label" in out and "x" in out and "y" in out + # Unchanged attr "active" must not appear + assert "Active" not in out + + +def test_render_diff_table_returns_false_when_nothing_changes(db_session_mock): + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + snap = {"name": "Same", "count": 3} + + rendered = render_diff_table( + doc, snap, snap, + attrs=("name", "count"), + labels={"name": "Name", "count": "Count"}, + ) + + assert rendered is False + assert doc.render() == "" # nothing appended + + +def test_render_diff_table_redacts_secret_attrs(db_session_mock): + """secret_attrs render as ``[secret]`` when present and em-dash when absent — value + is never echoed.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"name": "Demo", "api_key": None} + after = {"name": "Demo", "api_key": "super-secret-value"} + + rendered = render_diff_table( + doc, before, after, + attrs=("name", "api_key"), + labels={"name": "Name", "api_key": "API key"}, + secret_attrs=frozenset({"api_key"}), + ) + + assert rendered is True + out = doc.render() + assert "API key" in out + assert "[secret]" in out + assert "super-secret-value" not in out + + +def test_render_diff_table_honors_attr_ordering(db_session_mock): + """Row order matches the supplied ``attrs`` tuple, not dict insertion order.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"b": 1, "a": 1, "c": 1} + after = {"b": 2, "a": 2, "c": 2} + + render_diff_table( + doc, before, after, + attrs=("a", "b", "c"), + labels={"a": "A", "b": "B", "c": "C"}, + ) + + out = doc.render() + # Body rows render in attrs order. Label cells are code-wrapped (column 0 uses + # ``code=[0]`` in ``MdDoc.table``). + pos_a = out.find("| `A` |") + pos_b = out.find("| `B` |") + pos_c = out.find("| `C` |") + assert 0 < pos_a < pos_b < pos_c + + +def test_render_diff_table_custom_value_renderer(db_session_mock): + """``value_renderer`` overrides the default Yes/No/em-dash formatting.""" + from testgen.mcp.tools.common import render_diff_table + from testgen.mcp.tools.markdown import MdDoc + + doc = MdDoc() + before = {"n": 1} + after = {"n": 2} + + rendered = render_diff_table( + doc, before, after, + attrs=("n",), + labels={"n": "N"}, + value_renderer=lambda v: f"#{v}", + ) + + assert rendered is True + out = doc.render() + assert "#1" in out + assert "#2" in out + + +def test_default_render_diff_value_bool_yes_no(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value(True) == "Yes" + assert _default_render_diff_value(False) == "No" + + +def test_default_render_diff_value_none_and_empty(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value(None) is None + assert _default_render_diff_value("") is None + + +def test_default_render_diff_value_str(): + from testgen.mcp.tools.common import _default_render_diff_value + + assert _default_render_diff_value("hello") == "hello" + assert _default_render_diff_value(42) == "42" diff --git a/tests/unit/mcp/test_tools_connections.py b/tests/unit/mcp/test_tools_connections.py new file mode 100644 index 00000000..b324eb42 --- /dev/null +++ b/tests/unit/mcp/test_tools_connections.py @@ -0,0 +1,419 @@ +"""Tests for the core MCP connection tools — list / get / test. + +The tools take a flavor-shaped ``connection_params`` dict (keyed by UI label) +plus an explicit ``connection_mode``; mapping + validation are delegated to +``connection_service``. +""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.common.database.connection_service import ConnectionStatus +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.connections" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="administer"): + """Inject a ProjectPermissions for the given access set.""" + memberships = memberships or dict.fromkeys(allowed, "role_a") + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_connection(**overrides) -> MagicMock: + """Build a MagicMock matching the Connection model surface used by the tools.""" + conn = MagicMock() + conn.connection_id = overrides.get("connection_id", 42) + conn.project_code = overrides.get("project_code", "demo") + conn.connection_name = overrides.get("connection_name", "Local PG") + conn.sql_flavor = overrides.get("sql_flavor", "postgresql") + conn.sql_flavor_code = overrides.get("sql_flavor_code", "postgresql") + conn.project_host = overrides.get("project_host", "localhost") + conn.project_port = overrides.get("project_port", "5432") + conn.project_db = overrides.get("project_db", "testgen_local") + conn.project_user = overrides.get("project_user", "testgen") + conn.project_pw_encrypted = overrides.get("project_pw_encrypted", "stored_pw") + conn.connect_by_url = overrides.get("connect_by_url", False) + conn.url = overrides.get("url", None) + conn.connect_by_key = overrides.get("connect_by_key", False) + conn.private_key = overrides.get("private_key", None) + conn.private_key_passphrase = overrides.get("private_key_passphrase", None) + conn.connect_with_identity = overrides.get("connect_with_identity", False) + conn.http_path = overrides.get("http_path", None) + conn.warehouse = overrides.get("warehouse", None) + conn.service_account_key = overrides.get("service_account_key", None) + conn.max_threads = overrides.get("max_threads", 4) + conn.max_query_chars = overrides.get("max_query_chars", None) + return conn + + +# --------------------------------------------------------------------------- +# test_connection +# --------------------------------------------------------------------------- + + +def test_test_connection_neither_id_nor_flavor(db_session_mock): + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection() + msg = str(exc.value) + assert "`connection_id`" in msg + assert "`sql_flavor`" in msg + + +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_rejects_sql_flavor_with_id(mock_resolve, db_session_mock): + """sql_flavor override is meaningless on a stored connection — must reject loud.""" + mock_resolve.return_value = _mock_connection() + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection(connection_id=7, sql_flavor="Snowflake") + assert "sql_flavor" in str(exc.value) + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_with_id_only(mock_resolve, mock_runner, db_session_mock): + conn = _mock_connection(connection_id=7, connection_name="Local PG", project_host="localhost") + mock_resolve.return_value = conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection(connection_id=7) + + assert "Connection test succeeded" in out + assert "**ID:** `7`" in out + assert "**Name:** `Local PG`" in out + assert "**Type:** PostgreSQL" in out + mock_runner.assert_called_once() + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_with_id_and_overrides(mock_resolve, mock_runner, db_session_mock): + """Overrides win: assigned to the loaded connection before status runs.""" + conn = _mock_connection(connection_id=7, project_host="old.host") + mock_resolve.return_value = conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection(connection_id=7, connection_params={"Host": "new.host"}) + + assert conn.project_host == "new.host" + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="The connection was successful.", successful=True)) +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_only(mock_conn_cls, mock_runner, db_session_mock): + """No connection_id supplied → builds an inline Connection from args.""" + inline_conn = _mock_connection() + mock_conn_cls.return_value = inline_conn + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection( + sql_flavor="PostgreSQL", + connection_params={ + "Host": "localhost", + "Port": 5432, + "Database": "d", + "Username": "u", + "Password": "p", + }, + ) + + assert "Connection test succeeded" in out + # No ID/Name lines on inline tests (no entity). + assert "**ID:**" not in out + assert "**Name:**" not in out + mock_runner.assert_called_once() + + +@patch(f"{MODULE}.test_connection_status", return_value=ConnectionStatus(message="OK", successful=True)) +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_applies_defaults(mock_conn_cls, mock_runner, db_session_mock): + """Inline tests get the same flavor defaults as create.""" + inline = _mock_connection(max_query_chars=None) + mock_conn_cls.return_value = inline + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection( + sql_flavor="PostgreSQL", + connection_params={"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + + assert inline.max_query_chars == 20000 + + +@patch(f"{MODULE}.test_connection_status") +@patch(f"{MODULE}.resolve_connection") +def test_test_connection_failure_renders_details(mock_resolve, mock_runner, db_session_mock): + """Failure with detail string renders the detail verbatim in a code block.""" + mock_resolve.return_value = _mock_connection() + driver_text = "FATAL: password authentication failed for user 'dq'" + mock_runner.return_value = ConnectionStatus( + message="Error attempting the connection.", + successful=False, + details=driver_text, + ) + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + out = test_connection(connection_id=7) + + assert "Connection test failed" in out + assert driver_text in out # verbatim, no scrubbing + + +def test_test_connection_inline_validation_failure(db_session_mock): + """Inline test with missing required field rejects before opening any DB connection.""" + from testgen.mcp.tools.connections import test_connection + + # PostgreSQL inline with no params -> uses a real Connection (no patch) with empty fields. + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection(sql_flavor="PostgreSQL", connection_params={}) + msg = str(exc.value) + assert "Cannot test connection" in msg + assert "`Host` is required for PostgreSQL." in msg + + +def test_test_connection_inline_multi_mode_requires_mode(db_session_mock): + """Inline test of a multi-mode flavor without connection_mode is rejected (not defaulted).""" + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + test_connection( + sql_flavor="Salesforce Data 360", + connection_params={"Login URL": "https://my.salesforce.com", "Consumer Key": "ck"}, + ) + msg = str(exc.value) + assert "requires a connection_mode" in msg + assert "JWT Bearer Flow" in msg and "Client Credentials Flow" in msg + + +@patch(f"{MODULE}.test_connection_status") +@patch(f"{MODULE}.Connection") +def test_test_connection_inline_does_not_save(mock_conn_cls, mock_runner, db_session_mock): + """Inline test never persists — the constructed connection is not saved.""" + inline = _mock_connection() + mock_conn_cls.return_value = inline + mock_runner.return_value = ConnectionStatus(message="OK", successful=True) + + from testgen.mcp.tools.connections import test_connection + + with _patch_perms(): + test_connection( + sql_flavor="PostgreSQL", + connection_params={"Host": "h", "Port": 5432, "Database": "d", "Username": "u", "Password": "p"}, + ) + + inline.save.assert_not_called() + + +# --------------------------------------------------------------------------- +# list_connections +# --------------------------------------------------------------------------- + + +def _list_item(**overrides): + from testgen.common.models.connection import ConnectionListItem + + return ConnectionListItem( + connection_id=overrides.get("connection_id", 1), + connection_name=overrides.get("connection_name", "warehouse_prod"), + project_code=overrides.get("project_code", "demo"), + sql_flavor_code=overrides.get("sql_flavor_code", "snowflake"), + project_host=overrides.get("project_host", "abc.snowflakecomputing.com"), + project_db=overrides.get("project_db", "ANALYTICS"), + table_group_count=overrides.get("table_group_count", 4), + ) + + +@patch(f"{MODULE}.Connection") +def test_list_connections_returns_rows_with_display_label(mock_conn_cls, db_session_mock): + mock_conn_cls.list_for_project.return_value = ( + [_list_item(project_host="warehouse-host.example.com")], + 1, + ) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo") + + assert "Connections for `demo`" in out + assert "warehouse_prod" in out + # Display label is rendered, NOT the raw sql_flavor_code: + assert "Snowflake" in out + assert "snowflake" not in out + assert "warehouse-host.example.com" in out + assert "ANALYTICS" in out + assert "4" in out # table_group_count + + +@patch(f"{MODULE}.Connection") +def test_list_connections_pagination_footer(mock_conn_cls, db_session_mock): + rows = [_list_item(connection_id=i, connection_name=f"c{i}") for i in range(1, 3)] + mock_conn_cls.list_for_project.return_value = (rows, 25) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo", page=1, limit=2) + + assert "Showing 1\u20132 of 25" in out + assert "Use `page=2`" in out + + +@patch(f"{MODULE}.Connection") +def test_list_connections_empty(mock_conn_cls, db_session_mock): + mock_conn_cls.list_for_project.return_value = ([], 0) + + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"): + out = list_connections("demo") + + assert "No connections" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_list_connections_rejects_inaccessible_project(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.connections import list_connections + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + list_connections("secret") + + +def test_list_connections_invalid_limit_raises(db_session_mock): + from testgen.mcp.tools.connections import list_connections + + with _patch_perms(permission="view"), pytest.raises(MCPUserError, match="Invalid limit"): + list_connections("demo", limit=500) + + +# --------------------------------------------------------------------------- +# get_connection +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_returns_config_with_display_label(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_connection( + sql_flavor_code="snowflake", + connection_name="warehouse_prod", + connection_id=12, + ) + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(12) + + assert "Connection `warehouse_prod`" in out + assert "**ID:** `12`" in out + assert "**Type:** Snowflake" in out + # The raw flavor code MUST NOT appear in the output + assert "snowflake" not in out + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_never_returns_secrets(mock_resolve, db_session_mock): + """Password, private key, passphrase, service-account key must not leak. + + Uses a real ``Connection`` ORM instance so the test would fail if the rendering + started reading the encrypted attributes (a MagicMock would silently return + a ``MagicMock`` for any attribute, masking the regression). + """ + from testgen.common.models.connection import Connection + + conn = Connection( + id=uuid4(), + connection_id=42, + project_code="demo", + connection_name="warehouse_prod", + sql_flavor="postgresql", + sql_flavor_code="postgresql", + project_host="localhost", + project_port="5432", + project_db="testgen", + project_user="dq_user", + project_pw_encrypted="sentinel-password-hunter2", + private_key="-----BEGIN PRIVATE KEY-----\nsentinel-private-key-body\n-----END PRIVATE KEY-----", + private_key_passphrase="sentinel-passphrase", # noqa: S106 + service_account_key={"client_email": "sentinel-account@example.com", "private_key": "sentinel-sak-key"}, + ) + mock_resolve.return_value = conn + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(42) + + # Every populated secret value must NOT appear in output. + for needle in ( + "sentinel-password-hunter2", + "sentinel-private-key-body", + "BEGIN PRIVATE KEY", + "sentinel-passphrase", + "sentinel-account@example.com", + "sentinel-sak-key", + ): + assert needle not in out, f"secret material `{needle}` leaked into get_connection output" + + # The non-secret content we DO render should be present — sanity check that the + # test is exercising the real rendering path (not silently bypassing it). + assert "warehouse_prod" in out + assert "localhost" in out + + +@patch(f"{MODULE}.resolve_connection") +def test_get_connection_authentication_password(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_connection( + connect_by_key=False, connect_by_url=False, connect_with_identity=False, + ) + + from testgen.mcp.tools.connections import get_connection + + with _patch_perms(permission="view"): + out = get_connection(42) + + assert "**Authentication:** Password" in out + + +@patch("testgen.mcp.tools.common.Connection") +def test_get_connection_raises_not_found_for_inaccessible(mock_conn_cls, db_session_mock): + mock_conn_cls.get.return_value = None + + from testgen.mcp.tools.connections import get_connection + + with pytest.raises(MCPResourceNotAccessible, match="Connection .* not found or not accessible"): + get_connection(999) diff --git a/tests/unit/mcp/test_tools_data_catalog.py b/tests/unit/mcp/test_tools_data_catalog.py new file mode 100644 index 00000000..b622cb17 --- /dev/null +++ b/tests/unit/mcp/test_tools_data_catalog.py @@ -0,0 +1,294 @@ +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def grant_perms(): + """Grant role_a disposition + view_pii on demo (overrides the conftest matrix).""" + with patch("testgen.mcp.permissions.PluginHook") as hook: + hook.instance.return_value.rbac.get_roles_with_permission.side_effect = ( + lambda perm: ["role_a"] if perm in ("disposition", "view_pii", "catalog", "view") else [] + ) + yield hook + + +def _mock_table_group(): + tg = MagicMock() + tg.id = uuid4() + tg.project_code = "demo" + tg.table_groups_name = "sales_group" + tg.profile_flag_cdes = True + tg.profile_flag_pii = True + return tg + + +def _mock_entity(): + """A stand-in table/column row that records attribute writes.""" + return MagicMock() + + +def _patch(tg=None, table_rows=None, column_rows=None): + """Patch the resolution boundary: TableGroup.get, DataTable/DataColumnChars.select_where.""" + tg = tg or _mock_table_group() + mock_tg_cls = patch("testgen.mcp.tools.common.TableGroup").start() + mock_tg_cls.get.return_value = tg + mock_dt = patch("testgen.mcp.tools.data_catalog.DataTable").start() + mock_dt.select_where.return_value = table_rows if table_rows is not None else [] + mock_dcc = patch("testgen.mcp.tools.data_catalog.DataColumnChars").start() + mock_dcc.select_where.return_value = column_rows if column_rows is not None else [] + return tg, mock_dt, mock_dcc + + +@pytest.fixture(autouse=True) +def _cleanup_patches(): + yield + patch.stopall() + + +# ---------------------------------------------------------------------- +# Happy paths +# ---------------------------------------------------------------------- + +def test_table_description_update(db_session_mock): + table = _mock_entity() + _patch(table_rows=[table]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "Order header table"}, + ]) + + assert table.description == "Order header table" + assert "Updated" in result + assert "1" in result + + +def test_column_pii_true_stores_manual(db_session_mock): + column = _mock_entity() + _patch(column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": True}, + ]) + + assert column.pii_flag == "MANUAL" + + +def test_column_pii_false_clears(db_session_mock): + column = _mock_entity() + _patch(column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": False}, + ]) + + assert column.pii_flag is None + + +# ---------------------------------------------------------------------- +# Per-row isolation +# ---------------------------------------------------------------------- + +def test_one_bad_row_does_not_block_others(db_session_mock): + good = _mock_entity() + _tg, _mock_dt, mock_dcc = _patch() + + # First column resolves to the good row; the second is not found in the catalog. + calls = {"n": 0} + + def _resolve(*_a, **_k): + calls["n"] += 1 + return [good] if calls["n"] == 1 else [] + + mock_dcc.select_where.side_effect = _resolve + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + tg_id = str(uuid4()) + result = update_catalog_metadata([ + {"table_group_id": tg_id, "table_name": "orders", "column_name": "known", "description": "ok"}, + {"table_group_id": tg_id, "table_name": "orders", "column_name": "missing", "description": "x"}, + ]) + + assert good.description == "ok" + assert "Updated" in result and "Failed" in result + assert "not found" in result.lower() + + +# ---------------------------------------------------------------------- +# Validation / rejection +# ---------------------------------------------------------------------- + +def test_xde_on_table_row_rejected(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "xde": True}, + ]) + + assert "Failed" in result + assert "xde" in result.lower() and "column" in result.lower() + + +def test_pii_without_view_pii_rejected(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + with patch("testgen.mcp.tools.data_catalog.get_project_permissions") as mock_perms: + perms = MagicMock() + perms.has_permission.return_value = False # no view_pii + mock_perms.return_value = perms + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": True}, + ]) + + assert "Failed" in result + assert "view_pii" not in result + assert "permission to view PII" in result + + +def test_description_too_long_rejected(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "x" * 1001}, + ]) + + assert "Failed" in result + assert "description" in result.lower() + + +def test_non_boolean_pii_rejected(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": "high"}, + ]) + + assert "Failed" in result + + +def test_unknown_table_not_found(db_session_mock): + _patch(table_rows=[]) # select_where finds nothing + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "ghost", "description": "x"}, + ]) + + assert "Failed" in result + assert "not found" in result.lower() + + +def test_inaccessible_table_group(db_session_mock): + tg, _dt, _dcc = _patch() + # TableGroup.get returns None → resolve_table_group raises not-found-or-not-accessible + patch.stopall() + mock_tg_cls = patch("testgen.mcp.tools.common.TableGroup").start() + mock_tg_cls.get.return_value = None + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "description": "x"}, + ]) + + assert "Failed" in result + assert "not found or not accessible" in result.lower() + + +# ---------------------------------------------------------------------- +# Side-effect notices +# ---------------------------------------------------------------------- + +def test_cde_write_auto_disables_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "amount", "cde": True}, + ]) + + assert tg.profile_flag_cdes is False + assert "Auto-disabled profile_flag_cdes" in result + + +def test_cde_noop_does_not_disable_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + column.critical_data_element = False # already false; cde: false is a no-op + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "amount", "cde": False}, + ]) + + assert tg.profile_flag_cdes is True + assert "Auto-disabled profile_flag_cdes" not in result + + +def test_pii_noop_does_not_disable_flag(db_session_mock): + tg = _mock_table_group() + column = _mock_entity() + column.pii_flag = None # already cleared; pii: false is a no-op + _patch(tg=tg, column_rows=[column]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "ssn", "pii": False}, + ]) + + assert tg.profile_flag_pii is True + assert "Auto-disabled profile_flag_pii" not in result + + +def test_table_cde_inheritance_notice(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "cde": True}, + ]) + + assert "affects all columns" in result.lower() + + +def test_xde_exclusion_notice(db_session_mock): + _patch(column_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders", "column_name": "scratch", "xde": True}, + ]) + + assert "excluded" in result.lower() + assert "next profiling run" in result.lower() + assert "skip" in result.lower() + + +def test_empty_fields_skipped(db_session_mock): + _patch(table_rows=[_mock_entity()]) + + from testgen.mcp.tools.data_catalog import update_catalog_metadata + result = update_catalog_metadata([ + {"table_group_id": str(uuid4()), "table_name": "orders"}, + ]) + + assert "Skipped" in result + + +def test_empty_updates_raises(db_session_mock): + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.data_catalog import update_catalog_metadata + with pytest.raises(MCPUserError): + update_catalog_metadata([]) diff --git a/tests/unit/mcp/test_tools_discovery.py b/tests/unit/mcp/test_tools_discovery.py index 40146b21..3ca70913 100644 --- a/tests/unit/mcp/test_tools_discovery.py +++ b/tests/unit/mcp/test_tools_discovery.py @@ -116,11 +116,9 @@ def test_list_projects_filters_for_scoped_user(mock_compute, mock_project, db_se assert "Secret" not in result -@patch("testgen.mcp.tools.discovery.TestRun") @patch("testgen.mcp.tools.discovery.TestSuite") -def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mock): +def test_list_test_suites_returns_stats(mock_suite, db_session_mock): run_id = uuid4() - job_exec_id = uuid4() summary = MagicMock() summary.id = uuid4() summary.test_suite = "Quality Suite" @@ -137,7 +135,6 @@ def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mo summary.last_run_error_ct = 0 summary.last_run_dismissed_ct = 0 mock_suite.select_summary.return_value = [summary] - mock_test_run.get_job_execution_ids.return_value = {run_id: job_exec_id} from testgen.mcp.tools.discovery import list_test_suites @@ -146,7 +143,8 @@ def test_list_test_suites_returns_stats(mock_suite, mock_test_run, db_session_mo assert "Quality Suite" in result assert "45 passed" in result assert "3 failed" in result - assert str(job_exec_id) in result + # latest_run_id is already the job execution id — surfaced directly. + assert str(run_id) in result @patch("testgen.mcp.tools.discovery.TestSuite") @@ -181,7 +179,7 @@ def test_list_test_suites_raises_not_found_for_inaccessible_project( from testgen.mcp.tools.discovery import list_test_suites - with pytest.raises(MCPPermissionDenied, match="No test suites found for project `secret_project`"): + with pytest.raises(MCPResourceNotAccessible, match="Project `secret_project` not found or not accessible"): list_test_suites("secret_project") @@ -229,3 +227,219 @@ def test_list_tables_scopes_data_lookup_to_resolved_tg_project(mock_tg_cls, mock call_kwargs = mock_dt.select_table_names.call_args assert call_kwargs.kwargs["project_codes"] == ["proj_a"] + + +# --------------------------------------------------------------------------- +# get_project +# --------------------------------------------------------------------------- + + +def _mock_project(**overrides): + project = MagicMock() + project.project_code = overrides.get("project_code", "demo") + project.project_name = overrides.get("project_name", "Demo Project") + project.use_dq_score_weights = overrides.get("use_dq_score_weights", True) + project.data_retention_enabled = overrides.get("data_retention_enabled", True) + project.data_retention_days = overrides.get("data_retention_days", 180) + project.observability_api_url = overrides.get("observability_api_url", None) + return project + + +def _mock_project_summary(**overrides): + summary = MagicMock() + summary.connection_count = overrides.get("connection_count", 3) + summary.table_group_count = overrides.get("table_group_count", 5) + summary.test_suite_count = overrides.get("test_suite_count", 7) + summary.test_definition_count = overrides.get("test_definition_count", 142) + summary.profiling_run_count = overrides.get("profiling_run_count", 12) + summary.test_run_count = overrides.get("test_run_count", 38) + summary.can_export_to_observability = overrides.get("can_export_to_observability", False) + return summary + + +@patch("testgen.mcp.tools.discovery.Project") +def test_get_project_returns_counts_and_config(mock_project_cls, db_session_mock): + mock_project_cls.get.return_value = _mock_project() + mock_project_cls.get_summary.return_value = _mock_project_summary() + + from testgen.mcp.tools.discovery import get_project + + out = get_project("demo") + + assert "Project `demo`" in out + assert "Demo Project" in out + assert "**Connections:** 3" in out + assert "**Table groups:** 5" in out + assert "**Test suites:** 7" in out + assert "**Test definitions:** 142" in out + assert "**Profiling runs:** 12" in out + assert "**Test runs:** 38" in out + assert "**Weighted data quality scoring:** Yes" in out + assert "## Observability Integration" in out + assert "**Configured:** No" in out # default summary has can_export_to_observability=False + assert "## Data Retention" in out + assert "**Automatically delete old profiling and test history:** Yes" in out + assert "**Delete history older than (days):** 180" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_project_raises_not_found_for_inaccessible(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.discovery import get_project + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + get_project("secret") + + +@patch("testgen.mcp.tools.discovery.Project") +def test_get_project_raises_not_found_when_missing(mock_project_cls, db_session_mock): + mock_project_cls.get.return_value = None + mock_project_cls.get_summary.return_value = None + + from testgen.mcp.tools.discovery import get_project + + with pytest.raises(MCPResourceNotAccessible, match="Project `demo` not found or not accessible"): + get_project("demo") + + +# --------------------------------------------------------------------------- +# get_test_suite +# --------------------------------------------------------------------------- + + +def _mock_test_suite(**overrides): + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.test_suite = overrides.get("test_suite", "qa_checks") + suite.project_code = overrides.get("project_code", "demo") + suite.connection_id = overrides.get("connection_id", 42) + suite.table_groups_id = overrides.get("table_groups_id", uuid4()) + suite.test_suite_description = overrides.get("test_suite_description", "Daily QA") + suite.severity = overrides.get("severity", "Warning") + suite.export_to_observability = overrides.get("export_to_observability", True) + suite.is_monitor = overrides.get("is_monitor", False) + return suite + + +_DEFAULT_TYPE_COUNTS = {"Aggregate Balance": 5, "Row Count": 12} + + +def _stats(total=47, locked=3, counts_by_type=None): + from testgen.common.models.test_suite import TestDefinitionStats + + return TestDefinitionStats( + total=total, + locked=locked, + counts_by_type=_DEFAULT_TYPE_COUNTS if counts_by_type is None else counts_by_type, + ) + + +@patch("testgen.mcp.tools.discovery.TestSuite") +@patch("testgen.mcp.tools.discovery.resolve_table_group") +@patch("testgen.mcp.tools.discovery.resolve_connection") +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_returns_full_config( + mock_common_suite, mock_resolve_conn, mock_resolve_tg, mock_suite_cls, db_session_mock, +): + suite = _mock_test_suite() + mock_common_suite.get.return_value = suite + conn = MagicMock(connection_id=42, connection_name="warehouse_prod", sql_flavor_code="snowflake") + mock_resolve_conn.return_value = conn + tg = MagicMock(table_groups_name="curated_payments") + tg.id = suite.table_groups_id + mock_resolve_tg.return_value = tg + mock_suite_cls.test_definition_stats.return_value = _stats() + + from testgen.mcp.tools.discovery import get_test_suite + + out = get_test_suite(str(suite.id)) + + assert "Test Suite `qa_checks`" in out + assert f"**ID:** `{suite.id}`" in out + assert "**Project:** `demo`" in out + assert "warehouse_prod" in out + assert "Snowflake" in out + assert "curated_payments" in out + assert "**Default severity:** Warning" in out + assert "**Export to observability:** Yes" in out + assert "**Total tests:** 47" in out + assert "**Locked tests:** 3" in out + assert "Tests by type" in out + # Renders short_name labels (e.g. "Aggregate Balance"), NOT raw test_type codes (e.g. "Aggregate_Balance") + assert "Aggregate Balance" in out + assert "Row Count" in out + assert "Aggregate_Balance" not in out # the raw test_type code must NOT appear + + +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_inaccessible_id(mock_common_suite, db_session_mock): + """A genuinely missing / inaccessible id (TestSuite.get returns None) raises the unified error.""" + mock_common_suite.get.return_value = None + + from testgen.mcp.tools.discovery import get_test_suite + + bogus_id = str(uuid4()) + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_test_suite(bogus_id) + + +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_actual_monitor_suite(mock_common_suite, db_session_mock): + """An existing ``is_monitor=True`` suite is rejected because resolve_test_suite + passes ``TestSuite.is_monitor.isnot(True)`` as a filter clause. + + Simulates the real DB filter behaviour: ``TestSuite.get`` returns the monitor + suite when called without the filter clause, ``None`` when the clause is present. + """ + monitor_suite = _mock_test_suite(is_monitor=True) + + def filtered_get(_uuid, *clauses): + # The resolver's contract: pass an `is_monitor.isnot(True)` clause to TestSuite.get. + # If the clause is present, a DB query against it would not match this monitor row. + for clause in clauses: + clause_str = str(clause).lower() + if "is_monitor" in clause_str and "not" in clause_str: + return None + return monitor_suite + + mock_common_suite.get.side_effect = filtered_get + + from testgen.mcp.tools.discovery import get_test_suite + + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_test_suite(str(monitor_suite.id)) + + +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_rejects_invalid_uuid(mock_common_suite, db_session_mock): + from testgen.mcp.exceptions import MCPUserError as _MCPUserError + from testgen.mcp.tools.discovery import get_test_suite + + with pytest.raises(_MCPUserError, match="not a valid UUID"): + get_test_suite("not-a-uuid") + + +@patch("testgen.mcp.tools.discovery.TestSuite") +@patch("testgen.mcp.tools.discovery.resolve_table_group") +@patch("testgen.mcp.tools.discovery.resolve_connection") +@patch("testgen.mcp.tools.common.TestSuite") +def test_get_test_suite_no_severity_renders_inherit( + mock_common_suite, mock_resolve_conn, mock_resolve_tg, mock_suite_cls, db_session_mock, +): + suite = _mock_test_suite(severity=None, connection_id=None, table_groups_id=None) + mock_common_suite.get.return_value = suite + # connection_id / table_groups_id are None, so resolvers should not be called + mock_suite_cls.test_definition_stats.return_value = _stats(total=0, locked=0, counts_by_type={}) + + from testgen.mcp.tools.discovery import get_test_suite + + out = get_test_suite(str(suite.id)) + + assert "Inherit from test type" in out + # No "Tests by type" table when there are no test definitions + assert "Tests by type" not in out + mock_resolve_conn.assert_not_called() + mock_resolve_tg.assert_not_called() diff --git a/tests/unit/mcp/test_tools_hygiene_issues.py b/tests/unit/mcp/test_tools_hygiene_issues.py index 741c9b4b..343aab73 100644 --- a/tests/unit/mcp/test_tools_hygiene_issues.py +++ b/tests/unit/mcp/test_tools_hygiene_issues.py @@ -252,7 +252,7 @@ def test_resolve_je_id_table_group_no_completed_runs(mock_latest, mock_resolve_t @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") def test_resolve_je_id_je_branch_unknown_run(mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import _resolve_profile_run_je_id @@ -267,7 +267,7 @@ def test_resolve_je_id_je_branch_unknown_run(mock_get, mock_tg_cls, db_session_m @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") def test_resolve_je_id_je_branch_inaccessible_tg(mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import _resolve_profile_run_je_id @@ -309,7 +309,7 @@ def test_list_hygiene_issues_invalid_je_uuid(db_session_mock): @patch("testgen.mcp.tools.hygiene_issues.TableGroup") -@patch.object(ProfilingRun, "get_by_id_or_job") +@patch.object(ProfilingRun, "get") @patch.object(HygieneIssue, "list_for_run") def test_list_hygiene_issues_resolves_via_je_id(mock_list, mock_get, mock_tg_cls, db_session_mock): from testgen.mcp.tools.hygiene_issues import list_hygiene_issues @@ -321,7 +321,8 @@ def test_list_hygiene_issues_resolves_via_je_id(mock_list, mock_get, mock_tg_cls list_hygiene_issues(job_execution_id=str(uuid4())) - assert mock_list.call_args.args[0] == run.job_execution_id + # The run id is the job execution id — list_for_run is called with run.id. + assert mock_list.call_args.args[0] == run.id @patch("testgen.mcp.tools.hygiene_issues.resolve_table_group") @@ -727,22 +728,35 @@ def test_update_hygiene_issue_invalid_disposition(db_session_mock, disposition_p update_hygiene_issue(issue_id=str(uuid4()), disposition="Bogus") -@patch.object(HygieneIssue, "update_disposition") -def test_update_hygiene_issue_muted_maps_to_inactive(mock_update, db_session_mock, disposition_perms): +@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue") +def test_update_hygiene_issue_muted_maps_to_inactive(mock_resolve, db_session_mock, disposition_perms): from testgen.mcp.tools.hygiene_issues import update_hygiene_issue - mock_update.return_value = True + issue = MagicMock() + mock_resolve.return_value = issue update_hygiene_issue(issue_id=str(uuid4()), disposition="Muted") - args = mock_update.call_args.args - assert args[1] == Disposition.INACTIVE + assert issue.disposition == Disposition.INACTIVE -@patch.object(HygieneIssue, "update_disposition") -def test_update_hygiene_issue_returns_success_markdown(mock_update, db_session_mock, disposition_perms): +@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue") +def test_update_hygiene_issue_sets_disposition_on_resolved_issue( + mock_resolve, db_session_mock, disposition_perms, +): + from testgen.mcp.tools.hygiene_issues import update_hygiene_issue + + issue = MagicMock() + mock_resolve.return_value = issue + update_hygiene_issue(issue_id=str(uuid4()), disposition="Dismissed") + + assert issue.disposition == Disposition.DISMISSED + + +@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue") +def test_update_hygiene_issue_returns_success_markdown(mock_resolve, db_session_mock, disposition_perms): from testgen.mcp.tools.hygiene_issues import update_hygiene_issue - mock_update.return_value = True + mock_resolve.return_value = MagicMock() issue_id = str(uuid4()) result = update_hygiene_issue(issue_id=issue_id, disposition="Dismissed") @@ -751,30 +765,30 @@ def test_update_hygiene_issue_returns_success_markdown(mock_update, db_session_m assert "Dismissed" in result -@patch.object(HygieneIssue, "update_disposition") -def test_update_hygiene_issue_not_updated_collapses_to_not_accessible( - mock_update, db_session_mock, disposition_perms, +@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue") +def test_update_hygiene_issue_not_accessible_propagates( + mock_resolve, db_session_mock, disposition_perms, ): from testgen.mcp.tools.hygiene_issues import update_hygiene_issue - mock_update.return_value = False + mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x") with pytest.raises(MCPResourceNotAccessible): update_hygiene_issue(issue_id=str(uuid4()), disposition="Confirmed") -@patch.object(HygieneIssue, "update_disposition") -def test_update_hygiene_issue_passes_project_scope_clause( - mock_update, db_session_mock, disposition_perms, +@patch("testgen.mcp.tools.hygiene_issues.resolve_hygiene_issue") +def test_update_hygiene_issue_delegates_scope_to_resolver( + mock_resolve, db_session_mock, disposition_perms, ): + """Project scoping lives in resolve_hygiene_issue (covered in test_tools_common); + the tool must route the issue id through it.""" from testgen.mcp.tools.hygiene_issues import update_hygiene_issue - mock_update.return_value = True - update_hygiene_issue(issue_id=str(uuid4()), disposition="Confirmed") + mock_resolve.return_value = MagicMock() + issue_id = str(uuid4()) + update_hygiene_issue(issue_id=issue_id, disposition="Confirmed") - # Trailing args after (issue_uuid, db_disposition) are the *clauses - clauses = mock_update.call_args.args[2:] - sql = "\n".join(str(c.compile(dialect=postgresql.dialect())) for c in clauses) - assert "profile_anomaly_results.project_code IN" in sql + mock_resolve.assert_called_once_with(issue_id) def test_update_hygiene_issue_uses_disposition_permission(): diff --git a/tests/unit/mcp/test_tools_monitors.py b/tests/unit/mcp/test_tools_monitors.py new file mode 100644 index 00000000..591bd74f --- /dev/null +++ b/tests/unit/mcp/test_tools_monitors.py @@ -0,0 +1,1543 @@ +"""Tests for the MCP monitor tools — read (``get_monitor_summary`` / ``list_monitored_tables`` / +``list_monitor_events`` / ``list_monitors`` / ``list_monitor_schema_changes``) and +lifecycle/settings (``enable_monitors`` / ``get_monitor_settings`` / ``update_monitor_settings`` +/ ``disable_monitors``).""" + +from datetime import UTC, datetime +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.common.models.data_structure_log import DataStructureLog +from testgen.common.models.table_group import MonitorGroupSummary, MonitorTableSummary +from testgen.common.models.test_suite import PredictSensitivity +from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.monitors" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="view"): + memberships = memberships or dict.fromkeys(allowed, "role_a") + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_table_group(**overrides) -> MagicMock: + tg = MagicMock() + tg.id = overrides.get("id", uuid4()) + tg.project_code = overrides.get("project_code", "demo") + tg.table_groups_name = overrides.get("table_groups_name", "Sales") + tg.monitor_test_suite_id = overrides.get("monitor_test_suite_id", uuid4()) + return tg + + +def _mock_monitor_suite(**overrides) -> MagicMock: + from testgen.common.models.test_suite import PredictSensitivity + + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.is_monitor = True + suite.monitor_lookback = overrides.get("monitor_lookback", 7) + suite.predict_sensitivity = overrides.get("predict_sensitivity", PredictSensitivity.medium) + return suite + + +def _group_summary(**overrides) -> MonitorGroupSummary: + defaults: dict = { + "lookback": 7, + "lookback_start": datetime(2026, 5, 25, 12, 0, tzinfo=UTC), + "lookback_end": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "total_monitored_tables": 12, + "freshness_anomalies": 2, + "volume_anomalies": 0, + "schema_anomalies": 1, + "metric_anomalies": 0, + "freshness_has_errors": False, + "volume_has_errors": False, + "schema_has_errors": False, + "metric_has_errors": False, + "freshness_is_training": False, + "volume_is_training": False, + "metric_is_training": True, + "freshness_is_pending": False, + "volume_is_pending": False, + "schema_is_pending": False, + "metric_is_pending": False, + } + defaults.update(overrides) + return MonitorGroupSummary(**defaults) + + +def _table_summary(**overrides) -> MonitorTableSummary: + defaults: dict = { + "table_name": "orders", + "lookback": 7, + "lookback_start": datetime(2026, 5, 25, 12, 0, tzinfo=UTC), + "lookback_end": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "freshness_anomalies": 2, + "volume_anomalies": 0, + "schema_anomalies": 0, + "metric_anomalies": 0, + "freshness_is_training": False, + "volume_is_training": False, + "metric_is_training": None, + "freshness_is_pending": False, + "volume_is_pending": False, + "schema_is_pending": False, + "metric_is_pending": True, + "freshness_error_message": None, + "volume_error_message": None, + "schema_error_message": None, + "metric_error_message": None, + "latest_update": datetime(2026, 6, 1, 8, 30, tzinfo=UTC), + "row_count": 1_234_567, + "previous_row_count": 1_200_000, + "column_adds": 0, + "column_drops": 0, + "column_mods": 0, + "table_state": None, + } + defaults.update(overrides) + return MonitorTableSummary(**defaults) + + +# --------------------------------------------------------------------------- +# get_monitor_summary +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_happy_path(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary() + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "# Monitor summary for `Sales`" in out + assert "**Project:** `demo`" in out + assert "**Monitored tables:** 12" in out + assert "**Lookback:** 7 runs" in out + assert "(override)" not in out + assert "**Next scheduled run:** not scheduled" in out + # Per-type rows + assert "| Freshness | 2 | Ok |" in out + assert "| Volume | 0 | Ok |" in out + assert "| Schema | 1 | Ok |" in out + assert "| Metric | 0 | Training |" in out + + +@patch(f"{MODULE}.next_scheduled_run", return_value=datetime(2026, 6, 2, 18, 0, tzinfo=UTC)) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_renders_next_scheduled(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary() + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "Next scheduled run" in out + assert "2026-06-02" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert out == "This table group is not monitored." + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_inaccessible(mock_resolve, db_session_mock): + bad_id = str(uuid4()) + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", bad_id) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible): + get_monitor_summary(bad_id) + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_lookback_override_applied(mock_resolve, mock_tg_cls, mock_next, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary(lookback=14) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id), lookback=14) + + mock_tg_cls.get_monitor_group_summary.assert_called_once_with(tg.id, lookback_override=14) + assert "**Lookback:** 14 runs (override)" in out + + +@pytest.mark.parametrize("bad_lookback", [0, 366, 500, -1]) +def test_get_monitor_summary_lookback_out_of_range(bad_lookback, db_session_mock): + """Both bounds (and beyond) are pinned — the model accepts 1..365 inclusive.""" + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + get_monitor_summary(str(uuid4()), lookback=bad_lookback) + assert "between 1 and 365" in str(exc.value) + assert f"`{bad_lookback}`" in str(exc.value) + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_empty_state_lookback_zero( + mock_resolve, mock_tg_cls, mock_next, db_session_mock, +): + """When a monitor suite is configured but no runs have happened yet, the model + method returns ``lookback=0`` (preserves the pre-refactor signal the dashboard + uses to render "No monitor runs yet"). The MCP output reflects the empty state + rather than fabricating a one-run window.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary( + lookback=0, + lookback_start=None, + lookback_end=None, + total_monitored_tables=0, + freshness_anomalies=0, volume_anomalies=0, + schema_anomalies=0, metric_anomalies=0, + freshness_is_pending=True, volume_is_pending=True, + schema_is_pending=True, metric_is_pending=True, + freshness_is_training=False, volume_is_training=False, + metric_is_training=False, + ) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "**Lookback:** 0 runs" in out, "must show 0, not a fabricated default like 1" + # Window start / end fields are absent in the empty case (None values render as em-dash + # at most, but the tool omits them when the value is falsy) + assert "**Window start:**" not in out + assert "**Window end:**" not in out + # All per-type status cells reflect "no results" + assert out.count("No results yet or not configured") == 4 + + +@patch(f"{MODULE}.next_scheduled_run", return_value=None) +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_summary_renders_error_and_pending_states( + mock_resolve, mock_tg_cls, mock_next, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.get_monitor_group_summary.return_value = _group_summary( + freshness_has_errors=True, + volume_is_pending=True, + schema_is_pending=True, + metric_is_pending=True, + ) + + from testgen.mcp.tools.monitors import get_monitor_summary + + with _patch_perms(): + out = get_monitor_summary(str(tg.id)) + + assert "| Freshness | 2 | Error |" in out + assert "| Volume | 0 | No results yet or not configured |" in out + assert "| Schema | 1 | No results yet or not configured |" in out + assert "| Metric | 0 | No results yet or not configured |" in out + + +# --------------------------------------------------------------------------- +# list_monitored_tables +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_happy_path(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary(table_name="orders", freshness_anomalies=2), + _table_summary(table_name="customers", freshness_anomalies=0, row_count=42_000), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert "# Monitored tables in `Sales`" in out + assert "Showing 1–2 of 2" in out # noqa: RUF001 — page-info formatter uses EN DASH + assert "`orders`" in out + assert "`customers`" in out + # "Row count change" column renders the signed delta (row_count - previous_row_count), + # not the raw current count. Defaults are row_count=1_234_567, previous=1_200_000. + assert "+34,567" in out + # customers overrides row_count to 42_000 — delta vs default previous of 1_200_000 is -1,158,000. + assert "-1,158,000" in out + # Absolute current count should not appear in the column. + assert "1,234,567" not in out + # Default sort_by is None → table_name asc + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=None, sort_by=None, page=1, limit=20, + ) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_anomaly_type_filter_translated( + mock_resolve, mock_tg_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + list_monitored_tables(str(tg.id), anomaly_type="freshness") + + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=["Freshness_Trend"], sort_by=None, page=1, limit=20, + ) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_sort_by_translated(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + list_monitored_tables(str(tg.id), sort_by="anomaly_count_desc") + + mock_tg_cls.list_monitor_table_summaries.assert_called_once_with( + tg.id, anomaly_types=None, sort_by="total_anomalies_desc", page=1, limit=20, + ) + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert out == "This table group is not monitored." + + +def test_list_monitored_tables_invalid_anomaly_type(db_session_mock): + """Error message must reference the caller's public arg name (``anomaly_type``) + rather than the helper's internal arg name (``monitor_type``).""" + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitored_tables(str(uuid4()), anomaly_type="bogus") + msg = str(exc.value) + assert "Invalid anomaly_type" in msg + assert "Invalid monitor_type" not in msg + + +def test_list_monitored_tables_invalid_sort_by(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitored_tables(str(uuid4()), sort_by="wat") + assert "Invalid sort_by" in str(exc.value) + + +def test_list_monitored_tables_invalid_page(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitored_tables(str(uuid4()), page=0) + + +def test_list_monitored_tables_limit_out_of_range(db_session_mock): + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitored_tables(str(uuid4()), limit=500) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_schema_change_column( + mock_resolve, mock_tg_cls, db_session_mock, +): + """The Schema column shows just the anomaly count; the new Schema change column + renders the verbose description (added with N columns / modified breakdown / + dropped with N columns).""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t_mod", schema_anomalies=2, table_state="modified", + column_adds=1, column_drops=2, column_mods=0, + ), + _table_summary( + table_name="t_add", schema_anomalies=1, table_state="added", + column_adds=5, column_drops=0, column_mods=0, + ), + _table_summary( + table_name="t_drop", schema_anomalies=1, table_state="dropped", + column_adds=0, column_drops=10, column_mods=0, + ), + _table_summary(table_name="t_quiet", schema_anomalies=0, table_state=None), + ], + 4, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + # New column appears in the header + assert "Schema change" in out + # Verbose strings appear in their respective rows + assert "Table added with 5 columns." in out + assert "Table dropped with 10 columns." in out + assert "1 column added. 2 columns dropped." in out + # No more parenthetical states on the Schema column + assert "(columns)" not in out + assert "(added)" not in out + assert "(dropped)" not in out + # Quiet row's Schema column is the raw count, Schema change is em-dash + quiet_row = next(line for line in out.splitlines() if "`t_quiet`" in line) + assert quiet_row.count(" 0 ") >= 1 # Schema = 0 for quiet + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_training_and_pending_cells( + mock_resolve, mock_tg_cls, db_session_mock, +): + """Status words render when the per-type count is zero.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t1", + freshness_anomalies=0, freshness_is_training=True, + volume_anomalies=0, volume_is_pending=True, + metric_anomalies=0, metric_error_message="boom", + ), + ], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + row = next(line for line in out.splitlines() if "`t1`" in line) + assert "Training" in row + assert "Pending" in row + assert "Error" in row + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_count_wins_over_training_and_pending( + mock_resolve, mock_tg_cls, db_session_mock, +): + """A positive anomaly count must surface even when the monitor is in training + or pending state — otherwise the cell hides the value that made the row match + an ``anomaly_type`` filter, and the table doesn't agree with itself. + + Precedence is error > positive count > pending > training > zero. Errors still + win (the latest measurement is suspect, so the historic count is misleading). + """ + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ( + [ + _table_summary( + table_name="t_busy_during_learn", + freshness_anomalies=5, freshness_is_training=True, + volume_anomalies=3, volume_is_pending=True, + metric_anomalies=2, metric_is_training=True, + ), + _table_summary( + table_name="t_error_with_count", + freshness_anomalies=4, freshness_error_message="db down", + ), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + # Row 1: counts visible despite training/pending + row_count = next(line for line in out.splitlines() if "`t_busy_during_learn`" in line) + assert " 5 " in row_count, "freshness count should render despite is_training" + assert " 3 " in row_count, "volume count should render despite is_pending" + assert " 2 " in row_count, "metric count should render despite is_training" + assert "Training" not in row_count + assert "Pending" not in row_count + # Row 2: error still wins over count + row_error = next(line for line in out.splitlines() if "`t_error_with_count`" in line) + assert "Error" in row_error + assert " 4 " not in row_error, "error must win over count (measurement is suspect)" + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_empty(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id)) + + assert "_No monitored tables match this filter._" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitored_tables_empty_beyond_last_page(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tg_cls.list_monitor_table_summaries.return_value = ([], 7) + + from testgen.mcp.tools.monitors import list_monitored_tables + + with _patch_perms(): + out = list_monitored_tables(str(tg.id), page=99) + + assert "No tables on page 99 (total: 7)." in out + + +# --------------------------------------------------------------------------- +# Lifecycle & settings helpers +# --------------------------------------------------------------------------- + + +def _mock_schedule(**overrides) -> MagicMock: + sched = MagicMock() + sched.id = overrides.get("id", uuid4()) + sched.cron_expr = overrides.get("cron_expr", "0 6 * * *") + sched.cron_tz = overrides.get("cron_tz", "UTC") + sched.active = overrides.get("active", True) + sched.get_sample_triggering_timestamps.return_value = [ + overrides.get("next_run", datetime(2026, 6, 10, 6, 0, tzinfo=UTC)) + ] + return sched + + +def _settings_suite(**overrides) -> MagicMock: + suite = _mock_monitor_suite(**overrides) + suite.predict_sensitivity = overrides.get("predict_sensitivity", PredictSensitivity.medium) + suite.monitor_lookback = overrides.get("monitor_lookback", 14) + suite.predict_min_lookback = overrides.get("predict_min_lookback", 30) + suite.predict_exclude_weekends = overrides.get("predict_exclude_weekends", False) + suite.monitor_regenerate_freshness = overrides.get("monitor_regenerate_freshness", True) + suite.holiday_codes_list = overrides.get("holiday_codes_list", None) + return suite + + +# --------------------------------------------------------------------------- +# enable_monitors +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_happy_path(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + mock_enable.return_value = (MagicMock(), 4) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"): + out = enable_monitors(str(tg.id), "0 6 * * *", "America/New_York") + + assert "# Monitoring enabled for `Sales`" in out + assert "**Initial monitors created:** 4" in out + assert "**Cron expression:** `0 6 * * *`" in out + assert "America/New_York" in out + mock_enable.assert_called_once_with(tg, "0 6 * * *", "America/New_York") + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_already_enabled(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="already enabled"): + enable_monitors(str(tg.id), "0 6 * * *") + + mock_enable.assert_not_called() + + +@patch(f"{MODULE}.enable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_enable_monitors_invalid_cron_rejected_before_side_effects(mock_resolve, mock_enable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import enable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError): + enable_monitors(str(tg.id), "not a cron") + + mock_enable.assert_not_called() + + +# --------------------------------------------------------------------------- +# get_monitor_settings +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_settings_happy_path(mock_resolve, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + suite = _settings_suite( + predict_sensitivity=PredictSensitivity.high, monitor_lookback=50, holiday_codes_list=["US", "NYSE"] + ) + mock_resolve.return_value = (tg, suite) + mock_js.get_for_monitor_suite.return_value = _mock_schedule(cron_expr="0 */12 * * *", cron_tz="UTC", active=True) + + from testgen.mcp.tools.monitors import get_monitor_settings + + with _patch_perms(): + out = get_monitor_settings(str(tg.id)) + + assert "# Monitor settings for `Sales`" in out + assert "**Sensitivity:** high" in out + assert "**Lookback runs:** 50" in out + assert "**Holiday codes:** US, NYSE" in out + assert "**Regenerate freshness:** Yes" in out + assert "## Schedule" in out + assert "**Cron expression:** `0 */12 * * *`" in out + assert "**Status:** Active" in out + assert "Next run" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_get_monitor_settings_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import get_monitor_settings + + with _patch_perms(): + out = get_monitor_settings(str(tg.id)) +# list_monitor_events (TG-1092) +# --------------------------------------------------------------------------- + + +def _monitor_event(test_type="Volume_Trend", **overrides): + from testgen.common.models.test_result import MonitorEvent + + defaults: dict = { + "monitor_id": uuid4(), + "test_type": test_type, + "test_time": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "is_anomaly": False, + "is_training": False, + "is_pending": False, + "is_error": False, + "message": None, + "signal": "1000", + "lower_bound": "900", + "upper_bound": "1100", + "schema_change_kind": None, + "column_adds": None, + "column_drops": None, + "column_mods": None, + "metric_name": None, + } + defaults.update(overrides) + return MonitorEvent(**defaults) + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_happy_volume(mock_resolve, mock_tr_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event(signal="1000", is_anomaly=True), + _monitor_event(signal="800", is_training=True), + ], + 2, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume") + + assert "Monitor events: `orders` — `volume` in `Sales`" in out + assert "| Time | Status | Row count | Lower bound | Upper bound |" in out + assert "Anomaly" in out + assert "Training" in out + # Volume rendering must not pull in Metric's name header + assert "Metric name" not in out + mock_tr_cls.list_monitor_events_for_table.assert_called_once_with( + mock_resolve.return_value[1].id, + "orders", + monitor_type="Volume_Trend", + page=1, + limit=20, + ) + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_parses_message_into_columns( + mock_resolve, mock_tr_cls, db_session_mock, +): + """Freshness events carry a structured message from the SQL template: + ``"Table update detected: {Yes|No}[. {Detail}]"``. Render it as two + distinct columns so the LLM consumer doesn't have to re-parse the text.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. On time."), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: No. Late.", is_anomaly=True), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. Later than expected.", is_anomaly=True), + _monitor_event(test_type="Freshness_Trend", message="Table update detected: No"), + ], + 4, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "freshness") + + assert "| Time | Status | Update detected | Detail |" in out + # Freshness rendering should not surface raw bound columns (they map to + # the static-mode upper_tolerance, which is configuration, not per-event data). + assert "Lower bound" not in out + assert "Row count" not in out + # Each row's structured fields land in their own columns. + assert "| Yes | On time |" in out + assert "| No | Late |" in out + assert "| Yes | Later than expected |" in out + # Missing detail renders as em-dash. + assert "| No | — |" in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_renders_name_in_heading( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """For Metric monitors the metric name belongs in the heading (one metric + per call — see the ``monitor_id`` requirement). The data table itself + only carries Time / Status / Value / bounds.""" + monitor_id = str(uuid4()) + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_metric_monitor_events.return_value = ( + [_monitor_event(test_type="Metric_Trend", metric_name="total_amount")], + 1, + ) + # Scoped lookup returns the metric's summary (select_where, not the bare get). + mock_td_cls.select_where.return_value = [ + SimpleNamespace(test_type="Metric_Trend", column_name="total_amount"), + ] + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "metric", monitor_id=monitor_id) + + assert "Monitor events: Metric `total_amount` on `orders` in `Sales`" in out + assert "| Time | Status | Value | Lower bound | Upper bound |" in out + # metric_name lives in the heading, not as a column + assert "Metric name" not in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_requires_monitor_id(mock_resolve, db_session_mock): + """Metric is the only multi-instance monitor type — without ``monitor_id`` + the query would interleave every metric on the table.""" + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="monitor_id.*required.*metric"): + list_monitor_events(str(tg.id), "orders", "metric") + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_monitor_id_rejected_for_non_metric(mock_resolve, db_session_mock): + """Singleton monitor types (freshness/volume/schema) are uniquely + identified by table + type; ``monitor_id`` is meaningless for them and + must be rejected so the caller gets a clear discovery path.""" + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="monitor_id.*only applies.*metric"): + list_monitor_events(str(tg.id), "orders", "volume", monitor_id=str(uuid4())) + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_invalid_monitor_id(mock_resolve, db_session_mock): + """A malformed ``monitor_id`` is rejected with a clean MCPUserError before + it reaches the query — not surfaced as a raw Postgres UUID-cast error.""" + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + with _patch_perms(), pytest.raises(MCPUserError, match="Invalid monitor_id"): + list_monitor_events(str(tg.id), "orders", "metric", monitor_id="not-a-uuid") + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_metric_monitor_id_scoped_and_not_found( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """The metric lookup is scoped to this suite + table + Metric_Trend, so a + ``monitor_id`` that doesn't resolve within scope (wrong table, wrong suite, + or not a metric) yields a discovery-hint error rather than another table's + events under this table's heading.""" + from testgen.mcp.tools.monitors import list_monitor_events + + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_td_cls.select_where.return_value = [] # no metric matches the scoped clauses + + with _patch_perms(), pytest.raises(MCPUserError, match=r"No metric monitor.*list_monitors"): + list_monitor_events(str(tg.id), "orders", "metric", monitor_id=str(uuid4())) + + # The events query never runs when the scoped definition lookup fails. + mock_tr_cls.list_metric_monitor_events.assert_not_called() + # Lookup was scoped (id + suite + table + type clauses), not a bare PK fetch. + assert len(mock_td_cls.select_where.call_args[0]) == 4 + + +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_schema_uses_dedicated_columns(mock_resolve, mock_tr_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [ + _monitor_event( + test_type="Schema_Drift", + signal="A|3|1|0|2026-05-01T00:00:00", + schema_change_kind="A", + column_adds=3, column_drops=1, column_mods=0, + is_anomaly=True, + ), + ], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "schema") + + assert "| Time | Status | Table change | Columns added | Columns dropped | Columns modified |" in out + assert "added" in out + # Internal codes never leak + assert "Schema_Drift" not in out + assert "| A |" not in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_forecast_renders_separate_section( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """``include_predictions=True`` against a Prediction-Model monitor appends + a ``## Forecast`` section listing forecast points (future timestamps with + predicted bounds). Predictions never bleed into the historical events + table — the LLM consumer can distinguish observed from predicted at a + glance.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event()], + 1, + ) + # Build a prediction JSONB the helper will read at render time. Keys are + # epoch-ms strings — same format the dashboard uses. + def _epoch_ms(ts): + return str(int(ts.timestamp() * 1000)) + + forecast_t1 = datetime(2026, 6, 22, 12, 0, tzinfo=UTC) + forecast_t2 = datetime(2026, 6, 29, 12, 0, tzinfo=UTC) + monitor_def = MagicMock() + monitor_def.history_calculation = "PREDICT" + monitor_def.prediction = { + "lower_tolerance|medium": {_epoch_ms(forecast_t1): 500.4, _epoch_ms(forecast_t2): 500.3}, + "upper_tolerance|medium": {_epoch_ms(forecast_t1): 503.6, _epoch_ms(forecast_t2): 503.7}, + } + mock_td_cls.get_singleton_monitor.return_value = monitor_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + # Forecast section is separate from the historical table + assert "## Forecast" in out + assert "**Sensitivity:** medium" in out + assert "| Time | Predicted lower | Predicted upper |" in out + assert "500.4" in out and "503.6" in out + # Singleton lookup ran for the non-metric forecast path + mock_td_cls.get_singleton_monitor.assert_called_once() + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_predictions_unavailable_note( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """include_predictions=True against a non-Prediction-Model monitor surfaces + a note under the forecast heading explaining why no forecast follows. The + historical events table itself stays clean.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ([_monitor_event()], 1) + monitor_def = MagicMock() + monitor_def.history_calculation = None # Static or Historical, not Prediction + mock_td_cls.get_singleton_monitor.return_value = monitor_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + assert "Predictions not available" in out + assert "Static or Historical Calculation" in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_predictions_not_applicable_for_schema( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event(test_type="Schema_Drift", schema_change_kind="M", column_mods=1)], + 1, + ) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "schema", include_predictions=True) + + assert "## Forecast" in out + assert "Predictions not applicable to Schema monitors" in out + # Schema short-circuits before any TestDefinition lookup — predictions + # never apply, so don't waste a DB round-trip. + mock_td_cls.get_singleton_monitor.assert_not_called() + + +@patch(f"{MODULE}.next_update_window") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_prediction_shows_window( + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_window, db_session_mock, +): + """A Freshness monitor in Prediction Model mode forecasts a next-update time + window (the same one the dashboard computes), not a value band — so the + forecast section presents that window rather than a band table or a note.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ( + [_monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes. On time.")], + 1, + ) + monitor_def = MagicMock() + monitor_def.history_calculation = "PREDICT" + mock_td_cls.get_singleton_monitor.return_value = monitor_def + mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") + start_ms = int(datetime(2026, 7, 1, 9, 0, tzinfo=UTC).timestamp() * 1000) + end_ms = int(datetime(2026, 7, 1, 17, 0, tzinfo=UTC).timestamp() * 1000) + mock_window.return_value = {"start": start_ms, "end": end_ms} + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "freshness", include_predictions=True) + + assert "## Forecast" in out + assert "Next update expected" in out + assert "2026-07-01 09:00 to 2026-07-01 17:00 UTC" in out + # A window, not a value-band table. + assert "| Time | Predicted lower | Predicted upper |" not in out + + +@patch(f"{MODULE}.next_update_window") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_coupled_volume_shows_gated_band( + mock_resolve, mock_tr_cls, mock_td_cls, mock_sched, mock_window, db_session_mock, +): + """A Volume/Metric monitor coupled to a Freshness monitor holds at its + baseline until the next expected refresh — the forecast renders that coupled + band (the UI's), and the internal coupling is never named in the output.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + # First call: the volume monitor's own events; second: the freshness events + # the window is derived from. + mock_tr_cls.list_monitor_events_for_table.side_effect = [ + ([_monitor_event()], 1), + ([_monitor_event(test_type="Freshness_Trend", message="Table update detected: Yes.")], 1), + ] + volume_def = MagicMock() + volume_def.history_calculation = "PREDICT" + end_ms = int(datetime(2026, 6, 2, 12, 0, tzinfo=UTC).timestamp() * 1000) + volume_def.prediction = {"freshness_gated": True, "baseline_value": 500.0, "mean": {str(end_ms): 480.0}} + volume_def.lower_tolerance = "450" + volume_def.upper_tolerance = "550" + freshness_def = MagicMock() + # list_monitor_events looks up the volume singleton; _compute_forecast then + # looks up the table's freshness definition for the window. + mock_td_cls.get_singleton_monitor.side_effect = [volume_def, freshness_def] + mock_sched.get_for_monitor_suite.return_value = SimpleNamespace(cron_tz="UTC") + start_ms = int(datetime(2026, 6, 1, 18, 0, tzinfo=UTC).timestamp() * 1000) + mock_window.return_value = {"start": start_ms, "end": end_ms} + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + # The coupled band IS rendered (the step to the next-refresh tolerances). + assert "| Time | Predicted lower | Predicted upper |" in out + assert "450" in out and "550" in out + # Internal coupling terminology must never reach the client. + assert "gated" not in out.lower() + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.TestResult") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_freshness_coupled_volume_without_tolerances_shows_note( + mock_resolve, mock_tr_cls, mock_td_cls, db_session_mock, +): + """A coupled Volume/Metric monitor with no configured tolerance has no band on + the dashboard either, so the MCP shows a note (and skips the window queries), + rather than diverging by computing a band the UI never plots.""" + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_tr_cls.list_monitor_events_for_table.return_value = ([_monitor_event()], 1) + volume_def = MagicMock() + volume_def.history_calculation = "PREDICT" + volume_def.prediction = {"freshness_gated": True, "baseline_value": 500.0} + volume_def.lower_tolerance = None + volume_def.upper_tolerance = None + mock_td_cls.get_singleton_monitor.return_value = volume_def + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume", include_predictions=True) + + assert "## Forecast" in out + assert "No forecast available for this monitor right now" in out + # No band, and no terminology leak. + assert "| Time | Predicted lower | Predicted upper |" not in out + assert "gated" not in out.lower() + # The freshness window queries are skipped when there's no tolerance to band. + assert mock_td_cls.get_singleton_monitor.call_count == 1 # only the volume singleton lookup + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_events_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(): + out = list_monitor_events(str(tg.id), "orders", "volume") + + assert out == "This table group is not monitored." + + +# --------------------------------------------------------------------------- +# update_monitor_settings +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_partial_maps_args(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite(predict_sensitivity=PredictSensitivity.high)) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + out = update_monitor_settings(str(tg.id), sensitivity="high", lookback_runs=50, exclude_weekends=True) + + assert "# Monitor settings updated for `Sales`" in out + mock_update.assert_called_once() + suite_attrs = mock_update.call_args.kwargs["suite_attrs"] + assert suite_attrs["predict_sensitivity"] == PredictSensitivity.high + assert suite_attrs["monitor_lookback"] == 50 + assert suite_attrs["predict_exclude_weekends"] is True + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + out = update_monitor_settings(str(tg.id), sensitivity="high") + + assert out == "This table group is not monitored." + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_no_fields(mock_resolve, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="No fields supplied"): + update_monitor_settings(str(tg.id)) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_invalid_sensitivity(mock_resolve, mock_js, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="Invalid sensitivity"): + update_monitor_settings(str(tg.id), sensitivity="extreme") + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_lookback_out_of_range(mock_resolve, mock_js, mock_update, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="between 1 and 200"): + update_monitor_settings(str(tg.id), lookback_runs=5000) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_unknown_holiday_code(mock_resolve, mock_js, mock_update, db_session_mock): + # ``is_supported_holiday_code`` runs for real here — ``US_FEDERAL`` is not a valid calendar. + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="Unknown holiday codes"): + update_monitor_settings(str(tg.id), holiday_codes=["US_FEDERAL"]) + + mock_update.assert_not_called() + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_holiday_codes_serialized(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + update_monitor_settings(str(tg.id), holiday_codes=["US", "NYSE"]) + + assert mock_update.call_args.kwargs["suite_attrs"]["predict_holiday_codes"] == "US,NYSE" + + +@patch(f"{MODULE}._last_monitor_run", return_value=None) +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.update_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_update_monitor_settings_empty_holiday_codes_clears(mock_resolve, mock_update, mock_js, mock_last, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _settings_suite()) + mock_js.get_for_monitor_suite.return_value = _mock_schedule() + + from testgen.mcp.tools.monitors import update_monitor_settings + + with _patch_perms(permission="edit"): + update_monitor_settings(str(tg.id), holiday_codes=[]) + + assert mock_update.call_args.kwargs["suite_attrs"]["predict_holiday_codes"] is None + + +# --------------------------------------------------------------------------- +# disable_monitors +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.disable_monitoring", return_value={"monitors": 4, "events": 2, "runs": 1}) +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_disable_monitors_happy_path(mock_resolve, mock_disable, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + + from testgen.mcp.tools.monitors import disable_monitors + + with _patch_perms(permission="edit"): + out = disable_monitors(str(tg.id)) + + assert "# Monitoring disabled for `Sales`" in out + assert "**Monitors removed:** 4" in out + assert "**Events removed:** 2" in out + assert "**Runs removed:** 1" in out + mock_disable.assert_called_once_with(suite) + + +@patch(f"{MODULE}.disable_monitoring") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_disable_monitors_not_enabled(mock_resolve, mock_disable, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import disable_monitors + + with _patch_perms(permission="edit"), pytest.raises(MCPUserError, match="not enabled"): + disable_monitors(str(tg.id)) + + mock_disable.assert_not_called() +def test_list_monitor_events_invalid_monitor_type(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_events + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + list_monitor_events(str(uuid4()), "orders", "metrics") + assert "Invalid monitor_type" in str(exc.value) + + +# --------------------------------------------------------------------------- +# list_monitors (TG-1092) +# --------------------------------------------------------------------------- + + +def _monitor_config(**overrides): + from testgen.common.models.test_definition import ( + MonitorConfig, + ThresholdMode, + ) + + defaults: dict = { + "monitor_id": uuid4(), + "test_type": "Volume_Trend", + "table_name": "orders", + "metric_name": None, + "threshold_mode": ThresholdMode.STATIC, + "threshold_lower": "900", + "threshold_upper": "1100", + "custom_query": None, + } + defaults.update(overrides) + return MonitorConfig(**defaults) + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_happy_path(mock_resolve, mock_td_cls, db_session_mock): + tg = _mock_table_group() + suite = _mock_monitor_suite() + mock_resolve.return_value = (tg, suite) + mock_td_cls.list_monitor_configs_for_table.return_value = [ + _monitor_config(test_type="Freshness_Trend", threshold_lower=None, threshold_upper="60"), + _monitor_config( + test_type="Volume_Trend", + threshold_mode="Prediction Model", + threshold_lower=None, threshold_upper=None, + ), + _monitor_config( + test_type="Metric_Trend", + metric_name="total_amount", + custom_query="SUM(total_amount)", + ), + ] + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert "Monitors on `orders` in `Sales`" in out + assert "**Prediction model sensitivity:** medium" in out + # Type column uses Title Case (reuses _MONITOR_LABEL — no duplicate + # lowercase dict needed; parse_monitor_type accepts either casing). + assert "Freshness" in out + assert "Volume" in out + assert "Metric" in out + # Metric expression column only populated for Metric + assert "SUM(total_amount)" in out + # Configuration tool must not surface runtime prediction bands + assert "Prediction bands" not in out + # Internal codes never leak + assert "Volume_Trend" not in out + assert "Freshness_Trend" not in out + + +@patch(f"{MODULE}.TestDefinition") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_empty(mock_resolve, mock_td_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_td_cls.list_monitor_configs_for_table.return_value = [] + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert "_No monitors configured for this table._" in out + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitors_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitors + + with _patch_perms(): + out = list_monitors(str(tg.id), "orders") + + assert out == "This table group is not monitored." + + +# --------------------------------------------------------------------------- +# list_monitor_schema_changes (TG-1092) +# --------------------------------------------------------------------------- + + +def _schema_log_entry(**overrides): + from testgen.common.models.data_structure_log import DataStructureLogEntry + + defaults: dict = { + "log_id": uuid4(), + "table_groups_id": uuid4(), + "table_name": "orders", + "column_name": "customer_id", + "change_date": datetime(2026, 6, 1, 12, 0, tzinfo=UTC), + "change": "A", + "old_data_type": None, + "new_data_type": "INTEGER", + } + defaults.update(overrides) + return DataStructureLogEntry(**defaults) + + +@patch.object(DataStructureLog, "list_for_table_group") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_happy_path(mock_resolve, mock_list, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_list.return_value = ( + [ + _schema_log_entry(change="A", column_name="new_col", old_data_type=None, new_data_type="TEXT"), + _schema_log_entry(change="D", column_name="old_col", old_data_type="INTEGER", new_data_type=None), + _schema_log_entry(change="M", column_name="total", old_data_type="NUMERIC", new_data_type="DECIMAL(10,2)"), + ], + 3, + ) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id)) + + assert "Schema changes in `Sales`" in out + assert "| Time | Change | Table | Column | Old type | New type |" in out + # Codes are mapped to user-facing words; raw codes don't leak + assert "added" in out + assert "dropped" in out + assert "modified" in out + out_row_section = out.split("| --- |", 1)[1] + assert "| A |" not in out_row_section + assert "| D |" not in out_row_section + assert "| M |" not in out_row_section + + +@patch.object(DataStructureLog, "list_for_table_group") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_table_filter_in_heading( + mock_resolve, mock_list, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_list.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id), table_name="orders", since="7 days") + + assert "table `orders`" in out + assert "since `7 days`" in out + # table_name + since reach the model as WHERE clauses, not named kwargs. + sql = _compile_clauses(mock_list) + assert "data_structure_log.table_name = 'orders'" in sql + assert "data_structure_log.change_date >=" in sql + + +@patch.object(DataStructureLog, "list_for_table_group") +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_since_passed_to_model( + mock_resolve, mock_list, db_session_mock, +): + tg = _mock_table_group() + mock_resolve.return_value = (tg, _mock_monitor_suite()) + mock_list.return_value = ([], 0) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + list_monitor_schema_changes(str(tg.id), since="2026-05-01") + + # ``since`` is passed as a ``change_date >= `` WHERE clause, not a kwarg. + sql = _compile_clauses(mock_list) + assert "data_structure_log.change_date >=" in sql + assert "2026-05-01" in sql + + +def _compile_clauses(mock_method) -> str: + """Compile the ``*clauses`` positional args of a captured + ``list_for_table_group`` call into one SQL string.""" + clauses = mock_method.call_args[0][1:] # drop table_group_id (first positional) + return " ".join(str(c.compile(compile_kwargs={"literal_binds": True})) for c in clauses) + + +@patch(f"{MODULE}.resolve_monitored_table_group") +def test_list_monitor_schema_changes_not_monitored(mock_resolve, db_session_mock): + tg = _mock_table_group(monitor_test_suite_id=None) + mock_resolve.return_value = (tg, None) + + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(): + out = list_monitor_schema_changes(str(tg.id)) + + assert out == "This table group is not monitored." + + +def test_list_monitor_schema_changes_invalid_since(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitor_schema_changes(str(uuid4()), since="bogus") + + +def test_list_monitor_schema_changes_limit_out_of_range(db_session_mock): + from testgen.mcp.tools.monitors import list_monitor_schema_changes + + with _patch_perms(), pytest.raises(MCPUserError): + list_monitor_schema_changes(str(uuid4()), limit=500) diff --git a/tests/unit/mcp/test_tools_profile_history.py b/tests/unit/mcp/test_tools_profile_history.py index e82b5bf3..910b0764 100644 --- a/tests/unit/mcp/test_tools_profile_history.py +++ b/tests/unit/mcp/test_tools_profile_history.py @@ -96,7 +96,6 @@ def _profile_row( def _profiling_run( id_=None, - job_execution_id=None, table_groups_id=None, status="Complete", profiling_starttime=None, @@ -105,7 +104,6 @@ def _profiling_run( ): run = MagicMock() run.id = id_ or uuid4() - run.job_execution_id = job_execution_id or uuid4() run.table_groups_id = table_groups_id or uuid4() run.status = status run.profiling_starttime = profiling_starttime or datetime(2026, 5, 10, 12, 0) @@ -283,7 +281,7 @@ def test_compare_profiling_runs_auto_baseline( mock_iss_type.select_where.return_value = [] with _patch_session([_je(), _je()]): - result = compare_profiling_runs(str(target_run.job_execution_id)) + result = compare_profiling_runs(str(target_run.id)) assert "Profiling Run Comparison" in result assert "Target" in result and "Baseline" in result @@ -298,7 +296,7 @@ def test_compare_profiling_runs_rejects_non_completed_target(mock_resolve, db_se with _patch_session([_je(status=JobStatus.RUNNING)]): with pytest.raises(MCPUserError, match="Target run is in `Running` state"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.resolve_profiling_run") @@ -308,7 +306,7 @@ def test_compare_profiling_runs_rejects_canceled_target(mock_resolve, db_session with _patch_session([_je(status=JobStatus.CANCELED)]): with pytest.raises(MCPUserError, match="`Canceled`"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.resolve_profiling_run") @@ -320,8 +318,8 @@ def test_compare_profiling_runs_rejects_cross_table_group(mock_resolve, db_sessi with _patch_session([_je()]): with pytest.raises(MCPUserError, match="same table group"): compare_profiling_runs( - str(target_run.job_execution_id), - str(baseline_run.job_execution_id), + str(target_run.id), + str(baseline_run.id), ) @@ -338,7 +336,7 @@ def test_compare_profiling_runs_auto_baseline_first_run(mock_resolve, db_session with _patch_session([_je()]): with pytest.raises(MCPUserError, match="no earlier completed profiling run"): - compare_profiling_runs(str(target_run.job_execution_id)) + compare_profiling_runs(str(target_run.id)) @patch("testgen.mcp.tools.profile_history.HygieneIssue") @@ -361,7 +359,7 @@ def test_compare_profiling_runs_identical_runs_renders_no_changes( mock_iss_type.select_where.return_value = [] with _patch_session([_je(), _je()]): - result = compare_profiling_runs(str(target_run.job_execution_id)) + result = compare_profiling_runs(str(target_run.id)) assert "No changes between target and baseline" in result diff --git a/tests/unit/mcp/test_tools_profiling.py b/tests/unit/mcp/test_tools_profiling.py index 588bd42d..b0136ba1 100644 --- a/tests/unit/mcp/test_tools_profiling.py +++ b/tests/unit/mcp/test_tools_profiling.py @@ -5,6 +5,7 @@ import pytest from testgen.common.models.data_column import ColumnProfileDetail, ColumnProfileSummary, DataColumnChars +from testgen.common.models.hygiene_issue import HygieneIssueCounts, IssueCounts, PotentialPiiCounts from testgen.common.pii_masking import PII_REDACTED from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -267,7 +268,7 @@ def test_list_column_profiles_with_valid_job_execution_id( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.list_for_table_group.return_value = ([_column_summary()], 1) from testgen.mcp.tools.profiling import list_column_profiles @@ -289,7 +290,7 @@ def test_list_column_profiles_rejects_je_from_different_tg( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr from testgen.mcp.tools.profiling import list_column_profiles with pytest.raises(MCPResourceNotAccessible, match="Profiling run .* not found or not accessible"): @@ -300,7 +301,7 @@ def test_list_column_profiles_rejects_je_from_different_tg( @patch("testgen.mcp.tools.common.TableGroup") def test_list_column_profiles_rejects_unknown_je(mock_tg_cls, mock_pr_cls, db_session_mock): mock_tg_cls.get.return_value = _mock_table_group() - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None from testgen.mcp.tools.profiling import list_column_profiles with pytest.raises(MCPResourceNotAccessible, match="Profiling run .* not found or not accessible"): @@ -583,20 +584,73 @@ def test_list_profiling_runs_invalid_status(mock_tg_cls, mock_run_cls, mock_next list_profiling_runs(table_group_id=str(uuid4()), status="Bogus") +@patch("testgen.mcp.tools.profiling.JobExecution") +@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.profiling.ProfilingRun") +@patch("testgen.mcp.tools.common.TableGroup") +def test_list_profiling_runs_schedule_filter(mock_tg_cls, mock_run_cls, mock_next, mock_je, db_session_mock): + mock_je.select_active_by_kwargs.return_value = [] + mock_tg_cls.get.return_value = _mock_table_group() + mock_run_cls.select_summary.return_value = ([], 0) + schedule_id = str(uuid4()) + + from testgen.mcp.tools.profiling import list_profiling_runs + list_profiling_runs(table_group_id=str(uuid4()), schedule_id=schedule_id, status="Completed") + + call_kwargs = mock_run_cls.select_summary.call_args.kwargs + assert call_kwargs["schedule_id"] == schedule_id + assert call_kwargs["statuses"] == [JobStatus.COMPLETED] + # The schedule clause is forwarded to the pending-JE query (positional *clauses arg). + assert mock_je.select_active_by_kwargs.call_args.args + + +@patch("testgen.mcp.tools.profiling.JobExecution") +@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.profiling.ProfilingRun") +@patch("testgen.mcp.tools.common.TableGroup") +def test_list_profiling_runs_unknown_schedule_returns_empty_envelope(mock_tg_cls, mock_run_cls, mock_next, mock_je, db_session_mock): + mock_je.select_active_by_kwargs.return_value = [] + mock_tg_cls.get.return_value = _mock_table_group() + mock_run_cls.select_summary.return_value = ([], 0) + + from testgen.mcp.tools.profiling import list_profiling_runs + result = list_profiling_runs(table_group_id=str(uuid4()), schedule_id=str(uuid4())) + + assert "No profiling runs" in result + + +@patch("testgen.mcp.tools.profiling.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.profiling.ProfilingRun") +@patch("testgen.mcp.tools.common.TableGroup") +def test_list_profiling_runs_malformed_schedule_raises(mock_tg_cls, mock_run_cls, mock_next, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + + from testgen.mcp.tools.profiling import list_profiling_runs + with pytest.raises(MCPUserError): + list_profiling_runs(table_group_id=str(uuid4()), schedule_id="not-a-uuid") + mock_run_cls.select_summary.assert_not_called() + + # ---------------------------------------------------------------------- # get_profiling_run # ---------------------------------------------------------------------- +@patch("testgen.mcp.tools.profiling.HygieneIssue") @patch("testgen.mcp.tools.profiling.ProfilingRun") -def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): +def test_get_profiling_run_returns_detail(mock_run_cls, mock_hygiene_cls, db_session_mock): summary = _mock_profiling_run() mock_run_cls.select_summary.return_value = ([summary], 1) mock_run = MagicMock(project_code="demo") - mock_run_cls.get_by_id_or_job.return_value = mock_run + mock_run_cls.get.return_value = mock_run mock_run_cls.select_table_breakdown.return_value = [ MagicMock(schema_name="demo", table_name="orders", record_ct=1000, column_ct=5, anomaly_ct=2), ] + mock_hygiene_cls.count_for_run.return_value = IssueCounts( + hygiene_issues=HygieneIssueCounts(definite=2, likely=2, possible=4), + potential_pii=PotentialPiiCounts(high=0, moderate=8), + dismissed=0, + ) with patch("testgen.mcp.permissions._compute_project_permissions") as mock_compute: mock_compute.return_value = ProjectPermissions( @@ -612,6 +666,10 @@ def test_get_profiling_run_returns_detail(mock_run_cls, db_session_mock): assert "Completed" in result assert "Per-table breakdown" in result assert "orders" in result + # Hygiene breakdown comes from count_for_run, keeping Potential PII separate from "possible". + mock_hygiene_cls.count_for_run.assert_called_once_with(summary.profiling_run_id) + assert "8 total — 2 definite, 2 likely, 4 possible" in result + assert "Potential PII:** 0 high, 8 moderate" in result @patch("testgen.mcp.tools.profiling.ProfilingRun") @@ -624,7 +682,7 @@ def test_get_profiling_run_pending_no_breakdown(mock_run_cls, db_session_mock): anomalies_possible_ct=None, dq_score_profiling=None, ) mock_run_cls.select_summary.return_value = ([summary], 1) - mock_run_cls.get_by_id_or_job.return_value = MagicMock(project_code="demo") + mock_run_cls.get.return_value = MagicMock(project_code="demo") with patch("testgen.mcp.permissions._compute_project_permissions") as mock_compute: mock_compute.return_value = ProjectPermissions( @@ -760,7 +818,7 @@ def _column_detail(**overrides) -> ColumnProfileDetail: # Run identity "profile_run_id": uuid4(), "profile_run_je_id": uuid4(), - "profile_run_status": "Complete", + "profile_run_status": JobStatus.COMPLETED, "profile_run_started_at": datetime(2026, 5, 1, 12, 0, 0), "profile_run_ended_at": datetime(2026, 5, 1, 12, 5, 0), "profile_run_log_message": None, @@ -1026,7 +1084,7 @@ def test_get_column_profile_detail_pinned_run_without_column_rejects( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.get_column_detail.return_value = _column_detail( profile_run_id=None, profile_run_je_id=None, @@ -1105,7 +1163,7 @@ def test_get_column_profile_detail_pinned_run_passes_id_to_model( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr mock_dcc_cls.get_column_detail.return_value = _column_detail() from testgen.mcp.tools.profiling import get_column_profile_detail @@ -1126,7 +1184,7 @@ def test_get_column_profile_detail_pinned_run_from_different_tg_unified_error( pr.project_code = tg.project_code mock_tg_cls.get.return_value = tg - mock_pr_cls.get_by_id_or_job.return_value = pr + mock_pr_cls.get.return_value = pr from testgen.mcp.tools.profiling import get_column_profile_detail with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -1141,7 +1199,7 @@ def test_get_column_profile_detail_pinned_run_unknown_unified_error( mock_tg_cls, mock_pr_cls, db_session_mock, ): mock_tg_cls.get.return_value = _mock_table_group() - mock_pr_cls.get_by_id_or_job.return_value = None + mock_pr_cls.get.return_value = None from testgen.mcp.tools.profiling import get_column_profile_detail with pytest.raises(MCPResourceNotAccessible, match=r"Profiling run .* not found or not accessible"): @@ -1161,7 +1219,7 @@ def test_get_column_profile_detail_running_run_rejects_with_status( mock_tg_cls.get.return_value = _mock_table_group() je_id = uuid4() mock_dcc_cls.get_column_detail.return_value = _column_detail( - profile_run_status="Running", + profile_run_status=JobStatus.RUNNING, profile_run_je_id=je_id, profile_run_ended_at=None, ) @@ -1183,7 +1241,7 @@ def test_get_column_profile_detail_error_run_includes_log_message( mock_tg_cls.get.return_value = _mock_table_group() je_id = uuid4() mock_dcc_cls.get_column_detail.return_value = _column_detail( - profile_run_status="Error", + profile_run_status=JobStatus.ERROR, profile_run_je_id=je_id, profile_run_log_message="connection timed out", ) @@ -1621,8 +1679,8 @@ def test_get_column_frequent_values_surfaces_job_execution_id_not_profile_run_id from testgen.mcp.tools.profiling import get_column_frequent_values result = get_column_frequent_values(str(uuid4()), "customers", "country") - # The internal profile_run_id PK must not leak; only the job_execution_id is followable. - assert str(run.job_execution_id) in result + # The internal profile_run_id PK must not leak; only the run id (the job execution id) is followable. + assert str(run.id) in result assert str(profile.profile_run_id) not in result @@ -1898,3 +1956,33 @@ def test_search_columns_table_group_scope_skips_per_project_summary( assert "Matches by project" not in result mock_dcc_cls.summarize_matches_by_project.assert_not_called() + + +# ---------------------------------------------------------------------- +# generate_create_table_script +# ---------------------------------------------------------------------- + +@patch("testgen.mcp.tools.profiling.build_create_table_script") +@patch("testgen.mcp.tools.common.TableGroup") +def test_generate_create_table_script_happy_path(mock_tg_cls, mock_build, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_build.return_value = 'CREATE TABLE "demo"."orders" (\n "id" INTEGER\n);' + + from testgen.mcp.tools.profiling import generate_create_table_script + result = generate_create_table_script(str(uuid4()), "orders") + + assert "CREATE TABLE" in result + assert '"id"' in result + assert "-- WAS" not in result + # MCP tool requests clean DDL (no change annotations) + assert mock_build.call_args.kwargs.get("annotate_changes", False) is False + + +@patch("testgen.mcp.tools.profiling.build_create_table_script", return_value=None) +@patch("testgen.mcp.tools.common.TableGroup") +def test_generate_create_table_script_unknown_table(mock_tg_cls, _mock_build, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + + from testgen.mcp.tools.profiling import generate_create_table_script + with pytest.raises(MCPResourceNotAccessible): + generate_create_table_script(str(uuid4()), "missing") diff --git a/tests/unit/mcp/test_tools_projects.py b/tests/unit/mcp/test_tools_projects.py new file mode 100644 index 00000000..7000e766 --- /dev/null +++ b/tests/unit/mcp/test_tools_projects.py @@ -0,0 +1,435 @@ +"""Tests for the MCP project write tools — currently just update_project. + +Project create/delete live in the enterprise plugin (gated on ``global_admin``); +this module owns the per-project ``administer`` slice that ships in core. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.projects" + + +def _patch_perms(allowed=("demo",), memberships=None, permission="administer", role="role_a"): + memberships = memberships or dict.fromkeys(allowed, role) + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions(memberships=memberships, permission=permission, username="test_user"), + ) + + +def _mock_project(**overrides) -> MagicMock: + project = MagicMock() + project.project_code = overrides.get("project_code", "demo") + project.project_name = overrides.get("project_name", "Demo Project") + project.use_dq_score_weights = overrides.get("use_dq_score_weights", True) + project.observability_api_url = overrides.get("observability_api_url", None) + project.observability_api_key = overrides.get("observability_api_key", None) + project.data_retention_enabled = overrides.get("data_retention_enabled", True) + project.data_retention_days = overrides.get("data_retention_days", 180) + return project + + +def _mock_schedule(cron_expr: str = "0 1 * * *", cron_tz: str = "UTC") -> MagicMock: + schedule = MagicMock() + schedule.cron_expr = cron_expr + schedule.cron_tz = cron_tz + return schedule + + +# --------------------------------------------------------------------------- +# update_project — guards +# --------------------------------------------------------------------------- + + +def test_update_project_no_fields_supplied(db_session_mock): + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo") + assert str(exc.value) == "No fields supplied to update." + + +@patch(f"{MODULE}.resolve_project") +def test_update_project_inaccessible_uses_unified_wording(mock_resolve, db_session_mock): + """When resolve_project raises (out-of-scope or missing), the tool re-raises with the + unified wording. resolve_project's own behaviour is covered in test_tools_common.py.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Project", "secret") + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(allowed=("demo",)), pytest.raises(MCPResourceNotAccessible) as exc: + update_project(project_code="secret", project_name="Anything") + assert "Project `secret` not found or not accessible" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_not_found_uses_unified_wording(mock_resolve, mock_schedule_cls, db_session_mock): + """resolve_project collapses 'no row' into the unified error.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Project", "demo") + mock_schedule_cls.get.return_value = None + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible) as exc: + update_project(project_code="demo", project_name="Anything") + assert "Project `demo` not found or not accessible" in str(exc.value) + + +def test_update_project_requires_administer(db_session_mock): + """role_d has edit but NOT administer (per conftest matrix).""" + from testgen.mcp.tools.projects import update_project + + with _patch_perms(memberships={"demo": "role_d"}), pytest.raises(MCPPermissionDenied): + update_project(project_code="demo", project_name="anything") + + +# --------------------------------------------------------------------------- +# update_project — rename +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_renames(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", project_name="Demo Renamed") + + project.save.assert_called_once() + assert "Project `demo` updated" in out + assert "| Field | Before | After |" in out + assert "Demo Project" in out + assert "Demo Renamed" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_no_op_when_value_unchanged(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", project_name="Demo Project") + + project.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_empty_name_rejected(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo") + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", project_name=" ") + assert "project_name: must not be empty" in str(exc.value) + project.save.assert_not_called() + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_strips_whitespace(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(project_name="Demo Project") + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", project_name=" Demo Renamed ") + + assert project.project_name == "Demo Renamed" + + +# --------------------------------------------------------------------------- +# update_project — weights toggle (recalc-scores side effect) +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobExecution") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_weights_toggle_submits_recalc( + mock_resolve, mock_schedule_cls, mock_job_exec_cls, db_session_mock, +): + project = _mock_project(use_dq_score_weights=True) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", use_dq_score_weights=False) + + project.save.assert_called_once() + # Background score recalc is submitted, matching the UI behaviour. + mock_job_exec_cls.submit.assert_called_once() + _, kwargs = mock_job_exec_cls.submit.call_args + assert kwargs["project_code"] == "demo" + + +@patch(f"{MODULE}.JobExecution") +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_weights_unchanged_no_recalc( + mock_resolve, mock_schedule_cls, mock_job_exec_cls, db_session_mock, +): + project = _mock_project(use_dq_score_weights=True) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + # Supplied but unchanged — no side effect. + update_project(project_code="demo", use_dq_score_weights=True) + + mock_job_exec_cls.submit.assert_not_called() + + +# --------------------------------------------------------------------------- +# update_project — observability fields +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_observability_url_diff_visible(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(observability_api_url=None) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", observability_api_url="https://obs.example/api") + + project.save.assert_called_once() + assert "DataOps Observability API URL" in out + assert "https://obs.example/api" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_observability_key_redacted_in_diff(mock_resolve, mock_schedule_cls, db_session_mock): + """Per mcp-patterns 'Secrets in inputs': the key is consumed but never echoed back.""" + project = _mock_project(observability_api_key=None) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", observability_api_key="super-secret-key-value") + + project.save.assert_called_once() + assert "DataOps Observability API key" in out + assert "[secret]" in out + assert "super-secret-key-value" not in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_observability_url_empty_string_clears( + mock_resolve, mock_schedule_cls, db_session_mock, +): + """Empty string clears the field (NullIfEmptyString column).""" + project = _mock_project(observability_api_url="https://existing.example") + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", observability_api_url="") + + # After normalization "" → None on the way in. + assert project.observability_api_url is None + + +# --------------------------------------------------------------------------- +# update_project — data retention +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_disable_retention_deletes_schedule( + mock_resolve, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_enabled=False) + + project.save.assert_called_once() + mock_schedule_cls.delete_for_retention.assert_called_once_with("demo") + mock_schedule_cls.upsert_for_retention.assert_not_called() + assert project.data_retention_enabled is False + # Days cleared on disable (matches UI behaviour). + assert project.data_retention_days is None + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_change_retention_days_upserts_schedule( + mock_resolve, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_days=30) + + project.save.assert_called_once() + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["project_code"] == "demo" + assert kwargs["retention_days"] == 30 + # Cron is preserved from existing schedule when not supplied. + assert kwargs["cron_expr"] == "0 1 * * *" + assert kwargs["cron_tz"] == "UTC" + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_enable_retention_with_defaults(mock_resolve, mock_schedule_cls, db_session_mock): + """Re-enabling retention without explicit days/cron falls back to the system defaults.""" + project = _mock_project(data_retention_enabled=False, data_retention_days=None) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = None # no current schedule + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", data_retention_enabled=True) + + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["retention_days"] == 180 # default + assert kwargs["cron_expr"] == "0 1 * * *" # DEFAULT_DATA_CLEANUP_CRON + assert kwargs["cron_tz"] == "UTC" # DEFAULT_RETENTION_CRON_TZ + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_change_cron_only(mock_resolve, mock_schedule_cls, db_session_mock): + project = _mock_project(data_retention_enabled=True, data_retention_days=180) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule(cron_expr="0 1 * * *", cron_tz="UTC") + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + out = update_project(project_code="demo", retention_cron_expr="0 2 * * *") + + mock_schedule_cls.upsert_for_retention.assert_called_once() + _, kwargs = mock_schedule_cls.upsert_for_retention.call_args + assert kwargs["cron_expr"] == "0 2 * * *" + # Days carried over from the project's current value. + assert kwargs["retention_days"] == 180 + assert "Retention cron expression" in out + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_retention_days_rejected_when_disabled( + mock_resolve, mock_schedule_cls, db_session_mock, +): + """Setting days while disabling retention is inconsistent — reject loudly instead of silently dropping.""" + project = _mock_project(data_retention_enabled=True) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project( + project_code="demo", + data_retention_enabled=False, + data_retention_days=90, + ) + msg = str(exc.value) + assert "data_retention_days: cannot be set when data_retention_enabled is False" in msg + project.save.assert_not_called() + mock_schedule_cls.upsert_for_retention.assert_not_called() + mock_schedule_cls.delete_for_retention.assert_not_called() + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_cron_rejected_when_retention_currently_disabled( + mock_resolve, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=False, data_retention_days=None) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = None + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", retention_cron_expr="0 2 * * *") + assert "retention_cron_expr: cannot be set when data_retention_enabled is False" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_retention_days_must_be_positive( + mock_resolve, mock_schedule_cls, db_session_mock, +): + project = _mock_project(data_retention_enabled=True) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_project(project_code="demo", data_retention_days=0) + assert "data_retention_days: must be a positive integer" in str(exc.value) + + +@patch(f"{MODULE}.JobSchedule") +@patch(f"{MODULE}.resolve_project") +def test_update_project_name_only_does_not_touch_schedule( + mock_resolve, mock_schedule_cls, db_session_mock, +): + """Renaming a project must NOT incidentally upsert / delete the retention schedule.""" + project = _mock_project(project_name="Old", data_retention_enabled=True) + mock_resolve.return_value = project + mock_schedule_cls.get.return_value = _mock_schedule() + + from testgen.mcp.tools.projects import update_project + + with _patch_perms(): + update_project(project_code="demo", project_name="New") + + project.save.assert_called_once() + mock_schedule_cls.upsert_for_retention.assert_not_called() + mock_schedule_cls.delete_for_retention.assert_not_called() diff --git a/tests/unit/mcp/test_tools_reference.py b/tests/unit/mcp/test_tools_reference.py index 96ead51f..97bb03bd 100644 --- a/tests/unit/mcp/test_tools_reference.py +++ b/tests/unit/mcp/test_tools_reference.py @@ -1,5 +1,7 @@ from unittest.mock import MagicMock, patch +import pytest + @patch("testgen.mcp.tools.reference.TestType") def test_get_test_type_found(mock_tt_cls, db_session_mock): @@ -263,3 +265,92 @@ def test_server_instructions_reference_column_profile_fields_resource(): # Sanity check the existing references are still present. assert "testgen://test-types" in SERVER_INSTRUCTIONS assert "testgen://hygiene-issue-types" in SERVER_INSTRUCTIONS + + +# --------------------------------------------------------------------------- +# connection_parameters_resource +# --------------------------------------------------------------------------- + + +def test_connection_parameters_index_lists_flavors(): + from testgen.mcp.tools.reference import connection_parameters_index_resource + + out = connection_parameters_index_resource() + # Accepted sql_flavor labels. + assert "PostgreSQL" in out and "Azure SQL Database" in out and "Salesforce Data 360" in out + # Each links to its per-flavor resource (keyed by code). + assert "testgen://connection-parameters/postgresql" in out + assert "testgen://connection-parameters/salesforce_data360" in out + + +def test_connection_parameters_resource_unknown_flavor(): + from testgen.mcp.exceptions import MCPUserError + from testgen.mcp.tools.reference import connection_parameters_resource + + with pytest.raises(MCPUserError) as exc: + connection_parameters_resource("not_a_flavor") + msg = str(exc.value) + assert "snowflake" in msg and "salesforce_data360" in msg + + +def test_connection_parameters_resource_postgresql_single_mode(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("postgresql") + assert "PostgreSQL Connection Parameters" in out + assert "Host" in out and "Username" in out + assert "Required (host mode)" in out # Host/Port/Database + # URL alternative is advertised. + assert "connect by URL" in out + # Single-mode flavor: no connection_mode instruction. + assert "Set `connection_mode`" not in out + + +def test_connection_parameters_resource_snowflake_both_modes(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("snowflake") + assert "Mode: Key-Pair" in out + assert "Mode: Password" in out + assert "Private Key" in out + assert "Warehouse" in out + assert "Set `connection_mode`" in out + assert 'connection_mode="Key-Pair"' in out + + +def test_connection_parameters_resource_databricks_pat_fields(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("databricks") + assert "Mode: Access Token" in out + assert "Mode: Service Principal (OAuth)" in out + assert "Catalog" in out + assert "HTTP Path" in out + assert "Client ID" in out + + +def test_connection_parameters_resource_port_default_note(): + """The Port row documents the flavor's conventional default port (doc-only — + the field stays required; the LLM supplies the default when the user doesn't).""" + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("postgresql") + assert "Default for PostgreSQL is 5432" in out + assert "Required (host mode)" in out # Port requirement unchanged + + +def test_connection_parameters_resource_port_default_per_flavor(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("sap_hana") + assert "Default for SAP HANA is 39015" in out + + +def test_connection_parameters_resource_marks_secrets(): + from testgen.mcp.tools.reference import connection_parameters_resource + + out = connection_parameters_resource("bigquery") + assert "Service Account Key" in out + assert "Secret" in out + # No URL alternative for BigQuery. + assert "connect by URL" not in out diff --git a/tests/unit/mcp/test_tools_source_data.py b/tests/unit/mcp/test_tools_source_data.py index 0a888b46..a7370982 100644 --- a/tests/unit/mcp/test_tools_source_data.py +++ b/tests/unit/mcp/test_tools_source_data.py @@ -1,9 +1,10 @@ -from unittest.mock import patch +from unittest.mock import MagicMock, patch from uuid import uuid4 import pandas as pd import pytest +from testgen.common.data_catalog_service import TableSampleResult from testgen.common.source_data_service import SourceDataResult from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -21,6 +22,21 @@ def _make_context(**overrides): return base +def _make_hygiene_issue(**overrides): + issue = MagicMock() + issue.table_groups_id = uuid4() + issue.type_id = "1001" + issue.detail = "Empty String: 12" + issue.schema_name = "public" + issue.table_name = "users" + issue.column_name = "address" + issue.profile_run_id = uuid4() + issue.project_code = "demo" + for key, value in overrides.items(): + setattr(issue, key, value) + return issue + + # --- get_source_data_query --- @@ -281,3 +297,305 @@ def test_get_source_data_passes_project_codes(mock_compute, mock_td, mock_fetch, call_kwargs = mock_td.get_source_data_context.call_args.kwargs assert call_kwargs["project_codes"] == ["proj_a"] + + +# --- argument mutual-exclusion + reference_date rule --- + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_both_ids_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="Provide exactly one of test_definition_id or issue_id"): + tool(test_definition_id=str(uuid4()), issue_id=str(uuid4())) + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_neither_id_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="Provide exactly one of test_definition_id or issue_id"): + tool() + + +@pytest.mark.parametrize("tool_name", ["get_source_data", "get_source_data_query"]) +def test_source_data_reference_date_with_issue_id_rejected(tool_name, db_session_mock): + import testgen.mcp.tools.source_data as mod + + tool = getattr(mod, tool_name) + with pytest.raises(MCPUserError, match="reference_date"): + tool(issue_id=str(uuid4()), reference_date="2026-01-01") + + +# --- get_source_data_query: hygiene issue --- + + +@patch("testgen.mcp.tools.source_data.build_hygiene_query") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_basic(mock_resolve, mock_pr, mock_build, db_session_mock): + issue_id = str(uuid4()) + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_build.return_value = "SELECT * FROM users WHERE address = ''" + + from testgen.mcp.tools.source_data import get_source_data_query + + result = get_source_data_query(issue_id=issue_id) + + assert f"# Source Data Query for Hygiene Issue `{issue_id}`" in result + assert "public.users" in result + assert "`address`" in result + assert "SELECT * FROM users" in result + assert "Test Definition" not in result + + +@patch("testgen.mcp.tools.source_data.build_hygiene_query") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_no_query_available(mock_resolve, mock_pr, mock_build, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_build.return_value = None + + from testgen.mcp.tools.source_data import get_source_data_query + + result = get_source_data_query(issue_id=str(uuid4())) + + assert "not available" in result + + +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_query_hygiene_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x") + + from testgen.mcp.tools.source_data import get_source_data_query + + with pytest.raises(MCPResourceNotAccessible, match="Hygiene issue"): + get_source_data_query(issue_id=str(uuid4())) + + +def test_get_source_data_query_hygiene_invalid_uuid(db_session_mock): + from testgen.mcp.tools.source_data import get_source_data_query + + with pytest.raises(MCPUserError, match="not a valid UUID"): + get_source_data_query(issue_id="bad-uuid") + + +# --- get_source_data: hygiene issue --- + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_ok(mock_resolve, mock_pr, mock_fetch, db_session_mock): + issue_id = str(uuid4()) + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + df = pd.DataFrame({"address": ["", ""], "count": [12, 3]}) + mock_fetch.return_value = SourceDataResult(status="OK", message=None, query="SELECT ...", df=df) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=issue_id) + + assert f"# Source Data for Hygiene Issue `{issue_id}`" in result + assert "**Rows returned:** 2" in result + assert "public.users" in result + assert "`address`" in result + assert "SELECT ..." in result + assert "Test Definition" not in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_na(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="NA", message="Source data lookup is not available for this hygiene issue.", query=None, df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "not available for this hygiene issue" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_nd(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="ND", message="No matching data.", query="SELECT * FROM users WHERE 1=0", df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "No matching data." in result + assert "SELECT * FROM users WHERE 1=0" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_err(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + mock_fetch.return_value = SourceDataResult( + status="ERR", message="Connection refused", query="SELECT 1", df=None, + ) + + from testgen.mcp.tools.source_data import get_source_data + + result = get_source_data(issue_id=str(uuid4())) + + assert "**Error:** Connection refused" in result + assert "SELECT 1" in result + + +@patch("testgen.mcp.tools.source_data.fetch_hygiene_source_data") +@patch("testgen.mcp.tools.source_data.ProfilingRun") +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_mask_pii_passed(mock_resolve, mock_pr, mock_fetch, db_session_mock): + mock_resolve.return_value = _make_hygiene_issue() + mock_pr.get.return_value = MagicMock(profiling_starttime="2026-05-12 00:00:00") + df = pd.DataFrame({"address": [""]}) + mock_fetch.return_value = SourceDataResult(status="OK", message=None, query=None, df=df) + + from testgen.mcp.tools.source_data import get_source_data + + get_source_data(issue_id=str(uuid4())) + + # mask_pii is the third positional arg + assert isinstance(mock_fetch.call_args[0][2], bool) + + +@patch("testgen.mcp.tools.source_data.resolve_hygiene_issue") +def test_get_source_data_hygiene_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Hygiene issue", "x") + + from testgen.mcp.tools.source_data import get_source_data + + with pytest.raises(MCPResourceNotAccessible, match="Hygiene issue"): + get_source_data(issue_id=str(uuid4())) + + +def test_get_source_data_hygiene_invalid_uuid(db_session_mock): + from testgen.mcp.tools.source_data import get_source_data + + with pytest.raises(MCPUserError, match="not a valid UUID"): + get_source_data(issue_id="bad-uuid") + + +# ---------------------------------------------------------------------- +# get_table_sample +# ---------------------------------------------------------------------- + +def _mock_table_group(tg_id=None, project_code="demo"): + tg = MagicMock() + tg.id = tg_id or uuid4() + tg.project_code = project_code + return tg + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_happy_path(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult( + "OK", df=pd.DataFrame([{"id": 1, "name": "a"}]), pii_redacted=False, + ) + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert "id" in result + assert "name" in result + assert "PII" not in result + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_empty_table(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult("ND") + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert result.strip() == "Table has no rows." + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_redaction_note(mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="main") + mock_fetch.return_value = TableSampleResult( + "OK", df=pd.DataFrame([{"id": 1}]), pii_redacted=True, + ) + + from testgen.mcp.tools.source_data import get_table_sample + result = get_table_sample(str(uuid4()), "orders") + + assert "redacted" in result.lower() + + +@patch("testgen.mcp.tools.source_data.fetch_table_sample") +@patch("testgen.mcp.tools.source_data.Connection") +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_connection_failure_names_connection( + mock_tg_cls, mock_dc, mock_conn, mock_fetch, db_session_mock, +): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = ("demo", []) + mock_conn.get_by_table_group.return_value = MagicMock(connection_name="prod-warehouse") + mock_fetch.return_value = TableSampleResult("ERR", message="boom") + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPUserError) as exc: + get_table_sample(str(uuid4()), "orders") + + assert "prod-warehouse" in str(exc.value) + + +@patch("testgen.mcp.tools.source_data.DataColumnChars") +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_unknown_table(mock_tg_cls, mock_dc, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + mock_dc.list_for_create_script.return_value = (None, []) + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPResourceNotAccessible): + get_table_sample(str(uuid4()), "missing") + + +@pytest.mark.parametrize("bad_limit", [0, 501]) +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_sample_rejects_out_of_range_limit(mock_tg_cls, bad_limit, db_session_mock): + mock_tg_cls.get.return_value = _mock_table_group() + + from testgen.mcp.tools.source_data import get_table_sample + with pytest.raises(MCPUserError): + get_table_sample(str(uuid4()), "orders", limit=bad_limit) diff --git a/tests/unit/mcp/test_tools_table_groups.py b/tests/unit/mcp/test_tools_table_groups.py new file mode 100644 index 00000000..febbb853 --- /dev/null +++ b/tests/unit/mcp/test_tools_table_groups.py @@ -0,0 +1,1094 @@ +"""Tests for the MCP table-group CRUD tools — create / update / preview.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from sqlalchemy.exc import IntegrityError + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.table_groups" + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _patch_perms(allowed=("demo",), memberships=None, permission="edit", role="role_a"): + # role_a has edit but NOT view_pii; role_d has edit + view_pii (see conftest matrix). + memberships = memberships or dict.fromkeys(allowed, role) + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions( + memberships=memberships, permission=permission, username="test_user", + ), + ) + + +def _mock_connection(**overrides) -> MagicMock: + conn = MagicMock() + conn.connection_id = overrides.get("connection_id", 42) + conn.project_code = overrides.get("project_code", "demo") + conn.connection_name = overrides.get("connection_name", "Local PG") + conn.sql_flavor = overrides.get("sql_flavor", "postgresql") + conn.sql_flavor_code = overrides.get("sql_flavor_code", "postgresql") + return conn + + +def _mock_table_group(**overrides) -> MagicMock: + """Build a MagicMock matching the TableGroup model surface used by the tools.""" + tg_id = overrides.get("id", uuid4()) + tg = MagicMock() + tg.id = tg_id + tg.project_code = overrides.get("project_code", "demo") + tg.connection_id = overrides.get("connection_id", 42) + tg.table_groups_name = overrides.get("table_groups_name", "Sample TG") + tg.table_group_schema = overrides.get("table_group_schema", "public") + tg.profiling_table_set = overrides.get("profiling_table_set", None) + tg.profiling_include_mask = overrides.get("profiling_include_mask", None) + tg.profiling_exclude_mask = overrides.get("profiling_exclude_mask", None) + tg.profile_id_column_mask = overrides.get("profile_id_column_mask", "%id") + tg.profile_sk_column_mask = overrides.get("profile_sk_column_mask", "%_sk") + tg.profile_use_sampling = overrides.get("profile_use_sampling", False) + tg.profile_sample_percent = overrides.get("profile_sample_percent", "30") + tg.profile_sample_min_count = overrides.get("profile_sample_min_count", 100000) + tg.profiling_delay_days = overrides.get("profiling_delay_days", "0") + tg.profile_flag_cdes = overrides.get("profile_flag_cdes", True) + tg.profile_flag_pii = overrides.get("profile_flag_pii", True) + tg.profile_exclude_xde = overrides.get("profile_exclude_xde", True) + tg.include_in_dashboard = overrides.get("include_in_dashboard", True) + tg.description = overrides.get("description", None) + tg.data_source = overrides.get("data_source", None) + tg.source_system = overrides.get("source_system", None) + tg.source_process = overrides.get("source_process", None) + tg.data_location = overrides.get("data_location", None) + tg.business_domain = overrides.get("business_domain", None) + tg.stakeholder_group = overrides.get("stakeholder_group", None) + tg.transform_level = overrides.get("transform_level", None) + tg.data_product = overrides.get("data_product", None) + tg.data_classification = overrides.get("data_classification", None) + return tg + + +def _integrity_error(message: str) -> IntegrityError: + orig = Exception(message) + return IntegrityError("stmt", {}, orig) + + +# --------------------------------------------------------------------------- +# create_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_happy_path(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(table_groups_name="Sample TG") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + instance.save.assert_called_once() + assert "Table Group `Sample TG` created" in out + assert "**Project:** `demo`" in out + assert "**Schema:** `public`" in out + + +def test_model_defaults_normalize_ynstring_columns_to_python_bool(): + """``YNString`` columns store ``"Y"``/``"N"`` but expose ``bool``. + + The defaults dict mirrors ``Column(default=...)`` raw values, so without + normalization a ``"N"`` default would render as truthy (non-empty string) + in the create-tool output until the row is reloaded from the DB. + """ + from testgen.mcp.tools.table_groups import _MODEL_DEFAULTS + + assert _MODEL_DEFAULTS["profile_use_sampling"] is False + assert _MODEL_DEFAULTS["profile_do_pair_rules"] is False + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_applies_model_defaults_when_optional_args_omitted( + mock_resolve, mock_tg_cls, db_session_mock, +): + """Optional profiling fields must seed from the model's ``Column(default=...)`` values. + + SQLAlchemy column defaults only fire at flush time, but validation runs + before flush — so the create tool must populate them in-memory or + validation rejects every call that omits them. + """ + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + kwargs = mock_tg_cls.call_args.kwargs + assert kwargs["profile_sample_percent"] == "30" + assert kwargs["profile_sample_min_count"] == 100000 + assert kwargs["profiling_delay_days"] == "0" + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_with_table_set_joins_to_string(mock_resolve, mock_tg_cls, db_session_mock): + """``table_set: list[str]`` is comma-joined into the model's ``profiling_table_set`` column.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profiling_table_set="film,actor,customer") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + table_set=["film", "actor", "customer"], + ) + + assert instance.profiling_table_set == "film,actor,customer" + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_renders_sample_settings_when_sampling_on( + mock_resolve, mock_tg_cls, db_session_mock, +): + """Sample % and Sample min rows only matter when sampling is on; render them then.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group( + profile_use_sampling=True, + profile_sample_percent="50", + profile_sample_min_count=5000, + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + profile_use_sampling=True, + profile_sample_percent=50, + profile_sample_min_count=5000, + ) + + assert "Sample %" in out + assert "50" in out + assert "Sample min rows" in out + assert "5000" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_skips_sample_settings_when_sampling_off( + mock_resolve, mock_tg_cls, db_session_mock, +): + """When sampling is off, Sample % / min rows are irrelevant — omit them.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profile_use_sampling=False) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + assert "Sample %" not in out + assert "Sample min rows" not in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_with_catalog_tags_rendered(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group( + data_source="Postgres", + business_domain="Sales", + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + out = create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + data_source="Postgres", + business_domain="Sales", + ) + + assert "## Catalog" in out + assert "Data source" in out + assert "Postgres" in out + assert "Business domain" in out + assert "Sales" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_validation_error_no_save(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(table_groups_name="ab") + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_table_group( + connection_id=42, + table_group_name="ab", + schema="public", + ) + msg = str(exc.value) + assert "Table group creation rejected" in msg + assert "must be between 3 and 40 characters" in msg + instance.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_duplicate_name_maps_to_user_error(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + instance.save.side_effect = _integrity_error( + 'duplicate key value violates unique constraint "table_groups_name_unique"' + ) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + assert str(exc.value) == "A Table Group with the same name already exists." + + +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_connection_not_accessible(mock_resolve, db_session_mock): + """Connection's project not in allowed_codes → MCPResourceNotAccessible from resolve_connection.""" + mock_resolve.side_effect = MCPResourceNotAccessible("Connection", "42") + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible) as exc: + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + assert "Connection" in str(exc.value) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_adds_scorecard_by_default(mock_resolve, mock_tg_cls, db_session_mock): + """Mirrors the UI checkbox default — a new table group gets a scorecard unless opted out.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + instance.save.assert_called_once_with(add_scorecard_definition=True) + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_add_scorecard_false_skips_scorecard(mock_resolve, mock_tg_cls, db_session_mock): + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group() + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + add_scorecard=False, + ) + + instance.save.assert_called_once_with(add_scorecard_definition=False) + + +def test_create_table_group_requires_edit(db_session_mock): + """Role without 'edit' permission → MCPPermissionDenied.""" + from testgen.mcp.tools.table_groups import create_table_group + + # role_c lacks 'edit' + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + ) + + +# --------------------------------------------------------------------------- +# create_table_group — PII flag is not gated on create +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_connection") +def test_create_table_group_pii_on_allowed_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """A new table group has no manually-marked PII to overwrite, so creating with the + flag on is allowed even without view_pii — the gate applies only to editing it.""" + mock_resolve.return_value = _mock_connection() + instance = _mock_table_group(profile_flag_pii=False) + mock_tg_cls.return_value = instance + + from testgen.mcp.tools.table_groups import create_table_group + + with _patch_perms(): # role_a: edit, no view_pii + create_table_group( + connection_id=42, + table_group_name="Sample TG", + schema="public", + profile_flag_pii=True, + ) + + instance.save.assert_called_once() + assert instance.profile_flag_pii is True + + +# --------------------------------------------------------------------------- +# update_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_no_fields_supplied(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_table_group() + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(uuid4())) + assert str(exc.value) == "No fields supplied to update." + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_description_diff_table(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(description=None) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), description="Pulled from postgres tutorial DB") + + tg.save.assert_called_once() + assert "| Field | Before | After |" in out + assert "Description" in out + assert "Pulled from postgres tutorial DB" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_no_op(mock_resolve, mock_tg_cls, db_session_mock): + """Supplying the current value → no-op message, no save call.""" + tg = _mock_table_group(description="existing") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), description="existing") + + tg.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_schema_locked_when_in_use(mock_resolve, mock_tg_cls, db_session_mock): + """``is_in_use=True`` + new schema → MCPUserError, no save.""" + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = True + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), schema="staging") + assert "Schema cannot be changed" in str(exc.value) + assert "Delete and recreate" in str(exc.value) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_schema_unlocked_when_not_in_use(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), schema="staging") + + tg.save.assert_called_once() + assert "Schema" in out + assert "staging" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_same_schema_on_in_use_group_is_noop(mock_resolve, mock_tg_cls, db_session_mock): + """Re-supplying the current schema on an in-use group is a no-op, not a lock violation.""" + tg = _mock_table_group(table_group_schema="public") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = True + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + out = update_table_group(table_group_id=str(tg.id), schema="public") + + tg.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_validation_error_no_save(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), table_group_name="ab") + msg = str(exc.value) + assert "Update rejected" in msg + assert "must be between 3 and 40 characters" in msg + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_duplicate_name_maps_to_user_error(mock_resolve, mock_tg_cls, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG") + tg.save.side_effect = _integrity_error( + 'duplicate key value violates unique constraint "table_groups_name_unique"' + ) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_table_group(table_group_id=str(tg.id), table_group_name="Existing Name") + assert str(exc.value) == "A Table Group with the same name already exists." + + +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", "abc") + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + update_table_group(table_group_id=str(uuid4()), description="any") + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_delay_days_int_cast_to_str(mock_resolve, mock_tg_cls, db_session_mock): + """``profiling_delay_days: int`` from the caller gets cast to ``str`` to match the model column.""" + tg = _mock_table_group(profiling_delay_days="0") + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(): + update_table_group(table_group_id=str(tg.id), profiling_delay_days=3) + + assert tg.profiling_delay_days == "3" + + +# --------------------------------------------------------------------------- +# update_table_group — PII flag gating (view_pii permission) +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_enable_pii_denied_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """role_a has edit but not view_pii (real ProjectPermissions) — enabling PII is denied. + + role_a *does* hold administer, so this also proves the gate checks view_pii + specifically, not some broader permission. + """ + tg = _mock_table_group(profile_flag_pii=False) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"), pytest.raises(MCPPermissionDenied): + update_table_group(table_group_id=str(tg.id), profile_flag_pii=True) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_disable_pii_denied_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """Change-detection mirrors the disabled checkbox — the value can't be touched either way.""" + tg = _mock_table_group(profile_flag_pii=True) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"), pytest.raises(MCPPermissionDenied): + update_table_group(table_group_id=str(tg.id), profile_flag_pii=False) + tg.save.assert_not_called() + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_unchanged_pii_allowed_without_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """Re-sending the current PII value (as the disabled UI checkbox does) is not a change.""" + tg = _mock_table_group(profile_flag_pii=True, description=None) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_a"): + out = update_table_group( + table_group_id=str(tg.id), + profile_flag_pii=True, + description="Edited elsewhere", + ) + + tg.save.assert_called_once() + assert tg.profile_flag_pii is True + assert "Description" in out + + +@patch(f"{MODULE}.TableGroup") +@patch(f"{MODULE}.resolve_table_group") +def test_update_table_group_enable_pii_allowed_with_view_pii(mock_resolve, mock_tg_cls, db_session_mock): + """role_d holds edit + view_pii (real ProjectPermissions) — enabling PII is allowed.""" + tg = _mock_table_group(profile_flag_pii=False) + mock_resolve.return_value = tg + mock_tg_cls.is_in_use.return_value = False + + from testgen.mcp.tools.table_groups import update_table_group + + with _patch_perms(role="role_d"): + out = update_table_group(table_group_id=str(tg.id), profile_flag_pii=True) + + tg.save.assert_called_once() + assert tg.profile_flag_pii is True + assert "Flag PII" in out + + +# --------------------------------------------------------------------------- +# preview_table_group +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_success_renders_table(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 2, + "column_ct": 5, + "approx_record_ct": 150, + "approx_data_point_ct": 350, + }, + "tables": { + "customer": { + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + "can_access": True, + }, + "rental": { + "column_ct": 2, + "approx_record_ct": 50, + "approx_data_point_ct": 50, + "can_access": True, + }, + }, + "success": True, + "message": None, + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "Preview for table group" in out + assert "Sample TG" in out + assert "customer" in out + assert "rental" in out + # verify_access defaults to False — no Read Access column + assert "| Table | Columns | Approx Rows | Approx Data Points |" in out + assert "Read Access" not in out + _, kwargs = mock_preview_svc.call_args + assert kwargs["verify_access"] is False + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_verify_access_adds_column(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 1, + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + }, + "tables": { + "customer": { + "column_ct": 3, + "approx_record_ct": 100, + "approx_data_point_ct": 300, + "can_access": True, + }, + }, + "success": True, + "message": None, + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id), verify_access=True) + + assert "| Table | Columns | Approx Rows | Approx Data Points | Read Access |" in out + _, kwargs = mock_preview_svc.call_args + assert kwargs["verify_access"] is True + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_partial_inaccessible_appends_footer(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 2, + "column_ct": 5, + "approx_record_ct": 150, + "approx_data_point_ct": 350, + }, + "tables": { + "customer": { + "column_ct": 3, "approx_record_ct": 100, + "approx_data_point_ct": 300, "can_access": True, + }, + "rental": { + "column_ct": 2, "approx_record_ct": 50, + "approx_data_point_ct": 50, "can_access": False, + }, + }, + "success": True, + "message": "Some tables were not accessible. Please the check the database permissions.", + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id), verify_access=True) + + assert "Some tables were not accessible" in out + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_no_match(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + tg = _mock_table_group(table_groups_name="Sample TG", table_group_schema="public") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, + "table_groups_name": "Sample TG", + "table_group_schema": "public", + "table_ct": 0, + "column_ct": 0, + "approx_record_ct": None, + "approx_data_point_ct": None, + }, + "tables": {}, + "success": False, + "message": ( + "No tables found matching the criteria. Please check the Table Group configuration" + " or the database permissions." + ), + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "returned no tables" in out + assert "No tables found matching the criteria" in out + + +@patch(f"{MODULE}.preview_table_group_service") +@patch(f"{MODULE}.Connection") +@patch(f"{MODULE}.resolve_table_group") +def test_preview_failed_returns_text_not_raises(mock_resolve, mock_conn_cls, mock_preview_svc, db_session_mock): + """A failed preview surfaces as a text response — no exception raised.""" + tg = _mock_table_group(table_groups_name="Sample TG") + mock_resolve.return_value = tg + mock_conn_cls.get_by_table_group.return_value = _mock_connection() + mock_preview_svc.return_value = ( + { + "stats": { + "id": tg.id, "table_groups_name": "Sample TG", "table_group_schema": "public", + "table_ct": 0, "column_ct": 0, "approx_record_ct": None, "approx_data_point_ct": None, + }, + "tables": {}, + "success": False, + "message": "Could not connect to target DB: connection refused", + }, + None, + None, + ) + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(): + out = preview_table_group(table_group_id=str(tg.id)) + + assert "Preview failed" in out + assert "Could not connect" in out + + +@patch(f"{MODULE}.resolve_table_group") +def test_preview_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", "abc") + + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + preview_table_group(table_group_id=str(uuid4())) + + +def test_preview_requires_edit(db_session_mock): + from testgen.mcp.tools.table_groups import preview_table_group + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + preview_table_group(table_group_id=str(uuid4())) + + +# --------------------------------------------------------------------------- +# list_table_groups +# --------------------------------------------------------------------------- + + +def _list_item(**overrides): + from testgen.common.models.table_group import TableGroupListItem + + return TableGroupListItem( + id=overrides.get("id", uuid4()), + table_groups_name=overrides.get("table_groups_name", "core_tables"), + table_group_schema=overrides.get("table_group_schema", "public"), + project_code=overrides.get("project_code", "demo"), + connection_name=overrides.get("connection_name", "warehouse_prod"), + table_count=overrides.get("table_count", 12), + column_count=overrides.get("column_count", 84), + row_count=overrides.get("row_count", 100_000), + last_profiled_date=overrides.get("last_profiled_date", None), + last_tested_date=overrides.get("last_tested_date", None), + profiling_score=overrides.get("profiling_score", 0.95), + testing_score=overrides.get("testing_score", 0.97), + quality_score=overrides.get("quality_score", 0.92), + ) + + +def test_list_table_groups_requires_one_arg(db_session_mock): + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"), pytest.raises( + MCPUserError, match="Pass either `project_code` or `connection_id`", + ): + list_table_groups() + + +def test_list_table_groups_rejects_both_args(db_session_mock): + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"), pytest.raises( + MCPUserError, match="Pass either `project_code` or `connection_id`", + ): + list_table_groups(project_code="demo", connection_id=12) + + +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_by_project_renders_rows(mock_tg_cls, db_session_mock): + mock_tg_cls.list_for_project.return_value = ([_list_item()], 1) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(project_code="demo") + + assert "Table groups for project `demo`" in out + assert "core_tables" in out + assert "warehouse_prod" in out + assert "public" in out + assert "Quality Score" in out # column header (#8) + assert "Columns" in out # column header (#10) + assert "Rows" in out # column header (#10) + assert "12" in out # table count + # quality_score 0.92 rendered via friendly_score → "92.0" (percentage form, mirrors UI). + assert "92.0" in out + mock_tg_cls.list_for_project.assert_called_once_with("demo", page=1, limit=20) + + +@patch(f"{MODULE}.resolve_connection") +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_by_connection_renders_rows(mock_tg_cls, mock_resolve, db_session_mock): + conn = _mock_connection() + conn.connection_id = 42 + conn.connection_name = "warehouse_prod" + mock_resolve.return_value = conn + mock_tg_cls.list_for_connection.return_value = ([_list_item()], 1) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(connection_id=42) + + assert "Table groups on connection `warehouse_prod` (`42`)" in out + mock_tg_cls.list_for_connection.assert_called_once_with(42, page=1, limit=20) + + +@patch(f"{MODULE}.TableGroup") +def test_list_table_groups_empty(mock_tg_cls, db_session_mock): + mock_tg_cls.list_for_project.return_value = ([], 0) + + from testgen.mcp.tools.table_groups import list_table_groups + + with _patch_perms(permission="view"): + out = list_table_groups(project_code="demo") + + assert "none found" in out + + +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_list_table_groups_rejects_inaccessible_project(mock_compute, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"other": "role_a"}, permission="view", username="test_user", + ) + + from testgen.mcp.tools.table_groups import list_table_groups + + with pytest.raises(MCPResourceNotAccessible, match="Project `secret` not found or not accessible"): + list_table_groups(project_code="secret") + + +# --------------------------------------------------------------------------- +# get_table_group +# --------------------------------------------------------------------------- + + +def _read_mock_table_group(**overrides): + """Variant of _mock_table_group that also sets dq score fields used by get_table_group.""" + from testgen.utils import score + + tg = _mock_table_group(**overrides) + testing = overrides.get("dq_score_testing", None) + profiling = overrides.get("dq_score_profiling", None) + tg.dq_score_testing = testing + tg.dq_score_profiling = profiling + # Mirror TableGroup.quality_score property — MagicMock would return a MagicMock for + # the attribute, so wire it explicitly through the same `score` helper the property uses. + tg.quality_score = score(profiling, testing) + return tg + + +@patch(f"{MODULE}.resolve_connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_renders_all_dialog_sections(mock_resolve, mock_resolve_conn, db_session_mock): + tg = _read_mock_table_group( + description="Curated payments tables", + profiling_include_mask="payments_%", + profiling_exclude_mask="tmp_%", + profiling_table_set="payments,refunds", + profile_use_sampling=True, + profile_sample_percent="50", + data_source="DataKitchen", + business_domain="Finance", + dq_score_testing=0.91, + dq_score_profiling=0.95, + ) + mock_resolve.return_value = tg + conn = _mock_connection(connection_name="warehouse_prod", sql_flavor_code="snowflake") + mock_resolve_conn.return_value = conn + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + # Identity + assert "Table group `Sample TG`" in out + assert "warehouse_prod" in out + assert "Snowflake" in out + assert "`public`" in out + assert "Curated payments tables" in out + # Criteria — both table-name and column-name masks rendered (labels via _DIFF_LABELS) + assert "## Criteria" in out + assert "payments_%" in out + assert "tmp_%" in out + assert "payments,refunds" in out + assert "ID column mask" in out + assert "SK column mask" in out + # Settings + assert "## Settings" in out + assert "Flag CDEs" in out # _DIFF_LABELS["profile_flag_cdes"] + # Sampling enabled → percent + min count rendered (labels via _DIFF_LABELS) + assert "## Sampling parameters" in out + assert "**Sample %:** 50" in out + assert "Sample min rows" in out + # Catalog tags only render when set + assert "## Catalog tags" in out + assert "DataKitchen" in out + assert "Finance" in out + # Latest activity — scores rendered via friendly_score (percentage form, mirrors UI). + assert "## Latest activity" in out + assert "**Profiling Score:** 95.0" in out + assert "**Testing Score:** 91.0" in out + # 0.95 * 0.91 = 0.8645 → friendly_score → "86.4" (Python banker's rounding of 86.45). + assert "**Quality Score:** 86.4" in out + + +@patch(f"{MODULE}.resolve_connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_skips_catalog_when_no_tags(mock_resolve, mock_resolve_conn, db_session_mock): + tg = _read_mock_table_group() # all catalog tags None + mock_resolve.return_value = tg + mock_resolve_conn.return_value = _mock_connection() + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + assert "## Catalog tags" not in out + + +@patch(f"{MODULE}.resolve_connection") +@patch(f"{MODULE}.resolve_table_group") +def test_get_table_group_skips_sample_details_when_sampling_off( + mock_resolve, mock_resolve_conn, db_session_mock, +): + tg = _read_mock_table_group(profile_use_sampling=False) + mock_resolve.return_value = tg + mock_resolve_conn.return_value = _mock_connection() + + from testgen.mcp.tools.table_groups import get_table_group + + with _patch_perms(permission="view"): + out = get_table_group(str(tg.id)) + + # Sampling section header still present (with the toggle), but percent/min not shown + assert "**Sampling:** No" in out # _DIFF_LABELS["profile_use_sampling"] + assert "**Sample %:**" not in out + assert "Sample min rows" not in out + + +@patch("testgen.mcp.tools.common.TableGroup") +def test_get_table_group_raises_not_found_for_inaccessible(mock_tg_cls, db_session_mock): + mock_tg_cls.get.return_value = None + + from testgen.mcp.tools.table_groups import get_table_group + + with pytest.raises(MCPResourceNotAccessible, match="Table group .* not found or not accessible"): + get_table_group(str(uuid4())) diff --git a/tests/unit/mcp/test_tools_test_definitions.py b/tests/unit/mcp/test_tools_test_definitions.py index 46488143..ebd814a6 100644 --- a/tests/unit/mcp/test_tools_test_definitions.py +++ b/tests/unit/mcp/test_tools_test_definitions.py @@ -1099,6 +1099,36 @@ def test_update_test_multi_field(mock_resolve_td, mock_tt_model, db_session_mock td.save.assert_called_once() +@patch("testgen.mcp.tools.test_definitions.TestType") +@patch("testgen.mcp.tools.test_definitions.resolve_test_definition") +def test_update_test_custom_metadata_json_string_coerced(mock_resolve_td, mock_tt_model, db_session_mock): + td = _make_td_orm() + td.editable_fields.return_value = td.editable_fields.return_value | {"custom_metadata"} + mock_resolve_td.return_value = td + mock_tt_model.get.return_value = _make_test_type() + + from testgen.mcp.tools.test_definitions import update_test + + update_test(str(td.id), fields={"custom_metadata": '{"pipeline": "p1", "task": "t1"}'}) + assert td.custom_metadata == {"pipeline": "p1", "task": "t1"} + td.save.assert_called_once() + + +@patch("testgen.mcp.tools.test_definitions.TestType") +@patch("testgen.mcp.tools.test_definitions.resolve_test_definition") +def test_update_test_custom_metadata_invalid_json_rejected(mock_resolve_td, mock_tt_model, db_session_mock): + td = _make_td_orm() + td.editable_fields.return_value = td.editable_fields.return_value | {"custom_metadata"} + mock_resolve_td.return_value = td + mock_tt_model.get.return_value = _make_test_type() + + from testgen.mcp.tools.test_definitions import update_test + + with pytest.raises(MCPUserError, match="custom_metadata"): + update_test(str(td.id), fields={"custom_metadata": "{not valid json"}) + td.save.assert_not_called() + + # -- validate_custom_test ----------------------------------------------------- @@ -1291,7 +1321,7 @@ def test_bulk_update_tests_invalid_action(mock_resolve_suite, mock_session, db_s from testgen.mcp.tools.test_definitions import bulk_update_tests - with pytest.raises(MCPUserError, match="`action`"): + with pytest.raises(MCPUserError, match="Invalid action `toggle`"): bulk_update_tests(test_suite_id=str(uuid4()), action="toggle") # Suite resolution happens before action validation in current code path? @@ -1313,3 +1343,205 @@ def test_bulk_update_tests_no_match(mock_resolve_suite, mock_session, db_session assert "No tests matched" in result assert "nonexistent" in result + + +# -- export_tests / import_tests ---------------------------------------------- + + +def _make_export_document(definitions=None): + from datetime import UTC, datetime + + from testgen.common.test_definition_export_import_service import ( + ExportDocument, + ExportSource, + TestDefinitionExport, + ) + + return ExportDocument( + version=1, + source=ExportSource( + project_code="demo", + test_suite="demo_suite", + table_group="tg", + table_group_schema="public", + exported_at=datetime(2026, 1, 1, tzinfo=UTC), + ), + definitions=definitions + if definitions is not None + else [TestDefinitionExport(test_type="Row_Ct", table_name="orders")], + ) + + +def _make_import_response(*, created=0, updated=0, skipped=0, deleted=0, items=None): + from testgen.common.test_definition_export_import_service import ImportResponse, ImportSummary + + return ImportResponse( + summary=ImportSummary(created=created, updated=updated, skipped=skipped, deleted=deleted), + items=items or [], + ) + + +@patch("testgen.mcp.tools.test_definitions.export_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_basic(mock_resolve_suite, mock_export, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_export.return_value = _make_export_document() + + from testgen.mcp.tools.test_definitions import export_tests + + result = export_tests(test_suite_id=str(uuid4())) + + assert "Test Definition Export from suite `demo_suite`" in result + assert "Tests Exported" in result + assert "```json" in result # payload rendered as a fenced JSON block + assert "Row_Ct" in result + # default origin=both → no Filters line + assert "Filters" not in result + + +@patch("testgen.mcp.tools.test_definitions.export_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_type") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_with_filters(mock_resolve_suite, mock_resolve_tt, mock_export, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_resolve_tt.return_value = "Row_Ct" + mock_export.return_value = _make_export_document() + + from testgen.mcp.tools.test_definitions import export_tests + + result = export_tests(test_suite_id=str(uuid4()), origin="manual", table_name="orders", test_type="Row Count") + + assert "Filters" in result + assert "manual tests only" in result + assert "table `orders`" in result + mock_resolve_tt.assert_called_once_with("Row Count") + # origin code, not label, is forwarded to the service + assert mock_export.call_args.args[1].value == "manual" + + +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_export_tests_invalid_origin(mock_resolve_suite, db_session_mock): + from testgen.mcp.tools.test_definitions import export_tests + + with pytest.raises(MCPUserError, match="Invalid origin `sideways`"): + export_tests(test_suite_id=str(uuid4()), origin="sideways") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_preview(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportItem, + ImportItemTD, + ImportReason, + ) + + mock_resolve_suite.return_value = _make_suite() + mock_import.return_value = _make_import_response( + created=1, + updated=1, + items=[ + ImportItem(action=ImportAction.create, reason=ImportReason.no_match, tds=[ImportItemTD(idx=0, target_id=uuid4())]), + ImportItem(action=ImportAction.update, reason=ImportReason.matched, tds=[ImportItemTD(idx=1, target_id=uuid4())]), + ], + ) + + from testgen.mcp.tools.test_definitions import import_tests + + result = import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}') + + assert "Import Preview" in result + assert "no changes were persisted" in result + assert "Created (projected)" in result + assert "Create (1)" in result + assert "Update (1)" in result + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_apply(mock_resolve_suite, mock_import, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + mock_import.return_value = _make_import_response(created=2) + + from testgen.mcp.tools.test_definitions import import_tests + + result = import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply") + + assert "Import into suite `demo_suite`" in result + assert "no changes were persisted" not in result + assert "Created" in result + # applied mode → no "(projected)" suffix + assert "(projected)" not in result + + +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_invalid_mode(mock_resolve_suite, db_session_mock): + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="Invalid mode `yolo`"): + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="yolo") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_malformed_payload(mock_resolve_suite, mock_import, db_session_mock): + mock_resolve_suite.return_value = _make_suite() + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="not a valid export document"): + import_tests(test_suite_id=str(uuid4()), payload="{not json") + mock_import.assert_not_called() + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_domain_error_surfaced(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ImportErrorCode, InvalidImportPayload + + mock_resolve_suite.return_value = _make_suite() + mock_import.side_effect = InvalidImportPayload( + ImportErrorCode.duplicate_natural_key, "Duplicate auto-gen key at index 1" + ) + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError, match="Duplicate auto-gen key at index 1"): + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply") + + +@patch("testgen.mcp.tools.test_definitions.import_definitions") +@patch("testgen.mcp.tools.test_definitions.resolve_test_suite") +def test_import_tests_strict_violation_embeds_breakdown(mock_resolve_suite, mock_import, db_session_mock): + from testgen.common.test_definition_export_import_service import ( + ImportAction, + ImportItem, + ImportItemTD, + ImportReason, + ImportStrictViolation, + ) + + mock_resolve_suite.return_value = _make_suite() + result = _make_import_response( + created=1, + skipped=1, + items=[ + ImportItem(action=ImportAction.create, reason=ImportReason.no_match, tds=[ImportItemTD(idx=0)]), + ImportItem(action=ImportAction.skip, reason=ImportReason.invalid_test_type, tds=[ImportItemTD(idx=1)]), + ], + ) + mock_import.side_effect = ImportStrictViolation(result) + + from testgen.mcp.tools.test_definitions import import_tests + + with pytest.raises(MCPUserError) as exc_info: + import_tests(test_suite_id=str(uuid4()), payload='{"definitions": []}', mode="apply_strict") + + msg = str(exc_info.value) + assert "Strict import failed" in msg + assert "1 tests would be skipped" in msg + assert "Nothing was applied" in msg + # the projected breakdown is embedded + assert "Skip (1)" in msg + assert "invalid_test_type" in msg diff --git a/tests/unit/mcp/test_tools_test_results.py b/tests/unit/mcp/test_tools_test_results.py index 7b3d08e4..64991e20 100644 --- a/tests/unit/mcp/test_tools_test_results.py +++ b/tests/unit/mcp/test_tools_test_results.py @@ -4,8 +4,9 @@ import pytest -from testgen.common.enums import JobStatus +from testgen.common.enums import Disposition, JobStatus from testgen.common.models.test_result import TestResultStatus +from testgen.common.test_result_disposition_service import DispositionUpdate from testgen.mcp.exceptions import MCPResourceNotAccessible, MCPUserError from testgen.mcp.permissions import ProjectPermissions @@ -20,10 +21,11 @@ def _mock_test_run(test_run_id=None): @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") -def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): +def test_list_test_results_basic(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): job_id = str(uuid4()) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() r1 = MagicMock() @@ -36,6 +38,7 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo r1.threshold_value = "10.0" r1.message = "Truncation detected" mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -56,9 +59,46 @@ def test_list_test_results_basic(mock_result, mock_tt_cls, mock_test_run_cls, mo @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") -def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() +def test_list_test_results_emits_test_result_id(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): + mock_test_run_cls.get.return_value = _mock_test_run() + mock_suite_cls.get_regular.return_value = _mock_test_suite() + + result_id = uuid4() + r1 = MagicMock() + r1.id = result_id + r1.status = TestResultStatus.Failed + r1.test_type = "Alpha_Trunc" + r1.test_definition_id = uuid4() + r1.table_name = "orders" + r1.column_names = "customer_name" + r1.result_measure = "15.3" + r1.threshold_value = "10.0" + r1.message = "Truncation detected" + mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] + + tt = MagicMock() + tt.test_type = "Alpha_Trunc" + tt.test_name_short = "Alpha Truncation" + mock_tt_cls.select_where.return_value = [tt] + + from testgen.mcp.tools.test_results import list_test_results + + result = list_test_results(str(uuid4())) + + assert "Test result" in result + assert str(result_id) in result + + +@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.test_results.TestRun") +@patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") +@patch("testgen.mcp.tools.test_results.TestResult") +def test_list_test_results_table_level_title(mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock): + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() r1 = MagicMock() @@ -71,6 +111,7 @@ def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test r1.threshold_value = "500" r1.message = None mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Row_Ct" @@ -89,7 +130,7 @@ def test_list_test_results_table_level_title(mock_result, mock_tt_cls, mock_test @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_empty(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() mock_result.select_results.return_value = [] @@ -108,7 +149,7 @@ def test_list_test_results_empty(mock_result, mock_test_run_cls, mock_suite_cls, def test_list_test_results_with_filters( mock_result, mock_tt_cls, mock_test_run_cls, mock_tt_common, mock_suite_cls, db_session_mock ): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -136,7 +177,7 @@ def test_list_test_results_invalid_uuid(db_session_mock): @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_invalid_status(mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite() from testgen.mcp.tools.test_results import list_test_results @@ -147,7 +188,7 @@ def test_list_test_results_invalid_status(mock_test_run_cls, mock_suite_cls, db_ @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_run_not_found(mock_test_run_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import list_test_results @@ -159,7 +200,7 @@ def test_list_test_results_run_not_found(mock_test_run_cls, db_session_mock): @patch("testgen.mcp.tools.test_results.TestRun") def test_list_test_results_run_in_monitor_suite_rejected(mock_test_run_cls, mock_suite_cls, db_session_mock): # Run exists, but the resolved suite is monitor → TestSuite.get_regular returns None. - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import list_test_results @@ -175,7 +216,7 @@ def test_list_test_results_run_in_forbidden_project( mock_compute, mock_test_run_cls, mock_suite_cls, db_session_mock ): mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") from testgen.mcp.tools.test_results import list_test_results @@ -196,7 +237,7 @@ def test_list_test_results_passes_project_codes( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_results.return_value = [] @@ -212,12 +253,12 @@ def test_list_test_results_passes_project_codes( @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.tools.test_results.TestType") -def test_list_test_results_resolves_via_get_by_id_or_job( +def test_list_test_results_resolves_via_get( mock_tt_cls, mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock ): """Verify the resolved test_run.id is passed to select_results.""" resolved_run_id = uuid4() - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run(resolved_run_id) + mock_test_run_cls.get.return_value = _mock_test_run(resolved_run_id) mock_suite_cls.get_regular.return_value = _mock_test_suite() mock_result.select_results.return_value = [] @@ -289,18 +330,17 @@ def test_list_test_results_by_suite_id_no_completed_runs(mock_suite_cls, db_sess @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestType") +@patch("testgen.mcp.tools.test_results.TestDefinition") @patch("testgen.mcp.tools.test_results.TestResult") def test_list_test_results_by_suite_id_resolves_latest_run( - mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock + mock_result, mock_td_cls, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock ): last_run_id = uuid4() mock_suite_cls.get_regular.return_value = _mock_test_suite(last_complete_test_run_id=last_run_id) resolved_run_id = uuid4() - resolved_je_id = uuid4() resolved_run = _mock_test_run(resolved_run_id) - resolved_run.job_execution_id = resolved_je_id - mock_test_run_cls.get_by_id_or_job.return_value = resolved_run + mock_test_run_cls.get.return_value = resolved_run r1 = MagicMock() r1.status = TestResultStatus.Failed @@ -312,6 +352,7 @@ def test_list_test_results_by_suite_id_resolves_latest_run( r1.threshold_value = "1" r1.message = None mock_result.select_results.return_value = [r1] + mock_td_cls.select_where.return_value = [] tt = MagicMock() tt.test_type = "Alpha_Trunc" @@ -323,11 +364,11 @@ def test_list_test_results_by_suite_id_resolves_latest_run( suite_id = str(uuid4()) result = list_test_results(test_suite_id=suite_id) - # Resolution chain: suite.last_complete_test_run_id → TestRun.get_by_id_or_job → test_run.id → select_results - mock_test_run_cls.get_by_id_or_job.assert_called_once_with(last_run_id) + # Resolution chain: suite.last_complete_test_run_id → TestRun.get → test_run.id → select_results + mock_test_run_cls.get.assert_called_once_with(last_run_id) assert mock_result.select_results.call_args.kwargs["test_run_id"] == resolved_run_id - # Output indicates which run the suite was resolved to. - assert str(resolved_je_id) in result + # Output indicates which run the suite was resolved to (run id is the job execution id). + assert str(resolved_run_id) in result assert f"Latest completed run of test suite `{suite_id}`" in result @@ -338,7 +379,7 @@ def test_list_test_results_by_suite_id_resolves_latest_run( def test_get_failure_summary_by_test_type( mock_result, mock_tt_cls, mock_test_run_cls, mock_suite_cls, db_session_mock, ): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [ ("Alpha_Trunc", TestResultStatus.Failed, 5), @@ -370,7 +411,7 @@ def test_get_failure_summary_by_test_type( @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_empty(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [] @@ -385,7 +426,7 @@ def test_get_failure_summary_empty(mock_result, mock_test_run_cls, mock_suite_cl @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_by_table(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [("orders", 10)] @@ -402,7 +443,7 @@ def test_get_failure_summary_by_table(mock_result, mock_test_run_cls, mock_suite @patch("testgen.mcp.tools.test_results.TestRun") @patch("testgen.mcp.tools.test_results.TestResult") def test_get_failure_summary_by_column(mock_result, mock_test_run_cls, mock_suite_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="demo") mock_result.select_failures.return_value = [("orders", "total_value", 34), ("orders", None, 2)] @@ -425,7 +466,7 @@ def test_get_failure_summary_invalid_uuid(db_session_mock): @patch("testgen.mcp.tools.test_results.TestRun") def test_get_failure_summary_run_not_found(mock_test_run_cls, db_session_mock): - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -440,7 +481,7 @@ def test_get_failure_summary_run_in_forbidden_project( mock_compute, mock_test_run_cls, mock_suite_cls, db_session_mock, ): mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") from testgen.mcp.tools.test_results import get_failure_summary @@ -453,7 +494,7 @@ def test_get_failure_summary_run_in_forbidden_project( @patch("testgen.mcp.tools.test_results.TestRun") def test_get_failure_summary_run_in_monitor_suite_rejected(mock_test_run_cls, mock_suite_cls, db_session_mock): # Run exists, but the resolved suite is monitor → TestSuite.get_regular returns None. - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -474,7 +515,7 @@ def test_get_failure_summary_passes_project_codes( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = _mock_test_run() + mock_test_run_cls.get.return_value = _mock_test_run() mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_failures.return_value = [] @@ -625,7 +666,7 @@ def test_get_failure_summary_cross_run_by_project(mock_compute, mock_result, db_ assert call_kwargs["since"] is not None -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.permissions._compute_project_permissions") def test_get_failure_summary_cross_run_by_project_and_suite( @@ -636,7 +677,7 @@ def test_get_failure_summary_cross_run_by_project_and_suite( permission="view", username="test_user", ) - mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="proj_a") + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_a") mock_result.select_failures.return_value = [] from testgen.mcp.tools.test_results import get_failure_summary @@ -664,12 +705,13 @@ def test_get_failure_summary_rejects_inaccessible_project(mock_compute, db_sessi get_failure_summary(project_code="proj_b", since="7 days") -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") @patch("testgen.mcp.permissions._compute_project_permissions") def test_get_failure_summary_rejects_inaccessible_test_suite(mock_compute, mock_suite_cls, db_session_mock): - """test_suite_id branch validates suite access — same contract as list_test_results.""" + """test_suite_id branch validates suite access — inaccessible suites are filtered out by the + access-scoped query in resolve_test_suite, which returns None.""" mock_compute.return_value = ProjectPermissions(memberships={"proj_a": "role_a"}, permission="view", username="test_user") - mock_suite_cls.get_regular.return_value = _mock_test_suite(project_code="forbidden_project") + mock_suite_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -677,10 +719,10 @@ def test_get_failure_summary_rejects_inaccessible_test_suite(mock_compute, mock_ get_failure_summary(test_suite_id=str(uuid4())) -@patch("testgen.mcp.tools.test_results.TestSuite") +@patch("testgen.mcp.tools.common.TestSuite") def test_get_failure_summary_rejects_unknown_or_monitor_test_suite(mock_suite_cls, db_session_mock): - # TestSuite.get_regular returns None for monitor suites and unknown ids alike. - mock_suite_cls.get_regular.return_value = None + # resolve_test_suite's access + is_monitor scoped query returns None for monitor suites and unknown ids alike. + mock_suite_cls.get.return_value = None from testgen.mcp.tools.test_results import get_failure_summary @@ -688,6 +730,29 @@ def test_get_failure_summary_rejects_unknown_or_monitor_test_suite(mock_suite_cl get_failure_summary(test_suite_id=str(uuid4())) +def test_get_failure_summary_rejects_invalid_group_by(db_session_mock): + from testgen.mcp.tools.test_results import get_failure_summary + + with pytest.raises(MCPUserError, match="Invalid group_by"): + get_failure_summary(job_execution_id=str(uuid4()), group_by="bogus") + + +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_summary_rejects_cross_project_suite(mock_compute, mock_suite_cls, db_session_mock): + """A suite the caller can access but in a different project than project_code is rejected, + not silently scoped away to an empty result.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_summary + + with pytest.raises(MCPUserError, match="belongs to project `proj_b`, not `proj_a`"): + get_failure_summary(project_code="proj_a", test_suite_id=str(uuid4())) + + # ---------------------------------------------------------------------- # search_test_results # ---------------------------------------------------------------------- @@ -866,6 +931,55 @@ def test_get_failure_trend_exclude_today_shifts_end_date(mock_compute, mock_fail assert mock_failure_trend.call_args.kwargs["end_date"] == real_today +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_cross_project_suite(mock_compute, mock_suite_cls, db_session_mock): + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPUserError, match="belongs to project `proj_b`, not `proj_a`"): + get_failure_trend(project_code="proj_a", test_suite_id=str(uuid4())) + + +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_inaccessible_suite(mock_compute, mock_suite_cls, db_session_mock): + """An inaccessible (or unknown/monitor) suite errors instead of silently returning an empty trend.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = None + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPResourceNotAccessible, match="Test suite .* not found or not accessible"): + get_failure_trend(test_suite_id=str(uuid4())) + + +@patch("testgen.mcp.tools.common.TableGroup") +@patch("testgen.mcp.tools.common.TestSuite") +@patch("testgen.mcp.permissions._compute_project_permissions") +def test_get_failure_trend_rejects_suite_and_table_group_in_different_projects( + mock_compute, mock_suite_cls, mock_tg_cls, db_session_mock +): + """No project_code, but the suite and table group resolve to different accessible projects: + the two filters would AND to an empty result, so reject instead of silently returning empty.""" + mock_compute.return_value = ProjectPermissions( + memberships={"proj_a": "role_a", "proj_b": "role_a"}, permission="view", username="test_user", + ) + mock_suite_cls.get.return_value = _mock_test_suite(project_code="proj_a") + mock_tg_cls.get.return_value = MagicMock(project_code="proj_b") + + from testgen.mcp.tools.test_results import get_failure_trend + + with pytest.raises(MCPUserError, match="different projects"): + get_failure_trend(test_suite_id=str(uuid4()), table_group_id=str(uuid4())) + + # ---------------------------------------------------------------------- # compare_test_runs # ---------------------------------------------------------------------- @@ -926,12 +1040,13 @@ def test_compare_test_runs_happy_path( baseline_run = _mock_run(suite_id) target_run = _mock_run(suite_id) # Tool resolves target first, then baseline. - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") diff = MagicMock() diff.total_baseline = 100 diff.total_target = 100 + diff.stable_passes = 98 diff.regressions = [ _mock_diff_row( TestResultStatus.Passed, @@ -952,6 +1067,8 @@ def test_compare_test_runs_happy_path( out = compare_test_runs(str(uuid4()), str(uuid4())) assert "Test Run Comparison" in out + assert "Stable passes (Baseline passed → Target passed)" in out + assert "| Stable passes (Baseline passed → Target passed) | 98 |" in out assert "Regressions" in out assert "Pattern Match" in out assert "Passed → Failed" in out @@ -961,6 +1078,51 @@ def test_compare_test_runs_happy_path( mock_result.diff_with_details.assert_called_once_with(baseline_run.id, target_run.id) +def _fetch_row(test_definition_id, status): + row = MagicMock() + row.test_definition_id = test_definition_id + row.test_type = "Pattern_Match" + row.test_name_short = "Pattern Match" + row.table_name = "orders" + row.column_names = "customer_id" + row.status = status + row.result_measure = "0" + row.threshold_value = "0" + return row + + +def test_diff_with_details_counts_stable_passes(): + from testgen.common.models.test_result import TestResult + + stable_1, stable_2, regressed, improved = uuid4(), uuid4(), uuid4(), uuid4() + baseline_rows = [ + _fetch_row(stable_1, TestResultStatus.Passed), + _fetch_row(stable_2, TestResultStatus.Passed), + _fetch_row(regressed, TestResultStatus.Passed), + _fetch_row(improved, TestResultStatus.Failed), + ] + target_rows = [ + _fetch_row(stable_1, TestResultStatus.Passed), + _fetch_row(stable_2, TestResultStatus.Passed), + _fetch_row(regressed, TestResultStatus.Failed), + _fetch_row(improved, TestResultStatus.Passed), + ] + session = MagicMock() + session.execute.side_effect = [baseline_rows, target_rows] + + with patch("testgen.common.models.test_result.get_current_session", return_value=session): + diff = TestResult.diff_with_details(uuid4(), uuid4()) + + assert diff.stable_passes == 2 + assert len(diff.regressions) == 1 + assert len(diff.improvements) == 1 + assert len(diff.persistent_failures) == 0 + # Internal consistency: with no out-of-bucket statuses (Error/Log) in the fixture, the four + # named buckets account for every shared test definition. + shared = diff.total_target - len(diff.new_tests) + assert diff.stable_passes + len(diff.regressions) + len(diff.improvements) + len(diff.persistent_failures) == shared + + @patch("testgen.mcp.tools.test_results.TestSuite") @patch("testgen.mcp.tools.test_results.TestResult") @patch("testgen.mcp.tools.test_results.TestRun") @@ -978,7 +1140,7 @@ def test_compare_test_runs_single_arg_resolves_previous( target_run = _mock_run(suite_id) baseline_run = _mock_run(suite_id) target_run.get_previous.return_value = baseline_run - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") diff = MagicMock( @@ -994,8 +1156,8 @@ def test_compare_test_runs_single_arg_resolves_previous( target_run.get_previous.assert_called_once_with() mock_result.diff_with_details.assert_called_once_with(baseline_run.id, target_run.id) - # Rendered Baseline cell shows the resolved JE ID, not an input string. - assert str(baseline_run.job_execution_id) in out + # Rendered Baseline cell shows the resolved run id (the job execution id), not an input string. + assert str(baseline_run.id) in out @patch("testgen.mcp.tools.test_results.TestSuite") @@ -1013,7 +1175,7 @@ def test_compare_test_runs_single_arg_no_previous_raises( suite_id = uuid4() target_run = _mock_run(suite_id) target_run.get_previous.return_value = None - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1036,7 +1198,7 @@ def test_compare_test_runs_single_arg_inaccessible_target( ) suite_id = uuid4() target_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run # Monitor suite or inaccessible project — get_regular returns None either way. mock_test_suite_cls.get_regular.return_value = None @@ -1059,7 +1221,7 @@ def test_compare_test_runs_run_not_found( permission="view", username="test_user", ) - mock_test_run_cls.get_by_id_or_job.return_value = None + mock_test_run_cls.get.return_value = None from testgen.mcp.tools.test_results import compare_test_runs @@ -1081,7 +1243,7 @@ def test_compare_test_runs_rejects_inaccessible_project( ) suite_id = uuid4() run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = run + mock_test_run_cls.get.return_value = run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_forbidden") from testgen.mcp.tools.test_results import compare_test_runs @@ -1106,7 +1268,7 @@ def test_compare_test_runs_rejects_different_suites( suite_id_baseline = uuid4() target_run = _mock_run(suite_id_target) baseline_run = _mock_run(suite_id_baseline) - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.side_effect = [ _mock_test_suite(suite_id=suite_id_target, project_code="proj_a"), _mock_test_suite(suite_id=suite_id_baseline, project_code="proj_a"), @@ -1139,7 +1301,7 @@ def test_compare_test_runs_rejects_monitor_suite( ) suite_id = uuid4() run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = run + mock_test_run_cls.get.return_value = run mock_test_suite_cls.get_regular.return_value = None from testgen.mcp.tools.test_results import compare_test_runs @@ -1162,7 +1324,7 @@ def test_compare_test_runs_rejects_target_not_completed( ) suite_id = uuid4() target_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.return_value = target_run + mock_test_run_cls.get.return_value = target_run mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1188,7 +1350,7 @@ def test_compare_test_runs_rejects_baseline_not_completed( suite_id = uuid4() target_run = _mock_run(suite_id) baseline_run = _mock_run(suite_id) - mock_test_run_cls.get_by_id_or_job.side_effect = [target_run, baseline_run] + mock_test_run_cls.get.side_effect = [target_run, baseline_run] mock_test_suite_cls.get_regular.return_value = _mock_test_suite(suite_id=suite_id, project_code="proj_a") from testgen.mcp.tools.test_results import compare_test_runs @@ -1196,3 +1358,217 @@ def test_compare_test_runs_rejects_baseline_not_completed( with _patch_test_results_session([_je(), _je(status=JobStatus.ERROR)]), \ pytest.raises(MCPUserError, match=r"Baseline run is in `Error` state"): compare_test_runs(str(uuid4()), str(uuid4())) + + +# --------------------------------------------------------------------------- +# update_test_result +# --------------------------------------------------------------------------- + + +@pytest.fixture +def disposition_perms(): + """Grant 'disposition' permission on demo (the conftest matrix omits it).""" + perms = MagicMock(spec=ProjectPermissions) + perms.memberships = {"demo": "role_a"} + perms.permission = "disposition" + perms.username = "test_user" + perms.allowed_codes = ["demo"] + perms.codes_allowed_to.return_value = ["demo"] + perms.has_access.side_effect = lambda code: code in ["demo"] + with patch("testgen.mcp.permissions._compute_project_permissions", return_value=perms): + yield perms + + +def test_update_test_result_invalid_uuid(db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + with pytest.raises(MCPUserError, match="not a valid UUID"): + update_test_result(test_result_id="bogus", disposition="Confirmed") + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_muted_maps_to_inactive(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + update_test_result(test_result_id=str(uuid4()), disposition="Muted") + + assert mock_set.call_args.args[1] == Disposition.INACTIVE + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_no_decision_clears(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + update_test_result(test_result_id=str(uuid4()), disposition="No Decision") + + assert mock_set.call_args.args[1] is None + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_success_message(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + rid = str(uuid4()) + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + out = update_test_result(test_result_id=rid, disposition="Dismissed") + + assert rid in out and "Dismissed" in out + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.resolve_test_result") +def test_update_test_result_passed_row_is_noop(mock_resolve, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import update_test_result + + mock_resolve.return_value = MagicMock(id=uuid4()) + mock_set.return_value = DispositionUpdate(matched=0, passed_skipped=1) + out = update_test_result(test_result_id=str(uuid4()), disposition="Confirmed") + + assert "passed" in out and "No change" in out + + +def test_update_test_result_uses_disposition_permission(): + import testgen.mcp.tools.test_results as mod + + closure = {c.cell_contents for c in mod.update_test_result.__wrapped__.__closure__} + assert "disposition" in closure + + +# --------------------------------------------------------------------------- +# bulk_update_test_results +# --------------------------------------------------------------------------- + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_uses_latest_run_when_run_omitted( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + matched_ids = [uuid4(), uuid4()] + mock_session.return_value.scalars.return_value.all.return_value = matched_ids + mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=0) + + with patch("testgen.mcp.tools.test_results.TestRun.get", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Dismissed") + + assert mock_set.call_args.args[0] == matched_ids + assert mock_set.call_args.args[1] == Disposition.DISMISSED + assert "2" in out + + +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_no_completed_run_errors(mock_resolve_suite, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import bulk_update_test_results + + mock_resolve_suite.return_value = MagicMock(last_complete_test_run_id=None, test_suite="suite_a") + with pytest.raises(MCPUserError, match="No completed test runs"): + bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Confirmed") + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_reports_passed_exclusions( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + mock_session.return_value.scalars.return_value.all.return_value = [uuid4(), uuid4(), uuid4()] + mock_set.return_value = DispositionUpdate(matched=2, passed_skipped=1) + + with patch("testgen.mcp.tools.test_results.TestRun.get", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Muted") + + assert "2" in out # matched + assert "1" in out and "passed" in out # exclusions surfaced + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_no_matches(mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms): + from testgen.mcp.tools.test_results import bulk_update_test_results + + run_id = uuid4() + suite = MagicMock(id=uuid4(), test_suite="suite_a", last_complete_test_run_id=run_id) + mock_resolve_suite.return_value = suite + mock_session.return_value.scalars.return_value.all.return_value = [] + mock_set.return_value = DispositionUpdate(matched=0, passed_skipped=0) + + with patch("testgen.mcp.tools.test_results.TestRun.get", + return_value=MagicMock(id=run_id, job_execution_id=run_id, test_suite_id=suite.id)): + out = bulk_update_test_results(test_suite_id=str(uuid4()), disposition="Confirmed") + + assert "No test results matched" in out + + +def test_bulk_update_uses_disposition_permission(): + import testgen.mcp.tools.test_results as mod + + closure = {c.cell_contents for c in mod.bulk_update_test_results.__wrapped__.__closure__} + assert "disposition" in closure + + +@patch("testgen.mcp.tools.test_results.set_test_results_disposition") +@patch("testgen.mcp.tools.test_results.get_current_session") +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_explicit_run_in_suite( + mock_resolve_suite, mock_session, mock_set, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + suite = MagicMock(id=uuid4(), test_suite="suite_a") + mock_resolve_suite.return_value = suite + run_id = uuid4() + matched_ids = [uuid4()] + mock_session.return_value.scalars.return_value.all.return_value = matched_ids + mock_set.return_value = DispositionUpdate(matched=1, passed_skipped=0) + + with patch( + "testgen.mcp.tools.test_results.TestRun.get", + return_value=MagicMock(id=uuid4(), job_execution_id=run_id, test_suite_id=suite.id), + ): + out = bulk_update_test_results( + test_suite_id=str(uuid4()), disposition="Confirmed", job_execution_id=str(run_id) + ) + + assert mock_set.call_args.args[0] == matched_ids + assert "1" in out + + +@patch("testgen.mcp.tools.test_results.resolve_test_suite") +def test_bulk_update_explicit_run_from_other_suite_rejected( + mock_resolve_suite, db_session_mock, disposition_perms +): + from testgen.mcp.tools.test_results import bulk_update_test_results + + suite = MagicMock(id=uuid4(), test_suite="suite_a") + mock_resolve_suite.return_value = suite + + with patch( + "testgen.mcp.tools.test_results.TestRun.get", + return_value=MagicMock(test_suite_id=uuid4()), # different suite + ): + with pytest.raises(MCPResourceNotAccessible, match=r"Test run .* not found or not accessible"): + bulk_update_test_results( + test_suite_id=str(uuid4()), disposition="Confirmed", job_execution_id=str(uuid4()) + ) diff --git a/tests/unit/mcp/test_tools_test_runs.py b/tests/unit/mcp/test_tools_test_runs.py index 7a71380e..e37f6df9 100644 --- a/tests/unit/mcp/test_tools_test_runs.py +++ b/tests/unit/mcp/test_tools_test_runs.py @@ -43,6 +43,7 @@ def test_list_test_runs_default(mock_suite, mock_run, mock_next, db_session_mock project_code="demo", table_group_id=None, test_suite_id=None, + schedule_id=None, statuses=None, page=1, page_size=10, @@ -67,6 +68,66 @@ def test_list_test_runs_with_status_filter(mock_suite, mock_run, mock_next, db_s assert call_kwargs["statuses"] == [JobStatus.PENDING, JobStatus.CLAIMED] +@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.test_runs.TestRun") +@patch("testgen.mcp.tools.test_runs.TestSuite") +def test_list_test_runs_with_schedule_filter(mock_suite, mock_run, mock_next, db_session_mock): + mock_run.select_summary.return_value = ([], 0) + schedule_id = str(uuid4()) + + from testgen.mcp.tools.test_runs import list_test_runs + + list_test_runs(project_code="demo", schedule_id=schedule_id, status="Completed") + + call_kwargs = mock_run.select_summary.call_args.kwargs + assert call_kwargs["schedule_id"] == schedule_id + assert call_kwargs["statuses"] == [JobStatus.COMPLETED] + + +@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.test_runs.TestRun") +@patch("testgen.mcp.tools.test_runs.TestSuite") +def test_list_test_runs_unknown_schedule_returns_empty_envelope(mock_suite, mock_run, mock_next, db_session_mock): + # Unknown/inaccessible/wrong-kind schedule yields no rows — standard empty envelope, not an error. + mock_run.select_summary.return_value = ([], 0) + + from testgen.mcp.tools.test_runs import list_test_runs + + result = list_test_runs(project_code="demo", schedule_id=str(uuid4())) + + assert "No test runs" in result + + +@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.test_runs.TestRun") +@patch("testgen.mcp.tools.test_runs.TestSuite") +def test_list_test_runs_malformed_schedule_raises(mock_suite, mock_run, mock_next, db_session_mock): + from testgen.mcp.tools.test_runs import list_test_runs + + with pytest.raises(MCPUserError): + list_test_runs(project_code="demo", schedule_id="not-a-uuid") + mock_run.select_summary.assert_not_called() + + +@patch("testgen.mcp.tools.test_runs.JobExecution") +@patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None) +@patch("testgen.mcp.tools.test_runs.TestRun") +@patch("testgen.mcp.tools.test_runs.TestSuite") +def test_list_test_runs_schedule_filters_pending(mock_suite, mock_run, mock_next, mock_je, db_session_mock): + mock_je.select_active_by_kwargs.return_value = [] + suite_id = uuid4() + mock_suite.select_minimal_where.return_value = [MagicMock(id=suite_id)] + mock_run.select_summary.return_value = ([], 0) + schedule_id = str(uuid4()) + + from testgen.mcp.tools.test_runs import list_test_runs + + list_test_runs(project_code="demo", test_suite="Quality", schedule_id=schedule_id) + + # A schedule clause is forwarded to the pending-JE query (positional *clauses arg). + assert mock_je.select_active_by_kwargs.call_args.args + + @patch("testgen.mcp.tools.test_runs.JobExecution") @patch("testgen.mcp.tools.test_runs.next_scheduled_run", return_value=None) @patch("testgen.mcp.tools.test_runs.TestRun") diff --git a/tests/unit/mcp/test_tools_test_suites.py b/tests/unit/mcp/test_tools_test_suites.py new file mode 100644 index 00000000..09907608 --- /dev/null +++ b/tests/unit/mcp/test_tools_test_suites.py @@ -0,0 +1,479 @@ +"""Tests for the MCP test-suite write tools — create / update for regular (non-monitor) suites.""" + +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from testgen.mcp.exceptions import MCPPermissionDenied, MCPResourceNotAccessible, MCPUserError +from testgen.mcp.permissions import ProjectPermissions + +pytestmark = pytest.mark.unit + +MODULE = "testgen.mcp.tools.test_suites" + + +def _patch_perms(allowed=("demo",), memberships=None, permission="edit", role="role_a"): + memberships = memberships or dict.fromkeys(allowed, role) + return patch( + "testgen.mcp.permissions._compute_project_permissions", + return_value=ProjectPermissions(memberships=memberships, permission=permission, username="test_user"), + ) + + +def _mock_table_group(**overrides) -> MagicMock: + tg_id = overrides.get("id", uuid4()) + tg = MagicMock() + tg.id = tg_id + tg.project_code = overrides.get("project_code", "demo") + tg.connection_id = overrides.get("connection_id", 42) + tg.table_groups_name = overrides.get("table_groups_name", "Sample TG") + return tg + + +def _mock_test_suite(**overrides) -> MagicMock: + suite = MagicMock() + suite.id = overrides.get("id", uuid4()) + suite.project_code = overrides.get("project_code", "demo") + suite.test_suite = overrides.get("test_suite", "Sample Suite") + suite.connection_id = overrides.get("connection_id", 42) + suite.table_groups_id = overrides.get("table_groups_id", uuid4()) + suite.test_suite_description = overrides.get("test_suite_description", None) + suite.severity = overrides.get("severity", None) + suite.export_to_observability = overrides.get("export_to_observability", False) + suite.dq_score_exclude = overrides.get("dq_score_exclude", False) + suite.component_key = overrides.get("component_key", None) + suite.component_type = overrides.get("component_type", None) + suite.component_name = overrides.get("component_name", None) + suite.is_monitor = overrides.get("is_monitor", False) + return suite + + +# --------------------------------------------------------------------------- +# create_test_suite +# --------------------------------------------------------------------------- + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_happy_path(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite( + test_suite="Sales Tests", + project_code="demo", + table_groups_id=tg.id, + ) + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + out = create_test_suite(table_group_id=str(tg.id), test_suite_name="Sales Tests") + + # S2 review feedback: the tool no longer calls instance.save(); it relies on + # session.add + flush. The mock TestSuite constructor was invoked and the + # rendered response is returned — that's enough at the unit layer; the smoke + # test exercises the actual DB persistence. + mock_suite_cls.assert_called_once() + instance.save.assert_not_called() + assert "Test Suite `Sales Tests` created" in out + assert "**Project:** `demo`" in out + assert "Sample TG" in out + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_default_export_matches_ui(mock_resolve, mock_suite_cls, db_session_mock): + """S1 review feedback: export_to_observability defaults to False (matches the UI dialog), + not the model's legacy Y default.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite(table_group_id=str(tg.id), test_suite_name="Anything") + + _, kwargs = mock_suite_cls.call_args + assert kwargs["export_to_observability"] is False + # S3 review feedback: component_type defaults to "dataset" on create + assert kwargs["component_type"] == "dataset" + assert kwargs["dq_score_exclude"] is False + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_passes_all_new_args_to_model(mock_resolve, mock_suite_cls, db_session_mock): + """S9: create_test_suite accepts the full UI-editable surface (export, dq_score_exclude, component_*) + so the LLM can configure the suite in one round-trip.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite( + table_group_id=str(tg.id), + test_suite_name="Full", + description="d", + severity_default="Warning", + export_to_observability=True, + component_key="ck", + component_type="dataset", + component_name="cn", + dq_score_exclude=True, + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite_description"] == "d" + assert kwargs["severity"] == "Warning" + assert kwargs["export_to_observability"] is True + assert kwargs["component_key"] == "ck" + assert kwargs["component_type"] == "dataset" + assert kwargs["component_name"] == "cn" + assert kwargs["dq_score_exclude"] is True + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_normalizes_empty_text_args_to_none(mock_resolve, mock_suite_cls, db_session_mock): + """S4: NullIfEmptyString writes "" -> NULL; pre-normalize on the way in so the model + receives the canonical value.""" + tg = _mock_table_group() + mock_resolve.return_value = tg + mock_suite_cls.return_value = _mock_test_suite() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite( + table_group_id=str(tg.id), + test_suite_name="x", + description="", + component_key="", + component_type="", + component_name="", + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite_description"] is None + assert kwargs["component_key"] is None + assert kwargs["component_type"] is None + assert kwargs["component_name"] is None + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_with_severity(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="Sev Suite", severity="Fail") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + out = create_test_suite( + table_group_id=str(tg.id), + test_suite_name="Sev Suite", + severity_default="Fail", + ) + + _, kwargs = mock_suite_cls.call_args + assert kwargs["severity"] == "Fail" + assert kwargs["is_monitor"] is False + assert "Fail" in out + + +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_invalid_severity(mock_resolve, db_session_mock): + mock_resolve.return_value = _mock_table_group() + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_test_suite( + table_group_id=str(uuid4()), + test_suite_name="Anything", + severity_default="Critical", + ) + msg = str(exc.value) + # S6: error is field-scoped and enumerates valid values without double-labeling. + assert "severity_default" in msg + assert "Fail" in msg and "Warning" in msg + assert "Critical" in msg + assert "Invalid severity" not in msg # the double-labeled phrasing is gone + + +def test_create_test_suite_empty_name_rejected(db_session_mock): + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + create_test_suite(table_group_id=str(uuid4()), test_suite_name=" ") + assert "test_suite_name: must not be empty" in str(exc.value) + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_strips_whitespace(mock_resolve, mock_suite_cls, db_session_mock): + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="Sales Tests") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(): + create_test_suite(table_group_id=str(tg.id), test_suite_name=" Sales Tests ") + + _, kwargs = mock_suite_cls.call_args + assert kwargs["test_suite"] == "Sales Tests" + + +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_table_group_not_accessible(mock_resolve, db_session_mock): + mock_resolve.side_effect = MCPResourceNotAccessible("Table group", str(uuid4())) + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + create_test_suite(table_group_id=str(uuid4()), test_suite_name="Anything") + + +def test_create_test_suite_requires_edit(db_session_mock): + """role_c lacks edit (per conftest matrix).""" + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + create_test_suite(table_group_id=str(uuid4()), test_suite_name="Anything") + + +@patch(f"{MODULE}.TestSuite") +@patch(f"{MODULE}.resolve_table_group") +def test_create_test_suite_allows_role_with_edit_but_not_administer( + mock_resolve, mock_suite_cls, db_session_mock, +): + """role_d has edit but NOT administer (per conftest matrix) — must still be allowed. + + Mirrors the production data_quality role: write access to suites without + project-level administer rights. + """ + tg = _mock_table_group() + mock_resolve.return_value = tg + instance = _mock_test_suite(test_suite="DQ Suite") + mock_suite_cls.return_value = instance + + from testgen.mcp.tools.test_suites import create_test_suite + + with _patch_perms(memberships={"demo": "role_d"}): + out = create_test_suite(table_group_id=str(tg.id), test_suite_name="DQ Suite") + + mock_suite_cls.assert_called_once() + instance.save.assert_not_called() + assert "Test Suite `DQ Suite` created" in out + + +# --------------------------------------------------------------------------- +# update_test_suite +# --------------------------------------------------------------------------- + + +def test_update_test_suite_no_fields_supplied(db_session_mock): + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite(test_suite_id=str(uuid4())) + assert str(exc.value) == "No fields supplied to update." + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_renames(mock_resolve, db_session_mock): + suite = _mock_test_suite(test_suite="Original Name") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite(test_suite_id=str(suite.id), test_suite_name="New Name") + + suite.save.assert_called_once() + assert "Test Suite `New Name` updated" in out + assert "| Field | Before | After |" in out + assert "Original Name" in out + assert "New Name" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_multi_field_diff(mock_resolve, db_session_mock): + suite = _mock_test_suite( + test_suite="Suite", + severity=None, + export_to_observability=True, + dq_score_exclude=False, + ) + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + severity_default="Warning", + export_to_observability=False, + dq_score_exclude=True, + ) + + suite.save.assert_called_once() + assert "Default severity" in out + assert "Warning" in out + # S13/S15: column header capitalises "Observability" (product name). + assert "Export to Observability" in out + assert "Exclude from quality scoring" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_no_op(mock_resolve, db_session_mock): + suite = _mock_test_suite(test_suite_description="same", severity="Fail") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + description="same", + severity_default="Fail", + ) + + suite.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_text_arg_on_null_field_is_noop(mock_resolve, db_session_mock): + """S4 phantom-diff regression: an "" arg on a currently-NULL NullIfEmptyString column must + NOT show up as a "changed" diff row (the DB would read back identical to before).""" + suite = _mock_test_suite( + test_suite_description=None, + component_key=None, + component_type=None, + component_name=None, + ) + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite( + test_suite_id=str(suite.id), + description="", + component_key="", + component_type="", + component_name="", + ) + + suite.save.assert_not_called() + assert "No fields changed" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_text_arg_clears_populated_field(mock_resolve, db_session_mock): + """S4 complement: "" on a currently-populated field clears it to NULL and shows in the diff.""" + suite = _mock_test_suite(component_key="existing-key") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(): + out = update_test_suite(test_suite_id=str(suite.id), component_key="") + + suite.save.assert_called_once() + assert "Component key" in out + assert "existing-key" in out + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_empty_name_rejected(mock_resolve, db_session_mock): + suite = _mock_test_suite() + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite(test_suite_id=str(suite.id), test_suite_name=" ") + assert "test_suite_name: must not be empty" in str(exc.value) + suite.save.assert_not_called() + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_invalid_severity_collected(mock_resolve, db_session_mock): + """Field-level validation collected, then raised together — name + severity errors in one message.""" + suite = _mock_test_suite() + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPUserError) as exc: + update_test_suite( + test_suite_id=str(suite.id), + test_suite_name="", + severity_default="Critical", + ) + msg = str(exc.value) + assert "Update rejected" in msg + assert "test_suite_name: must not be empty" in msg + # S6: field-scoped phrasing, no double-labeled "Invalid severity" + assert "severity_default" in msg + assert "Fail" in msg and "Warning" in msg + assert "Invalid severity" not in msg + suite.save.assert_not_called() + + +def test_update_test_suite_monitor_suite_unified_wording(db_session_mock): + """Per TG-1053 review feedback: exercising an is_monitor=True suite must surface + the unified missing-or-inaccessible error, not a distinct rejection. The filter + side-effect is in ``resolve_test_suite``; this test exercises it end-to-end with + the real TestSuite.get patched to behave as it would in DB (filter applied → None).""" + with patch("testgen.mcp.tools.common.TestSuite") as mock_ts: + mock_ts.get.return_value = None # filter excluded the monitor suite + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(), pytest.raises(MCPResourceNotAccessible) as exc: + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + assert "Test suite" in str(exc.value) + assert "not found or not accessible" in str(exc.value) + + +def test_update_test_suite_not_accessible(db_session_mock): + """Suite in a project the user can't access → unified wording.""" + with patch("testgen.mcp.tools.common.TestSuite") as mock_ts: + mock_ts.get.return_value = None + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(allowed=("other",)), pytest.raises(MCPResourceNotAccessible): + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + + +def test_update_test_suite_requires_edit(db_session_mock): + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(memberships={"demo": "role_c"}), pytest.raises(MCPPermissionDenied): + update_test_suite(test_suite_id=str(uuid4()), test_suite_name="Anything") + + +@patch(f"{MODULE}.resolve_test_suite") +def test_update_test_suite_allows_role_with_edit_but_not_administer(mock_resolve, db_session_mock): + """role_d has edit but NOT administer — must be allowed (mirrors production data_quality).""" + suite = _mock_test_suite(test_suite="DQ Suite") + mock_resolve.return_value = suite + + from testgen.mcp.tools.test_suites import update_test_suite + + with _patch_perms(memberships={"demo": "role_d"}): + out = update_test_suite(test_suite_id=str(suite.id), description="updated by data_quality") + + suite.save.assert_called_once() + assert "Test Suite `DQ Suite` updated" in out diff --git a/tests/unit/mcp/test_transport_security.py b/tests/unit/mcp/test_transport_security.py index 22b101ec..a2b6f048 100644 --- a/tests/unit/mcp/test_transport_security.py +++ b/tests/unit/mcp/test_transport_security.py @@ -30,14 +30,34 @@ def test_loopback_and_base_url_always_present(): assert "https://localhost:*" in settings.allowed_origins -def test_extra_host_without_port_gets_wildcard(): - """An extras entry without `:` gets `:*` automatically appended.""" +def test_ported_base_url_also_allows_bare_host_and_origin(): + """A BASE_URL with a non-default port also allows the bare (port-less) host and origin. + + Clients such as hosted MCP gateways may send a port-less Host/Origin for the server's + own host; the `:*` wildcard requires a port, so the bare forms must be present too. + """ + settings = _build_with("https://tg.example.com:8530") + + assert "tg.example.com:8530" in settings.allowed_hosts + assert "tg.example.com" in settings.allowed_hosts + assert "https://tg.example.com:8530" in settings.allowed_origins + assert "https://tg.example.com" in settings.allowed_origins + + +def test_extra_host_without_port_gets_wildcard_and_bare(): + """An extras entry without `:` is allowed both with a `:*` port wildcard and bare. + + The bare (port-less) form is required because some MCP gateways (e.g. Databricks) + send an Origin with no port, which the `:*` wildcard does not match. + """ settings = _build_with("http://localhost:8530", extras=["tg.example.com"]) assert "tg.example.com:*" in settings.allowed_hosts - assert "tg.example.com" not in settings.allowed_hosts # bare entry should NOT be present + assert "tg.example.com" in settings.allowed_hosts assert "http://tg.example.com:*" in settings.allowed_origins assert "https://tg.example.com:*" in settings.allowed_origins + assert "http://tg.example.com" in settings.allowed_origins + assert "https://tg.example.com" in settings.allowed_origins def test_extra_host_with_explicit_port_preserved_literally(): diff --git a/tests/unit/test_test_criteria.py b/tests/unit/test_test_criteria.py new file mode 100644 index 00000000..582c7d94 --- /dev/null +++ b/tests/unit/test_test_criteria.py @@ -0,0 +1,104 @@ +from pathlib import Path + +import pytest +from yaml import safe_load + +import testgen +from testgen.common.models.test_definition import TestCriteria, derive_test_criteria + +YAML_DIR = Path(testgen.__file__).parent / "template" / "dbsetup_test_types" + +# Authoritative test-type -> Criteria mapping. The derivation is pure Python (no column), +# shared by the UI lookup and MCP. This dict pins the expected facet value for every test type; +# redline here and the derivation together. +EXPECTED_CRITERIA = { + # Custom Criteria: user authors the test in SQL (custom scope) + "CUSTOM": TestCriteria.CUSTOM_CRITERIA, + "Condition_Flag": TestCriteria.CUSTOM_CRITERIA, + # Reference Dataset: compared against another dataset (referential scope) + "Aggregate_Balance": TestCriteria.REFERENCE_DATASET, + "Aggregate_Balance_Percent": TestCriteria.REFERENCE_DATASET, + "Aggregate_Balance_Range": TestCriteria.REFERENCE_DATASET, + "Aggregate_Minimum": TestCriteria.REFERENCE_DATASET, + "Combo_Match": TestCriteria.REFERENCE_DATASET, + "Distribution_Shift": TestCriteria.REFERENCE_DATASET, + "Timeframe_Combo_Gain": TestCriteria.REFERENCE_DATASET, + "Timeframe_Combo_Match": TestCriteria.REFERENCE_DATASET, + # List of Values: tested against a set/lookup of allowed values (Set / lookup algorithm) + "LOV_All": TestCriteria.LIST_OF_VALUES, + "LOV_Match": TestCriteria.LIST_OF_VALUES, + "US_State": TestCriteria.LIST_OF_VALUES, + # Defined Rule: predefined validity/integrity rule, no user-supplied value (enumerated) + "Dupe_Rows": TestCriteria.DEFINED_RULE, + "Unique": TestCriteria.DEFINED_RULE, + "Email_Format": TestCriteria.DEFINED_RULE, + "Street_Addr_Pattern": TestCriteria.DEFINED_RULE, + "Valid_Characters": TestCriteria.DEFINED_RULE, + "Valid_Month": TestCriteria.DEFINED_RULE, + "Valid_US_Zip": TestCriteria.DEFINED_RULE, + "Valid_US_Zip3": TestCriteria.DEFINED_RULE, + # Defined Value: user asserts the expected value/pattern (enumerated) + "Constant": TestCriteria.DEFINED_VALUE, + "Pattern_Match": TestCriteria.DEFINED_VALUE, + "Required": TestCriteria.DEFINED_VALUE, + # Defined Threshold: user tunes a numeric threshold/tolerance (fallthrough) + "Alpha_Trunc": TestCriteria.DEFINED_THRESHOLD, + "Avg_Shift": TestCriteria.DEFINED_THRESHOLD, + "Daily_Record_Ct": TestCriteria.DEFINED_THRESHOLD, + "Dec_Trunc": TestCriteria.DEFINED_THRESHOLD, + "Distinct_Date_Ct": TestCriteria.DEFINED_THRESHOLD, + "Distinct_Value_Ct": TestCriteria.DEFINED_THRESHOLD, + "Freshness_Trend": TestCriteria.DEFINED_THRESHOLD, + "Future_Date": TestCriteria.DEFINED_THRESHOLD, + "Future_Date_1Y": TestCriteria.DEFINED_THRESHOLD, + "Incr_Avg_Shift": TestCriteria.DEFINED_THRESHOLD, + "Metric_Trend": TestCriteria.DEFINED_THRESHOLD, + "Min_Date": TestCriteria.DEFINED_THRESHOLD, + "Min_Val": TestCriteria.DEFINED_THRESHOLD, + "Missing_Pct": TestCriteria.DEFINED_THRESHOLD, + "Monthly_Rec_Ct": TestCriteria.DEFINED_THRESHOLD, + "Outlier_Pct_Above": TestCriteria.DEFINED_THRESHOLD, + "Outlier_Pct_Below": TestCriteria.DEFINED_THRESHOLD, + "Recency": TestCriteria.DEFINED_THRESHOLD, + "Row_Ct": TestCriteria.DEFINED_THRESHOLD, + "Row_Ct_Pct": TestCriteria.DEFINED_THRESHOLD, + "Table_Freshness": TestCriteria.DEFINED_THRESHOLD, + "Unique_Pct": TestCriteria.DEFINED_THRESHOLD, + "Variability_Increase": TestCriteria.DEFINED_THRESHOLD, + "Variability_Decrease": TestCriteria.DEFINED_THRESHOLD, + "Volume_Trend": TestCriteria.DEFINED_THRESHOLD, + "Weekly_Rec_Ct": TestCriteria.DEFINED_THRESHOLD, +} + +# Test types that never surface in the Add-Test picker, so they carry no Criteria facet at all. +NON_PUBLIC_TEST_TYPES = frozenset({"Schema_Drift"}) + + +def _test_type_docs(): + for path in sorted([*YAML_DIR.glob("*.yaml"), *YAML_DIR.glob("*.yml")]): + yield safe_load(path.read_text())["test_types"] + + +@pytest.mark.unit +def test_every_test_type_has_expected_criteria(): + seen = set() + for tt in _test_type_docs(): + test_type = tt["test_type"] + if test_type in NON_PUBLIC_TEST_TYPES: + continue + seen.add(test_type) + criteria = derive_test_criteria(test_type, tt.get("test_scope"), tt.get("algorithm")) + assert criteria == EXPECTED_CRITERIA[test_type], ( + f"{test_type}: derived {criteria!r}, expected {EXPECTED_CRITERIA[test_type]!r}" + ) + missing = set(EXPECTED_CRITERIA) - seen + assert not missing, f"expected-criteria mapping references unknown test types: {sorted(missing)}" + + +@pytest.mark.unit +def test_derive_returns_a_valid_criteria_for_every_test_type(): + for tt in _test_type_docs(): + if tt["test_type"] in NON_PUBLIC_TEST_TYPES: + continue + criteria = derive_test_criteria(tt["test_type"], tt.get("test_scope"), tt.get("algorithm")) + assert isinstance(criteria, TestCriteria) diff --git a/tests/unit/test_test_type_taxonomy.py b/tests/unit/test_test_type_taxonomy.py new file mode 100644 index 00000000..123eed95 --- /dev/null +++ b/tests/unit/test_test_type_taxonomy.py @@ -0,0 +1,45 @@ +from pathlib import Path + +import pytest +from yaml import safe_load + +import testgen +from testgen.common.models.test_definition import StatisticalTechnique, TestAlgorithm + +YAML_DIR = Path(testgen.__file__).parent / "template" / "dbsetup_test_types" + +ALGORITHMS = {a.value for a in TestAlgorithm} +TECHNIQUES = {t.value for t in StatisticalTechnique} +HEALTH_DIMENSIONS = {"Schema Drift", "Data Drift", "Statistical Drift", "Volume", "Freshness"} + + +def _test_type_docs(): + for path in sorted([*YAML_DIR.glob("*.yaml"), *YAML_DIR.glob("*.yml")]): + data = safe_load(path.read_text())["test_types"] + yield path.name, data + + +@pytest.mark.unit +def test_every_test_type_declares_a_valid_algorithm(): + for filename, tt in _test_type_docs(): + algorithm = tt.get("algorithm") + assert algorithm in ALGORITHMS, f"{filename}: invalid/missing algorithm {algorithm!r}" + + +@pytest.mark.unit +def test_statistical_technique_is_valid_when_present(): + for filename, tt in _test_type_docs(): + algorithm = tt.get("algorithm") + technique = tt.get("statistical_technique") + if technique is not None: + assert technique in TECHNIQUES, f"{filename}: invalid statistical_technique {technique!r}" + if algorithm == TestAlgorithm.STATISTICAL_DRIFT.value: + assert technique in TECHNIQUES, f"{filename}: drift test needs a valid technique, got {technique!r}" + + +@pytest.mark.unit +def test_health_dimension_uses_freshness_not_recency(): + for filename, tt in _test_type_docs(): + health = tt.get("health_dimension") + assert health != "Recency", f"{filename}: health_dimension must be 'Freshness', not 'Recency'" + assert health is None or health in HEALTH_DIMENSIONS, f"{filename}: unexpected health_dimension {health!r}" diff --git a/tests/unit/ui/conftest.py b/tests/unit/ui/conftest.py index 0e70ee57..cd74174c 100644 --- a/tests/unit/ui/conftest.py +++ b/tests/unit/ui/conftest.py @@ -9,3 +9,4 @@ # widgets/__init__.py calls components_v2.component() at module level; stub it out so # views that import widgets don't fail collection without a Streamlit runtime. sys.modules.setdefault("testgen.ui.components.widgets", MagicMock()) +sys.modules.setdefault("testgen.ui.components.widgets.download_dialog", MagicMock()) diff --git a/tests/unit/ui/test_data_catalog.py b/tests/unit/ui/test_data_catalog.py new file mode 100644 index 00000000..51e15e93 --- /dev/null +++ b/tests/unit/ui/test_data_catalog.py @@ -0,0 +1,39 @@ +import pytest + +from testgen.ui.views.data_catalog import merge_tag_defaults + +pytestmark = pytest.mark.unit + + +class Test_MergeTagDefaults: + def test_db_value_beats_matching_seed_case_insensitively(self): + values = {"data_classification": ["confidential", "public"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["confidential", "Internal", "public", "Restricted"] + + def test_seed_added_when_no_db_match(self): + values = {"data_classification": []} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential", "Internal", "Public", "Restricted"] + + def test_result_sorted_case_insensitively(self): + values = {"data_classification": ["RESTRICTED"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential", "Internal", "Public", "RESTRICTED"] + + def test_non_classified_tags_unaffected(self): + values = {"data_product": ["prod-a"], "data_classification": ["internal"]} + defaults = {"data_classification": ["Confidential", "Internal", "Public", "Restricted"]} + result = merge_tag_defaults(values, defaults) + assert result["data_product"] == ["prod-a"] + assert "Internal" not in result["data_classification"] + assert "internal" in result["data_classification"] + + def test_empty_db_values_for_tagged_field(self): + values = {} + defaults = {"data_classification": ["Confidential"]} + result = merge_tag_defaults(values, defaults) + assert result["data_classification"] == ["Confidential"] diff --git a/tests/unit/ui/test_import_metadata.py b/tests/unit/ui/test_import_metadata.py index e5f7f718..627958f8 100644 --- a/tests/unit/ui/test_import_metadata.py +++ b/tests/unit/ui/test_import_metadata.py @@ -1,4 +1,5 @@ import base64 +from unittest.mock import patch import pandas as pd import pytest @@ -6,6 +7,7 @@ from testgen.ui.views.dialogs.import_metadata_dialog import ( DESCRIPTION_MAX_LENGTH, TAG_MAX_LENGTH, + _build_metadata_fields, _extract_metadata_fields, _parse_csv, _set_row_status, @@ -326,3 +328,43 @@ def test_preview_props_unmatched_preserved(): } result = build_import_preview_props(preview) assert result["preview_rows"][0]["_status"] == "unmatched" + + +# --- _build_metadata_fields --- + +_ALL_COLUMNS = ["description", "critical_data_element", "excluded_data_element", "pii_flag", "data_source"] + + +def test_build_fields_table_excludes_column_only_fields(): + row = { + "description": "Orders table", + "critical_data_element": True, + "excluded_data_element": True, + "pii_flag": "MANUAL", + "data_source": "Salesforce", + } + fields = _build_metadata_fields(row, _ALL_COLUMNS, is_column=False) + assert fields == {"description": "Orders table", "critical_data_element": True, "data_source": "Salesforce"} + + +def test_build_fields_blank_text_clears(): + row = {"description": "", "data_source": ""} + fields = _build_metadata_fields(row, ["description", "data_source"], is_column=False) + assert fields == {"description": None, "data_source": None} + + +def test_build_fields_column_includes_pii_when_permitted(): + row = {"pii_flag": "MANUAL", "excluded_data_element": True} + with patch("testgen.ui.views.dialogs.import_metadata_dialog.session") as mock_session: + mock_session.auth.user_has_permission.return_value = True + fields = _build_metadata_fields(row, ["pii_flag", "excluded_data_element"], is_column=True) + assert fields == {"pii_flag": "MANUAL", "excluded_data_element": True} + + +def test_build_fields_column_drops_pii_without_permission(): + row = {"pii_flag": "MANUAL", "excluded_data_element": True} + with patch("testgen.ui.views.dialogs.import_metadata_dialog.session") as mock_session: + mock_session.auth.user_has_permission.return_value = False + fields = _build_metadata_fields(row, ["pii_flag", "excluded_data_element"], is_column=True) + assert fields == {"excluded_data_element": True} + assert "pii_flag" not in fields diff --git a/tests/unit/ui/test_monitors_dashboard.py b/tests/unit/ui/test_monitors_dashboard.py new file mode 100644 index 00000000..75f06cf0 --- /dev/null +++ b/tests/unit/ui/test_monitors_dashboard.py @@ -0,0 +1,157 @@ +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +from testgen.common.monitor_forecast import gated_forecast_prediction +from testgen.ui.views.monitors_dashboard import _freshness_next_update_window + +pytestmark = pytest.mark.unit + +# The business-time helpers live in (and are called from) the shared forecast +# module, so patch them there even when exercising the dashboard wrapper. +MODULE = "testgen.common.monitor_forecast" + + +def _freshness_def(history_calculation="PREDICT", prediction=None, upper="1000", lower="500"): + d = MagicMock() + d.history_calculation = history_calculation + d.prediction = {"schedule_stage": "active"} if prediction is None else prediction + d.upper_tolerance = upper + d.lower_tolerance = lower + return d + + +def _events(changed=True, is_training=False, is_pending=False, time=datetime(2026, 6, 22, 16, 0)): + return {"freshness_events": [{"changed": changed, "is_training": is_training, "is_pending": is_pending, "time": time}]} + + +def _suite(): + s = MagicMock() + s.predict_exclude_weekends = False + s.holiday_codes_list = None + return s + + +def _schedule(): + s = MagicMock() + s.cron_tz = "UTC" + return s + + +# --- _freshness_next_update_window: guard branches return None --- + + +def test_window_none_when_no_freshness_definition(): + assert _freshness_next_update_window(None, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_not_predict(): + d = _freshness_def(history_calculation=None) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_prediction_lacks_schedule_stage(): + # Non-empty prediction without schedule_stage (e.g. still learning) → no window + d = _freshness_def(prediction={"frequency": "daily"}) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_no_upper_tolerance(): + d = _freshness_def(upper=None) + assert _freshness_next_update_window(d, _events(), _suite(), _schedule()) is None + + +def test_window_none_when_no_qualifying_update_events(): + # Only training / pending / unchanged events → nothing to anchor the window to + events = {"freshness_events": [ + {"changed": True, "is_training": True, "is_pending": False, "time": datetime(2026, 6, 22)}, + {"changed": False, "is_training": False, "is_pending": False, "time": datetime(2026, 6, 22)}, + ]} + assert _freshness_next_update_window(_freshness_def(), events, _suite(), _schedule()) is None + + +@patch(f"{MODULE}.get_schedule_params", return_value=MagicMock(excluded_days=[])) +@patch(f"{MODULE}.add_business_minutes") +def test_window_returns_start_end_when_valid(mock_abm, _mock_sched): + # add_business_minutes is called for the end (upper_tolerance) then the start (lower_tolerance) + mock_abm.side_effect = [pd.Timestamp("2026-06-23 19:00"), pd.Timestamp("2026-06-23 04:00")] + window = _freshness_next_update_window(_freshness_def(), _events(), _suite(), _schedule()) + assert window == { + "start": int(pd.Timestamp("2026-06-23 04:00").timestamp() * 1000), + "end": int(pd.Timestamp("2026-06-23 19:00").timestamp() * 1000), + } + + +@patch(f"{MODULE}.get_schedule_params", return_value=MagicMock(excluded_days=[])) +@patch(f"{MODULE}.add_business_minutes") +def test_window_start_is_none_when_no_lower_tolerance(mock_abm, _mock_sched): + mock_abm.return_value = pd.Timestamp("2026-06-23 19:00") + window = _freshness_next_update_window(_freshness_def(lower=None), _events(), _suite(), _schedule()) + assert window["start"] is None + assert window["end"] == int(pd.Timestamp("2026-06-23 19:00").timestamp() * 1000) + + +# --- gated_forecast_prediction --- + +LAST_RUN = datetime(2026, 6, 23, 16, 0) +NOW_MS = int(pd.Timestamp(LAST_RUN).timestamp() * 1000) +HOUR = 3600 * 1000 + + +def _gated_def(baseline=1000.0, mean=None, lower="950", upper="1400"): + d = MagicMock() + d.prediction = {"freshness_gated": True, "baseline_value": baseline} + if mean is not None: + d.prediction["mean"] = mean + d.lower_tolerance = lower + d.upper_tolerance = upper + return d + + +def test_gated_prediction_none_without_window(): + assert gated_forecast_prediction(_gated_def(), None, LAST_RUN) is None + + +def test_gated_prediction_none_when_window_elapsed(): + window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS - HOUR} # window_end already in the past + assert gated_forecast_prediction(_gated_def(), window, LAST_RUN) is None + + +def test_gated_prediction_none_without_baseline(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + assert gated_forecast_prediction(_gated_def(baseline=None), window, LAST_RUN) is None + + +def test_gated_prediction_none_without_last_run(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + assert gated_forecast_prediction(_gated_def(), window, None) is None + + +def test_gated_prediction_anchors_at_now_when_window_started(): + # window_start is before the latest run → anchor clamps to now (forecast never draws backward) + window = {"start": NOW_MS - 12 * HOUR, "end": NOW_MS + 3 * HOUR} + mean = {str(NOW_MS + 3 * HOUR): 1200.0, str(NOW_MS + 27 * HOUR): 1300.0} + result = gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + assert result["method"] == "predict" + # mean holds flat at baseline then steps to the next-refresh forecast + assert result["mean"] == {NOW_MS: 1000.0, NOW_MS + 3 * HOUR: 1200.0} + # band carries the next-refresh tolerance across the whole forecast (continuous width, no + # zero-width pinch at the flat anchor); coerced from VARCHAR to float + assert result["lower_tolerance"] == {NOW_MS: 950.0, NOW_MS + 3 * HOUR: 950.0} + assert result["upper_tolerance"] == {NOW_MS: 1400.0, NOW_MS + 3 * HOUR: 1400.0} + + +def test_gated_prediction_anchors_at_window_start_when_future(): + # window opens after the latest run → flat segment runs out to window_start + window = {"start": NOW_MS + HOUR, "end": NOW_MS + 3 * HOUR} + mean = {str(NOW_MS + 3 * HOUR): 1200.0} + result = gated_forecast_prediction(_gated_def(mean=mean), window, LAST_RUN) + assert set(result["mean"].keys()) == {NOW_MS + HOUR, NOW_MS + 3 * HOUR} + + +def test_gated_prediction_next_mean_falls_back_to_baseline_without_forecast(): + window = {"start": NOW_MS - HOUR, "end": NOW_MS + 3 * HOUR} + result = gated_forecast_prediction(_gated_def(mean=None), window, LAST_RUN) + assert result["mean"][NOW_MS + 3 * HOUR] == 1000.0 # baseline used as the step value
{{profiling_run.log_message}}
View {{format_number issue_count}} issues > @@ -241,6 +242,8 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct if not notifications: return + job_execution = profiling_run.job_execution + previous_run = profiling_run.get_previous() issues = list( HygieneIssue.select_with_diff( @@ -251,7 +254,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct ) triggers = {ProfilingRunNotificationTrigger.always} - if profiling_run.status in ("Error", "Cancelled") or {None, True} & {is_new for _, is_new in issues}: + if job_execution.status in (JobStatus.ERROR, JobStatus.CANCELED) or {None, True} & {is_new for _, is_new in issues}: triggers.add(ProfilingRunNotificationTrigger.on_changes) notifications = [ns for ns in notifications if ns.trigger in triggers] @@ -263,7 +266,7 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct settings.UI_BASE_URL, "/profiling-runs:hygiene?project_code=", str(profiling_run.project_code), - "&run_id=", str(profiling_run.job_execution_id), + "&run_id=", str(profiling_run.id), "&source=email" ) ) @@ -316,14 +319,15 @@ def send_profiling_run_notifications(profiling_run: ProfilingRun, result_list_ct "/profiling-runs:results?project_code=", str(profiling_run.project_code), "&run_id=", - str(profiling_run.job_execution_id), + str(profiling_run.id), "&source=email" ) ), - "start_time": profiling_run.profiling_starttime, - "end_time": profiling_run.profiling_endtime, - "status": profiling_run.status, - "log_message": profiling_run.log_message, + "start_time": job_execution.started_at, + "end_time": job_execution.completed_at, + "status": job_execution.status, + "status_label": JOB_STATUS_LABEL.get(job_execution.status, job_execution.status), + "log_message": job_execution.error_message, "table_ct": profiling_run.table_ct, "column_ct": profiling_run.column_ct, }, diff --git a/testgen/common/notifications/test_run.py b/testgen/common/notifications/test_run.py index 4f5985a8..72ad5140 100644 --- a/testgen/common/notifications/test_run.py +++ b/testgen/common/notifications/test_run.py @@ -3,6 +3,7 @@ from sqlalchemy import case, literal, select from testgen import settings +from testgen.common.enums import JobStatus from testgen.common.models import get_current_session, with_database_session from testgen.common.models.notification_settings import ( TestRunNotificationSettings, @@ -136,7 +137,7 @@ def get_main_content_template(self): {{/if}} {{#if (eq test_run.status 'error')}}
{{test_run.log_message}}
{{test_run.error_message}}