relationalai 1.0.0a3__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

v0/relationalai/tools/cli_helpers.py

@@ -1,6 +1,7 @@
 #pyright: reportPrivateImportUsage=false
 from __future__ import annotations
 import io
+import json
 import os
 import re
 import sys
@@ -9,10 +10,12 @@ import rich
 import click
 import functools
 import pytz
+from typing import NoReturn
 
 from v0.relationalai.util.constants import TOP_LEVEL_PROFILE_NAME
+from v0.relationalai.errors import RAIException
 from rich.table import Table
-from typing import Callable, Dict, Any
+from typing import Callable, Dict, Any, List, cast
 from ..clients import config
 from click.core import Context
 from rich.console import Console
@@ -23,9 +26,10 @@ from ..clients.config import ConfigFile
 from datetime import datetime, timedelta
 from click.formatting import HelpFormatter
 from ..clients.client import ResourcesBase
-from v0.relationalai.tools.constants import GlobalProfile
+from v0.relationalai.tools.constants import GlobalProfile, SHOW_FULL_TRACES
 from v0.relationalai.tools.cli_controls import divider
 from v0.relationalai.util.format import humanized_bytes, humanized_duration
+from InquirerPy.base.control import Choice
 
 #--------------------------------------------------
 # Helpers
@@ -152,6 +156,291 @@ def validate_engine_name(name:str) -> tuple[bool, str|None]:
         return False, ENGINE_NAME_ERROR
     return True, None
 
+#--------------------------------------------------
+# Engine types & selection helpers (Snowflake)
+#--------------------------------------------------
+
+def format_state_with_color(state: str) -> str:
+    """Format engine state with colors for display."""
+    if not state:
+        return ""
+    state_upper = state.upper()
+    if state_upper == "READY":
+        return f"[green]{state_upper}[/green]"
+    if state_upper == "SUSPENDED":
+        return f"[yellow]{state_upper}[/yellow]"
+    if state_upper in ("PENDING", "SYNCING", "PROCESSING"):
+        return f"[bold yellow]{state_upper}[/bold yellow]"
+    if state_upper in ("ABORTED", "QUARANTINED", "GONE"):
+        return f"[red]{state_upper}[/red]"
+    return state_upper
+
+def _get_engine_type_api():
+    # Local import to avoid importing snowflake modules for non-snowflake usage
+    from v0.relationalai.clients.resources.snowflake import EngineType
+    return EngineType
+
+def _get_internal_engine_sizes() -> list[str]:
+    # Local import to avoid snowflake imports for non-snowflake usage
+    from v0.relationalai.clients.resources.snowflake import INTERNAL_ENGINE_SIZES
+    return list(INTERNAL_ENGINE_SIZES)
+
+def get_engine_type_choices(cfg: config.Config, exclude_types: List[str] | None = None) -> List[Choice]:
+    """Get sorted list of engine type choices for interactive selection."""
+    EngineType = _get_engine_type_api()
+    if exclude_types is None:
+        exclude_types = []
+    exclude_types_upper = [et.upper() for et in exclude_types]
+    engine_types_list = [EngineType.LOGIC, EngineType.ML, EngineType.SOLVER]
+    engine_types_list = [et for et in engine_types_list if et.upper() not in exclude_types_upper]
+    engine_types_list.sort(key=lambda et: EngineType.get_label(et))
+    return [
+        Choice(value=et, name=f"{EngineType.get_label(et)}: {EngineType.get_description(et)}")
+        for et in engine_types_list
+    ]
+
+def select_engine_type_interactive(cfg: config.Config) -> str:
+    """Show interactive engine type selection and return the selected type."""
+    from . import cli_controls as controls
+    rich.print("")
+    engine_type_choices = get_engine_type_choices(cfg)
+    return controls.select("Engine type:", cast("list[str | Choice]", engine_type_choices), None, newline=True)
+
+def select_engine_interactive(
+    provider: ResourcesBase,
+    prompt: str = "Select an engine:",
+    engine_name: str | None = None,
+    engines: List[Dict[str, Any]] | None = None,
+) -> tuple[str, str | None] | None:
+    """Interactive engine picker returning (name, type)."""
+    from . import cli_controls as controls
+
+    engine_map: Dict[str, tuple[str, str | None]] = {}
+
+    def get_engines():
+        engine_list = engines if engines is not None else provider.list_engines()
+        engine_map.clear()
+        items: List[str] = []
+        EngineType = _get_engine_type_api()
+        for engine in engine_list:
+            eng_name = engine.get("name", "")
+            if engine_name and eng_name.upper() != engine_name.upper():
+                continue
+            eng_type = engine.get("type", "")
+            eng_size = engine.get("size", "")
+            if eng_type:
+                label = f"{EngineType.get_label(eng_type)} ({eng_type})" if EngineType.is_valid(eng_type) else f"{eng_type} ({eng_type})"
+                display = f"{eng_name}, {label}, {eng_size}"
+            else:
+                display = f"{eng_name}, {eng_size}" if eng_size else eng_name
+            engine_map[display] = (eng_name, eng_type or None)
+            items.append(display)
+        return items
+
+    # Auto-select when engine_name uniquely identifies an engine
+    if engine_name:
+        engine_list = engines if engines is not None else provider.list_engines()
+        matches = [e for e in engine_list if e.get("name", "").upper() == engine_name.upper()]
+        if len(matches) == 1:
+            e = matches[0]
+            return (e.get("name", ""), e.get("type"))
+        if len(matches) == 0:
+            return None
+
+    selected = controls.fuzzy_with_refetch(prompt, "engines", get_engines)
+    if not selected or isinstance(selected, Exception):
+        return None
+    return engine_map.get(selected)
+
+def select_engine_with_state_filter(
+    provider: ResourcesBase,
+    engine_name: str | None,
+    engine_type: str | None,
+    state_filter: str,
+    prompt_no_name: str,
+    prompt_with_name: str,
+    error_no_engines: str,
+    error_no_matching: str,
+) -> tuple[str, str | None] | None:
+    """Select an engine with optional state filtering + optional name filtering."""
+    EngineType = _get_engine_type_api()
+
+    if not engine_name:
+        filtered = provider.list_engines(state_filter)
+        if not filtered:
+            exit_with_error(error_no_engines)
+        return select_engine_interactive(provider, prompt_no_name, engines=filtered)
+
+    # If type provided and valid, return directly
+    if engine_type and EngineType.is_valid(engine_type):
+        return (engine_name, engine_type)
+
+    # Filter by name + state; selection handles (name,type)
+    filtered = provider.list_engines(state_filter, name=engine_name)
+    if not filtered:
+        exit_with_error(error_no_matching)
+    return select_engine_interactive(provider, prompt_with_name, engine_name=engine_name, engines=filtered)
+
+def ensure_engine_type_for_snowflake(
+    provider: ResourcesBase,
+    engine_name: str,
+    engine_type: str | None,
+    error_message: str,
+) -> str:
+    """Ensure engine_type is provided and valid; default to LOGIC if omitted."""
+    EngineType = _get_engine_type_api()
+    # If --type was omitted, default to LOGIC for backwards compatibility
+    if engine_type is None:
+        return EngineType.LOGIC
+    assert isinstance(engine_type, str)
+    if engine_type == "" or not EngineType.is_valid(engine_type):
+        cfg = get_config()
+        if engine_type == "":
+            rich.print(f"[yellow]Empty engine type provided for engine '{engine_name}'.")
+        else:
+            rich.print(f"[yellow]Invalid engine type '{engine_type}' for engine '{engine_name}'.")
+        return select_engine_type_interactive(cfg)
+    return engine_type
+
+def build_engine_operation_messages(
+    provider: ResourcesBase,
+    engine_name: str,
+    engine_type: str | None,
+    action: str,
+    action_past: str,
+) -> tuple[str, str]:
+    EngineType = _get_engine_type_api()
+    if engine_type:
+        label = EngineType.get_label(engine_type) if EngineType.is_valid(engine_type) else engine_type
+        return (f"{action} {label} engine '{engine_name}'", f"{label} Engine '{engine_name}' {action_past.lower()}")
+    return (f"{action} '{engine_name}' engine", f"Engine '{engine_name}' {action_past.lower()}")
+
+def prompt_and_validate_engine_name(name: str | None) -> str:
+    """Prompt for engine name if missing; validate using ENGINE_NAME_REGEX."""
+    from . import cli_controls as controls
+    if not name:
+        name = controls.prompt(
+            "Engine name:",
+            name,
+            validator=ENGINE_NAME_REGEX.match,
+            invalid_message=ENGINE_NAME_ERROR,
+            newline=True,
+        )
+    assert isinstance(name, str)
+    return name
+
+def validate_auto_suspend_mins(auto_suspend_mins: int | str | None) -> int | None:
+    if auto_suspend_mins is None:
+        return None
+    if isinstance(auto_suspend_mins, int):
+        return auto_suspend_mins
+    error_msg = f"[yellow]Error: auto_suspend_mins must be an integer instead of {type(auto_suspend_mins)}"
+    try:
+        return int(auto_suspend_mins)
+    except ValueError:
+        exit_with_error(error_msg)
+    return None
+
+def get_engine_type_for_creation(provider: ResourcesBase, cfg: config.Config, engine_type: str | None) -> str | None:
+    """Get engine type for engine creation; defaults to LOGIC when omitted."""
+    EngineType = _get_engine_type_api()
+    if engine_type is None:
+        return EngineType.LOGIC
+    if engine_type == "" or not EngineType.is_valid(engine_type):
+        if engine_type == "":
+            rich.print("[yellow]Empty engine type provided.")
+        else:
+            valid_types_display = ", ".join(EngineType.get_all_types())
+            rich.print(f"[yellow]Invalid engine type '{engine_type}'. Valid types: {valid_types_display}")
+        return select_engine_type_interactive(cfg)
+    return engine_type
+
+def get_and_validate_engine_size(
+    provider: ResourcesBase,
+    cfg: config.Config,
+    size: str | None,
+    engine_type: str | None = None,
+) -> str:
+    from . import cli_controls as controls
+    EngineType = _get_engine_type_api()
+    internal_sizes = set(_get_internal_engine_sizes())
+
+    cloud_provider = provider.get_cloud_provider()
+    valid_sizes = provider.get_engine_sizes(cloud_provider)
+
+    # Engine-type-aware filtering:
+    # Internal sizes (XS/S/M/L) are only valid for LOGIC engines (and only on some accounts).
+    # For ML/SOLVER, hide them to avoid presenting invalid options.
+    if engine_type and EngineType.is_valid(engine_type) and engine_type != EngineType.LOGIC:
+        valid_sizes = [s for s in valid_sizes if s not in internal_sizes]
+
+    # Ask if missing and not in config
+    if not size and not cfg.get("engine_size", None):
+        rich.print("")
+        # This refers to the cloud backing your Snowflake account (AWS/Azure), not the engine "platform".
+        size = controls.fuzzy(f"Engine size (Snowflake cloud: {cloud_provider.upper()}):", choices=valid_sizes)
+    elif size is None and cfg.get("engine_size", None):
+        size = cfg.get("engine_size", None)
+    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]")
+    return size
+
+def create_engine_with_spinner(
+    provider: ResourcesBase,
+    engine_name: str,
+    engine_size: str,
+    engine_type: str | None,
+    auto_suspend_mins: int | None,
+) -> None:
+    """Create an engine with appropriate spinner messages and error handling."""
+    from .cli_controls import Spinner
+
+    EngineType = _get_engine_type_api()
+    # Build creation message with engine type when available
+    if engine_type:
+        creation_message = (
+            f"Creating {EngineType.get_label(engine_type)} engine '{engine_name}' with size {engine_size}... "
+            f"(this may take several minutes)"
+        )
+    else:
+        creation_message = (
+            f"Creating engine '{engine_name}' with size {engine_size}... "
+            f"(this may take several minutes)"
+        )
+
+    with Spinner(
+        creation_message,
+        f"Engine '{engine_name}' created!",
+        failed_message=None,  # We handle error display ourselves below
+    ):
+        try:
+            provider.create_engine(
+                engine_name,
+                type=engine_type,
+                size=engine_size,
+                auto_suspend_mins=auto_suspend_mins,
+            )
+        except Exception as e:
+            # Prefer richer error messages when available
+            error_msg = None
+
+            # EngineProvisioningFailed has a format_message() method that provides better error details
+            if hasattr(e, "format_message"):
+                try:
+                    error_msg = getattr(e, "format_message")()
+                except Exception:
+                    pass
+
+            # Try content/message fallbacks
+            if not error_msg and hasattr(e, "content"):
+                error_msg = getattr(e, "content", None)
+            if not error_msg and hasattr(e, "message"):
+                error_msg = str(getattr(e, "message", ""))
+            if not error_msg:
+                error_msg = str(e)
+
+            raise Exception(error_msg)
+
 #--------------------------------------------------
 # Tables
 #--------------------------------------------------
@@ -190,7 +479,7 @@ def format_value(value) -> str:
 def format_row(key: str, value) -> dict:
     result = {}
     result[key] = value
-    if "status" or "state" in key.lower():
+    if "status" in key.lower() or "state" in key.lower():
         result["style"] = get_color_by_state(value)
     if key == "query_size" and isinstance(value, int):
         result[key] = humanized_bytes(value)
@@ -216,6 +505,28 @@ def show_dictionary_table(dict, format_fn:Callable|None=None):
 
 
 class RichGroup(click.Group):
+    def invoke(self, ctx: Context) -> Any:
+        """Invoke the CLI command, suppressing tracebacks for handled RAIExceptions.
+
+        Any `RAIException` is expected to already know how to render itself nicely via
+        `pprint()`. When such an exception bubbles up to the top-level Click runner,
+        Click will otherwise print a full Python traceback, which is noisy for users.
+        """
+        try:
+            return super().invoke(ctx)
+        except RAIException as exc:
+            # Respect config-based full-trace setting when available.
+            try:
+                show_full_traces = get_config().get("show_full_traces", SHOW_FULL_TRACES)
+            except Exception:
+                show_full_traces = SHOW_FULL_TRACES
+
+            if show_full_traces:
+                raise
+
+            exc.pprint()
+            raise click.exceptions.Exit(1) from None
+
     def format_help(self, ctx: Context, formatter: HelpFormatter) -> None:
         is_latest, current_ver, latest_ver = is_latest_cli_version()
 
@@ -383,16 +694,109 @@ def show_engines(engines):
         table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
         table.add_column("#")
         table.add_column("Name")
+        # Show type column if present
+        table.add_column("Type")
         table.add_column("Size")
         table.add_column("State")
+        table.add_column("Created By")
+        table.add_column("Created On")
+        EngineType = _get_engine_type_api()
         for index, engine in enumerate(engines):
-            table.add_row(f"{index+1}", engine.get("name"), engine.get("size"), engine.get("state"))
+            engine_type = engine.get("type", "")
+            type_display = EngineType.get_label_with_value(engine_type) if engine_type and EngineType.is_valid(engine_type) else (engine_type or "")
+            created_on = format_value(engine.get("created_on"))
+            state_display = format_state_with_color(engine.get("state", ""))
+            table.add_row(
+                f"{index+1}",
+                engine.get("name"),
+                type_display,
+                engine.get("size"),
+                state_display,
+                engine.get("created_by", ""),
+                created_on,
+            )
         rich.print(table)
 
-def exit_with_error(message:str):
+
+def show_engine_details(engine: dict[str, Any]) -> None:
+    """Print a vertical table of engine details (one field per row)."""
+    from v0.relationalai.clients.resources.snowflake import EngineType as _EngineType
+
+    table = Table(
+        show_header=True,
+        border_style="dim",
+        header_style="bold",
+        box=rich_box.SIMPLE_HEAD,
+    )
+    table.add_column("Field")
+    table.add_column("Value", overflow="fold")
+
+    engine_type_from_db = engine.get("type", "")
+    type_display = (
+        _EngineType.get_label_with_value(engine_type_from_db)
+        if engine_type_from_db and _EngineType.is_valid(engine_type_from_db)
+        else engine_type_from_db
+    )
+
+    rows: list[tuple[str, str]] = [
+        ("Name", str(engine.get("name", ""))),
+        ("Type", str(type_display)),
+        ("Size", str(engine.get("size", ""))),
+        ("State", str(format_state_with_color(engine.get("state", "")))),
+        ("Created By", str(engine.get("created_by", ""))),
+        ("Created On", str(format_value(engine.get("created_on")))),
+    ]
+
+    # Optional fields (may not exist on older backends).
+    for key, label, formatter in [
+        ("version", "Version", lambda v: "" if v is None else str(v)),
+        ("updated_on", "Updated On", format_value),
+        ("suspends_at", "Suspends At", format_value),
+    ]:
+        if key in engine:
+            try:
+                val = engine.get(key)
+                rows.append((label, str(formatter(val))))
+            except Exception:
+                rows.append((label, str(engine.get(key))))
+
+    # Auto-suspend minutes is represented differently across backends:
+    # - list_engines tends to return auto_suspend_mins
+    # - get_engine (Snowflake EngineServiceSQL) returns auto_suspend
+    if "auto_suspend_mins" in engine or "auto_suspend" in engine:
+        auto_suspend_val = engine.get("auto_suspend_mins", engine.get("auto_suspend"))
+        rows.append(("Auto Suspend (mins)", "" if auto_suspend_val is None else str(auto_suspend_val)))
+
+    settings = engine.get("settings")
+    if settings in (None, {}, ""):
+        settings_str = ""
+    elif isinstance(settings, dict):
+        settings_str = json.dumps(settings, indent=2, sort_keys=True)
+    else:
+        settings_str = str(settings)
+    rows.append(("Settings", settings_str))
+
+    for field, value in rows:
+        table.add_row(field, value)
+
+    rich.print(table)
+
+def exit_with_error(message: str) -> NoReturn:
     rich.print(message, file=sys.stderr)
     exit_with_divider(1)
 
-def exit_with_divider(exit_code:int=0):
+def exit_with_handled_exception(context: str, exc: Exception) -> NoReturn:
+    """Exit with a nicely formatted message for handled RAIExceptions.
+
+    - If `exc` is a RAIException (i.e. produced by an error handler), print its rich
+        formatted content via `exc.pprint()`.
+    - Otherwise, print a raw one-line error including the context and exception string.
+    """
+    if isinstance(exc, RAIException):
+        exc.pprint()
+        sys.exit(1)
+    exit_with_error(f"\n\n[yellow]{context}: {exc}")
+
+def exit_with_divider(exit_code: int = 0) -> NoReturn:
     divider()
     sys.exit(exit_code)

v0/relationalai/util/format.py

@@ -49,7 +49,7 @@ def humanized_duration(ms: int) -> str:
     return " ".join(parts[:3])
 
 
-def format_duration(duration_seconds: float) -> str:
+def format_duration(duration_seconds: float, seconds_decimals: bool = True) -> str:
     """Format duration in seconds to human-readable format.
 
     Args:
@@ -78,7 +78,10 @@ def format_duration(duration_seconds: float) -> str:
 
     # For durations less than 1 minute, show seconds with one decimal
     elif duration_seconds < 60:
-        return f"{round(duration_seconds, 1)}s"
+        if seconds_decimals:
+            return f"{round(duration_seconds, 1)}s"
+        else:
+            return f"{int(duration_seconds)}s"
 
     # For durations less than 1 hour, show minutes and seconds
     elif duration_seconds < 3600: