Repository: TobikoData/sqlmesh
Pull request #3240 (2024-10-10 21:20:24)

Fix: use parentheses for databricks' CLUSTER BY clause
Context: https://tobiko-data.slack.com/archives/C044BRE5W4S/p1728573925367219?thread_ts=1728490239.247289&cid=C044BRE5W4S


Patch:
diff --git a/sqlmesh/core/config/gateway.py b/sqlmesh/core/config/gateway.py
index 5c330a16..6184d500 100644
--- a/sqlmesh/core/config/gateway.py
+++ b/sqlmesh/core/config/gateway.py
@@ -9,7 +9,7 @@ from sqlmesh.core.config.connection import (
     SerializableConnectionConfig,
     connection_config_validator,
 )
-from sqlmesh.core.config.scheduler import SchedulerConfig
+from sqlmesh.core.config.scheduler import SchedulerConfig, scheduler_config_validator
 
 
 class GatewayConfig(BaseConfig):
@@ -36,4 +36,5 @@ class GatewayConfig(BaseConfig):
     variables: t.Dict[str, t.Any] = {}
 
     _connection_config_validator = connection_config_validator
+    _scheduler_config_validator = scheduler_config_validator
     _variables_validator = variables_validator
diff --git a/sqlmesh/core/config/root.py b/sqlmesh/core/config/root.py
index 80edfd7e..a08b6af8 100644
--- a/sqlmesh/core/config/root.py
+++ b/sqlmesh/core/config/root.py
@@ -30,7 +30,11 @@ from sqlmesh.core.config.model import ModelDefaultsConfig
 from sqlmesh.core.config.naming import NameInferenceConfig as NameInferenceConfig
 from sqlmesh.core.config.plan import PlanConfig
 from sqlmesh.core.config.run import RunConfig
-from sqlmesh.core.config.scheduler import BuiltInSchedulerConfig, SchedulerConfig
+from sqlmesh.core.config.scheduler import (
+    BuiltInSchedulerConfig,
+    SchedulerConfig,
+    scheduler_config_validator,
+)
 from sqlmesh.core.config.ui import UIConfig
 from sqlmesh.core.loader import Loader, SqlMeshLoader
 from sqlmesh.core.notification_target import NotificationTarget
@@ -137,6 +141,7 @@ class Config(BaseConfig):
     }
 
     _connection_config_validator = connection_config_validator
+    _scheduler_config_validator = scheduler_config_validator
     _variables_validator = variables_validator
 
     @field_validator("gateways", mode="before", always=True)
diff --git a/sqlmesh/core/config/scheduler.py b/sqlmesh/core/config/scheduler.py
index de1da68f..708c7c4b 100644
--- a/sqlmesh/core/config/scheduler.py
+++ b/sqlmesh/core/config/scheduler.py
@@ -8,6 +8,7 @@ import typing as t
 from pydantic import Field
 from requests import Session
 
+from sqlglot.helper import subclasses
 from sqlmesh.core.config.base import BaseConfig
 from sqlmesh.core.config.common import concurrent_tasks_validator
 from sqlmesh.core.console import Console
@@ -30,15 +31,15 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.context import GenericContext
 
 if sys.version_info >= (3, 9):
-    from typing import Annotated, Literal
+    from typing import Literal
 else:
-    from typing_extensions import Annotated, Literal
+    from typing_extensions import Literal
 
 
 logger = logging.getLogger(__name__)
 
 
-class _SchedulerConfig(abc.ABC):
+class SchedulerConfig(abc.ABC):
     """Abstract base class for Scheduler configurations."""
 
     @abc.abstractmethod
@@ -77,7 +78,7 @@ class _SchedulerConfig(abc.ABC):
         """
 
 
-class _EngineAdapterStateSyncSchedulerConfig(_SchedulerConfig):
+class _EngineAdapterStateSyncSchedulerConfig(SchedulerConfig):
     def create_state_sync(self, context: GenericContext) -> StateSync:
         state_connection = (
             context.config.get_state_connection(context.gateway) or context._connection_config
@@ -386,12 +387,31 @@ class MWAASchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig):
         return self.default_catalog_override or default_catalog
 
 
-SchedulerConfig = Annotated[
-    t.Union[
-        BuiltInSchedulerConfig,
-        AirflowSchedulerConfig,
-        CloudComposerSchedulerConfig,
-        MWAASchedulerConfig,
-    ],
-    Field(discriminator="type_"),
-]
+SCHEDULER_CONFIG_TO_TYPE = {
+    tpe.all_field_infos()["type_"].default: tpe
+    for tpe in subclasses(__name__, BaseConfig, exclude=(BaseConfig,))
+}
+
+
+def _scheduler_config_validator(
+    cls: t.Type, v: SchedulerConfig | t.Dict[str, t.Any] | None
+) -> SchedulerConfig | None:
+    if v is None or isinstance(v, SchedulerConfig):
+        return v
+
+    if "type" not in v:
+        raise ConfigError("Missing scheduler type.")
+
+    scheduler_type = v["type"]
+    if scheduler_type not in SCHEDULER_CONFIG_TO_TYPE:
+        raise ConfigError(f"Unknown scheduler type '{scheduler_type}'.")
+
+    return SCHEDULER_CONFIG_TO_TYPE[scheduler_type](**v)
+
+
+scheduler_config_validator = field_validator(
+    "scheduler",
+    "default_scheduler",
+    mode="before",
+    check_fields=False,
+)(_scheduler_config_validator)
diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py
index 4b715b2a..f259731e 100644
--- a/sqlmesh/core/engine_adapter/databricks.py
+++ b/sqlmesh/core/engine_adapter/databricks.py
@@ -280,7 +280,10 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
             table_kind=table_kind,
         )
         if clustered_by:
-            clustered_by_exp = exp.Cluster(expressions=[exp.column(col) for col in clustered_by])
+            cluster_key = exp.maybe_parse(
+                f"({','.join(clustered_by)})", into=exp.Ordered, dialect=self.dialect
+            )
+            clustered_by_exp = exp.Cluster(expressions=[cluster_key])
             expressions = properties.expressions if properties else []
             expressions.append(clustered_by_exp)
             properties = exp.Properties(expressions=expressions)


Test patch:
diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py
index df5ab3d3..8a63dad6 100644
--- a/tests/core/engine_adapter/test_databricks.py
+++ b/tests/core/engine_adapter/test_databricks.py
@@ -176,5 +176,5 @@ def test_create_table_clustered_by(make_mocked_engine_adapter: t.Callable):
 
     sql_calls = to_sql_calls(adapter)
     assert sql_calls == [
-        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING) CLUSTER BY `cola`",
+        "CREATE TABLE IF NOT EXISTS `test_table` (`cola` INT, `colb` STRING) CLUSTER BY (`cola`)",
     ]
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index ad4d25dc..6e533e7f 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -14,6 +14,9 @@ from sqlmesh.core.config import (
     ModelDefaultsConfig,
     BigQueryConnectionConfig,
     MotherDuckConnectionConfig,
+    BuiltInSchedulerConfig,
+    MWAASchedulerConfig,
+    AirflowSchedulerConfig,
 )
 from sqlmesh.core.config.connection import DuckDBAttachOptions
 from sqlmesh.core.config.feature_flag import DbtFeatureFlag, FeatureFlag
@@ -663,3 +666,39 @@ model_defaults:
     assert config.model_defaults.audits[1][1]["column"].this.this == "id"
     assert type(config.model_defaults.audits[1][1]["threshold"]) == exp.Literal
     assert config.model_defaults.audits[1][1]["threshold"].this == "1000"
+
+
+def test_scheduler_config(tmp_path_factory):
+    config_path = tmp_path_factory.mktemp("yaml_config") / "config.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+gateways:
+    airflow_gateway:
+        scheduler:
+            type: airflow
+            airflow_url: https://airflow.url
+    mwaa_gateway:
+        scheduler:
+            type: mwaa
+            environment: test_environment
+    builtin_gateway:
+        scheduler:
+            type: builtin
+
+default_scheduler:
+    type: builtin
+
+model_defaults:
+    dialect: bigquery
+        """
+        )
+
+    config = load_config_from_paths(
+        Config,
+        project_paths=[config_path],
+    )
+    assert isinstance(config.default_scheduler, BuiltInSchedulerConfig)
+    assert isinstance(config.get_gateway("airflow_gateway").scheduler, AirflowSchedulerConfig)
+    assert isinstance(config.get_gateway("mwaa_gateway").scheduler, MWAASchedulerConfig)
+    assert isinstance(config.get_gateway("builtin_gateway").scheduler, BuiltInSchedulerConfig)