sqlsaber 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlsaber/cli/onboarding.py

@@ -0,0 +1,325 @@
+"""Interactive onboarding flow for first-time SQLSaber users."""
+
+import sys
+
+from rich.panel import Panel
+
+from sqlsaber.cli.models import ModelManager
+from sqlsaber.config.api_keys import APIKeyManager
+from sqlsaber.config.auth import AuthConfigManager
+from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.theme.manager import create_console
+
+console = create_console()
+
+
+def needs_onboarding(database_arg: str | None = None) -> bool:
+    """Check if user needs onboarding.
+
+    Onboarding is needed if:
+    - No database is configured AND no database connection string provided via CLI
+    """
+    # If user provided a database argument, skip onboarding
+    if database_arg:
+        return False
+
+    # Check if databases are configured
+    db_manager = DatabaseConfigManager()
+    has_db = db_manager.has_databases()
+
+    return not has_db
+
+
+def welcome_screen() -> None:
+    """Display welcome screen to new users."""
+    banner = """
+███████  ██████  ██      ███████  █████  ██████  ███████ ██████
+██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
+███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
+     ██ ██ ▄▄ ██ ██           ██ ██   ██ ██   ██ ██      ██   ██
+███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
+           ▀▀
+    """
+
+    console.print(Panel.fit(banner, style="bold blue"))
+    console.print()
+
+    welcome_message = """
+[bold]Welcome to SQLsaber! 🎉[/bold]
+
+SQLsaber is an agentic SQL assistant that lets you query your database using natural language.
+
+Let's get you set up in just a few steps.
+    """
+
+    console.print(Panel(welcome_message.strip(), border_style="blue", padding=(1, 2)))
+    console.print()
+
+
+async def setup_database_guided() -> str | None:
+    """Guide user through database setup.
+
+    Returns the name of the configured database or None if cancelled.
+    """
+    from sqlsaber.application.db_setup import (
+        build_config,
+        collect_db_input,
+        save_database,
+        test_connection,
+    )
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    console.print("━" * 80, style="dim")
+    console.print("[bold cyan]Step 1 of 2: Database Connection[/bold cyan]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    try:
+        # Ask for connection name
+        prompter = AsyncPrompter()
+        name = await prompter.text(
+            "What would you like to name this connection?",
+            default="mydb",
+            validate=lambda x: bool(x.strip()) or "Name cannot be empty",
+        )
+
+        if name is None:
+            return None
+
+        name = name.strip()
+
+        # Check if name already exists
+        db_manager = DatabaseConfigManager()
+        if db_manager.get_database(name):
+            console.print(
+                f"[yellow]Database connection '{name}' already exists.[/yellow]"
+            )
+            return name
+
+        # Collect database input (simplified - no SSL in onboarding)
+        db_input = await collect_db_input(
+            prompter=prompter, name=name, db_type="postgresql", include_ssl=False
+        )
+
+        if db_input is None:
+            return None
+
+        # Build config
+        db_config = build_config(db_input)
+
+        # Test the connection
+        console.print(f"[dim]Testing connection to '{name}'...[/dim]")
+        connection_success = await test_connection(db_config, db_input.password)
+
+        if not connection_success:
+            retry = await prompter.confirm(
+                "Would you like to try again with different settings?", default=True
+            )
+            if retry:
+                return await setup_database_guided()
+            else:
+                console.print(
+                    "[yellow]You can add a database later using 'saber db add'[/yellow]"
+                )
+                return None
+
+        # Save the configuration
+        try:
+            save_database(db_manager, db_config, db_input.password)
+            console.print(f"[green]✓ Connection to '{name}' successful![/green]")
+            console.print()
+            return name
+        except Exception as e:
+            console.print(f"[bold red]Error saving database:[/bold red] {e}")
+            return None
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        return None
+    except Exception as e:
+        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        return None
+
+
+async def select_model_for_provider(provider: str) -> str | None:
+    """Fetch and let user select a model for the given provider.
+
+    Returns the selected model ID or None if cancelled/failed.
+    """
+    from sqlsaber.application.model_selection import choose_model, fetch_models
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    try:
+        console.print()
+        console.print(f"[dim]Fetching available {provider.title()} models...[/dim]")
+
+        model_manager = ModelManager()
+        models = await fetch_models(model_manager, providers=[provider])
+
+        if not models:
+            console.print(
+                f"[yellow]Could not fetch models for {provider}. Using default.[/yellow]"
+            )
+            # Use provider-specific default or fallback to Anthropic
+            default_model_id = ModelManager.RECOMMENDED_MODELS.get(
+                provider, ModelManager.DEFAULT_MODEL
+            )
+            # Format it properly if we have a recommended model for this provider
+            if provider in ModelManager.RECOMMENDED_MODELS:
+                return f"{provider}:{ModelManager.RECOMMENDED_MODELS[provider]}"
+            return default_model_id
+
+        prompter = AsyncPrompter()
+        console.print()
+        selected_model = await choose_model(
+            prompter, models, restrict_provider=provider, use_search_filter=True
+        )
+
+        return selected_model
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Model selection cancelled.[/yellow]")
+        return None
+    except Exception as e:
+        console.print(f"[yellow]Error selecting model: {e}. Using default.[/yellow]")
+        # Fallback to provider default
+        if provider in ModelManager.RECOMMENDED_MODELS:
+            return f"{provider}:{ModelManager.RECOMMENDED_MODELS[provider]}"
+        return ModelManager.DEFAULT_MODEL
+
+
+async def setup_auth_guided() -> tuple[bool, str | None]:
+    """Guide user through auth setup.
+
+    Returns tuple of (success: bool, selected_model: str | None).
+    """
+    from sqlsaber.application.auth_setup import setup_auth
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    console.print("━" * 80, style="dim")
+    console.print("[bold cyan]Step 2 of 2: Authentication[/bold cyan]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    try:
+        # Run auth setup
+        prompter = AsyncPrompter()
+        auth_manager = AuthConfigManager()
+        api_key_manager = APIKeyManager()
+
+        success, provider = await setup_auth(
+            prompter=prompter,
+            auth_manager=auth_manager,
+            api_key_manager=api_key_manager,
+            allow_oauth=True,
+            default_provider="anthropic",
+            run_oauth_in_thread=True,
+        )
+
+        if not success:
+            console.print(
+                "[yellow]You can set it up later using 'saber auth setup'[/yellow]"
+            )
+            console.print()
+            return False, None
+
+        # If auth configured but we don't know the provider (already configured case)
+        if provider is None:
+            console.print()
+            return True, None
+
+        # Select model for this provider
+        selected_model = await select_model_for_provider(provider)
+        if selected_model:
+            model_manager = ModelManager()
+            model_manager.set_model(selected_model)
+            console.print(f"[green]✓ Model set to: {selected_model}[/green]")
+        console.print()
+        return True, selected_model
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        console.print()
+        return False, None
+    except Exception as e:
+        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        console.print()
+        return False, None
+
+
+def success_screen(
+    database_name: str | None, auth_configured: bool, model_name: str | None = None
+) -> None:
+    """Display success screen after onboarding."""
+    console.print("━" * 80, style="dim")
+    console.print("[bold green]You're all set! 🚀[/bold green]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    if database_name and auth_configured:
+        console.print(
+            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+        )
+        console.print("[green]✓ Authentication configured[/green]")
+        if model_name:
+            console.print(f"[green]✓ Model: {model_name}[/green]")
+    elif database_name:
+        console.print(
+            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+        )
+        console.print(
+            "[yellow]⚠ AI authentication not configured - you'll be prompted when needed[/yellow]"
+        )
+    elif auth_configured:
+        console.print("[green]✓ AI authentication configured[/green]")
+        if model_name:
+            console.print(f"[green]✓ Model: {model_name}[/green]")
+        console.print(
+            "[yellow]⚠ No database configured - you'll need to provide one via -d flag[/yellow]"
+        )
+
+    console.print()
+    console.print("[dim]Starting interactive session...[/dim]")
+    console.print()
+
+
+async def run_onboarding() -> bool:
+    """Run the complete onboarding flow.
+
+    Returns True if onboarding completed successfully (at least database configured),
+    False if user cancelled or onboarding failed.
+    """
+    try:
+        # Welcome screen
+        welcome_screen()
+
+        # Database setup
+        database_name = await setup_database_guided()
+
+        # If user cancelled database setup, exit
+        if database_name is None:
+            console.print("[yellow]Database setup is required to continue.[/yellow]")
+            console.print(
+                "[dim]You can also provide a connection string using: saber -d <connection-string>[/dim]"
+            )
+            return False
+
+        # Auth setup
+        auth_configured, model_name = await setup_auth_guided()
+
+        # Show success screen
+        success_screen(database_name, auth_configured, model_name)
+
+        return True
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Onboarding cancelled.[/yellow]")
+        console.print(
+            "[dim]You can run setup commands manually:[/dim]\n"
+            "[dim]  - saber db add <name>  # Add database connection[/dim]\n"
+            "[dim]  - saber auth setup     # Configure authentication[/dim]"
+        )
+        sys.exit(0)
+    except Exception as e:
+        console.print(f"[bold red]Onboarding failed:[/bold red] {e}")
+        return False

sqlsaber/cli/streaming.py

@@ -170,7 +170,7 @@ class StreamingQueryHandler:
         except asyncio.CancelledError:
             # Show interruption message outside of Live
             self.display.show_newline()
-            self.console.print("[yellow]Query interrupted[/yellow]")
+            self.console.print("[warning]Query interrupted[/warning]")
             return None
         finally:
             # End any active status and live markdown segments

sqlsaber/cli/threads.py

@@ -12,10 +12,12 @@ from rich.markdown import Markdown
 from rich.panel import Panel
 from rich.table import Table
 
+from sqlsaber.theme.manager import create_console, get_theme_manager
 from sqlsaber.threads import ThreadStorage
 
 # Globals consistent with other CLI modules
-console = Console()
+console = create_console()
+tm = get_theme_manager()
 
 
 threads_app = cyclopts.App(
@@ -84,13 +86,23 @@ def _render_transcript(
                         console.print(f"**User:**\n\n{text}\n")
                     else:
                         console.print(
-                            Panel.fit(Markdown(text), title="User", border_style="cyan")
+                            Panel.fit(
+                                Markdown(text, code_theme=tm.pygments_style_name),
+                                title="User",
+                                border_style=tm.style("panel.border.user"),
+                            )
                         )
                     return
         if is_redirected:
             console.print("**User:** (no content)\n")
         else:
-            console.print(Panel.fit("(no content)", title="User", border_style="cyan"))
+            console.print(
+                Panel.fit(
+                    "(no content)",
+                    title="User",
+                    border_style=tm.style("panel.border.user"),
+                )
+            )
 
     def _render_response(message: ModelMessage) -> None:
         for part in getattr(message, "parts", []):
@@ -103,7 +115,9 @@ def _render_transcript(
                     else:
                         console.print(
                             Panel.fit(
-                                Markdown(text), title="Assistant", border_style="green"
+                                Markdown(text, code_theme=tm.pygments_style_name),
+                                title="Assistant",
+                                border_style=tm.style("panel.border.assistant"),
                             )
                         )
             elif kind in ("tool-call", "builtin-tool-call"):
@@ -211,11 +225,11 @@ def list_threads(
         console.print("No threads found.")
         return
     table = Table(title="Threads")
-    table.add_column("ID", style="cyan")
-    table.add_column("Database", style="magenta")
-    table.add_column("Title", style="green")
-    table.add_column("Last Activity", style="dim")
-    table.add_column("Model", style="yellow")
+    table.add_column("ID", style=tm.style("info"))
+    table.add_column("Database", style=tm.style("accent"))
+    table.add_column("Title", style=tm.style("success"))
+    table.add_column("Last Activity", style=tm.style("muted"))
+    table.add_column("Model", style=tm.style("warning"))
     for t in threads:
         table.add_row(
             t.id,
@@ -235,7 +249,7 @@ def show(
     store = ThreadStorage()
     thread = asyncio.run(store.get_thread(thread_id))
     if not thread:
-        console.print(f"[red]Thread not found:[/red] {thread_id}")
+        console.print(f"[error]Thread not found:[/error] {thread_id}")
         return
     msgs = asyncio.run(store.get_thread_messages(thread_id))
     console.print(f"[bold]Thread: {thread.id}[/bold]")
@@ -273,12 +287,12 @@ def resume(
 
         thread = await store.get_thread(thread_id)
         if not thread:
-            console.print(f"[red]Thread not found:[/red] {thread_id}")
+            console.print(f"[error]Thread not found:[/error] {thread_id}")
             return
         db_selector = database or thread.database_name
         if not db_selector:
             console.print(
-                "[red]No database specified or stored with this thread.[/red]"
+                "[error]No database specified or stored with this thread.[/error]"
             )
             return
         try:
@@ -287,7 +301,7 @@ def resume(
             connection_string = resolved.connection_string
             db_name = resolved.name
         except DatabaseResolutionError as e:
-            console.print(f"[red]Database resolution error:[/red] {e}")
+            console.print(f"[error]Database resolution error:[/error] {e}")
             return
 
         db_conn = DatabaseConnection(connection_string)
@@ -295,7 +309,12 @@ def resume(
             sqlsaber_agent = SQLSaberAgent(db_conn, db_name)
             history = await store.get_thread_messages(thread_id)
             if console.is_terminal:
-                console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
+                console.print(
+                    Panel.fit(
+                        f"Thread: {thread.id}",
+                        border_style=tm.style("panel.border.thread"),
+                    )
+                )
             else:
                 console.print(f"# Thread: {thread.id}\n")
             _render_transcript(console, history, None)
@@ -310,7 +329,7 @@ def resume(
             await session.run()
         finally:
             await db_conn.close()
-            console.print("\n[green]Goodbye![/green]")
+            console.print("\n[success]Goodbye![/success]")
 
     asyncio.run(_run())
 
@@ -329,7 +348,7 @@ def prune(
 
     async def _run() -> None:
         deleted = await store.prune_threads(older_than_days=days)
-        console.print(f"[green]✓ Pruned {deleted} thread(s).[/green]")
+        console.print(f"[success]✓ Pruned {deleted} thread(s).[/success]")
 
     asyncio.run(_run())
 

sqlsaber/config/api_keys.py

@@ -4,11 +4,11 @@ import getpass
 import os
 
 import keyring
-from rich.console import Console
 
 from sqlsaber.config import providers
+from sqlsaber.theme.manager import create_console
 
-console = Console()
+console = create_console()
 
 
 class APIKeyManager:
@@ -19,7 +19,7 @@ class APIKeyManager:
 
     def get_api_key(self, provider: str) -> str | None:
         """Get API key for the specified provider using cascading logic."""
-        env_var_name = self._get_env_var_name(provider)
+        env_var_name = self.get_env_var_name(provider)
         service_name = self._get_service_name(provider)
 
         # 1. Check environment variable first
@@ -41,7 +41,7 @@ class APIKeyManager:
         # 3. Prompt user for API key
         return self._prompt_and_store_key(provider, env_var_name, service_name)
 
-    def _get_env_var_name(self, provider: str) -> str:
+    def get_env_var_name(self, provider: str) -> str:
         """Get the expected environment variable name for a provider."""
         # Normalize aliases to canonical provider keys
         key = providers.canonical(provider) or provider

sqlsaber/config/oauth_flow.py

@@ -10,12 +10,13 @@ from datetime import datetime, timezone
 
 import httpx
 import questionary
-from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 
+from sqlsaber.theme.manager import create_console
+
 from .oauth_tokens import OAuthToken, OAuthTokenManager
 
-console = Console()
+console = create_console()
 logger = logging.getLogger(__name__)
 
 

sqlsaber/config/oauth_tokens.py

@@ -6,9 +6,10 @@ from datetime import datetime, timedelta, timezone
 from typing import Any
 
 import keyring
-from rich.console import Console
 
-console = Console()
+from sqlsaber.theme.manager import create_console
+
+console = create_console()
 logger = logging.getLogger(__name__)
 
 
@@ -158,9 +159,6 @@ class OAuthTokenManager:
             keyring.delete_password(service_name, provider)
             console.print(f"OAuth token for {provider} removed", style="green")
             return True
-        except keyring.errors.PasswordDeleteError:
-            # Token doesn't exist
-            return True
         except Exception as e:
             logger.error(f"Failed to remove OAuth token for {provider}: {e}")
             console.print(f"Warning: Could not remove OAuth token: {e}", style="yellow")

sqlsaber/database/base.py

@@ -61,6 +61,12 @@ class BaseDatabaseConnection(ABC):
         self.connection_string = connection_string
         self._pool = None
 
+    @property
+    @abstractmethod
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name for this database."""
+        pass
+
     @abstractmethod
     async def get_pool(self):
         """Get or create connection pool."""

sqlsaber/database/csv.py

@@ -58,6 +58,11 @@ class CSVConnection(BaseDatabaseConnection):
 
         self.table_name = Path(self.csv_path).stem or "csv_table"
 
+    @property
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name."""
+        return "duckdb"
+
     async def get_pool(self):
         """CSV connections do not maintain a pool."""
         return None

sqlsaber/database/duckdb.py

@@ -52,6 +52,11 @@ class DuckDBConnection(BaseDatabaseConnection):
 
         self.database_path = db_path or ":memory:"
 
+    @property
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name."""
+        return "duckdb"
+
     async def get_pool(self):
         """DuckDB creates connections per query, return database path."""
         return self.database_path

sqlsaber/database/mysql.py

@@ -23,6 +23,11 @@ class MySQLConnection(BaseDatabaseConnection):
         self._pool: aiomysql.Pool | None = None
         self._parse_connection_string()
 
+    @property
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name."""
+        return "mysql"
+
     def _parse_connection_string(self):
         """Parse MySQL connection string into components."""
         parsed = urlparse(self.connection_string)

sqlsaber/database/postgresql.py

@@ -23,6 +23,11 @@ class PostgreSQLConnection(BaseDatabaseConnection):
         self._pool: asyncpg.Pool | None = None
         self._ssl_context = self._create_ssl_context()
 
+    @property
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name."""
+        return "postgres"
+
     def _create_ssl_context(self) -> ssl.SSLContext | None:
         """Create SSL context from connection string parameters."""
         parsed = urlparse(self.connection_string)

sqlsaber/database/sqlite.py

@@ -21,6 +21,11 @@ class SQLiteConnection(BaseDatabaseConnection):
         # Extract database path from sqlite:///path format
         self.database_path = connection_string.replace("sqlite:///", "")
 
+    @property
+    def sqlglot_dialect(self) -> str:
+        """Return the sqlglot dialect name."""
+        return "sqlite"
+
     async def get_pool(self):
         """SQLite doesn't use connection pooling, return database path."""
         return self.database_path

sqlsaber/theme/__init__.py

@@ -0,0 +1,5 @@
+"""Theme management for SQLSaber."""
+
+from sqlsaber.theme.manager import ThemeManager, create_console, get_theme_manager
+
+__all__ = ["ThemeManager", "create_console", "get_theme_manager"]