Add dialect and display name to ConnectionConfig classes

Author: treysp

sqlmesh/cli/example_project.py

@@ -8,7 +8,7 @@
 from sqlmesh.utils.date import yesterday_ds
 from sqlmesh.utils.errors import SQLMeshError
 
-from sqlmesh.core.config.connection import CONNECTION_CONFIG_TO_TYPE
+from sqlmesh.core.config.connection import CONNECTION_CONFIG_TO_TYPE, DIALECT_TO_TYPE
 
 
 PRIMITIVES = (str, int, bool, float)
@@ -27,7 +27,7 @@ class InitCliMode(Enum):
 
 
 def _gen_config(
-    dialect: t.Optional[str],
+    engine_type: t.Optional[str],
     settings: t.Optional[str],
     start: t.Optional[str],
     template: ProjectTemplate,
@@ -39,7 +39,7 @@ def _gen_config(
       database: db.db"""
     )
 
-    engine = "mssql" if dialect == "tsql" else dialect
+    engine = "mssql" if engine_type == "tsql" else engine_type
 
     if not settings and template != ProjectTemplate.DBT:
         doc_link = "https://sqlmesh.readthedocs.io/en/stable/integrations/engines{engine_link}"
@@ -51,6 +51,9 @@ def _gen_config(
 
             for name, field in CONNECTION_CONFIG_TO_TYPE[engine].model_fields.items():
                 field_name = field.alias or name
+                if field_name in ("dialect", "display_name"):
+                    continue
+
                 default_value = field.get_default()
 
                 if isinstance(default_value, Enum):
@@ -77,7 +80,7 @@ def _gen_config(
 
         connection_settings = (
             "      # For more information on configuring the connection to your execution engine, visit:\n"
-            "      # https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#connections\n"
+            "      # https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#connection\n"
             f"      # {doc_link.format(engine_link=engine_link)}\n{connection_settings}"
         )
 
@@ -93,7 +96,7 @@ def _gen_config(
 # https://sqlmesh.readthedocs.io/en/stable/reference/model_configuration/#model-defaults
 
 model_defaults:
-  dialect: {dialect}
+  dialect: {DIALECT_TO_TYPE[engine]}
   start: {start or yesterday_ds()} # Start date for backfill history
   cron: '@daily'    # Run models daily at 12am UTC (can override per model)
 

sqlmesh/cli/main.py

@@ -15,7 +15,7 @@
 from sqlmesh.core.console import configure_console, get_console
 from sqlmesh.utils import Verbosity
 from sqlmesh.core.config import load_configs
-from sqlmesh.core.config.connection import CONNECTION_CONFIG_TO_TYPE
+from sqlmesh.core.config.connection import CONNECTION_CONFIG_TO_TYPE, DISPLAY_NAME_TO_TYPE
 from sqlmesh.core.context import Context
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import MissingDependencyError, SQLMeshError
@@ -41,24 +41,24 @@
 SKIP_CONTEXT_COMMANDS = ("init", "ui")
 
 # These are ordered for user display - do not reorder
-ENGINE_DISPLAY_NAME_TO_CONNECTION_TYPE = {
-    "DuckDB": "duckdb",
-    "Snowflake": "snowflake",
-    "Databricks": "databricks",
-    "BigQuery": "bigquery",
-    "MotherDuck": "duckdb",
-    "ClickHouse": "clickhouse",
-    "Redshift": "redshift",
-    "Spark": "spark",
-    "Trino": "trino",
-    "Azure SQL": "azuresql",
-    "MSSQL": "tsql",
-    "Postgres": "postgres",
-    "GCP Postgres": "gcp_postgres",
-    "MySQL": "mysql",
-    "Athena": "athena",
-    "RisingWave": "risingwave",
-}
+ENGINE_TYPE_DISPLAY_ORDER = [
+    "duckdb",
+    "snowflake",
+    "databricks",
+    "bigquery",
+    "motherduck",
+    "clickhouse",
+    "redshift",
+    "spark",
+    "trino",
+    "azuresql",
+    "mssql",
+    "postgres",
+    "gcp_postgres",
+    "mysql",
+    "athena",
+    "risingwave",
+]
 
 
 def _sqlmesh_version() -> str:
@@ -231,7 +231,7 @@ def init(
     )
 
     next_step_text = {
-        ProjectTemplate.DEFAULT: f"• Update your gateway connection settings (e.g., username/password) in the project configuration file:\n    {config_path}\nRun command in CLI: sqlmesh plan\n(Optional) Explain a plan: sqlmesh plan --explain",
+        ProjectTemplate.DEFAULT: f"• Update your gateway connection settings (e.g., username/password) in the project configuration file:\n    {config_path}",
         ProjectTemplate.DBT: "",
     }
     next_step_text[ProjectTemplate.EMPTY] = next_step_text[ProjectTemplate.DEFAULT]
@@ -1376,9 +1376,9 @@ def _init_engine_prompt(console: Console) -> str:
     console.print("Choose your SQL engine:\n")
 
     display_num_to_engine = {}
-    for i, engine in enumerate(ENGINE_DISPLAY_NAME_TO_CONNECTION_TYPE.keys()):
-        console.print(f"    \\[{i + 1}] {' ' if i < 9 else ''}{engine}")
-        display_num_to_engine[i + 1] = engine
+    for i, engine_type in enumerate(ENGINE_TYPE_DISPLAY_ORDER):
+        console.print(f"    \\[{i + 1}] {' ' if i < 9 else ''}{DISPLAY_NAME_TO_TYPE[engine_type]}")
+        display_num_to_engine[i + 1] = engine_type
     console.print("")
 
     #         self._print("""Need another engine? See: https://sqlmesh.readthedocs.io/en/stable/integrations/overview/#execution-engines)
@@ -1388,10 +1388,10 @@ def _init_engine_prompt(console: Console) -> str:
     # """)
 
     engine_num = _init_integer_prompt(
-        console, "engine", len(ENGINE_DISPLAY_NAME_TO_CONNECTION_TYPE), _init_engine_prompt
+        console, "engine", len(ENGINE_TYPE_DISPLAY_ORDER), _init_engine_prompt
     )
 
-    return ENGINE_DISPLAY_NAME_TO_CONNECTION_TYPE[display_num_to_engine[engine_num]]
+    return display_num_to_engine[engine_num]
 
 
 def _init_cli_mode_prompt(console: Console) -> InitCliMode:
@@ -1426,5 +1426,5 @@ def _check_engine_installed(console: Console, engine_type: t.Optional[str] = Non
     except ModuleNotFoundError:
         install_command = f'pip install "sqlmesh[{engine_type}]"'
         raise SQLMeshError(
-            f"Unable to load required Python dependencies for the {engine_type.upper()} engine.\n\nPlease run `{install_command}` to install them before running `sqlmesh init` again."
+            f"Unable to load required Python dependencies for the {DISPLAY_NAME_TO_TYPE[engine_type]} engine.\n\nPlease run `{install_command}` to install them before running `sqlmesh init` again."
         )

sqlmesh/core/config/connection.py

@@ -88,6 +88,8 @@ def validate(cls: t.Any, data: t.Any) -> t.Any:
 
 class ConnectionConfig(abc.ABC, BaseConfig):
     type_: str
+    dialect: str
+    display_name: str
     concurrent_tasks: int
     register_comments: bool
     pre_ping: bool
@@ -464,6 +466,8 @@ class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
     """Configuration for the MotherDuck connection."""
 
     type_: t.Literal["motherduck"] = Field(alias="type", default="motherduck")
+    dialect: t.Literal["duckdb"] = "duckdb"
+    display_name: t.Literal["MotherDuck"] = "MotherDuck"
 
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
@@ -488,6 +492,8 @@ class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
     """Configuration for the DuckDB connection."""
 
     type_: t.Literal["duckdb"] = Field(alias="type", default="duckdb")
+    dialect: t.Literal["duckdb"] = "duckdb"
+    display_name: t.Literal["DuckDB"] = "DuckDB"
 
 
 class SnowflakeConnectionConfig(ConnectionConfig):
@@ -538,6 +544,8 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     session_parameters: t.Optional[dict] = None
 
     type_: t.Literal["snowflake"] = Field(alias="type", default="snowflake")
+    dialect: t.Literal["snowflake"] = "snowflake"
+    display_name: t.Literal["Snowflake"] = "Snowflake"
 
     _concurrent_tasks_validator = concurrent_tasks_validator
     _engine_import_validator = _get_engine_import_validator("snowflake", "snowflake")
@@ -731,6 +739,8 @@ class DatabricksConnectionConfig(ConnectionConfig):
     pre_ping: t.Literal[False] = False
 
     type_: t.Literal["databricks"] = Field(alias="type", default="databricks")
+    dialect: t.Literal["databricks"] = "databricks"
+    display_name: t.Literal["Databricks"] = "Databricks"
 
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
@@ -986,6 +996,8 @@ class BigQueryConnectionConfig(ConnectionConfig):
     pre_ping: t.Literal[False] = False
 
     type_: t.Literal["bigquery"] = Field(alias="type", default="bigquery")
+    dialect: t.Literal["bigquery"] = "bigquery"
+    display_name: t.Literal["BigQuery"] = "BigQuery"
 
     _engine_import_validator = _get_engine_import_validator("google.cloud.bigquery", "bigquery")
 
@@ -1127,6 +1139,9 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
     scopes: t.Tuple[str, ...] = ("https://www.googleapis.com/auth/sqlservice.admin",)
     driver: str = "pg8000"
     type_: t.Literal["gcp_postgres"] = Field(alias="type", default="gcp_postgres")
+    dialect: t.Literal["postgres"] = "postgres"
+    display_name: t.Literal["GCP Postgres"] = "GCP Postgres"
+
     concurrent_tasks: int = 4
     register_comments: bool = True
     pre_ping: bool = True
@@ -1261,6 +1276,8 @@ class RedshiftConnectionConfig(ConnectionConfig):
     pre_ping: bool = False
 
     type_: t.Literal["redshift"] = Field(alias="type", default="redshift")
+    dialect: t.Literal["redshift"] = "redshift"
+    display_name: t.Literal["Redshift"] = "Redshift"
 
     _engine_import_validator = _get_engine_import_validator("redshift_connector", "redshift")
 
@@ -1322,6 +1339,8 @@ class PostgresConnectionConfig(ConnectionConfig):
     pre_ping: bool = True
 
     type_: t.Literal["postgres"] = Field(alias="type", default="postgres")
+    dialect: t.Literal["postgres"] = "postgres"
+    display_name: t.Literal["Postgres"] = "Postgres"
 
     _engine_import_validator = _get_engine_import_validator("psycopg2", "postgres")
 
@@ -1375,6 +1394,8 @@ class MySQLConnectionConfig(ConnectionConfig):
     pre_ping: bool = True
 
     type_: t.Literal["mysql"] = Field(alias="type", default="mysql")
+    dialect: t.Literal["mysql"] = "mysql"
+    display_name: t.Literal["MySQL"] = "MySQL"
 
     _engine_import_validator = _get_engine_import_validator("pymysql", "mysql")
 
@@ -1427,6 +1448,8 @@ class MSSQLConnectionConfig(ConnectionConfig):
     pre_ping: bool = True
 
     type_: t.Literal["mssql"] = Field(alias="type", default="mssql")
+    dialect: t.Literal["tsql"] = "tsql"
+    display_name: t.Literal["MSSQL"] = "MSSQL"
 
     _engine_import_validator = _get_engine_import_validator("pymssql", "mssql")
 
@@ -1464,6 +1487,8 @@ def _extra_engine_config(self) -> t.Dict[str, t.Any]:
 
 class AzureSQLConnectionConfig(MSSQLConnectionConfig):
     type_: t.Literal["azuresql"] = Field(alias="type", default="azuresql")  # type: ignore
+    dialect: t.Literal["tsql"] = "tsql"
+    display_name: t.Literal["Azure SQL"] = "Azure SQL"  # type: ignore
 
     @property
     def _extra_engine_config(self) -> t.Dict[str, t.Any]:
@@ -1484,6 +1509,8 @@ class SparkConnectionConfig(ConnectionConfig):
     pre_ping: t.Literal[False] = False
 
     type_: t.Literal["spark"] = Field(alias="type", default="spark")
+    dialect: t.Literal["spark"] = "spark"
+    display_name: t.Literal["Spark"] = "Spark"
 
     _engine_import_validator = _get_engine_import_validator("pyspark", "spark")
 
@@ -1602,6 +1629,8 @@ class TrinoConnectionConfig(ConnectionConfig):
     pre_ping: t.Literal[False] = False
 
     type_: t.Literal["trino"] = Field(alias="type", default="trino")
+    dialect: t.Literal["trino"] = "trino"
+    display_name: t.Literal["Trino"] = "Trino"
 
     _engine_import_validator = _get_engine_import_validator("trino", "trino")
 
@@ -1762,6 +1791,8 @@ class ClickhouseConnectionConfig(ConnectionConfig):
     connection_pool_options: t.Optional[t.Dict[str, t.Any]] = None
 
     type_: t.Literal["clickhouse"] = Field(alias="type", default="clickhouse")
+    dialect: t.Literal["clickhouse"] = "clickhouse"
+    display_name: t.Literal["ClickHouse"] = "ClickHouse"
 
     _engine_import_validator = _get_engine_import_validator("clickhouse_connect", "clickhouse")
 
@@ -1886,6 +1917,8 @@ class AthenaConnectionConfig(ConnectionConfig):
     pre_ping: t.Literal[False] = False
 
     type_: t.Literal["athena"] = Field(alias="type", default="athena")
+    dialect: t.Literal["athena"] = "athena"
+    display_name: t.Literal["Athena"] = "Athena"
 
     _engine_import_validator = _get_engine_import_validator("pyathena", "athena")
 
@@ -1954,6 +1987,8 @@ class RisingwaveConnectionConfig(ConnectionConfig):
     pre_ping: bool = True
 
     type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
+    dialect: t.Literal["risingwave"] = "risingwave"
+    display_name: t.Literal["RisingWave"] = "RisingWave"
 
     _engine_import_validator = _get_engine_import_validator("psycopg2", "risingwave")
 
@@ -1998,6 +2033,34 @@ def init(cursor: t.Any) -> None:
     )
 }
 
+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),
+    )
+}
+
+DIALECT_TO_TYPE = {
+    tpe.all_field_infos()["type_"].default: tpe.all_field_infos()["dialect"].default
+    for tpe in subclasses(
+        __name__,
+        ConnectionConfig,
+        exclude=(ConnectionConfig, BaseDuckDBConnectionConfig),
+    )
+}
+
+DISPLAY_NAME_TO_TYPE = {
+    tpe.all_field_infos()["type_"].default: tpe.all_field_infos()["display_name"].default
+    for tpe in subclasses(
+        __name__,
+        ConnectionConfig,
+        exclude=(ConnectionConfig, BaseDuckDBConnectionConfig),
+    )
+}
+
 
 def parse_connection_config(v: t.Dict[str, t.Any]) -> ConnectionConfig:
     if "type" not in v: