feat: add context to linter rules (#4247)

Author: tobymao

examples/sushi/linter/user.py

@@ -2,6 +2,7 @@
 
 import typing as t
 
+from sqlmesh.core.context import Context
 from sqlmesh.core.linter.rule import Rule, RuleViolation
 from sqlmesh.core.model import Model
 
@@ -10,4 +11,6 @@ class NoMissingOwner(Rule):
     """All models should have an owner specified."""
 
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:
+        assert isinstance(self.context, Context)
+        assert len(self.context.models) > 10
         return self.violation() if not model.owner else None

sqlmesh/core/context.py

@@ -2537,7 +2537,7 @@ def lint_models(
             # Linter may be `None` if the context is not loaded yet
             if linter := self._linters.get(model.project):
                 lint_violation, violations = (
-                    linter.lint_model(model, console=self.console) or found_error
+                    linter.lint_model(model, self, console=self.console) or found_error
                 )
                 if lint_violation:
                     found_error = True

sqlmesh/core/linter/definition.py

@@ -11,6 +11,9 @@
 from sqlmesh.core.linter.rule import Rule, RuleViolation
 from sqlmesh.core.console import LinterConsole, get_console
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.context import GenericContext
+
 
 def select_rules(all_rules: RuleSet, rule_names: t.Set[str]) -> RuleSet:
     if "all" in rule_names:
@@ -52,7 +55,7 @@ def from_rules(cls, all_rules: RuleSet, config: LinterConfig) -> Linter:
         return Linter(config.enabled, all_rules, rules, warn_rules)
 
     def lint_model(
-        self, model: Model, console: LinterConsole = get_console()
+        self, model: Model, context: GenericContext, console: LinterConsole = get_console()
     ) -> t.Tuple[bool, t.List[AnnotatedRuleViolation]]:
         if not self.enabled:
             return False, []
@@ -62,8 +65,8 @@ def lint_model(
         rules = self.rules.difference(ignored_rules)
         warn_rules = self.warn_rules.difference(ignored_rules)
 
-        error_violations = rules.check_model(model)
-        warn_violations = warn_rules.check_model(model)
+        error_violations = rules.check_model(model, context)
+        warn_violations = warn_rules.check_model(model, context)
 
         all_violations: t.List[AnnotatedRuleViolation] = [
             AnnotatedRuleViolation(
@@ -96,11 +99,11 @@ class RuleSet(Mapping[str, type[Rule]]):
     def __init__(self, rules: Iterable[type[Rule]] = ()) -> None:
         self._underlying = {rule.name: rule for rule in rules}
 
-    def check_model(self, model: Model) -> t.List[RuleViolation]:
+    def check_model(self, model: Model, context: GenericContext) -> t.List[RuleViolation]:
         violations = []
 
         for rule in self._underlying.values():
-            violation = rule().check_model(model)
+            violation = rule(context).check_model(model)
 
             if violation:
                 violations.append(violation)

sqlmesh/core/linter/rule.py

@@ -9,6 +9,10 @@
 import typing as t
 
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.context import GenericContext
+
+
 class _Rule(abc.ABCMeta):
     def __new__(cls: Type[_Rule], clsname: str, bases: t.Tuple, attrs: t.Dict) -> _Rule:
         attrs["name"] = clsname.lower()
@@ -20,6 +24,9 @@ class Rule(abc.ABC, metaclass=_Rule):
 
     name = "rule"
 
+    def __init__(self, context: GenericContext):
+        self.context = context
+
     @abc.abstractmethod
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:
         """The evaluation function that'll check for a violation of this rule."""