From f224eb56fde54a18165db94397743b729d50f301 Mon Sep 17 00:00:00 2001 From: Bart Schuijt Date: Mon, 7 Sep 2026 22:23:20 +0200 Subject: [PATCH 1/3] feat(registry): unified CheckRegistry and granular rule execution (tff#146) --- README.md | 5 + packages/tff-core/src/tff/core/health.py | 112 +--- packages/tff-core/src/tff/core/registry.py | 635 +++++++++++++++++++ packages/tff-core/src/tff/core/report.py | 100 +-- packages/tff-core/src/tff/dataform/runner.py | 85 +-- packages/tff-core/src/tff/dbt/runner.py | 88 +-- packages/tff-core/src/tff/sqlmesh/runner.py | 127 ++-- packages/tff-core/tests/test_registry.py | 509 +++++++++++++++ 8 files changed, 1288 insertions(+), 373 deletions(-) create mode 100644 packages/tff-core/src/tff/core/registry.py create mode 100644 packages/tff-core/tests/test_registry.py diff --git a/README.md b/README.md index 2ea99f8..9e7d03e 100644 --- a/README.md +++ b/README.md @@ -127,6 +127,11 @@ Run linting on the current project: tff lint ``` +Run specific fitness checks or rules: +```bash +tff lint --checks no_missing_owner,ban_select_star +``` + Automatically fix simple linting violations (positional GROUP BY/ORDER BY, missing owner/description metadata): ```bash tff lint --fix diff --git a/packages/tff-core/src/tff/core/health.py b/packages/tff-core/src/tff/core/health.py index 3a809c7..4f32416 100644 --- a/packages/tff-core/src/tff/core/health.py +++ b/packages/tff-core/src/tff/core/health.py @@ -15,109 +15,17 @@ from tff.core.config import FitnessFunctionsConfig from tff.core.report import CHECK_LABELS, CONNASCENCE_CATEGORIES, LintFinding -PROJECT_LEVEL_CHECKS = { - "layer_integrity", - "custom_exclusions", - "schema_contracts", - "dependency_graph", - "materialization_depth", -} - -CATEGORIES = { - "Connascence of Name (CoN)": [ - "banselectstar", - "filenameequalsmodelname", - "columnnames", - "martmodelnamingconvention", - "ambiguousorinvalidcolumn", - "invalidselectstarexpansion", - ], - "Connascence of Type (CoT)": [ - "columntypes", - "schema_contracts", - ], - "Connascence of Position (CoP)": [ - "nopositionalgroupbyororderby", - ], - "Connascence of Meaning (CoM)": [ - "classificationmacros", - ], - "Connascence of Algorithm (CoA)": [ - "duplicate_ctes", - ], - "Connascence of Value (CoV)": [ - "connascence_of_value", - ], - "Dynamic Coupling & DAG Structure": [ - "layer_integrity", - "custom_exclusions", - "dependency_graph", - "materialization_depth", - "environmentagnosticreferences", - ], - "Quality & Metadata (Non-Connascence)": [ - "nomissingowner", - "nomissingdescription", - "nomissinggrain", - "nomissingnotnull", - "nomissinguniquevalues", - "sqlcomplexity", - ], -} - - -def is_check_enabled(config: FitnessFunctionsConfig, check_name: str, provider: str) -> bool: +from tff.core.registry import registry + +PROJECT_LEVEL_CHECKS: set[str] = registry.get_project_level_check_names() +CATEGORIES: dict[str, list[str]] = registry.get_categories() + + +def is_check_enabled( + config: FitnessFunctionsConfig, check_name: str, provider: str +) -> bool: """Determine if a check/rule is enabled in the configuration.""" - if check_name == "layer_integrity": - return config.checks.layer_integrity.enabled - if check_name == "custom_exclusions": - return config.checks.custom_exclusions.enabled - if check_name == "schema_contracts": - return config.checks.schema_contracts.enabled - if check_name == "dependency_graph": - return config.checks.dependency_graph.enabled - if check_name == "materialization_depth": - return config.checks.materialization_depth.enabled - if check_name == "duplicate_ctes": - return config.checks.duplicate_ctes.enabled - if check_name == "connascence_of_value": - return config.checks.connascence_of_value.enabled - if check_name == "classificationmacros": - return config.rules.classification_macros.enabled - if check_name == "sqlcomplexity": - return config.rules.sql_complexity.enabled - if check_name == "martmodelnamingconvention": - return config.rules.mart_naming.enabled - if check_name == "columnnames": - return config.rules.column_names.enabled - if check_name == "columntypes": - return config.rules.column_types.enabled - if check_name == "filenameequalsmodelname": - return config.rules.filename_equals_modelname.enabled - if check_name == "banselectstar": - return config.rules.ban_select_star.enabled - if check_name == "nopositionalgroupbyororderby": - return config.rules.no_positional_group_by_or_order_by.enabled - if check_name == "environmentagnosticreferences": - return config.rules.environment_agnostic_references.enabled - - # Metadata sub-rules - if check_name == "nomissingowner": - return config.rules.metadata.enabled and config.rules.metadata.owner - if check_name == "nomissingdescription": - return config.rules.metadata.enabled and config.rules.metadata.description - if check_name == "nomissinggrain": - return config.rules.metadata.enabled and config.rules.metadata.grain - if check_name == "nomissingnotnull": - return config.rules.metadata.enabled and config.rules.metadata.not_null - if check_name == "nomissinguniquevalues": - return config.rules.metadata.enabled and config.rules.metadata.unique_values - - # SQLMesh native rules - if check_name in {"ambiguousorinvalidcolumn", "invalidselectstarexpansion"}: - return provider == "sqlmesh" - - return False + return registry.is_check_enabled(config, check_name, provider) def _matches_scope(finding_path: str | None, scope: list[str]) -> bool: diff --git a/packages/tff-core/src/tff/core/registry.py b/packages/tff-core/src/tff/core/registry.py new file mode 100644 index 0000000..ddf6694 --- /dev/null +++ b/packages/tff-core/src/tff/core/registry.py @@ -0,0 +1,635 @@ +"""Centralized, declarative CheckRegistry for fitness checks and rules.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Literal + +if TYPE_CHECKING: + from tff.core.config import FitnessFunctionsConfig + from tff.core.model import ModelRepresentation + from tff.core.report import LintFinding, Severity + from tff.core.rules.base import Rule + +Scope = Literal["model", "dag"] + + +def normalize_check_name(name: str) -> str: + """Normalize a check name/alias for case-insensitive and separator-agnostic lookups.""" + return name.lower().replace("-", "").replace("_", "").replace(" ", "") + + +def run_model_rule( + rule_cls: type[Rule], + models: dict[str, ModelRepresentation], + severity: Severity = "error", + check_name: str | None = None, +) -> list[LintFinding]: + """Execute a single model-level Rule across all eligible models in a project.""" + from tff.core.report import LintFinding + from tff.core.utils.paths import model_path_relative + + rule = rule_cls() + findings: list[LintFinding] = [] + finding_check = check_name or getattr(rule, "name", rule_cls.__name__.lower()) + + for model in models.values(): + if model.is_external or model.is_symbolic: + continue + + violation = rule.check_model(model) + if violation: + msgs = violation.violation_msg + if isinstance(msgs, str): + msgs = [msgs] + for msg in msgs: + model_label = f"{model.name}: " + clean_msg = msg.removeprefix(model_label) + findings.append( + LintFinding( + check=finding_check, + severity=severity, + model=model.name, + path=model_path_relative(model), + message=clean_msg, + ) + ) + return findings + + +@dataclass(frozen=True) +class CheckDefinition: + """Metadata and execution specification for a single fitness check or rule.""" + + id: str + label: str + category: str + scope: Scope + default_severity: Severity = "error" + aliases: tuple[str, ...] = () + finding_check_id: str | None = None + rule_cls: type[Rule] | None = None + rule_module: str | None = None + rule_class_name: str | None = None + collector_fn: ( + Callable[[dict[str, ModelRepresentation], FitnessFunctionsConfig], list[LintFinding]] + | None + ) = None + collector_module: str | None = None + collector_func_name: str | None = None + is_enabled_fn: Callable[[FitnessFunctionsConfig, str], bool] | None = None + + @property + def canonical_id(self) -> str: + return self.id + + @property + def finding_id(self) -> str: + return self.finding_check_id or (self.aliases[0] if self.aliases else self.id) + + def get_rule_cls(self) -> type[Rule] | None: + if self.rule_cls is not None: + return self.rule_cls + if self.rule_module and self.rule_class_name: + import importlib + + mod = importlib.import_module(self.rule_module) + return getattr(mod, self.rule_class_name) + return None + + def get_collector_fn(self) -> Callable | None: + if self.collector_fn is not None: + return self.collector_fn + if self.collector_module and self.collector_func_name: + import importlib + + mod = importlib.import_module(self.collector_module) + return getattr(mod, self.collector_func_name) + return None + + def is_enabled(self, config: FitnessFunctionsConfig, provider: str = "dbt") -> bool: + if self.is_enabled_fn is not None: + return self.is_enabled_fn(config, provider) + return True + + def run( + self, + models: dict[str, ModelRepresentation], + config: FitnessFunctionsConfig, + ) -> list[LintFinding]: + if self.scope == "model": + rule_cls = self.get_rule_cls() + if rule_cls is not None: + return run_model_rule( + rule_cls, + models, + severity=self.default_severity, + check_name=self.finding_id, + ) + return [] + elif self.scope == "dag": + collector = self.get_collector_fn() + if collector is not None: + return collector(models, config) + return [] + return [] + + +class CheckRegistry: + """Registry holding definitions of all fitness checks and rules.""" + + def __init__(self) -> None: + self._checks: dict[str, CheckDefinition] = {} + self._lookup: dict[str, CheckDefinition] = {} + + def register(self, check: CheckDefinition) -> None: + self._checks[check.id] = check + self._lookup[normalize_check_name(check.id)] = check + if check.finding_check_id: + self._lookup[normalize_check_name(check.finding_check_id)] = check + for alias in check.aliases: + self._lookup[normalize_check_name(alias)] = check + + def get(self, name_or_alias: str) -> CheckDefinition | None: + return self._lookup.get(normalize_check_name(name_or_alias)) + + def get_or_raise(self, name_or_alias: str) -> CheckDefinition: + check = self.get(name_or_alias) + if check is None: + available = sorted(self._checks.keys()) + raise ValueError( + f"Unknown check or rule: '{name_or_alias}'. " + f"Available checks: {', '.join(available)}" + ) + return check + + def all_checks(self) -> list[CheckDefinition]: + return list(self._checks.values()) + + def model_rules(self) -> list[CheckDefinition]: + return [ + c + for c in self._checks.values() + if c.scope == "model" and (c.rule_cls is not None or c.rule_module is not None) + ] + + def dag_checks(self) -> list[CheckDefinition]: + return [c for c in self._checks.values() if c.scope == "dag"] + + def is_check_enabled( + self, + config: FitnessFunctionsConfig, + check_name: str, + provider: str = "dbt", + ) -> bool: + check = self.get(check_name) + if check is None: + return False + return check.is_enabled(config, provider) + + def resolve_checks( + self, + checks: list[str] | None, + config: FitnessFunctionsConfig, + provider: str = "dbt", + ) -> list[CheckDefinition]: + if checks is None: + return [c for c in self.all_checks() if c.is_enabled(config, provider)] + + resolved: list[CheckDefinition] = [] + seen: set[str] = set() + + for name in checks: + norm = normalize_check_name(name) + if norm == "rules": + for rule_def in self.model_rules(): + if rule_def.id not in seen: + seen.add(rule_def.id) + resolved.append(rule_def) + elif norm == "sqlmesh": + # Container key for SQLMesh linter + continue + else: + check_def = self.get(name) + if check_def is not None: + if check_def.id not in seen: + seen.add(check_def.id) + resolved.append(check_def) + else: + self.get_or_raise(name) + + return resolved + + def run_checks( + self, + models: dict[str, ModelRepresentation], + config: FitnessFunctionsConfig, + checks: list[str] | None = None, + provider: str = "dbt", + ) -> tuple[list[LintFinding], list[str]]: + resolved = self.resolve_checks(checks, config, provider=provider) + findings: list[LintFinding] = [] + for check_def in resolved: + findings.extend(check_def.run(models, config)) + + if checks is not None: + executed_names = checks + else: + executed_names = ["rules"] + [ + c.finding_id for c in resolved if c.scope == "dag" + ] + + return findings, executed_names + + def get_check_labels(self) -> dict[str, str]: + labels: dict[str, str] = {} + for c in self.all_checks(): + labels[c.id] = c.label + if c.finding_check_id: + labels[c.finding_check_id] = c.label + for alias in c.aliases: + labels[alias] = c.label + return labels + + def get_connascence_categories(self) -> dict[str, str]: + cats: dict[str, str] = {} + for c in self.all_checks(): + cats[c.id] = c.category + if c.finding_check_id: + cats[c.finding_check_id] = c.category + for alias in c.aliases: + cats[alias] = c.category + return cats + + def get_categories(self) -> dict[str, list[str]]: + cats: dict[str, list[str]] = { + "Connascence of Name (CoN)": [], + "Connascence of Type (CoT)": [], + "Connascence of Position (CoP)": [], + "Connascence of Meaning (CoM)": [], + "Connascence of Algorithm (CoA)": [], + "Connascence of Value (CoV)": [], + "Dynamic Coupling & DAG Structure": [], + "Quality & Metadata (Non-Connascence)": [], + } + for c in self.all_checks(): + key = c.finding_id + if c.category in cats and key not in cats[c.category]: + cats[c.category].append(key) + return cats + + def get_project_level_check_names(self) -> set[str]: + return { + "layer_integrity", + "custom_exclusions", + "schema_contracts", + "dependency_graph", + "materialization_depth", + } + + def get_architectural_check_names(self) -> frozenset[str]: + return frozenset( + { + "layer_integrity", + "custom_exclusions", + "schema_contracts", + "dependency_graph", + "duplicate_ctes", + "connascence_of_value", + } + ) + + +def create_default_registry() -> CheckRegistry: + """Create and populate the default CheckRegistry with all TFF checks and rules.""" + reg = CheckRegistry() + + # 1. Connascence of Name (CoN) + reg.register( + CheckDefinition( + id="ban_select_star", + label="No SELECT *", + category="Connascence of Name (CoN)", + scope="model", + aliases=("banselectstar",), + finding_check_id="banselectstar", + rule_module="tff.core.rules.ban_select_star", + rule_class_name="BanSelectStar", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.ban_select_star.enabled), + ) + ) + reg.register( + CheckDefinition( + id="filename_equals_modelname", + label="Filename equals model name", + category="Connascence of Name (CoN)", + scope="model", + aliases=("filenameequalsmodelname",), + finding_check_id="filenameequalsmodelname", + rule_module="tff.core.rules.filename_equals_modelname", + rule_class_name="FilenameEqualsModelname", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.filename_equals_modelname.enabled), + ) + ) + reg.register( + CheckDefinition( + id="column_names", + label="Column names", + category="Connascence of Name (CoN)", + scope="model", + aliases=("columnnames",), + finding_check_id="columnnames", + rule_module="tff.core.rules.column_names", + rule_class_name="ColumnNames", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.column_names.enabled), + ) + ) + reg.register( + CheckDefinition( + id="mart_model_naming_convention", + label="Mart naming convention", + category="Connascence of Name (CoN)", + scope="model", + aliases=("martmodelnamingconvention", "mart_naming"), + finding_check_id="martmodelnamingconvention", + rule_module="tff.core.rules.mart_naming", + rule_class_name="MartModelNamingConvention", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.mart_naming.enabled), + ) + ) + reg.register( + CheckDefinition( + id="ambiguous_or_invalid_column", + label="Ambiguous/invalid column", + category="Connascence of Name (CoN)", + scope="model", + aliases=("ambiguousorinvalidcolumn",), + finding_check_id="ambiguousorinvalidcolumn", + is_enabled_fn=lambda cfg, p: p == "sqlmesh", + ) + ) + reg.register( + CheckDefinition( + id="invalid_select_star_expansion", + label="Invalid SELECT * expansion", + category="Connascence of Name (CoN)", + scope="model", + aliases=("invalidselectstarexpansion",), + finding_check_id="invalidselectstarexpansion", + is_enabled_fn=lambda cfg, p: p == "sqlmesh", + ) + ) + + # 2. Connascence of Type (CoT) + reg.register( + CheckDefinition( + id="column_types", + label="Column types", + category="Connascence of Type (CoT)", + scope="model", + aliases=("columntypes",), + finding_check_id="columntypes", + rule_module="tff.core.rules.column_types", + rule_class_name="ColumnTypes", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.column_types.enabled), + ) + ) + reg.register( + CheckDefinition( + id="schema_contracts", + label="Schema contracts", + category="Connascence of Type (CoT)", + scope="dag", + collector_module="tff.core.checks.schema_contracts", + collector_func_name="collect_schema_contract_findings", + collector_fn=lambda _models, cfg: __import__( + "tff.core.checks.schema_contracts", fromlist=["collect_schema_contract_findings"] + ).collect_schema_contract_findings(cfg), + is_enabled_fn=lambda cfg, p: bool(cfg.checks.schema_contracts.enabled), + ) + ) + + # 3. Connascence of Position (CoP) + reg.register( + CheckDefinition( + id="no_positional_group_by_or_order_by", + label="No positional GROUP BY or ORDER BY", + category="Connascence of Position (CoP)", + scope="model", + aliases=("nopositionalgroupbyororderby",), + finding_check_id="nopositionalgroupbyororderby", + rule_module="tff.core.rules.no_positional_group_by_or_order_by", + rule_class_name="NoPositionalGroupByOrOrderBy", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.no_positional_group_by_or_order_by.enabled + ), + ) + ) + + # 4. Connascence of Meaning (CoM) + reg.register( + CheckDefinition( + id="classification_macros", + label="Classification macros", + category="Connascence of Meaning (CoM)", + scope="model", + aliases=("classificationmacros",), + finding_check_id="classificationmacros", + rule_module="tff.core.rules.classification_macros", + rule_class_name="ClassificationMacros", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.classification_macros.enabled), + ) + ) + + # 5. Connascence of Algorithm (CoA) + reg.register( + CheckDefinition( + id="duplicate_ctes", + label="Duplicate CTEs", + category="Connascence of Algorithm (CoA)", + scope="dag", + collector_module="tff.core.checks.duplicate_ctes", + collector_func_name="collect_duplicate_cte_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.duplicate_ctes.enabled), + ) + ) + + # 6. Connascence of Value (CoV) + reg.register( + CheckDefinition( + id="connascence_of_value", + label="Connascence of Value", + category="Connascence of Value (CoV)", + scope="dag", + collector_module="tff.core.checks.connascence_of_value", + collector_func_name="collect_connascence_of_value_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.connascence_of_value.enabled), + ) + ) + + # 7. Dynamic Coupling & DAG Structure + reg.register( + CheckDefinition( + id="layer_integrity", + label="Layer integrity", + category="Dynamic Coupling & DAG Structure", + scope="dag", + collector_module="tff.core.checks.layer_integrity", + collector_func_name="collect_layer_integrity_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.layer_integrity.enabled), + ) + ) + reg.register( + CheckDefinition( + id="custom_exclusions", + label="Custom exclusions", + category="Dynamic Coupling & DAG Structure", + scope="dag", + collector_module="tff.core.checks.custom_exclusions", + collector_func_name="collect_custom_exclusion_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.custom_exclusions.enabled), + ) + ) + reg.register( + CheckDefinition( + id="dependency_graph", + label="Dependency graph", + category="Dynamic Coupling & DAG Structure", + scope="dag", + collector_module="tff.core.checks.dependency_graph", + collector_func_name="collect_dependency_graph_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.dependency_graph.enabled), + ) + ) + reg.register( + CheckDefinition( + id="materialization_depth", + label="materialization_depth", + category="Dynamic Coupling & DAG Structure", + scope="dag", + collector_module="tff.core.checks.materialization_depth", + collector_func_name="collect_materialization_depth_findings", + is_enabled_fn=lambda cfg, p: bool(cfg.checks.materialization_depth.enabled), + ) + ) + reg.register( + CheckDefinition( + id="environment_agnostic_references", + label="Environment-agnostic references", + category="Dynamic Coupling & DAG Structure", + scope="model", + aliases=("environmentagnosticreferences",), + finding_check_id="environmentagnosticreferences", + rule_module="tff.core.rules.environment_agnostic_references", + rule_class_name="EnvironmentAgnosticReferences", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.environment_agnostic_references.enabled + ), + ) + ) + + # 8. Quality & Metadata (Non-Connascence) + reg.register( + CheckDefinition( + id="no_missing_owner", + label="Missing owner", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissingowner",), + finding_check_id="nomissingowner", + rule_module="tff.core.rules.metadata", + rule_class_name="NoMissingOwner", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.metadata.enabled and cfg.rules.metadata.owner + ), + ) + ) + reg.register( + CheckDefinition( + id="no_missing_description", + label="Missing description", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissingdescription",), + finding_check_id="nomissingdescription", + rule_module="tff.core.rules.metadata", + rule_class_name="NoMissingDescription", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.metadata.enabled and cfg.rules.metadata.description + ), + ) + ) + reg.register( + CheckDefinition( + id="no_missing_grain", + label="Missing grain", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissinggrain",), + finding_check_id="nomissinggrain", + rule_module="tff.core.rules.metadata", + rule_class_name="NoMissingGrain", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.metadata.enabled and cfg.rules.metadata.grain + ), + ) + ) + reg.register( + CheckDefinition( + id="no_missing_not_null", + label="Missing not_null audit", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissingnotnull",), + finding_check_id="nomissingnotnull", + rule_module="tff.core.rules.metadata", + rule_class_name="NoMissingNotNull", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.metadata.enabled and cfg.rules.metadata.not_null + ), + ) + ) + reg.register( + CheckDefinition( + id="no_missing_unique_values", + label="Missing unique_values audit", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissinguniquevalues",), + finding_check_id="nomissinguniquevalues", + rule_module="tff.core.rules.metadata", + rule_class_name="NoMissingUniqueValues", + is_enabled_fn=lambda cfg, p: bool( + cfg.rules.metadata.enabled and cfg.rules.metadata.unique_values + ), + ) + ) + reg.register( + CheckDefinition( + id="no_missing_audits", + label="Missing audits", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("nomissingaudits",), + finding_check_id="nomissingaudits", + is_enabled_fn=lambda cfg, p: False, + ) + ) + reg.register( + CheckDefinition( + id="sql_complexity", + label="SQL complexity", + category="Quality & Metadata (Non-Connascence)", + scope="model", + aliases=("sqlcomplexity",), + finding_check_id="sqlcomplexity", + rule_module="tff.core.rules.sql_complexity", + rule_class_name="SqlComplexity", + is_enabled_fn=lambda cfg, p: bool(cfg.rules.sql_complexity.enabled), + ) + ) + + return reg + + +registry: CheckRegistry = create_default_registry() diff --git a/packages/tff-core/src/tff/core/report.py b/packages/tff-core/src/tff/core/report.py index 8e73d81..701a43c 100644 --- a/packages/tff-core/src/tff/core/report.py +++ b/packages/tff-core/src/tff/core/report.py @@ -12,89 +12,9 @@ from rich.table import Table from rich.text import Text -Severity = Literal["error", "warning"] +from tff.core.registry import registry -CHECK_LABELS: dict[str, str] = { - "classificationmacros": "Classification macros", - "sqlcomplexity": "SQL complexity", - "layer_integrity": "Layer integrity", - "custom_exclusions": "Custom exclusions", - "schema_contracts": "Schema contracts", - "dependency_graph": "Dependency graph", - "nomissinggrain": "Missing grain", - "nomissingowner": "Missing owner", - "nomissingdescription": "Missing description", - "nomissingaudits": "Missing audits", - "nomissingnotnull": "Missing not_null audit", - "nomissinguniquevalues": "Missing unique_values audit", - "banselectstar": "No SELECT *", - "filenameequalsmodelname": "Filename equals model name", - "columntypes": "Column types", - "columnnames": "Column names", - "martmodelnamingconvention": "Mart naming convention", - "ambiguousorinvalidcolumn": "Ambiguous/invalid column", - "invalidselectstarexpansion": "Invalid SELECT * expansion", - "nopositionalgroupbyororderby": "No positional GROUP BY or ORDER BY", - "environmentagnosticreferences": "Environment-agnostic references", - "duplicate_ctes": "Duplicate CTEs", - "connascence_of_value": "Connascence of Value", -} - -CONNASCENCE_CATEGORIES: dict[str, str] = { - # Connascence of Name (CoN) - "banselectstar": "Connascence of Name (CoN)", - "filenameequalsmodelname": "Connascence of Name (CoN)", - "columnnames": "Connascence of Name (CoN)", - "martmodelnamingconvention": "Connascence of Name (CoN)", - "ambiguousorinvalidcolumn": "Connascence of Name (CoN)", - "invalidselectstarexpansion": "Connascence of Name (CoN)", - - # Connascence of Type (CoT) - "columntypes": "Connascence of Type (CoT)", - "schema_contracts": "Connascence of Type (CoT)", - - # Connascence of Position (CoP) - "nopositionalgroupbyororderby": "Connascence of Position (CoP)", - - # Connascence of Meaning (CoM) - "classificationmacros": "Connascence of Meaning (CoM)", - - # Dynamic Coupling - "layer_integrity": "Dynamic Coupling & DAG Structure", - "custom_exclusions": "Dynamic Coupling & DAG Structure", - "dependency_graph": "Dynamic Coupling & DAG Structure", - "materialization_depth": "Dynamic Coupling & DAG Structure", - "environmentagnosticreferences": "Dynamic Coupling & DAG Structure", - - - # Quality & Metadata - "nomissingowner": "Quality & Metadata (Non-Connascence)", - "nomissingdescription": "Quality & Metadata (Non-Connascence)", - "nomissinggrain": "Quality & Metadata (Non-Connascence)", - "nomissingaudits": "Quality & Metadata (Non-Connascence)", - "nomissingnotnull": "Quality & Metadata (Non-Connascence)", - "nomissinguniquevalues": "Quality & Metadata (Non-Connascence)", - "sqlcomplexity": "Quality & Metadata (Non-Connascence)", - "duplicate_ctes": "Connascence of Algorithm (CoA)", - "connascence_of_value": "Connascence of Value (CoV)", -} - -ARCHITECTURAL_CHECKS = frozenset( - { - "layer_integrity", - "custom_exclusions", - "schema_contracts", - "dependency_graph", - "duplicate_ctes", - "connascence_of_value", - } -) - -ALWAYS_VISIBLE_CHECKS = [ - *ARCHITECTURAL_CHECKS, - "sqlcomplexity", - "classificationmacros", -] +Severity = Literal["error", "warning"] @dataclass(frozen=True) @@ -106,6 +26,16 @@ class LintFinding: path: str | None = None +CHECK_LABELS: dict[str, str] = registry.get_check_labels() +CONNASCENCE_CATEGORIES: dict[str, str] = registry.get_connascence_categories() +ARCHITECTURAL_CHECKS: frozenset[str] = registry.get_architectural_check_names() +ALWAYS_VISIBLE_CHECKS: list[str] = [ + *ARCHITECTURAL_CHECKS, + "sqlcomplexity", + "classificationmacros", +] + + def normalize_model_name(name: str) -> str: parts = name.replace('"', "").split(".") if len(parts) >= 2: @@ -128,18 +58,20 @@ def _summary_check_names( names: list[str] = [] for check in executed_checks: - if check == "sqlmesh": + if check in ("sqlmesh", "rules"): from_findings = { name for name in by_check if name not in ARCHITECTURAL_CHECKS } if from_findings: names.extend(from_findings) - else: + elif check == "sqlmesh": names.extend( name for name in ALWAYS_VISIBLE_CHECKS if name not in ARCHITECTURAL_CHECKS ) + else: + names.append("rules") else: names.append(check) diff --git a/packages/tff-core/src/tff/dataform/runner.py b/packages/tff-core/src/tff/dataform/runner.py index 077b5d2..b6a2cc0 100644 --- a/packages/tff-core/src/tff/dataform/runner.py +++ b/packages/tff-core/src/tff/dataform/runner.py @@ -5,79 +5,34 @@ import logging from pathlib import Path -from tff.core.checks.connascence_of_value import collect_connascence_of_value_findings -from tff.core.checks.custom_exclusions import collect_custom_exclusion_findings -from tff.core.checks.dependency_graph import collect_dependency_graph_findings -from tff.core.checks.duplicate_ctes import collect_duplicate_cte_findings -from tff.core.checks.layer_integrity import collect_layer_integrity_findings -from tff.core.checks.materialization_depth import collect_materialization_depth_findings -from tff.core.checks.schema_contracts import collect_schema_contract_findings from tff.core.config import FitnessFunctionsConfig, load_fitness_config from tff.core.context import set_ff_config from tff.core.model import ModelRepresentation +from tff.core.registry import registry from tff.core.report import LintFinding -from tff.core.rules import ALL_RULES -from tff.core.utils.paths import model_path_relative from tff.dataform.manifest import load_dataform_models logger = logging.getLogger(__name__) CHECK_COLLECTORS = { - "layer_integrity": lambda models, cfg: collect_layer_integrity_findings( - models, cfg - ), - "custom_exclusions": lambda models, cfg: collect_custom_exclusion_findings( - models, cfg - ), - "schema_contracts": lambda _models, cfg: collect_schema_contract_findings(cfg), - "dependency_graph": lambda models, cfg: collect_dependency_graph_findings( - models, cfg - ), - "materialization_depth": lambda models, cfg: collect_materialization_depth_findings( - models, cfg - ), - "duplicate_ctes": lambda models, cfg: collect_duplicate_cte_findings(models, cfg), - "connascence_of_value": lambda models, cfg: collect_connascence_of_value_findings( - models, cfg - ), + c.finding_id: c.get_collector_fn() + for c in registry.dag_checks() + if c.get_collector_fn() is not None } def collect_dataform_rules_findings( models: dict[str, ModelRepresentation], ) -> list[LintFinding]: - findings = [] - rules = [rule_cls() for rule_cls in ALL_RULES] - - for model in models.values(): - if model.is_external or model.is_symbolic: - continue - - for rule in rules: - violation = rule.check_model(model) - if violation: - msgs = violation.violation_msg - if isinstance(msgs, str): - msgs = [msgs] - for msg in msgs: - model_label = f"{model.name}: " - clean_msg = msg.removeprefix(model_label) - - findings.append( - LintFinding( - check=rule.name, - severity="error", - model=model.name, - path=model_path_relative(model), - message=clean_msg, - ) - ) + """Collect findings for all registered model-level rules.""" + findings: list[LintFinding] = [] + for rule_def in registry.model_rules(): + findings.extend(rule_def.run(models, config=None)) return findings def _check_enabled(config: FitnessFunctionsConfig, check_name: str) -> bool: - check = getattr(config.checks, check_name, None) - return bool(getattr(check, "enabled", False)) + return registry.is_check_enabled(config, check_name, provider="dataform") def run_all_checks( @@ -99,22 +54,12 @@ def run_all_checks( project_root, manifest_path=manifest_path, dialect=dialect ) - if checks is None: - selected = ["rules"] + [ - name for name in CHECK_COLLECTORS if _check_enabled(config, name) - ] - else: - selected = checks - - findings: list[LintFinding] = [] - - if "rules" in selected: - findings.extend(collect_dataform_rules_findings(models)) - - for check_name, collector in CHECK_COLLECTORS.items(): - if check_name not in selected: - continue - findings.extend(collector(models, config)) + findings, selected = registry.run_checks( + models=models, + config=config, + checks=checks, + provider="dataform", + ) # Count of non-external, non-symbolic models checked models_checked = sum( diff --git a/packages/tff-core/src/tff/dbt/runner.py b/packages/tff-core/src/tff/dbt/runner.py index bc774d3..bc518dc 100644 --- a/packages/tff-core/src/tff/dbt/runner.py +++ b/packages/tff-core/src/tff/dbt/runner.py @@ -5,80 +5,34 @@ import logging from pathlib import Path -from tff.core.checks.custom_exclusions import collect_custom_exclusion_findings -from tff.core.checks.dependency_graph import collect_dependency_graph_findings -from tff.core.checks.layer_integrity import collect_layer_integrity_findings -from tff.core.checks.schema_contracts import collect_schema_contract_findings -from tff.core.checks.materialization_depth import collect_materialization_depth_findings -from tff.core.checks.duplicate_ctes import collect_duplicate_cte_findings -from tff.core.checks.connascence_of_value import collect_connascence_of_value_findings from tff.core.config import FitnessFunctionsConfig, load_fitness_config from tff.core.context import set_ff_config -from tff.core.report import LintFinding -from tff.core.rules import ALL_RULES -from tff.core.utils.paths import model_path_relative from tff.core.model import ModelRepresentation +from tff.core.registry import registry +from tff.core.report import LintFinding from tff.dbt.manifest import load_dbt_models logger = logging.getLogger(__name__) CHECK_COLLECTORS = { - "layer_integrity": lambda models, cfg: collect_layer_integrity_findings( - models, cfg - ), - "custom_exclusions": lambda models, cfg: collect_custom_exclusion_findings( - models, cfg - ), - "schema_contracts": lambda _models, cfg: collect_schema_contract_findings(cfg), - "dependency_graph": lambda models, cfg: collect_dependency_graph_findings( - models, cfg - ), - "materialization_depth": lambda models, cfg: collect_materialization_depth_findings( - models, cfg - ), - "duplicate_ctes": lambda models, cfg: collect_duplicate_cte_findings(models, cfg), - "connascence_of_value": lambda models, cfg: collect_connascence_of_value_findings( - models, cfg - ), + c.finding_id: c.get_collector_fn() + for c in registry.dag_checks() + if c.get_collector_fn() is not None } def collect_dbt_rules_findings( models: dict[str, ModelRepresentation], ) -> list[LintFinding]: - findings = [] - rules = [rule_cls() for rule_cls in ALL_RULES] - - for model in models.values(): - if model.is_external or model.is_symbolic: - continue - - for rule in rules: - violation = rule.check_model(model) - if violation: - msgs = violation.violation_msg - if isinstance(msgs, str): - msgs = [msgs] - for msg in msgs: - # Strip model name prefix from message if the rule prepended it - model_label = f"{model.name}: " - clean_msg = msg.removeprefix(model_label) - - findings.append( - LintFinding( - check=rule.name, - severity="error", - model=model.name, - path=model_path_relative(model), - message=clean_msg, - ) - ) + """Collect findings for all registered model-level rules.""" + findings: list[LintFinding] = [] + for rule_def in registry.model_rules(): + findings.extend(rule_def.run(models, config=None)) return findings def _check_enabled(config: FitnessFunctionsConfig, check_name: str) -> bool: - check = getattr(config.checks, check_name, None) - return bool(getattr(check, "enabled", False)) + return registry.is_check_enabled(config, check_name, provider="dbt") def run_all_checks( @@ -97,22 +51,12 @@ def run_all_checks( if models is None: models = load_dbt_models(project_root, dialect=dialect) - if checks is None: - selected = ["rules"] + [ - name for name in CHECK_COLLECTORS if _check_enabled(config, name) - ] - else: - selected = checks - - findings: list[LintFinding] = [] - - if "rules" in selected: - findings.extend(collect_dbt_rules_findings(models)) - - for check_name, collector in CHECK_COLLECTORS.items(): - if check_name not in selected: - continue - findings.extend(collector(models, config)) + findings, selected = registry.run_checks( + models=models, + config=config, + checks=checks, + provider="dbt", + ) models_checked = sum( 1 for m in models.values() if not m.is_external and not m.is_symbolic diff --git a/packages/tff-core/src/tff/sqlmesh/runner.py b/packages/tff-core/src/tff/sqlmesh/runner.py index 409b11d..703c3c6 100644 --- a/packages/tff-core/src/tff/sqlmesh/runner.py +++ b/packages/tff-core/src/tff/sqlmesh/runner.py @@ -8,40 +8,20 @@ from sqlmesh.core.context import Context from sqlmesh.core.linter.definition import AnnotatedRuleViolation -from tff.core.checks.custom_exclusions import collect_custom_exclusion_findings -from tff.core.checks.dependency_graph import collect_dependency_graph_findings -from tff.core.checks.layer_integrity import collect_layer_integrity_findings -from tff.core.checks.schema_contracts import collect_schema_contract_findings -from tff.core.checks.materialization_depth import collect_materialization_depth_findings -from tff.core.checks.duplicate_ctes import collect_duplicate_cte_findings -from tff.core.checks.connascence_of_value import collect_connascence_of_value_findings from tff.core.config import FitnessFunctionsConfig, load_fitness_config from tff.core.context import set_ff_config +from tff.core.model import ModelRepresentation +from tff.core.registry import normalize_check_name, registry from tff.core.report import LintFinding, format_message, normalize_model_name from tff.core.utils.paths import model_path_relative -from tff.core.model import ModelRepresentation from tff.sqlmesh.loader import FitnessLoader, map_sqlmesh_model logger = logging.getLogger(__name__) CHECK_COLLECTORS = { - "layer_integrity": lambda models, cfg: collect_layer_integrity_findings( - models, cfg - ), - "custom_exclusions": lambda models, cfg: collect_custom_exclusion_findings( - models, cfg - ), - "schema_contracts": lambda _models, cfg: collect_schema_contract_findings(cfg), - "dependency_graph": lambda models, cfg: collect_dependency_graph_findings( - models, cfg - ), - "materialization_depth": lambda models, cfg: collect_materialization_depth_findings( - models, cfg - ), - "duplicate_ctes": lambda models, cfg: collect_duplicate_cte_findings(models, cfg), - "connascence_of_value": lambda models, cfg: collect_connascence_of_value_findings( - models, cfg - ), + c.finding_id: c.get_collector_fn() + for c in registry.dag_checks() + if c.get_collector_fn() is not None } @@ -97,8 +77,7 @@ def count_models_checked(context: Context) -> int: def _check_enabled(config: FitnessFunctionsConfig, check_name: str) -> bool: - check = getattr(config.checks, check_name, None) - return bool(getattr(check, "enabled", False)) + return registry.is_check_enabled(config, check_name, provider="sqlmesh") def map_sqlmesh_context_models(context: Context) -> dict[str, ModelRepresentation]: @@ -120,34 +99,92 @@ def run_all_checks( config = load_fitness_config(project_root) set_ff_config(config) + findings: list[LintFinding] = [] + if checks is None: selected = ["sqlmesh"] + [ name for name in CHECK_COLLECTORS if _check_enabled(config, name) ] - else: - selected = checks + if context is None and models is None: + context = Context( + paths=[str(project_root)], + loader=FitnessLoader, + ) - findings: list[LintFinding] = [] + if context is not None: + findings.extend(collect_sqlmesh_findings(context)) - if "sqlmesh" in selected or models is None: - context = context or Context( - paths=[str(project_root)], - loader=FitnessLoader, + mapped_models = ( + models if models is not None else map_sqlmesh_context_models(context) ) - if "sqlmesh" in selected and context is not None: - findings.extend(collect_sqlmesh_findings(context)) + for check_name, collector in CHECK_COLLECTORS.items(): + if _check_enabled(config, check_name) and collector is not None: + findings.extend(collector(mapped_models, config)) + else: + selected = checks + # Validate checks or raise ValueError + for chk in checks: + norm = normalize_check_name(chk) + if norm not in ("sqlmesh", "rules"): + registry.get_or_raise(chk) + + model_rules_requested = any( + normalize_check_name(c) == "rules" + or (registry.get(c) is not None and registry.get(c).scope == "model") + for c in checks + ) - mapped_models = ( - models if models is not None else map_sqlmesh_context_models(context) - ) + if context is None and ("sqlmesh" in selected or model_rules_requested or models is None): + try: + context = Context( + paths=[str(project_root)], + loader=FitnessLoader, + ) + except Exception as e: + logger.debug("Could not initialize SQLMesh Context: %s", e) + context = None + + mapped_models = ( + models + if models is not None + else (map_sqlmesh_context_models(context) if context is not None else {}) + ) - for check_name, collector in CHECK_COLLECTORS.items(): - if check_name not in selected: - continue - if checks is None and not _check_enabled(config, check_name): - continue - findings.extend(collector(mapped_models, config)) + # Run SQLMesh linter / model rules + if "sqlmesh" in selected: + if context is not None: + findings.extend(collect_sqlmesh_findings(context)) + elif model_rules_requested: + if context is not None: + all_sqlmesh_findings = collect_sqlmesh_findings(context) + if any(normalize_check_name(c) == "rules" for c in checks): + findings.extend(all_sqlmesh_findings) + else: + target_norms: set[str] = set() + for chk in checks: + c_def = registry.get(chk) + if c_def is not None and c_def.scope == "model": + target_norms.add(normalize_check_name(c_def.id)) + target_norms.add(normalize_check_name(c_def.finding_id)) + for alias in c_def.aliases: + target_norms.add(normalize_check_name(alias)) + findings.extend( + [f for f in all_sqlmesh_findings if normalize_check_name(f.check) in target_norms] + ) + elif mapped_models: + for chk in checks: + c_def = registry.get(chk) + if c_def is not None and c_def.scope == "model": + findings.extend(c_def.run(mapped_models, config)) + + # Run DAG checks + for chk in checks: + c_def = registry.get(chk) + if c_def is not None and c_def.scope == "dag": + collector = c_def.get_collector_fn() + if collector is not None: + findings.extend(collector(mapped_models, config)) checked_count = ( count_models_checked(context) diff --git a/packages/tff-core/tests/test_registry.py b/packages/tff-core/tests/test_registry.py new file mode 100644 index 0000000..ac21e11 --- /dev/null +++ b/packages/tff-core/tests/test_registry.py @@ -0,0 +1,509 @@ +"""Tests for CheckRegistry, CheckDefinition, and granular rule execution.""" + +from pathlib import Path +import pytest + +from tff.core.config import FitnessFunctionsConfig +from tff.core.context import set_ff_config +from tff.core.model import ModelRepresentation +from tff.core.registry import ( + CheckDefinition, + CheckRegistry, + normalize_check_name, + registry, + run_model_rule, +) +from tff.core.report import LintFinding +from tff.core.rules.ban_select_star import BanSelectStar +from tff.core.rules.base import Rule, RuleViolation +import tff.dbt.runner as dbt_runner +import tff.dataform.runner as dataform_runner +import tff.sqlmesh.runner as sqlmesh_runner + + +def test_normalize_check_name() -> None: + assert normalize_check_name("no_missing_owner") == "nomissingowner" + assert normalize_check_name("Ban-Select-Star") == "banselectstar" + assert normalize_check_name(" COLUMN NAMES ") == "columnnames" + + +def test_check_definition_properties() -> None: + c1 = CheckDefinition( + id="test_check", + label="Test Check", + category="Test Category", + scope="model", + finding_check_id="custom_finding", + aliases=("alias1", "alias2"), + ) + assert c1.canonical_id == "test_check" + assert c1.finding_id == "custom_finding" + + c2 = CheckDefinition( + id="test_alias_fallback", + label="Alias Fallback", + category="Test", + scope="model", + aliases=("alias_first",), + ) + assert c2.finding_id == "alias_first" + + c3 = CheckDefinition( + id="test_id_fallback", + label="ID Fallback", + category="Test", + scope="dag", + ) + assert c3.finding_id == "test_id_fallback" + + +def test_check_definition_get_rule_cls() -> None: + # Direct rule_cls + c1 = CheckDefinition( + id="c1", + label="C1", + category="Cat", + scope="model", + rule_cls=BanSelectStar, + ) + assert c1.get_rule_cls() is BanSelectStar + + # Dynamic import via module and class name + c2 = CheckDefinition( + id="c2", + label="C2", + category="Cat", + scope="model", + rule_module="tff.core.rules.ban_select_star", + rule_class_name="BanSelectStar", + ) + assert c2.get_rule_cls() is BanSelectStar + + # Neither provided + c3 = CheckDefinition(id="c3", label="C3", category="Cat", scope="model") + assert c3.get_rule_cls() is None + + +def test_check_definition_get_collector_fn() -> None: + dummy_fn = lambda _m, _c: [] # noqa: E731 + c1 = CheckDefinition( + id="c1", + label="C1", + category="Cat", + scope="dag", + collector_fn=dummy_fn, + ) + assert c1.get_collector_fn() is dummy_fn + + # Dynamic import + c2 = CheckDefinition( + id="c2", + label="C2", + category="Cat", + scope="dag", + collector_module="tff.core.checks.duplicate_ctes", + collector_func_name="collect_duplicate_cte_findings", + ) + fn = c2.get_collector_fn() + assert fn is not None + assert callable(fn) + + # Neither provided + c3 = CheckDefinition(id="c3", label="C3", category="Cat", scope="dag") + assert c3.get_collector_fn() is None + + +def test_check_definition_is_enabled() -> None: + cfg = FitnessFunctionsConfig() + c1 = CheckDefinition( + id="c1", + label="C1", + category="Cat", + scope="model", + is_enabled_fn=lambda c, p: p == "dbt", + ) + assert c1.is_enabled(cfg, provider="dbt") is True + assert c1.is_enabled(cfg, provider="sqlmesh") is False + + # Default without is_enabled_fn is True + c2 = CheckDefinition(id="c2", label="C2", category="Cat", scope="model") + assert c2.is_enabled(cfg) is True + + +def test_check_definition_run(tmp_path: Path) -> None: + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + set_ff_config(cfg) + + sql_file = tmp_path / "models/marts/test_model.sql" + sql_file.parent.mkdir(parents=True, exist_ok=True) + sql_file.write_text("SELECT * FROM tbl", encoding="utf-8") + + model = ModelRepresentation( + name="test_model", + path=str(sql_file), + dialect="duckdb", + is_symbolic=False, + ) + models = {"test_model": model} + + # 1. Model scope with valid rule + c1 = CheckDefinition( + id="ban_select_star", + label="No SELECT *", + category="CoN", + scope="model", + rule_cls=BanSelectStar, + finding_check_id="banselectstar", + ) + findings = c1.run(models, cfg) + assert len(findings) == 1 + assert findings[0].check == "banselectstar" + + # 2. Model scope without rule_cls + c2 = CheckDefinition(id="empty_model", label="Empty", category="Cat", scope="model") + assert c2.run(models, cfg) == [] + + # 3. DAG scope with valid collector + c3 = CheckDefinition( + id="dag_check", + label="DAG Check", + category="DAG", + scope="dag", + collector_fn=lambda m, c: [ + LintFinding(check="dag_check", severity="error", message="dag violation") + ], + ) + findings_dag = c3.run(models, cfg) + assert len(findings_dag) == 1 + assert findings_dag[0].message == "dag violation" + + # 4. DAG scope without collector + c4 = CheckDefinition(id="empty_dag", label="Empty", category="Cat", scope="dag") + assert c4.run(models, cfg) == [] + + # 5. Unsupported scope + c5 = CheckDefinition(id="other", label="Other", category="Cat", scope="other") # type: ignore[arg-type] + assert c5.run(models, cfg) == [] + + +def test_run_model_rule_variations(tmp_path: Path) -> None: + class MockCustomRule(Rule): + name = "mock_rule" + + def check_model(self, model: ModelRepresentation): + if model.name == "single_msg": + return RuleViolation(f"{model.name}: single violation") + if model.name == "multi_msg": + return RuleViolation(["multi 1", f"{model.name}: multi 2"]) + return None + + sql_file = tmp_path / "model.sql" + sql_file.write_text("SELECT 1", encoding="utf-8") + + m_symbolic = ModelRepresentation(name="sym", path=str(sql_file), dialect="duckdb", is_symbolic=True) + m_external = ModelRepresentation(name="ext", path=str(sql_file), dialect="duckdb", is_external=True) + m_single = ModelRepresentation(name="single_msg", path=str(sql_file), dialect="duckdb") + m_multi = ModelRepresentation(name="multi_msg", path=str(sql_file), dialect="duckdb") + m_clean = ModelRepresentation(name="clean", path=str(sql_file), dialect="duckdb") + + models = { + "sym": m_symbolic, + "ext": m_external, + "single_msg": m_single, + "multi_msg": m_multi, + "clean": m_clean, + } + + findings = run_model_rule(MockCustomRule, models, severity="warning", check_name="custom_check") + assert len(findings) == 3 + for f in findings: + assert f.check == "custom_check" + assert f.severity == "warning" + assert not f.message.startswith(f"{f.model}: ") + + +def test_check_registry_lookups_and_errors() -> None: + reg = CheckRegistry() + chk = CheckDefinition( + id="no_missing_owner", + label="Missing owner", + category="Quality", + scope="model", + finding_check_id="nomissingowner", + aliases=("owner", "missing_owner"), + ) + reg.register(chk) + + # Lookup variations + assert reg.get("no_missing_owner") is chk + assert reg.get("nomissingowner") is chk + assert reg.get("No-Missing-Owner") is chk + assert reg.get("owner") is chk + assert reg.get("missing_owner") is chk + assert reg.get("nonexistent") is None + + # get_or_raise + assert reg.get_or_raise("owner") is chk + with pytest.raises(ValueError, match="Unknown check or rule: 'nonexistent'"): + reg.get_or_raise("nonexistent") + + +def test_check_registry_default_collections() -> None: + all_chks = registry.all_checks() + assert len(all_chks) >= 20 + + model_rules = registry.model_rules() + assert len(model_rules) >= 10 + for r in model_rules: + assert r.scope == "model" + + dag_checks = registry.dag_checks() + assert len(dag_checks) >= 5 + for d in dag_checks: + assert d.scope == "dag" + + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + assert registry.is_check_enabled(cfg, "ban_select_star") is True + assert registry.is_check_enabled(cfg, "nonexistent_check") is False + + +def test_check_registry_resolve_checks() -> None: + reg = CheckRegistry() + r1 = CheckDefinition(id="rule1", label="R1", category="C", scope="model", rule_module="m", rule_class_name="c") + r2 = CheckDefinition(id="rule2", label="R2", category="C", scope="model", rule_module="m", rule_class_name="c") + d1 = CheckDefinition(id="dag1", label="D1", category="C", scope="dag") + reg.register(r1) + reg.register(r2) + reg.register(d1) + + cfg = FitnessFunctionsConfig() + + # None -> returns all enabled + assert len(reg.resolve_checks(None, cfg)) == 3 + + # "rules" -> expands to all model rules + resolved_rules = reg.resolve_checks(["rules"], cfg) + assert set(resolved_rules) == {r1, r2} + + # "sqlmesh" container skipped + resolved_sqlmesh = reg.resolve_checks(["sqlmesh", "rule1"], cfg) + assert resolved_sqlmesh == [r1] + + # Specific list without duplicates + resolved_multi = reg.resolve_checks(["rule1", "rules", "dag1"], cfg) + assert resolved_multi == [r1, r2, d1] + + # Unknown check raises + with pytest.raises(ValueError, match="Unknown check or rule: 'bogus'"): + reg.resolve_checks(["bogus"], cfg) + + +def test_check_registry_run_checks_helper(tmp_path: Path) -> None: + reg = CheckRegistry() + r1 = CheckDefinition( + id="ban_select_star", + label="No SELECT *", + category="CoN", + scope="model", + rule_cls=BanSelectStar, + finding_check_id="banselectstar", + ) + d1 = CheckDefinition( + id="custom_dag", + label="Custom DAG", + category="DAG", + scope="dag", + collector_fn=lambda m, c: [LintFinding(check="custom_dag", severity="error", message="failed")], + ) + reg.register(r1) + reg.register(d1) + + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + set_ff_config(cfg) + + sql_file = tmp_path / "model.sql" + sql_file.write_text("SELECT * FROM t", encoding="utf-8") + model = ModelRepresentation(name="m", path=str(sql_file), dialect="duckdb") + models = {"m": model} + + # checks=None + findings, executed = reg.run_checks(models, cfg, checks=None) + assert len(findings) == 2 + assert executed == ["rules", "custom_dag"] + + # checks specified + findings, executed = reg.run_checks(models, cfg, checks=["ban_select_star"]) + assert len(findings) == 1 + assert executed == ["ban_select_star"] + + +def test_check_registry_metadata_dictionaries() -> None: + labels = registry.get_check_labels() + assert "ban_select_star" in labels + assert "banselectstar" in labels + + connascence = registry.get_connascence_categories() + assert connascence["ban_select_star"] == "Connascence of Name (CoN)" + assert connascence["banselectstar"] == "Connascence of Name (CoN)" + + categories = registry.get_categories() + assert "Connascence of Name (CoN)" in categories + assert "banselectstar" in categories["Connascence of Name (CoN)"] + + proj_level = registry.get_project_level_check_names() + assert "layer_integrity" in proj_level + + arch_checks = registry.get_architectural_check_names() + assert "layer_integrity" in arch_checks + + +def test_dbt_runner_granular_execution(tmp_path: Path) -> None: + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + set_ff_config(cfg) + + sql_file = tmp_path / "model.sql" + sql_file.write_text("SELECT * FROM t", encoding="utf-8") + model = ModelRepresentation(name="m", path=str(sql_file), dialect="duckdb") + models = {"m": model} + + # Test collect_dbt_rules_findings direct call + rule_findings = dbt_runner.collect_dbt_rules_findings(models) + assert len(rule_findings) >= 1 + + # Test _check_enabled + assert dbt_runner._check_enabled(cfg, "schema_contracts") is True + cfg.checks.schema_contracts.enabled = False + assert dbt_runner._check_enabled(cfg, "schema_contracts") is False + assert dbt_runner._check_enabled(cfg, "nonexistent") is False + cfg.checks.schema_contracts.enabled = True + + # Granular check: only ban_select_star + findings, count, selected = dbt_runner.run_all_checks( + project_root=tmp_path, + config=cfg, + checks=["ban_select_star"], + models=models, + ) + assert count == 1 + assert selected == ["ban_select_star"] + assert len(findings) == 1 + assert findings[0].check == "banselectstar" + + +def test_dataform_runner_granular_execution(tmp_path: Path) -> None: + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + set_ff_config(cfg) + + sql_file = tmp_path / "model.sqlx" + sql_file.write_text("SELECT * FROM t", encoding="utf-8") + model = ModelRepresentation(name="m", path=str(sql_file), dialect="bigquery") + models = {"m": model} + + # Test collect_dataform_rules_findings direct call + rule_findings = dataform_runner.collect_dataform_rules_findings(models) + assert len(rule_findings) >= 1 + + # Test _check_enabled + assert dataform_runner._check_enabled(cfg, "schema_contracts") is True + cfg.checks.schema_contracts.enabled = False + assert dataform_runner._check_enabled(cfg, "schema_contracts") is False + assert dataform_runner._check_enabled(cfg, "nonexistent") is False + cfg.checks.schema_contracts.enabled = True + + # Granular check: only ban_select_star + findings, count, selected = dataform_runner.run_all_checks( + project_root=tmp_path, + config=cfg, + checks=["ban_select_star"], + models=models, + ) + assert count == 1 + assert selected == ["ban_select_star"] + assert len(findings) == 1 + assert findings[0].check == "banselectstar" + + +def test_sqlmesh_runner_granular_execution() -> None: + fixture_path = Path(__file__).parent / "fixtures" / "sqlmesh_minimal_project" + cfg = FitnessFunctionsConfig() + set_ff_config(cfg) + + # Test _check_enabled + assert sqlmesh_runner._check_enabled(cfg, "layer_integrity") is True + assert sqlmesh_runner._check_enabled(cfg, "nonexistent") is False + + # Granular check: DAG check only (layer_integrity) + findings, count, selected = sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["layer_integrity"], + ) + assert count == 2 + assert selected == ["layer_integrity"] + assert all(f.check == "layer_integrity" for f in findings) + assert len(findings) == 1 + + # Granular check: "sqlmesh" container + findings_sqlmesh, count_s, selected_s = sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["sqlmesh"], + ) + assert count_s == 2 + assert selected_s == ["sqlmesh"] + + # Granular check: "rules" container + findings_rules, count_r, selected_r = sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["rules"], + ) + assert selected_r == ["rules"] + + # Granular check: model rule + DAG check together + findings_combo, count_c, selected_c = sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["ban_select_star", "layer_integrity"], + ) + assert selected_c == ["ban_select_star", "layer_integrity"] + assert any(f.check == "layer_integrity" for f in findings_combo) + + # Granular check: single model rule + findings_ban, _, selected_ban = sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["ban_select_star"], + ) + assert selected_ban == ["ban_select_star"] + assert all(f.check == "banselectstar" for f in findings_ban) + + # Unknown check raises ValueError + with pytest.raises(ValueError, match="Unknown check or rule: 'invalid_rule'"): + sqlmesh_runner.run_all_checks( + project_root=fixture_path, + checks=["invalid_rule"], + ) + + +def test_sqlmesh_runner_fallback_without_context(tmp_path: Path) -> None: + cfg = FitnessFunctionsConfig() + cfg.rules.ban_select_star.enabled = True + set_ff_config(cfg) + + sql_file = tmp_path / "model.sql" + sql_file.write_text("SELECT * FROM t", encoding="utf-8") + model = ModelRepresentation(name="m", path=str(sql_file), dialect="duckdb") + models = {"m": model} + + # Pass nonexistent project root so context is None, but provide models dict + findings, count, selected = sqlmesh_runner.run_all_checks( + project_root=tmp_path / "nonexistent", + config=cfg, + checks=["ban_select_star"], + models=models, + ) + assert count == 1 + assert selected == ["ban_select_star"] + assert len(findings) == 1 + assert findings[0].check == "banselectstar" From 2d2da59e708510be7b381fd0703528c8a689523d Mon Sep 17 00:00:00 2001 From: Bart Schuijt Date: Mon, 7 Sep 2026 22:31:32 +0200 Subject: [PATCH 2/3] fix(sqlmesh): optimize context reuse in granular test and clean context init (tff#146) --- packages/tff-core/src/tff/sqlmesh/runner.py | 14 +++++--------- packages/tff-core/tests/test_registry.py | 17 +++++++++++------ 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/packages/tff-core/src/tff/sqlmesh/runner.py b/packages/tff-core/src/tff/sqlmesh/runner.py index 703c3c6..0aa3aff 100644 --- a/packages/tff-core/src/tff/sqlmesh/runner.py +++ b/packages/tff-core/src/tff/sqlmesh/runner.py @@ -135,15 +135,11 @@ def run_all_checks( for c in checks ) - if context is None and ("sqlmesh" in selected or model_rules_requested or models is None): - try: - context = Context( - paths=[str(project_root)], - loader=FitnessLoader, - ) - except Exception as e: - logger.debug("Could not initialize SQLMesh Context: %s", e) - context = None + if context is None and (models is None or "sqlmesh" in selected): + context = Context( + paths=[str(project_root)], + loader=FitnessLoader, + ) mapped_models = ( models diff --git a/packages/tff-core/tests/test_registry.py b/packages/tff-core/tests/test_registry.py index ac21e11..92437dc 100644 --- a/packages/tff-core/tests/test_registry.py +++ b/packages/tff-core/tests/test_registry.py @@ -437,7 +437,7 @@ def test_sqlmesh_runner_granular_execution() -> None: assert sqlmesh_runner._check_enabled(cfg, "layer_integrity") is True assert sqlmesh_runner._check_enabled(cfg, "nonexistent") is False - # Granular check: DAG check only (layer_integrity) + # Granular check: DAG check only (layer_integrity) with automatic context initialization findings, count, selected = sqlmesh_runner.run_all_checks( project_root=fixture_path, checks=["layer_integrity"], @@ -447,9 +447,14 @@ def test_sqlmesh_runner_granular_execution() -> None: assert all(f.check == "layer_integrity" for f in findings) assert len(findings) == 1 + from sqlmesh.core.context import Context + from tff.sqlmesh.loader import FitnessLoader + + context = Context(paths=[str(fixture_path)], loader=FitnessLoader) + # Granular check: "sqlmesh" container findings_sqlmesh, count_s, selected_s = sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=context, checks=["sqlmesh"], ) assert count_s == 2 @@ -457,14 +462,14 @@ def test_sqlmesh_runner_granular_execution() -> None: # Granular check: "rules" container findings_rules, count_r, selected_r = sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=context, checks=["rules"], ) assert selected_r == ["rules"] # Granular check: model rule + DAG check together findings_combo, count_c, selected_c = sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=context, checks=["ban_select_star", "layer_integrity"], ) assert selected_c == ["ban_select_star", "layer_integrity"] @@ -472,7 +477,7 @@ def test_sqlmesh_runner_granular_execution() -> None: # Granular check: single model rule findings_ban, _, selected_ban = sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=context, checks=["ban_select_star"], ) assert selected_ban == ["ban_select_star"] @@ -481,7 +486,7 @@ def test_sqlmesh_runner_granular_execution() -> None: # Unknown check raises ValueError with pytest.raises(ValueError, match="Unknown check or rule: 'invalid_rule'"): sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=context, checks=["invalid_rule"], ) From 126a0e2d1699dadf958d9baad5aa0f2809e792f4 Mon Sep 17 00:00:00 2001 From: Bart Schuijt Date: Mon, 7 Sep 2026 22:37:59 +0200 Subject: [PATCH 3/3] test(registry): use mock context for sqlmesh granular execution tests (tff#146) --- packages/tff-core/tests/test_registry.py | 107 ++++++++++++++++++++--- 1 file changed, 93 insertions(+), 14 deletions(-) diff --git a/packages/tff-core/tests/test_registry.py b/packages/tff-core/tests/test_registry.py index 92437dc..d3ea442 100644 --- a/packages/tff-core/tests/test_registry.py +++ b/packages/tff-core/tests/test_registry.py @@ -1,6 +1,7 @@ """Tests for CheckRegistry, CheckDefinition, and granular rule execution.""" from pathlib import Path +from unittest.mock import MagicMock import pytest from tff.core.config import FitnessFunctionsConfig @@ -429,7 +430,8 @@ def test_dataform_runner_granular_execution(tmp_path: Path) -> None: def test_sqlmesh_runner_granular_execution() -> None: - fixture_path = Path(__file__).parent / "fixtures" / "sqlmesh_minimal_project" + from sqlmesh.core.linter.definition import AnnotatedRuleViolation + cfg = FitnessFunctionsConfig() set_ff_config(cfg) @@ -437,9 +439,60 @@ def test_sqlmesh_runner_granular_execution() -> None: assert sqlmesh_runner._check_enabled(cfg, "layer_integrity") is True assert sqlmesh_runner._check_enabled(cfg, "nonexistent") is False - # Granular check: DAG check only (layer_integrity) with automatic context initialization + # Create mock context and models for deterministic, process-pool-free testing + mock_context = MagicMock() + mock_model_1 = MagicMock() + mock_model_1.name = "sqlmesh_example.src_model" + mock_model_1.project = "default" + mock_model_1.kind = MagicMock(is_symbolic=False) + mock_model_1._path = Path("models/src_model.sql") + + mock_model_2 = MagicMock() + mock_model_2.name = "sqlmesh_example.violating_model" + mock_model_2.project = "default" + mock_model_2.kind = MagicMock(is_symbolic=False) + mock_model_2._path = Path("models/violating_model.sql") + + mock_context.models = { + "sqlmesh_example.src_model": mock_model_1, + "sqlmesh_example.violating_model": mock_model_2, + } + + mock_rule = MagicMock() + mock_rule.name = "banselectstar" + violation = AnnotatedRuleViolation( + mock_rule, "sqlmesh_example.src_model: SELECT * is prohibited", mock_model_1, "error" + ) + + mock_linter = MagicMock() + mock_linter.enabled = True + + def mock_lint(model, *args, **kwargs): + if model.name == "sqlmesh_example.src_model": + return (None, [violation]) + return (None, []) + + mock_linter.lint_model.side_effect = mock_lint + mock_context._linters = {"default": mock_linter} + + mapped_models = { + "sqlmesh_example.src_model": ModelRepresentation( + name="sqlmesh_example.src_model", + path="models/sources/src_model.sql", + dialect="duckdb", + depends_on={"sqlmesh_example.violating_model"}, + ), + "sqlmesh_example.violating_model": ModelRepresentation( + name="sqlmesh_example.violating_model", + path="models/derived/violating_model.sql", + dialect="duckdb", + ), + } + + # Granular check: DAG check only (layer_integrity) findings, count, selected = sqlmesh_runner.run_all_checks( - project_root=fixture_path, + context=mock_context, + models=mapped_models, checks=["layer_integrity"], ) assert count == 2 @@ -447,46 +500,53 @@ def test_sqlmesh_runner_granular_execution() -> None: assert all(f.check == "layer_integrity" for f in findings) assert len(findings) == 1 - from sqlmesh.core.context import Context - from tff.sqlmesh.loader import FitnessLoader - - context = Context(paths=[str(fixture_path)], loader=FitnessLoader) - # Granular check: "sqlmesh" container findings_sqlmesh, count_s, selected_s = sqlmesh_runner.run_all_checks( - context=context, + context=mock_context, + models=mapped_models, checks=["sqlmesh"], ) assert count_s == 2 assert selected_s == ["sqlmesh"] + assert len(findings_sqlmesh) == 1 + assert findings_sqlmesh[0].check == "banselectstar" # Granular check: "rules" container findings_rules, count_r, selected_r = sqlmesh_runner.run_all_checks( - context=context, + context=mock_context, + models=mapped_models, checks=["rules"], ) assert selected_r == ["rules"] + assert len(findings_rules) == 1 + assert findings_rules[0].check == "banselectstar" # Granular check: model rule + DAG check together findings_combo, count_c, selected_c = sqlmesh_runner.run_all_checks( - context=context, + context=mock_context, + models=mapped_models, checks=["ban_select_star", "layer_integrity"], ) assert selected_c == ["ban_select_star", "layer_integrity"] + assert len(findings_combo) == 2 assert any(f.check == "layer_integrity" for f in findings_combo) + assert any(f.check == "banselectstar" for f in findings_combo) # Granular check: single model rule findings_ban, _, selected_ban = sqlmesh_runner.run_all_checks( - context=context, + context=mock_context, + models=mapped_models, checks=["ban_select_star"], ) assert selected_ban == ["ban_select_star"] - assert all(f.check == "banselectstar" for f in findings_ban) + assert len(findings_ban) == 1 + assert findings_ban[0].check == "banselectstar" # Unknown check raises ValueError with pytest.raises(ValueError, match="Unknown check or rule: 'invalid_rule'"): sqlmesh_runner.run_all_checks( - context=context, + context=mock_context, + models=mapped_models, checks=["invalid_rule"], ) @@ -512,3 +572,22 @@ def test_sqlmesh_runner_fallback_without_context(tmp_path: Path) -> None: assert selected == ["ban_select_star"] assert len(findings) == 1 assert findings[0].check == "banselectstar" + + +def test_sqlmesh_runner_initializes_context_when_missing(tmp_path: Path) -> None: + from unittest.mock import patch + + cfg = FitnessFunctionsConfig() + set_ff_config(cfg) + + mock_context = MagicMock() + mock_context.models = {} + mock_context._linters = {} + + with patch("tff.sqlmesh.runner.Context", return_value=mock_context) as mock_ctx_cls: + findings, count, selected = sqlmesh_runner.run_all_checks( + project_root=tmp_path, + checks=["sqlmesh"], + ) + assert mock_ctx_cls.called + assert selected == ["sqlmesh"]