Merge branch 'main' into fix/databricks-oauth-shared-connection

Author: gabepesco

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot~=30.0.1",
+    "sqlglot~=30.2.1",
     "tenacity",
     "time-machine",
     "json-stream"

sqlmesh/core/config/connection.py

@@ -2343,35 +2343,28 @@ def init(cursor: t.Any) -> None:
         return init
 
 
+_CONNECTION_CONFIG_EXCLUDE: t.Set[t.Type[ConnectionConfig]] = {
+    ConnectionConfig,  # type: ignore[type-abstract]
+    BaseDuckDBConnectionConfig,  # type: ignore[type-abstract]
+}
+
 CONNECTION_CONFIG_TO_TYPE = {
     # Map all subclasses of ConnectionConfig to the value of their `type_` field.
     tpe.all_field_infos()["type_"].default: tpe
-    for tpe in subclasses(
-        __name__,
-        ConnectionConfig,
-        exclude={ConnectionConfig, BaseDuckDBConnectionConfig},
-    )
+    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
 }
 
 DIALECT_TO_TYPE = {
     tpe.all_field_infos()["type_"].default: tpe.DIALECT
-    for tpe in subclasses(
-        __name__,
-        ConnectionConfig,
-        exclude={ConnectionConfig, BaseDuckDBConnectionConfig},
-    )
+    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
 }
 
 INIT_DISPLAY_INFO_TO_TYPE = {
     tpe.all_field_infos()["type_"].default: (
         tpe.DISPLAY_ORDER,
         tpe.DISPLAY_NAME,
     )
-    for tpe in subclasses(
-        __name__,
-        ConnectionConfig,
-        exclude={ConnectionConfig, BaseDuckDBConnectionConfig},
-    )
+    for tpe in subclasses(__name__, ConnectionConfig, exclude=_CONNECTION_CONFIG_EXCLUDE)
 }
 
 

sqlmesh/core/config/scheduler.py

@@ -144,9 +144,10 @@ def get_default_catalog_per_gateway(self, context: GenericContext) -> t.Dict[str
         return default_catalogs_per_gateway
 
 
-SCHEDULER_CONFIG_TO_TYPE = {
+SCHEDULER_CONFIG_TO_TYPE: t.Dict[str, t.Type[SchedulerConfig]] = {
     tpe.all_field_infos()["type_"].default: tpe
     for tpe in subclasses(__name__, BaseConfig, exclude={BaseConfig})
+    if issubclass(tpe, SchedulerConfig)
 }
 
 

sqlmesh/core/dialect.py

@@ -14,6 +14,7 @@
 from sqlglot.dialects.dialect import DialectType
 from sqlglot.dialects import DuckDB, Snowflake, TSQL
 import sqlglot.dialects.athena as athena
+import sqlglot.generators.athena as athena_generators
 from sqlglot.parsers.athena import AthenaTrinoParser
 from sqlglot.helper import seq_get
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
@@ -1048,8 +1049,8 @@ def extend_sqlglot() -> None:
         if dialect == athena.Athena:
             tokenizers.add(athena._TrinoTokenizer)
             parsers.add(AthenaTrinoParser)
-            generators.add(athena._TrinoGenerator)
-            generators.add(athena._HiveGenerator)
+            generators.add(athena_generators.AthenaTrinoGenerator)
+            generators.add(athena_generators._HiveGenerator)
 
         if hasattr(dialect, "Tokenizer"):
             tokenizers.add(dialect.Tokenizer)

sqlmesh/core/linter/rules/builtin.py

@@ -318,4 +318,5 @@ def check_model(self, model: Model) -> t.Optional[RuleViolation]:
         return None
 
 
-BUILTIN_RULES = RuleSet(subclasses(__name__, Rule, exclude={Rule}))
+_RULE_EXCLUDE: t.Set[t.Type[Rule]] = {Rule}  # type: ignore[type-abstract]
+BUILTIN_RULES = RuleSet(subclasses(__name__, Rule, exclude=_RULE_EXCLUDE))

sqlmesh/core/loader.py

@@ -840,7 +840,8 @@ def _load_linting_rules(self) -> RuleSet:
             if os.path.getsize(path):
                 self._track_file(path)
                 module = import_python_file(path, self.config_path)
-                module_rules = subclasses(module.__name__, Rule, exclude={Rule})
+                _rule_exclude: t.Set[t.Type[Rule]] = {Rule}  # type: ignore[type-abstract]
+                module_rules = subclasses(module.__name__, Rule, exclude=_rule_exclude)
                 for user_rule in module_rules:
                     user_rules[user_rule.name] = user_rule
 

sqlmesh/core/macros.py

@@ -294,7 +294,9 @@ def evaluate_macros(
             return node
 
         transformed = exp.replace_tree(
-            expression.copy(), evaluate_macros, prune=lambda n: isinstance(n, exp.Lambda)
+            expression.copy(),
+            evaluate_macros,  # type: ignore[arg-type]
+            prune=lambda n: isinstance(n, exp.Lambda),
         )
 
         if changed:

sqlmesh/core/renderer.py

@@ -331,7 +331,7 @@ def _resolve_table(
         deployability_index: t.Optional[DeployabilityIndex] = None,
     ) -> exp.Table:
         table = exp.replace_tables(
-            exp.maybe_parse(table_name, into=exp.Table, dialect=self._dialect),
+            t.cast(exp.Table, exp.maybe_parse(table_name, into=exp.Table, dialect=self._dialect)),
             {
                 **self._to_table_mapping((snapshots or {}).values(), deployability_index),
                 **(table_mapping or {}),

sqlmesh/core/table_diff.py

@@ -255,12 +255,13 @@ def __init__(
         self.source_alias = source_alias
         self.target_alias = target_alias
 
+        cols: t.List[str] = ensure_list(skip_columns)
         self.skip_columns = {
             normalize_identifiers(
-                exp.parse_identifier(t.cast(str, col)),
+                exp.parse_identifier(col),
                 dialect=self.model_dialect or self.dialect,
             ).name
-            for col in ensure_list(skip_columns)
+            for col in cols
         }
 
         self._on = on

sqlmesh/core/test/definition.py

@@ -612,20 +612,27 @@ def _concurrent_render_context(self) -> t.Iterator[None]:
         - Globally patch the SQLGlot dialect so that any date/time nodes are evaluated at the `execution_time` during generation
         """
         import time_machine
+        from sqlglot.generator import _DISPATCH_CACHE
 
         lock_ctx: AbstractContextManager = (
             self.CONCURRENT_RENDER_LOCK if self.concurrency else nullcontext()
         )
         time_ctx: AbstractContextManager = nullcontext()
         dialect_patch_ctx: AbstractContextManager = nullcontext()
+        dispatch_patch_ctx: AbstractContextManager = nullcontext()
 
         if self._execution_time:
+            generator_class = self._test_adapter_dialect.generator_class
             time_ctx = time_machine.travel(self._execution_time, tick=False)
-            dialect_patch_ctx = patch.dict(
-                self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms
-            )
+            dialect_patch_ctx = patch.dict(generator_class.TRANSFORMS, self._transforms)
+
+            # sqlglot caches a dispatch table per generator class, so we need to patch
+            # it as well to ensure the overridden transforms are actually used
+            dispatch = _DISPATCH_CACHE.get(generator_class)
+            if dispatch is not None:
+                dispatch_patch_ctx = patch.dict(dispatch, self._transforms)
 
-        with lock_ctx, time_ctx, dialect_patch_ctx:
+        with lock_ctx, time_ctx, dialect_patch_ctx, dispatch_patch_ctx:
             yield
 
     def _execute(self, query: exp.Query | str) -> pd.DataFrame: