Move engine order to ConnectionConfig classes, remove engine check

Author: treysp

sqlmesh/cli/example_project.py

@@ -15,8 +15,8 @@
 
 class ProjectTemplate(Enum):
     DEFAULT = "default"
-    EMPTY = "empty"
     DBT = "dbt"
+    EMPTY = "empty"
     DLT = "dlt"
 
 
@@ -48,8 +48,6 @@ def _gen_config(
 
             for name, field in CONNECTION_CONFIG_TO_TYPE[engine_type].model_fields.items():
                 field_name = field.alias or name
-                if field_name in ("dialect", "display_name"):
-                    continue
 
                 default_value = field.get_default()
 
@@ -279,9 +277,7 @@ def init_example_project(
     dlt_path: t.Optional[str] = None,
     schema_name: str = "sqlmesh_example",
     cli_mode: InitCliMode = InitCliMode.DEFAULT,
-) -> t.Union[str, Path]:
-    from sqlmesh.cli.main import ENGINE_TYPE_DISPLAY_ORDER
-
+) -> Path:
     root_path = Path(path)
     config_extension = "py" if template == ProjectTemplate.DBT else "yaml"
     config_path = root_path / f"config.{config_extension}"
@@ -296,32 +292,35 @@ def init_example_project(
             f"Found an existing config file '{config_path}'.\n\nPlease change to another directory or remove the existing file."
         )
 
-    engine_types = "'" + "', '".join(ENGINE_TYPE_DISPLAY_ORDER) + "'"
+    if template == ProjectTemplate.DBT and not Path(root_path, "dbt_project.yml").exists():
+        raise SQLMeshError(
+            "Required dbt project file 'dbt_project.yml' not found in the current directory.\n\n Please add it or change directories before running `sqlmesh init` to set up your project."
+        )
+
+    engine_types = "', '".join(CONNECTION_CONFIG_TO_TYPE)
     if engine_type is None and template != ProjectTemplate.DBT:
         raise SQLMeshError(
             f"Missing `engine` argument to `sqlmesh init` - please specify a SQL engine for your project. Options: '{engine_types}'."
         )
 
-    if engine_type and engine_type not in ENGINE_TYPE_DISPLAY_ORDER:
+    if engine_type and engine_type not in CONNECTION_CONFIG_TO_TYPE:
         raise SQLMeshError(
             f"Invalid engine '{engine_type}'. Please specify one of '{engine_types}'."
         )
 
     models: t.Set[t.Tuple[str, str]] = set()
     settings = None
     start = None
-    if engine_type:
-        dialect = DIALECT_TO_TYPE[engine_type]
-
-        if template == ProjectTemplate.DLT:
-            if pipeline and dialect:
-                models, settings, start = generate_dlt_models_and_settings(
-                    pipeline_name=pipeline, dialect=dialect, dlt_path=dlt_path
-                )
-            else:
-                raise SQLMeshError(
-                    "Please provide a DLT pipeline with the `--dlt-pipeline` flag to generate a SQLMesh project from DLT."
-                )
+    if engine_type and template == ProjectTemplate.DLT:
+        dialect = DIALECT_TO_TYPE.get(engine_type)
+        if pipeline and dialect:
+            models, settings, start = generate_dlt_models_and_settings(
+                pipeline_name=pipeline, dialect=dialect, dlt_path=dlt_path
+            )
+        else:
+            raise SQLMeshError(
+                "Please provide a DLT pipeline with the `--dlt-pipeline` flag to generate a SQLMesh project from DLT."
+            )
 
     _create_config(config_path, engine_type, settings, start, template, cli_mode)
     if template == ProjectTemplate.DBT:

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, DISPLAY_NAME_TO_TYPE
+from sqlmesh.core.config.connection import INIT_DISPLAY_INFO_TO_TYPE
 from sqlmesh.core.context import Context
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import MissingDependencyError, SQLMeshError
@@ -40,26 +40,6 @@
 )
 SKIP_CONTEXT_COMMANDS = ("init", "ui")
 
-# These are ordered for user display - do not reorder
-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:
     try:
@@ -194,23 +174,16 @@ def init(
         try:
             project_template = ProjectTemplate(template.lower())
         except ValueError:
-            template_strings = (
-                "'" + "', '".join([template.value for template in ProjectTemplate]) + "'"
-            )
+            template_strings = "', '".join([template.value for template in ProjectTemplate])
             raise click.ClickException(
-                f"Invalid project template '{template}'. Please specify one of {template_strings}."
+                f"Invalid project template '{template}'. Please specify one of '{template_strings}'."
             )
 
-    if project_template == ProjectTemplate.DBT and not Path(ctx.obj, "dbt_project.yml").exists():
-        raise click.ClickException(
-            "Required dbt project file 'dbt_project.yml' not found in the current directory.\n\n Please add it or change directories before running `sqlmesh init` to set up your dbt project with SQLMesh."
-        )
-
     if engine or project_template == ProjectTemplate.DBT:
         init_example_project(
             path=ctx.obj,
-            engine_type=engine,
             template=project_template or ProjectTemplate.DEFAULT,
+            engine_type=engine,
             pipeline=dlt_pipeline,
             dlt_path=dlt_path,
         )
@@ -221,8 +194,6 @@ def init(
     console = srich.console
 
     project_template, engine_type, cli_mode = _interactive_init(ctx.obj, console, project_template)
-    if project_template != ProjectTemplate.DBT:
-        _check_engine_installed(console, engine_type)
 
     config_path = init_example_project(
         path=ctx.obj,
@@ -1317,10 +1288,6 @@ def _interactive_init(
     project_template = _init_template_prompt(console) if not project_template else project_template
 
     if project_template == ProjectTemplate.DBT:
-        if not Path(path, "dbt_project.yml").exists():
-            raise SQLMeshError(
-                "Required dbt project file 'dbt_project.yml' not found in the current directory.\n\n Please add it or change directories before running `sqlmesh init` to set up your dbt project with SQLMesh."
-            )
         return (project_template, None, None)
 
     engine_type = _init_engine_prompt(console)
@@ -1332,24 +1299,26 @@ def _interactive_init(
 def _init_integer_prompt(
     console: Console, err_msg_entity: str, num_options: int, retry_func: t.Callable[[t.Any], t.Any]
 ) -> int:
-    err_msg = "\nERROR: '{option_str}' is not a valid {err_msg_entity} number - please enter a number between 1 and {num_options} or exit with control+c"
-    option_str = Prompt.ask("Enter a number", console=console)
-    try:
-        option_num = int(option_str)
-        if option_num < 1 or option_num > num_options:
-            raise ValueError
-    except ValueError:
-        console.print(
-            err_msg.format(
-                option_str=option_str, err_msg_entity=err_msg_entity, num_options=num_options
-            ),
-            style="red",
-        )
-        return retry_func(console)
-    finally:
-        console.print("")
+    err_msg = "\nERROR: '{option_str}' is not a valid {err_msg_entity} number - please enter a number between 1 and {num_options} or exit with control+c\n"
+    while True:
+        option_str = Prompt.ask("Enter a number", console=console)
 
-    return option_num
+        value_error = False
+        try:
+            option_num = int(option_str)
+        except ValueError:
+            value_error = True
+
+        if value_error or option_num < 1 or option_num > num_options:
+            console.print(
+                err_msg.format(
+                    option_str=option_str, err_msg_entity=err_msg_entity, num_options=num_options
+                ),
+                style="red",
+            )
+            continue
+        console.print("")
+        return option_num
 
 
 def _init_template_prompt(console: Console) -> ProjectTemplate:
@@ -1380,17 +1349,22 @@ def _init_engine_prompt(console: Console) -> str:
     console.print("──────────────────────────────\n")
     console.print("Choose your SQL engine:\n")
 
-    display_num_to_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
+    # INIT_DISPLAY_INFO_TO_TYPE is a dict of {engine_type: (display_order, display_name)}
+    ordered_engine_display_names = [
+        info[1] for info in sorted(INIT_DISPLAY_INFO_TO_TYPE.values(), key=lambda x: x[0])
+    ]
+    display_num_to_display_name = {}
+    for i, display_name in enumerate(ordered_engine_display_names):
+        console.print(f"    \\[{i + 1}] {' ' if i < 9 else ''}{display_name}")
+        display_num_to_display_name[i + 1] = display_name
     console.print("")
 
     engine_num = _init_integer_prompt(
-        console, "engine", len(ENGINE_TYPE_DISPLAY_ORDER), _init_engine_prompt
+        console, "engine", len(ordered_engine_display_names), _init_engine_prompt
     )
 
-    return display_num_to_engine[engine_num]
+    DISPLAY_NAME_TO_TYPE = {v[1]: k for k, v in INIT_DISPLAY_INFO_TO_TYPE.items()}
+    return DISPLAY_NAME_TO_TYPE[display_num_to_display_name[engine_num]]
 
 
 def _init_cli_mode_prompt(console: Console) -> InitCliMode:
@@ -1413,17 +1387,3 @@ def _init_cli_mode_prompt(console: Console) -> InitCliMode:
     )
 
     return InitCliMode(display_num_to_cli_mode[cli_mode_num].lower())
-
-
-def _check_engine_installed(console: Console, engine_type: t.Optional[str] = None) -> None:
-    if not engine_type:
-        return
-    connection_config = CONNECTION_CONFIG_TO_TYPE[engine_type]
-
-    try:
-        connection_config._connection_factory.fget(None)
-    except ModuleNotFoundError:
-        install_command = f'pip install "sqlmesh[{engine_type}]"'
-        raise SQLMeshError(
-            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."
-        )