relationalai 0.13.1__py3-none-any.whl → 0.13.2__py3-none-any.whl

relationalai/tools/cli.py

@@ -30,10 +30,9 @@ from ..tools import debugger as deb, qb_debugger as qb_deb
 from ..clients import config
 from relationalai.tools.constants import RAI_APP_NAME
 from relationalai.clients.resources.snowflake.cli_resources import CLIResources
+from relationalai.clients.resources.snowflake import EngineType
 from relationalai.tools.cli_helpers import (
     EMPTY_STRING_REGEX,
-    ENGINE_NAME_ERROR,
-    ENGINE_NAME_REGEX,
     PASSCODE_REGEX,
     UUID,
     RichGroup,
@@ -42,14 +41,25 @@ from relationalai.tools.cli_helpers import (
     ensure_config,
     exit_with_divider,
     exit_with_error,
+    exit_with_handled_exception,
     filter_profiles_by_platform,
     format_row,
     get_config, get_resource_provider,
     is_latest_cli_version,
     issue_top_level_profile_warning,
     latest_version,
+    select_engine_interactive,
+    select_engine_with_state_filter,
+    ensure_engine_type_for_snowflake,
+    build_engine_operation_messages,
+    prompt_and_validate_engine_name,
+    validate_auto_suspend_mins,
+    get_engine_type_for_creation,
+    get_and_validate_engine_size,
+    create_engine_with_spinner,
     show_dictionary_table,
     show_engines,
+    show_engine_details,
     show_imports,
     show_transactions,
     supports_platform,
@@ -113,6 +123,50 @@ def cli(profile):
         rich.print(f"[yellow]RelationalAI version ({latest_ver}) is the latest. Please consider upgrading.[/yellow]")
     GlobalProfile.set(profile)
 
+
+#--------------------------------------------------
+# Engines helpers
+#--------------------------------------------------
+
+def _exit_engine_requires_type(name: str, available_types: list[str], cmd: str) -> None:
+    """Exit with a consistent 'available types' hint and example command."""
+    types_display = ", ".join(available_types) if available_types else "<unknown>"
+    example_type = available_types[0] if available_types else "<ENGINE TYPE>"
+    exit_with_error(
+        f"[yellow]Engine '{name}' has no LOGIC type engine. Available types: {types_display}. "
+        f"Please re-run with [cyan]--type[/cyan]. Example: \n\n"
+        f"[green]rai {cmd} --name {name} --type {example_type}[/green]"
+    )
+
+
+def _get_engine_types_for_name(provider: ResourcesBase, name: str) -> list[str] | None:
+    """Return available engine types for a given name, or None if name not found.
+
+    Errors are handled via exit_with_handled_exception, since this is used for user-facing
+    CLI diagnostics.
+    """
+    try:
+        engines_with_name = provider.list_engines(name=name)
+    except Exception as e:
+        exit_with_handled_exception("Error fetching engines", e)
+        raise Exception("unreachable")
+    if not engines_with_name:
+        return None
+    return sorted({(e.get("type") or "").upper() for e in engines_with_name if e.get("type")})
+
+
+def _require_type_if_no_logic(provider: ResourcesBase, name: str, cmd: str) -> str:
+    """Return LOGIC if present, otherwise exit with a helpful types message."""
+    available_types = _get_engine_types_for_name(provider, name)
+    if not available_types:
+        exit_with_error(f"[yellow]No engine found with name '{name}'.")
+        raise Exception("unreachable")
+    has_logic = any(t == EngineType.LOGIC for t in available_types)
+    if has_logic:
+        return EngineType.LOGIC
+    _exit_engine_requires_type(name, available_types, cmd)
+    raise Exception("unreachable")
+
 #--------------------------------------------------
 # Init
 #--------------------------------------------------
@@ -716,7 +770,7 @@ def config_check(all_profiles:bool=False):
             engine_name = cfg.get("engine")
             assert isinstance(engine_name, str), f"Engine name must be a string, not {type(engine_name)}"
             # This essentially checks if the profile is valid since we are connecting to get the engine
-            engine = provider.get_engine(engine_name)
+            engine = provider.get_engine(engine_name, EngineType.LOGIC)
             if not engine or (engine and not provider.is_valid_engine_state(engine.get("state"))):
                 provider.auto_create_engine_async(engine_name)
         except Exception as e:
@@ -823,16 +877,20 @@ def debugger(host, port, old, qb, profile):
 
 @cli.command(name="engines:list", help="List all engines")
 @click.option("--state", help="Filter by engine state")
-def engines_list(state:str|None=None):
+@click.option("--name", help="Filter by engine name (case-insensitive partial match)")
+@click.option("--type", help="Filter by engine type")
+@click.option("--size", help="Filter by engine size")
+@click.option("--created-by", help="Filter by creator (case-insensitive partial match)")
+def engines_list(state:str|None=None, name:str|None=None, type:str|None=None, size:str|None=None, created_by:str|None=None):
     divider(flush=True)
     ensure_config()
     rich.print("Note: [cyan]Engine names are case sensitive")
     rich.print("")
     with Spinner("Fetching engines"):
         try:
-            engines = get_resource_provider().list_engines(state)
+            engines = get_resource_provider().list_engines(state=state, name=name, type=type, size=size, created_by=created_by)
         except Exception as e:
-            return exit_with_error(f"\n\n[yellow]Error fetching engines: {e}")
+            return exit_with_handled_exception("Error fetching engines", e)
 
     if len(engines):
         show_engines(engines)
@@ -842,31 +900,64 @@ def engines_list(state:str|None=None):
 
 @cli.command(name="engines:get", help="Get engine details")
 @click.option("--name", help="Name of the engine")
-def engines_get(name:str|None=None):
+@click.option("--type", help="Type of the engine")
+def engines_get(name:str|None=None, type:str|None=None):
     divider(flush=True)
     ensure_config()
+    provider = get_resource_provider()
+
+    # Default to LOGIC for backwards compatibility when --type is not provided but --name is.
+    # If LOGIC doesn't exist but other types do, we'll show a targeted hint after probing.
+    if name and type is None:
+        type = EngineType.LOGIC
 
     rich.print("Note: [cyan]Engine names are case sensitive")
     rich.print("")
 
-    if not name:
-        name = controls.text("Engine name:", validator=ENGINE_NAME_REGEX.match, invalid_message=ENGINE_NAME_ERROR)
-        rich.print("")
+    engine = None
+    if not name or (not type or not EngineType.is_valid(type)):
+        if type and not EngineType.is_valid(type):
+            rich.print(f"[yellow]Invalid engine type '{type}'.")
 
-    with Spinner("Fetching engine"):
         try:
-            engine = get_resource_provider().get_engine(name)
+            engines_list = provider.list_engines(name=name if name else None, type=type if type else None)
         except Exception as e:
-            return exit_with_error(f"\n\n[yellow]Error fetching engine: {e}")
+            return exit_with_handled_exception("Error fetching engines", e)
+        result = select_engine_interactive(
+            provider,
+            "Select an engine:",
+            engine_name=name,
+            engines=engines_list,
+        )
+        if result is None:
+            return
+        name, type = result
+
+        for eng in engines_list:
+            if eng.get("name", "").upper() == name.upper() and (type is None or eng.get("type", "").upper() == type.upper()):
+                engine = eng
+                break
+
+    if engine is None:
+        with Spinner("Fetching engine"):
+            try:
+                engine_type = type or EngineType.LOGIC
+                engine = provider.get_engine(name, engine_type)
+            except Exception as e:
+                return exit_with_handled_exception("Error fetching engine", e)
 
     if engine:
-        table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
-        table.add_column("Name")
-        table.add_column("Size")
-        table.add_column("State")
-        table.add_row(engine.get("name"), engine.get("size"), engine.get("state"))
-        rich.print(table)
+        show_engine_details(cast(dict[str, Any], engine))
     else:
+        # If the user didn't specify --type, try to detect whether the engine exists
+        # under a non-LOGIC type and provide a helpful hint.
+        if name and type == EngineType.LOGIC:
+            try:
+                available_types = _get_engine_types_for_name(provider, name)
+            except Exception:
+                available_types = None
+            if available_types and EngineType.LOGIC not in available_types:
+                _exit_engine_requires_type(name, available_types, "engines:get")
         exit_with_error(f'[yellow]Engine "{name}" not found')
     divider()
 
@@ -874,33 +965,18 @@ def engines_get(name:str|None=None):
 # Engine create
 #--------------------------------------------------
 
-def create_engine_flow(cfg:config.Config, name=None, size=None, auto_suspend_mins=None):
+def create_engine_flow(cfg:config.Config, name=None, engine_type=None, size=None, auto_suspend_mins=None):
+    """Main flow for creating an engine interactively or programmatically."""
     provider = get_resource_provider(None, cfg)
-    engine = None
-    is_engine_present = False
-    is_engine_suspended = False
     is_interactive = name is None or size is None
-    auto_suspend_mins = cfg.get("auto_suspend_mins", None) if auto_suspend_mins is None else auto_suspend_mins
     if is_interactive:
         rich.print("Note: [cyan]Engine names are case sensitive")
         rich.print("")
 
-    if not name:
-        name = controls.prompt(
-            "Engine name:",
-            name,
-            validator=ENGINE_NAME_REGEX.match,
-            invalid_message=ENGINE_NAME_ERROR,
-            newline=True
-        )
+    auto_suspend_mins = cfg.get("auto_suspend_mins", None) if auto_suspend_mins is None else auto_suspend_mins
+    auto_suspend_mins = validate_auto_suspend_mins(auto_suspend_mins)
 
-    if auto_suspend_mins is not None:
-        error_msg = f"[yellow]Error: auto_suspend_mins must be an integer instead of {type(auto_suspend_mins)}"
-        try:
-            auto_suspend_mins = int(auto_suspend_mins)
-        except ValueError:
-            exit_with_error(error_msg)
-        assert isinstance(auto_suspend_mins, int), error_msg
+    name = prompt_and_validate_engine_name(name)
 
     is_name_valid, msg = validate_engine_name(name)
     if not is_name_valid:
@@ -912,66 +988,49 @@ def create_engine_flow(cfg:config.Config, name=None, size=None, auto_suspend_min
         else:
             exit_with_divider(1)
 
-    with Spinner(f"Validating engine '{name}' name", "Engine name validated", "Error:"):
-        try:
-            engine = provider.get_engine(name)
-            is_engine_present = engine and engine is not None
-            is_engine_suspended = engine and engine.get("state", None) == "SUSPENDED"
-        except Exception as e:
-            if "already exists" in f"{e}":
-                is_engine_present = True
-            elif "Not Found" in f"{e}":
-                is_engine_present = False
-            else:
-                raise e
-    if is_engine_suspended:
-        rich.print("")
-        with Spinner(f"Resuming engine '{name}'", f"Engine '{name}' resumed", "Error:"):
-            provider.resume_engine(name)
-            return name
-    if is_engine_present:
-        rich.print("")
-        if is_interactive:
-            rich.print(f"Engine '{name}' already exists")
-            rich.print("")
-            return create_engine_flow(cfg)
-        else:
-            exit_with_error(f"[yellow]Engine '{name}' already exists")
+    # Backwards-compatible behavior:
+    # - If --type is omitted, default to LOGIC (script-friendly; avoids interactive prompts).
+    # Interactive behavior:
+    # - Only prompt for engine type when the user didn't provide --name (fully interactive flow).
+    if name is None and engine_type is None:
+        engine_type = ""
+    engine_type = get_engine_type_for_creation(provider, cfg, engine_type)
 
-    # If no size is provided via the params and no size is set in the config,
-    if not size and not cfg.get("engine_size", None):
-        rich.print("")
-        cloud_provider = provider.get_cloud_provider()
-        choices = provider.get_engine_sizes(cloud_provider)
-        size = controls.fuzzy(f"Engine size ({cloud_provider.upper()}):", choices=choices)
-    # If a size is provided via the params or the config, validate it
-    elif size or cfg.get("engine_size", None):
-        # If size is None but config has engine_size, use config value
-        if size is None:
-            size = cfg.get("engine_size", None)
-        valid_sizes = provider.get_engine_sizes()
-        if not isinstance(size, str) or size not in valid_sizes:
-            exit_with_error(f"\nInvalid engine size [yellow]{size}[/yellow] provided. Please check your config.\n\nValid sizes: [green]{valid_sizes}[/green]")
-    assert isinstance(size, str), "engine_size must be a string"
+    # Simple existence check using the new get_engine API
+    try:
+        existing = provider.get_engine(name, engine_type or EngineType.LOGIC)
+        if existing:
+            engine_type_label = EngineType.get_label(engine_type or EngineType.LOGIC)
+            exit_with_error(f"[yellow]Engine '{name}' with type '{engine_type_label} ({engine_type})' already exists.")
+    except Exception:
+        # If get_engine fails, proceed to creation path; real errors will be surfaced by create_engine
+        pass
 
-    rich.print("")
+    size = get_and_validate_engine_size(provider, cfg, size, engine_type)
+
+    if is_interactive:
+        rich.print("")
 
-    with Spinner(
-        f"Creating '{name}' engine with size {size}... (this may take several minutes)",
-        f"Engine '{name}' created!",
-        failed_message="Error:"
-    ):
-        provider.create_engine(name, size, auto_suspend_mins)
+    create_engine_with_spinner(provider, name, size, engine_type, auto_suspend_mins)
     return name
 
 @cli.command(name="engines:create", help="Create a new engine")
 @click.option("--name", help="Name of the engine")
+@click.option("--type", help="Type of the engine")
 @click.option("--size", help="Size of the engine")
-@click.option("--auto_suspend_mins", help="Suspend the engine after this many minutes of inactivity", default=None)
-def engines_create(name, size, auto_suspend_mins):
+@click.option(
+    "--auto-suspend-mins",
+    "--auto_suspend_mins",
+    help="Suspend the engine after this many minutes of inactivity",
+    default=None,
+)
+def engines_create(name, type, size, auto_suspend_mins):
     divider(flush=True)
     cfg = ensure_config()
-    create_engine_flow(cfg, name, size, auto_suspend_mins)
+    try:
+        create_engine_flow(cfg, name, type, size, auto_suspend_mins)
+    except Exception as e:
+        return exit_with_handled_exception("Error creating engine", e)
     divider()
 
 #--------------------------------------------------
@@ -980,40 +1039,62 @@ def engines_create(name, size, auto_suspend_mins):
 
 @cli.command(name="engines:delete", help="Delete an engine")
 @click.option("--name", help="Name of the engine")
-@click.option("--force", help="Force delete the engine", is_flag=True)
-def engines_delete(name, force=False):
+@click.option("--type", help="Type of the engine")
+def engines_delete(name, type):
     divider(flush=True)
     ensure_config()
     provider = get_resource_provider()
-    if not name:
-        name = controls.fuzzy_with_refetch(
-            "Select an engine:",
-            "engines",
-            lambda: [engine["name"] for engine in provider.list_engines()],
-        )
-        if not name or isinstance(name, Exception):
-            return
-    else:
-        try:
-            engine = provider.get_engine(name)
-            if not engine:
-                exit_with_error(f"[yellow]Engine '{name}' not found")
-        except Exception as e:
-            exit_with_error(f"[yellow]Error fetching engine: {e}")
+    try:
+        _engines_delete(provider, name, type)
+    except Exception as e:
+        return exit_with_handled_exception("Error deleting engine", e)
+    divider()
+
+def _engines_delete(provider: ResourcesBase, name, type) -> None:
+    # If --type is omitted but --name is provided:
+    # - prefer LOGIC for backwards compatibility if that engine exists
+    # - otherwise, require explicit --type to avoid accidentally deleting the wrong engine type
+    if name and type is None:
+        # We only auto-select LOGIC; otherwise require explicit --type (avoid accidental deletes).
+        type = _require_type_if_no_logic(provider, name, "engines:delete")
 
-    del_err = None
-    with Spinner(f"Deleting '{name}' engine", f"Engine '{name}' deleted!", "Error:"):
+    # Select engine if name or type missing
+    if not name or not type:
         try:
-            provider.delete_engine(name, force)
+            result = select_engine_interactive(provider, "Select an engine to delete:", engine_name=name)
         except Exception as e:
-            if "SETUP_CDC" in f"{e}":
-                del_err = Exception("[yellow]Imports are setup to utilize this engine.\nUse '[cyan]rai engines:delete --force[/cyan]' to force delete engines.")
-            else:
-                del_err = e
-            raise del_err
-    if isinstance(del_err, Exception):
-        exit_with_divider(1)
-    divider()
+            return exit_with_handled_exception("Error fetching engines", e)
+        if result is None:
+            if name:
+                exit_with_error(f"[yellow]No engine found with name '{name}'.")
+            return
+        name, type = result
+
+    engine_type = ensure_engine_type_for_snowflake(
+        provider,
+        name,
+        type,
+        f"[yellow]Engine type is required for engine '{name}'. Please specify --type or select from the list.",
+    )
+
+    operation_msg, success_msg = build_engine_operation_messages(provider, name, engine_type, "Deleting", "Deleted")
+
+    try:
+        with Spinner(operation_msg, success_msg):
+            provider.delete_engine(name, engine_type)
+    except Exception as e:
+        error_str = str(e).lower()
+        if "setup_cdc" in str(e):
+            exc = Exception(
+                "Imports are setup to utilize this engine.\n"
+                "Use 'rai engines:delete --force' to force delete engines."
+            )
+        elif "engine not found" in error_str or ("not found" in error_str and "engine" in error_str):
+            engine_type_label = EngineType.get_label(engine_type) if EngineType.is_valid(engine_type) else engine_type
+            exc = Exception(f"Engine '{name}' with type '{engine_type_label} ({engine_type})' not found.")
+        else:
+            exc = e
+        exit_with_handled_exception("Error deleting engine", exc)
 
 #--------------------------------------------------
 # Engine resume
@@ -1021,35 +1102,67 @@ def engines_delete(name, force=False):
 
 @cli.command(name="engines:resume", help="Resume an engine")
 @click.option("--name", help="Name of the engine")
-def engines_resume(name):
+@click.option("--type", help="Type of the engine")
+def engines_resume(name, type):
     divider(flush=True)
     ensure_config()
     provider = get_resource_provider()
-    if not name:
-        name = controls.fuzzy_with_refetch(
+    try:
+        _engines_resume(provider, name, type)
+    except Exception as e:
+        return exit_with_handled_exception("Error resuming engine", e)
+    divider()
+
+def _engines_resume(provider: ResourcesBase, name, type) -> None:
+    type_was_omitted = type is None
+    if name and type is None:
+        type = EngineType.LOGIC
+
+    # Validate type early if provided
+    if type and not EngineType.is_valid(type):
+        exit_with_error(f"[yellow]Invalid engine type '{type}'. Valid types: LOGIC, SOLVER, ML")
+
+    try:
+        result = select_engine_with_state_filter(
+            provider,
+            name,
+            type,
+            "SUSPENDED",
+            "Select a suspended engine to resume:",
             "Select a suspended engine to resume:",
-            "engines",
-            lambda: [engine["name"] for engine in provider.list_engines('SUSPENDED')],
+            "[yellow]No suspended engines found",
+            f"[yellow]No suspended engines found with name '{name}'" if name else "[yellow]No suspended engines found",
         )
-        if not name or isinstance(name, Exception):
-            return
-    else:
-        try:
-            engine = provider.get_engine(name)
-            if not engine:
-                exit_with_error(f"[yellow]Engine '{name}' not found")
-            if engine and engine.get("state") != "SUSPENDED":
-                exit_with_error(f"[yellow]Engine '{name}' not in 'SUSPENDED' state")
-        except Exception as e:
-            exit_with_error(f"[yellow]Error fetching engine: {e}")
+    except Exception as e:
+        return exit_with_handled_exception("Error fetching engines", e)
+    if result is None:
+        return
+    name, type = result
 
-    with Spinner(f"Resuming '{name}' engine", f"Engine '{name}' resumed", "Error:"):
-        try:
-            provider.resume_engine(name)
-        except Exception as e:
-            rich.print(f"[yellow]Error resuming engine: {e}")
-            exit_with_divider(1)
-    divider()
+    engine_type = ensure_engine_type_for_snowflake(
+        provider,
+        name,
+        type,
+        f"[yellow]Engine type is required for engine '{name}'. Please specify --type or select from the list.",
+    )
+
+    operation_msg, success_msg = build_engine_operation_messages(provider, name, engine_type, "Resuming", "Resumed")
+    try:
+        with Spinner(operation_msg, success_msg):
+            provider.resume_engine(name, engine_type)
+    except Exception as e:
+        error_str = str(e).lower()
+        if "engine not found" in error_str or ("not found" in error_str and "engine" in error_str):
+            # If the user omitted --type and we defaulted to LOGIC, try to hint at other types.
+            if type_was_omitted and engine_type == EngineType.LOGIC:
+                available_types = _get_engine_types_for_name(provider, name)
+                if available_types and EngineType.LOGIC not in available_types:
+                    _exit_engine_requires_type(name, available_types, "engines:resume")
+            engine_type_label = EngineType.get_label(engine_type) if EngineType.is_valid(engine_type) else engine_type
+            exc = Exception(f"Engine '{name}' with type '{engine_type_label} ({engine_type})' not found.")
+        else:
+            exc = e
+        exit_with_handled_exception("Error resuming engine", exc)
 
 #--------------------------------------------------
 # Engine suspend
@@ -1057,35 +1170,66 @@ def engines_resume(name):
 
 @cli.command(name="engines:suspend", help="Suspend an engine")
 @click.option("--name", help="Name of the engine")
-def engines_suspend(name):
+@click.option("--type", help="Type of the engine")
+def engines_suspend(name, type):
     divider(flush=True)
     ensure_config()
     provider = get_resource_provider()
-    if not name:
-        name = controls.fuzzy_with_refetch(
-            "Select an active engine to suspend:",
-            "engines",
-            lambda: [engine["name"] for engine in provider.list_engines('READY')],
+    try:
+        _engines_suspend(provider, name, type)
+    except Exception as e:
+        return exit_with_handled_exception("Error suspending engine", e)
+    divider()
+
+def _engines_suspend(provider: ResourcesBase, name, type) -> None:
+    type_was_omitted = type is None
+    if name and type is None:
+        type = EngineType.LOGIC
+
+    if type and not EngineType.is_valid(type):
+        exit_with_error(f"[yellow]Invalid engine type '{type}'. Valid types: LOGIC, SOLVER, ML")
+
+    try:
+        result = select_engine_with_state_filter(
+            provider,
+            name,
+            type,
+            "READY",
+            "Select a ready engine to suspend:",
+            "Select a ready engine to suspend:",
+            "[yellow]No ready engines found",
+            f"[yellow]No ready engines found with name '{name}'" if name else "[yellow]No ready engines found",
         )
-        if not name or isinstance(name, Exception):
-            return
-    else:
-        try:
-            engine = provider.get_engine(name)
-            if not engine:
-                exit_with_error(f"[yellow]Engine '{name}' not found")
-            if engine and engine.get("state") != "READY":
-                exit_with_error(f"[yellow]Engine '{name}' must be in 'READY' state to be suspended")
-        except Exception as e:
-            exit_with_error(f"[yellow]Error fetching engine: {e}")
+    except Exception as e:
+        return exit_with_handled_exception("Error fetching engines", e)
+    if result is None:
+        return
+    name, type = result
 
-    with Spinner(f"Suspending '{name}' engine", f"Engine '{name}' suspended", "Error:"):
-        try:
-            provider.suspend_engine(name)
-        except Exception as e:
-            rich.print(f"[yellow]Error suspending engine: {e}")
-            exit_with_divider(1)
-    divider()
+    engine_type = ensure_engine_type_for_snowflake(
+        provider,
+        name,
+        type,
+        f"[yellow]Engine type is required for engine '{name}'. Please specify --type or select from the list.",
+    )
+
+    operation_msg, success_msg = build_engine_operation_messages(provider, name, engine_type, "Suspending", "Suspended")
+    try:
+        with Spinner(operation_msg, success_msg):
+            provider.suspend_engine(name, engine_type)
+    except Exception as e:
+        error_str = str(e).lower()
+        if "engine not found" in error_str or ("not found" in error_str and "engine" in error_str):
+            # If the user omitted --type and we defaulted to LOGIC, try to hint at other types.
+            if type_was_omitted and engine_type == EngineType.LOGIC:
+                available_types = _get_engine_types_for_name(provider, name)
+                if available_types and EngineType.LOGIC not in available_types:
+                    _exit_engine_requires_type(name, available_types, "engines:suspend")
+            engine_type_label = EngineType.get_label(engine_type) if EngineType.is_valid(engine_type) else engine_type
+            exc = Exception(f"Engine '{name}' with type '{engine_type_label} ({engine_type})' not found.")
+        else:
+            exc = e
+        exit_with_handled_exception("Error suspending engine", exc)
 
 #--------------------------------------------------
 # Engine alter engine pool
@@ -1105,7 +1249,10 @@ def engines_alter_pool(size:str|None=None, min:int|None=None, max:int|None=None)
 
     # Ask for engine size if not provided
     if not size:
-        valid_sizes = provider.get_engine_sizes()
+        try:
+            valid_sizes = provider.get_engine_sizes()
+        except Exception as e:
+            return exit_with_handled_exception("Error fetching engine sizes", e)
         size = controls.fuzzy(
             "Select engine size:",
             choices=valid_sizes,
@@ -1113,7 +1260,10 @@ def engines_alter_pool(size:str|None=None, min:int|None=None, max:int|None=None)
         )
 
     # Validate engine size
-    valid_sizes = provider.get_engine_sizes()
+    try:
+        valid_sizes = provider.get_engine_sizes()
+    except Exception as e:
+        return exit_with_handled_exception("Error fetching engine sizes", e)
     if size not in valid_sizes:
         exit_with_error(f"Invalid engine size '{size}'. Valid sizes: {valid_sizes}")
 
@@ -1148,9 +1298,12 @@ def engines_alter_pool(size:str|None=None, min:int|None=None, max:int|None=None)
     rich.print()
 
     # Call the API method
-    with Spinner("Altering engine pool", "Engine pool altered", "Error:"):
-        # Type cast to ensure type checker recognizes the method
-        cast(ResourcesBase, provider).alter_engine_pool(size, min, max)
+    try:
+        with Spinner("Altering engine pool", "Engine pool altered"):
+            # Type cast to ensure type checker recognizes the method
+            cast(ResourcesBase, provider).alter_engine_pool(size, min, max)
+    except Exception as e:
+        return exit_with_handled_exception("Error altering engine pool", e)
     divider()
 
 #--------------------------------------------------
@@ -1319,7 +1472,7 @@ def parse_source(provider: ResourcesBase, raw: str) -> ImportSource:
 
 @supports_platform("snowflake")
 @cli.command(name="imports:setup", help="Modify and view imports setup")
-@click.option("--engine_size", help="Engine size")
+@click.option("--engine-size", "--engine_size", help="Engine size")
 @click.option("--resume", help="Resume imports", is_flag=True)
 @click.option("--suspend", help="Suspend imports", is_flag=True)
 def imports_setup(engine_size:str|None=None, resume:bool=False, suspend:bool=False):
@@ -1380,7 +1533,7 @@ def imports_setup(engine_size:str|None=None, resume:bool=False, suspend:bool=Fal
                 lambda k, v: {k: str(v), "style": "red"} if k == "enabled" and not v else format_row(k, v)
             )
         except Exception as e:
-            exit_with_error(f"\n\n[yellow]Error fetching imports setup: {e}")
+            exit_with_handled_exception("Error fetching imports setup", e)
     divider()
 
 
@@ -1496,7 +1649,7 @@ def imports_wait(source: List[str], model: str):
             try:
                 models = [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
         if not models:
             return exit_with_error("[yellow]No models found")
         rich.print()
@@ -1592,7 +1745,7 @@ def imports_stream(
             try:
                 models = ["[CREATE MODEL]"] + [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
 
         rich.print()
         model = controls.fuzzy("Select a model:", models)
@@ -1614,7 +1767,7 @@ def imports_stream(
         else:
             sources = [parse_source(provider, source_) for source_ in source]
     except Exception as e:
-        return exit_with_error(f"[yellow]Error: {e}")
+        return exit_with_handled_exception("Error", e)
 
     for import_source in sources:
         try:
@@ -1640,7 +1793,7 @@ def imports_stream(
                 exit_with_error("\n\n[yellow]Stream engine not found. Please use '[cyan]rai imports:setup[/cyan]' to set up imports.")
             else:
                 rich.print()
-                exit_with_error(f"\n[yellow]Error creating stream: {e}")
+                exit_with_handled_exception("Error creating stream", e)
     wait = not no_wait
     if wait:
         poll_imports(provider, [source.name for source in sources], model, no_wait_notice=True)
@@ -1673,7 +1826,7 @@ def imports_snapshot(source:str|None, model:str|None, name:str|None, type:str|No
             try:
                 models = [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
         if len(models) == 0:
             exit_with_error("[yellow]No models found")
         rich.print()
@@ -1695,7 +1848,7 @@ def imports_snapshot(source:str|None, model:str|None, name:str|None, type:str|No
             exit_with_error("\n[yellow]Error creating snapshot, aborting.")
 
         except Exception as e:
-            exit_with_error(f"\n\n[yellow]Error creating snapshot: {e}")
+            exit_with_handled_exception("Error creating snapshot", e)
     divider()
 
 #--------------------------------------------------
@@ -1715,7 +1868,7 @@ def imports_delete(object, model, force):
             try:
                 models = [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
         if len(models) == 0:
             rich.print()
             exit_with_error("[yellow]No models found")
@@ -1727,7 +1880,7 @@ def imports_delete(object, model, force):
         try:
             imports = provider.list_imports(model=model)
         except Exception as e:
-            return exit_with_error(f"\n\n[yellow]Error fetching imports: {e}")
+            return exit_with_handled_exception("Error fetching imports", e)
 
     if not imports and not force:
         rich.print()
@@ -1751,7 +1904,7 @@ def imports_delete(object, model, force):
             try:
                 provider.delete_import(object, model, force)
             except Exception as e:
-                exit_with_error(f"\n\n[yellow]Error deleting import: {e}")
+                exit_with_handled_exception("Error deleting import", e)
     divider()
 
 #--------------------------------------------------
@@ -1771,7 +1924,7 @@ def exports_list(model):
             try:
                 models = [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
         if len(models) == 0:
             return exit_with_error("[yellow]No models found")
         rich.print()
@@ -1782,7 +1935,7 @@ def exports_list(model):
         try:
             exports = provider.list_exports(model, "")
         except Exception as e:
-            return exit_with_error(f"\n\n[yellow]Error fetching exports: {e}")
+            return exit_with_handled_exception("Error fetching exports", e)
 
     rich.print()
     if len(exports):
@@ -1813,7 +1966,7 @@ def exports_delete(export, model):
             try:
                 models = [model["name"] for model in provider.list_graphs()]
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching models: {e}")
+                return exit_with_handled_exception("Error fetching models", e)
         if len(models) == 0:
             exit_with_error("[yellow]No models found")
         rich.print()
@@ -1849,7 +2002,7 @@ def transactions_get(id):
         try:
             transaction = provider.get_transaction(id)
         except Exception as e:
-            exit_with_error(f"\n\n[yellow]Error fetching transaction: {e}")
+            exit_with_handled_exception("Error fetching transaction", e)
     rich.print()
     if transaction:
         show_dictionary_table(transaction, format_row)
@@ -1881,7 +2034,7 @@ def transactions_list(id, state, engine, limit, all_users):
             )
         except Exception as e:
             rich.print()
-            return exit_with_error(f"\n\n[yellow]Error fetching transactions: {e}\n")
+            return exit_with_handled_exception("Error fetching transactions", e)
 
     if len(transactions) == 0:
         rich.print()
@@ -1912,7 +2065,7 @@ def transactions_cancel(id, all_users):
                     created_by=cfg.get("user", None),
                 )
             except Exception as e:
-                return exit_with_error(f"\n\n[yellow]Error fetching transactions: {e}")
+                return exit_with_handled_exception("Error fetching transactions", e)
 
         if not transactions:
             exit_with_error("\n[yellow]No active transactions found")