sqlsaber 0.27.0__py3-none-any.whl → 0.29.0__py3-none-any.whl

sqlsaber/cli/theme.py

@@ -0,0 +1,146 @@
+"""Theme management CLI commands."""
+
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+
+import cyclopts
+import questionary
+from platformdirs import user_config_dir
+from pygments.styles import get_all_styles
+
+from sqlsaber.theme.manager import DEFAULT_THEME_NAME, create_console
+
+console = create_console()
+
+# Create the theme management CLI app
+theme_app = cyclopts.App(
+    name="theme",
+    help="Manage theme settings",
+)
+
+
+class ThemeManager:
+    """Manages theme configuration persistence."""
+
+    def __init__(self):
+        self.config_dir = Path(user_config_dir("sqlsaber"))
+        self.config_file = self.config_dir / "theme.json"
+
+    def _ensure_config_dir(self) -> None:
+        """Ensure config directory exists."""
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+
+    def _load_config(self) -> dict:
+        """Load theme configuration from file."""
+        if not self.config_file.exists():
+            return {}
+
+        try:
+            with open(self.config_file, "r") as f:
+                return json.load(f)
+        except Exception:
+            return {}
+
+    def _save_config(self, config: dict) -> None:
+        """Save theme configuration to file."""
+        self._ensure_config_dir()
+
+        with open(self.config_file, "w") as f:
+            json.dump(config, f, indent=2)
+
+    def get_current_theme(self) -> str:
+        """Get the currently configured theme."""
+        config = self._load_config()
+        env_theme = os.getenv("SQLSABER_THEME")
+        if env_theme:
+            return env_theme
+        return config.get("theme", {}).get("pygments_style") or DEFAULT_THEME_NAME
+
+    def set_theme(self, theme_name: str) -> bool:
+        """Set the current theme."""
+        try:
+            config = self._load_config()
+            if "theme" not in config:
+                config["theme"] = {}
+            config["theme"]["name"] = theme_name
+            config["theme"]["pygments_style"] = theme_name
+            self._save_config(config)
+            return True
+        except Exception as e:
+            console.print(f"[error]Error setting theme: {e}[/error]")
+            return False
+
+    def reset_theme(self) -> bool:
+        """Reset to default theme."""
+        try:
+            if self.config_file.exists():
+                self.config_file.unlink()
+            return True
+        except Exception as e:
+            console.print(f"[error]Error resetting theme: {e}[/error]")
+            return False
+
+    def get_available_themes(self) -> list[str]:
+        """Get list of available Pygments themes."""
+        return sorted(get_all_styles())
+
+
+theme_manager = ThemeManager()
+
+
+@theme_app.command
+def set():
+    """Set the theme to use for syntax highlighting."""
+
+    async def interactive_set():
+        themes = theme_manager.get_available_themes()
+        current_theme = theme_manager.get_current_theme()
+
+        # Create choices with current theme highlighted
+        choices = [
+            questionary.Choice(
+                title=f"{theme} (current)" if theme == current_theme else theme,
+                value=theme,
+            )
+            for theme in themes
+        ]
+
+        selected_theme = await questionary.select(
+            "Select a theme:",
+            choices=choices,
+            default=current_theme,
+            use_search_filter=True,
+            use_jk_keys=False,
+        ).ask_async()
+
+        if selected_theme:
+            if theme_manager.set_theme(selected_theme):
+                console.print(f"[success]✓ Theme set to: {selected_theme}[/success]")
+            else:
+                console.print("[error]✗ Failed to set theme[/error]")
+                sys.exit(1)
+        else:
+            console.print("[warning]Operation cancelled[/warning]")
+
+    asyncio.run(interactive_set())
+
+
+@theme_app.command
+def reset():
+    """Reset to the default theme."""
+
+    if theme_manager.reset_theme():
+        console.print(
+            f"[success]✓ Theme reset to default: {DEFAULT_THEME_NAME}[/success]"
+        )
+    else:
+        console.print("[error]✗ Failed to reset theme[/error]")
+        sys.exit(1)
+
+
+def create_theme_app() -> cyclopts.App:
+    """Return the theme management CLI app."""
+    return theme_app

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:

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"]

sqlsaber/theme/manager.py

@@ -0,0 +1,229 @@
+"""Theme management for unified theming across Rich and prompt_toolkit."""
+
+import json
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Dict
+
+from platformdirs import user_config_dir
+from prompt_toolkit.styles import Style as PTStyle
+from prompt_toolkit.styles.pygments import style_from_pygments_cls
+from pygments.styles import get_style_by_name
+from pygments.token import Token
+from pygments.util import ClassNotFound
+from rich.console import Console
+from rich.theme import Theme
+
+DEFAULT_THEME_NAME = "nord"
+
+DEFAULT_ROLE_PALETTE = {
+    # base roles
+    "primary": "cyan",
+    "accent": "magenta",
+    "success": "green",
+    "warning": "yellow",
+    "error": "red",
+    "info": "cyan",
+    "muted": "dim",
+    # components
+    "table.header": "bold $primary",
+    "panel.border.user": "$info",
+    "panel.border.assistant": "$success",
+    "panel.border.thread": "$primary",
+    "spinner": "$warning",
+    "status": "$warning",
+    # domain-specific
+    "key.primary": "bold $warning",
+    "key.foreign": "bold $accent",
+    "key.index": "bold $primary",
+    "column.schema": "$info",
+    "column.name": "white",
+    "column.type": "$warning",
+    "heading": "bold $primary",
+    "section": "bold $accent",
+    "title": "bold $success",
+}
+
+ROLE_TOKEN_PREFERENCES: dict[str, tuple] = {
+    "primary": (
+        Token.Keyword,
+        Token.Keyword.Namespace,
+        Token.Name.Tag,
+    ),
+    "accent": (
+        Token.Name.Tag,
+        Token.Keyword.Type,
+        Token.Literal.Number,
+        Token.Operator.Word,
+    ),
+    "success": (
+        Token.Literal.String,
+        Token.Generic.Inserted,
+        Token.Name.Attribute,
+    ),
+    "warning": (
+        Token.Literal.String.Escape,
+        Token.Name.Constant,
+        Token.Generic.Emph,
+    ),
+    "error": (
+        Token.Error,
+        Token.Generic.Error,
+        Token.Generic.Deleted,
+        Token.Name.Exception,
+    ),
+    "info": (
+        Token.Name.Function,
+        Token.Name.Builtin,
+        Token.Keyword.Type,
+    ),
+    "muted": (
+        Token.Comment,
+        Token.Generic.Subheading,
+        Token.Text,
+    ),
+}
+
+
+def _normalize_hex(color: str | None) -> str | None:
+    if not color:
+        return None
+    color = color.strip()
+    if not color:
+        return None
+    if color.startswith("#"):
+        color = color[1:]
+    if len(color) == 3:
+        color = "".join(ch * 2 for ch in color)
+    if len(color) != 6:
+        return None
+    return f"#{color.lower()}"
+
+
+def _build_role_palette_from_style(style_name: str) -> dict[str, str]:
+    try:
+        style_cls = get_style_by_name(style_name)
+    except ClassNotFound:
+        return {}
+
+    palette: dict[str, str] = {}
+    try:
+        base_color = _normalize_hex(style_cls.style_for_token(Token.Text).get("color"))
+    except KeyError:
+        base_color = None
+    for role, tokens in ROLE_TOKEN_PREFERENCES.items():
+        for token in tokens:
+            try:
+                style_def = style_cls.style_for_token(token)
+            except KeyError:
+                continue
+            color = _normalize_hex(style_def.get("color"))
+            if not color or color == base_color:
+                continue
+            if role == "accent" and color == palette.get("primary"):
+                continue
+            palette[role] = color
+            break
+    return palette
+
+
+def _load_user_theme_config() -> dict:
+    """Load theme configuration from user config directory."""
+    cfg_dir = user_config_dir("sqlsaber")
+    path = os.path.join(cfg_dir, "theme.json")
+    if not os.path.exists(path):
+        return {}
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def _resolve_refs(palette: dict[str, str]) -> dict[str, str]:
+    """Resolve $var references in palette values."""
+    out = {}
+    for k, v in palette.items():
+        if isinstance(v, str) and "$" in v:
+            parts = v.split()
+            resolved = []
+            for part in parts:
+                if part.startswith("$"):
+                    ref = part[1:]
+                    resolved.append(palette.get(ref, ""))
+                else:
+                    resolved.append(part)
+            out[k] = " ".join(p for p in resolved if p)
+        else:
+            out[k] = v
+    return out
+
+
+@dataclass(frozen=True)
+class ThemeConfig:
+    """Theme configuration."""
+
+    name: str
+    pygments_style: str
+    roles: Dict[str, str]
+
+
+class ThemeManager:
+    """Manages theme configuration and provides themed components."""
+
+    def __init__(self, cfg: ThemeConfig):
+        self._cfg = cfg
+        self._roles = _resolve_refs({**DEFAULT_ROLE_PALETTE, **cfg.roles})
+        self._rich_theme = Theme(self._roles)
+        self._pt_style = None
+
+    @property
+    def rich_theme(self) -> Theme:
+        """Get Rich theme with semantic role mappings."""
+        return self._rich_theme
+
+    @property
+    def pygments_style_name(self) -> str:
+        """Get pygments style name for syntax highlighting."""
+        return self._cfg.pygments_style
+
+    def pt_style(self) -> PTStyle:
+        """Get prompt_toolkit style derived from Pygments theme."""
+        if self._pt_style is None:
+            try:
+                # Try to use Pygments style directly
+                pygments_style = get_style_by_name(self._cfg.pygments_style)
+                self._pt_style = style_from_pygments_cls(pygments_style)
+            except Exception:
+                # Fallback to basic style if Pygments theme not found
+                self._pt_style = PTStyle.from_dict({})
+        return self._pt_style
+
+    def style(self, role: str) -> str:
+        """Get style string for a semantic role."""
+        return self._roles.get(role, "")
+
+
+@lru_cache(maxsize=1)
+def get_theme_manager() -> ThemeManager:
+    """Get the global theme manager instance."""
+    user_cfg = _load_user_theme_config()
+    env_name = os.getenv("SQLSABER_THEME")
+
+    name = (
+        env_name or user_cfg.get("theme", {}).get("name") or DEFAULT_THEME_NAME
+    ).lower()
+    pygments_style = user_cfg.get("theme", {}).get("pygments_style") or name
+
+    roles = dict(DEFAULT_ROLE_PALETTE)
+    roles.update(_build_role_palette_from_style(pygments_style))
+    roles.update(user_cfg.get("roles", {}))
+
+    cfg = ThemeConfig(name=name, pygments_style=pygments_style, roles=roles)
+    return ThemeManager(cfg)
+
+
+def create_console(**kwargs):
+    """Create a Rich Console with theme applied."""
+    # from rich.console import Console
+
+    tm = get_theme_manager()
+    return Console(theme=tm.rich_theme, **kwargs)