sqlsaber 0.34.0__py3-none-any.whl → 0.36.0__py3-none-any.whl

sqlsaber/application/db_setup.py

@@ -1,7 +1,7 @@
 """Shared database setup logic for onboarding and CLI."""
 
 import getpass
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 
 from sqlsaber.application.prompts import Prompter
@@ -11,6 +11,21 @@ from sqlsaber.theme.manager import create_console
 console = create_console()
 
 
+def _normalize_schemas(schemas: list[str]) -> list[str]:
+    """Deduplicate schema list while preserving order and case."""
+    normalized: list[str] = []
+    seen: set[str] = set()
+    for schema in schemas:
+        name = schema.strip()
+        if not name:
+            continue
+        if name in seen:
+            continue
+        seen.add(name)
+        normalized.append(name)
+    return normalized
+
+
 @dataclass
 class DatabaseInput:
     """Input data for database configuration."""
@@ -26,6 +41,7 @@ class DatabaseInput:
     ssl_ca: str | None = None
     ssl_cert: str | None = None
     ssl_key: str | None = None
+    exclude_schemas: list[str] = field(default_factory=list)
 
 
 async def collect_db_input(
@@ -69,11 +85,20 @@ async def collect_db_input(
         port = 0
         username = db_type
         password = ""
+        exclude_schemas: list[str] = []
         ssl_mode = None
         ssl_ca = None
         ssl_cert = None
         ssl_key = None
 
+        if db_type == "duckdb":
+            exclude_prompt = await prompter.text(
+                "Schemas to exclude (comma separated, optional):", default=""
+            )
+            if exclude_prompt is None:
+                return None
+            exclude_schemas = _normalize_schemas(exclude_prompt.split(","))
+
     else:
         # PostgreSQL/MySQL need connection details
         host = await prompter.text("Host:", default="localhost")
@@ -155,6 +180,13 @@ async def collect_db_input(
                                 "SSL client private key file:"
                             )
 
+        exclude_prompt = await prompter.text(
+            "Schemas to exclude (comma separated, optional):", default=""
+        )
+        if exclude_prompt is None:
+            return None
+        exclude_schemas = _normalize_schemas(exclude_prompt.split(","))
+
     return DatabaseInput(
         name=name,
         type=db_type,
@@ -167,6 +199,7 @@ async def collect_db_input(
         ssl_ca=ssl_ca,
         ssl_cert=ssl_cert,
         ssl_key=ssl_key,
+        exclude_schemas=exclude_schemas,
     )
 
 
@@ -183,6 +216,7 @@ def build_config(db_input: DatabaseInput) -> DatabaseConfig:
         ssl_ca=db_input.ssl_ca,
         ssl_cert=db_input.ssl_cert,
         ssl_key=db_input.ssl_key,
+        exclude_schemas=_normalize_schemas(db_input.exclude_schemas),
     )
 
 
@@ -200,7 +234,9 @@ async def test_connection(config: DatabaseConfig, password: str | None) -> bool:
 
     try:
         connection_string = config.to_connection_string()
-        db_conn = DatabaseConnection(connection_string)
+        db_conn = DatabaseConnection(
+            connection_string, excluded_schemas=config.exclude_schemas
+        )
         await db_conn.execute_query("SELECT 1 as test")
         await db_conn.close()
         return True

sqlsaber/cli/commands.py

@@ -214,7 +214,9 @@ def query(
 
         # Create database connection
         try:
-            db_conn = DatabaseConnection(connection_string)
+            db_conn = DatabaseConnection(
+                connection_string, excluded_schemas=resolved.excluded_schemas
+            )
             log.info("db.connection.created", db_type=type(db_conn).__name__)
         except Exception as e:
             log.exception("db.connection.error", error=str(e))
@@ -229,8 +231,10 @@ def query(
                 # Single query mode with streaming
                 streaming_handler = StreamingQueryHandler(console)
                 db_type = sqlsaber_agent.db_type
+                model_name = sqlsaber_agent.agent.model.model_name
                 console.print(
                     f"[primary]Connected to:[/primary] {db_name} ({db_type})\n"
+                    f"[primary]Model:[/primary] {model_name}\n"
                 )
                 log.info("query.execute.start", db_name=db_name, db_type=db_type)
                 run = await streaming_handler.execute_streaming_query(

sqlsaber/cli/database.py

@@ -26,6 +26,28 @@ db_app = cyclopts.App(
 )
 
 
+def _normalize_schema_list(raw_schemas: list[str]) -> list[str]:
+    """Deduplicate schemas while preserving order and case."""
+    schemas: list[str] = []
+    seen: set[str] = set()
+    for schema in raw_schemas:
+        item = schema.strip()
+        if not item:
+            continue
+        if item in seen:
+            continue
+        seen.add(item)
+        schemas.append(item)
+    return schemas
+
+
+def _parse_schema_list(raw: str | None) -> list[str]:
+    """Parse comma-separated schema list into cleaned list."""
+    if not raw:
+        return []
+    return _normalize_schema_list(raw.split(","))
+
+
 @db_app.command
 def add(
     name: Annotated[str, cyclopts.Parameter(help="Name for the database connection")],
@@ -71,6 +93,13 @@ def add(
         str | None,
         cyclopts.Parameter(["--ssl-key"], help="SSL client private key file path"),
     ] = None,
+    exclude_schemas: Annotated[
+        str | None,
+        cyclopts.Parameter(
+            ["--exclude-schemas"],
+            help="Comma-separated list of schemas to exclude from introspection",
+        ),
+    ] = None,
     interactive: Annotated[
         bool,
         cyclopts.Parameter(
@@ -119,6 +148,7 @@ def add(
         ssl_ca = db_input.ssl_ca
         ssl_cert = db_input.ssl_cert
         ssl_key = db_input.ssl_key
+        exclude_schema_list = _normalize_schema_list(db_input.exclude_schemas)
     else:
         # Non-interactive mode - use provided values or defaults
         if type == "sqlite":
@@ -160,6 +190,7 @@ def add(
                 if questionary.confirm("Enter password?").ask()
                 else ""
             )
+        exclude_schema_list = _parse_schema_list(exclude_schemas)
 
     # Create database config
     # At this point, all required values should be set
@@ -180,6 +211,7 @@ def add(
         ssl_ca=ssl_ca,
         ssl_cert=ssl_cert,
         ssl_key=ssl_key,
+        exclude_schemas=exclude_schema_list,
     )
 
     try:
@@ -219,6 +251,7 @@ def list():
     table.add_column("Port", style="warning")
     table.add_column("Database", style="info")
     table.add_column("Username", style="info")
+    table.add_column("Excluded Schemas", style="muted")
     table.add_column("SSL", style="success")
     table.add_column("Default", style="error")
 
@@ -241,6 +274,7 @@ def list():
             str(db.port) if db.port else "",
             db.database,
             db.username,
+            ", ".join(db.exclude_schemas) if db.exclude_schemas else "",
             ssl_status,
             is_default,
         )
@@ -249,6 +283,116 @@ def list():
     logger.info("db.list.complete", count=len(databases))
 
 
+@db_app.command
+def exclude(
+    name: Annotated[
+        str,
+        cyclopts.Parameter(help="Name of the database connection to update"),
+    ],
+    set_schemas: Annotated[
+        str | None,
+        cyclopts.Parameter(
+            ["--set"],
+            help="Replace excluded schemas with this comma-separated list",
+        ),
+    ] = None,
+    add_schemas: Annotated[
+        str | None,
+        cyclopts.Parameter(
+            ["--add"],
+            help="Add comma-separated schemas to the existing exclude list",
+        ),
+    ] = None,
+    remove_schemas: Annotated[
+        str | None,
+        cyclopts.Parameter(
+            ["--remove"],
+            help="Remove comma-separated schemas from the existing exclude list",
+        ),
+    ] = None,
+    clear: Annotated[
+        bool,
+        cyclopts.Parameter(
+            ["--clear", "--no-clear"],
+            help="Clear all excluded schemas",
+        ),
+    ] = False,
+):
+    """Update excluded schemas for a database connection."""
+    logger.info(
+        "db.exclude.start",
+        name=name,
+        set=bool(set_schemas),
+        add=bool(add_schemas),
+        remove=bool(remove_schemas),
+        clear=clear,
+    )
+    db_config = config_manager.get_database(name)
+    if not db_config:
+        console.print(
+            f"[bold error]Error: Database connection '{name}' not found[/bold error]"
+        )
+        logger.error("db.exclude.not_found", name=name)
+        sys.exit(1)
+
+    actions_selected = sum(
+        bool(flag)
+        for flag in [
+            set_schemas is not None,
+            add_schemas is not None,
+            remove_schemas is not None,
+            clear,
+        ]
+    )
+    if actions_selected > 1:
+        console.print(
+            "[bold error]Error: Specify only one of --set, --add, --remove, or --clear[/bold error]"
+        )
+        logger.error("db.exclude.multiple_actions", name=name)
+        sys.exit(1)
+
+    current = [*(db_config.exclude_schemas or [])]
+
+    if clear:
+        updated = []
+    elif set_schemas is not None:
+        updated = _parse_schema_list(set_schemas)
+    elif add_schemas is not None:
+        additions = _parse_schema_list(add_schemas)
+        updated = [*current]
+        current_set = set(current)
+        for schema in additions:
+            if schema not in current_set:
+                updated.append(schema)
+                current_set.add(schema)
+    elif remove_schemas is not None:
+        removals = set(_parse_schema_list(remove_schemas))
+        updated = [schema for schema in current if schema not in removals]
+    else:
+        console.print(
+            "[info]Update excluded schemas for "
+            f"[primary]{name}[/primary] (leave blank to clear)[/info]"
+        )
+        default_value = ", ".join(current)
+        response = questionary.text(
+            "Schemas to exclude (comma separated):", default=default_value
+        ).ask()
+        if response is None:
+            console.print("[warning]Operation cancelled[/warning]")
+            logger.info("db.exclude.cancelled", name=name)
+            return
+        updated = _parse_schema_list(response)
+
+    db_config.exclude_schemas = _normalize_schema_list(updated)
+    config_manager.update_database(db_config)
+
+    console.print(
+        f"[success]Updated excluded schemas for '{name}':[/success] "
+        f"{', '.join(db_config.exclude_schemas) if db_config.exclude_schemas else '(none)'}"
+    )
+    logger.info("db.exclude.success", name=name, count=len(db_config.exclude_schemas))
+
+
 @db_app.command
 def remove(
     name: Annotated[
@@ -259,7 +403,7 @@ def remove(
     logger.info("db.remove.start", name=name)
     if not config_manager.get_database(name):
         console.print(
-            f"[bold error]Error:[/bold error] Database connection '{name}' not found"
+            f"[bold error]Error: Database connection '{name}' not found[/bold error]"
         )
         logger.error("db.remove.not_found", name=name)
         sys.exit(1)
@@ -269,17 +413,17 @@ def remove(
     ).ask():
         if config_manager.remove_database(name):
             console.print(
-                f"[green]Successfully removed database connection '{name}'[/green]"
+                f"[success]Successfully removed database connection '{name}'[/success]"
             )
             logger.info("db.remove.success", name=name)
         else:
             console.print(
-                f"[bold error]Error:[/bold error] Failed to remove database connection '{name}'"
+                f"[bold error]Error: Failed to remove database connection '{name}'[/bold error]"
             )
             logger.error("db.remove.failed", name=name)
             sys.exit(1)
     else:
-        console.print("Operation cancelled")
+        console.print("[warning]Operation cancelled[/warning]")
         logger.info("db.remove.cancelled", name=name)
 
 
@@ -294,17 +438,19 @@ def set_default(
     logger.info("db.default.start", name=name)
     if not config_manager.get_database(name):
         console.print(
-            f"[bold error]Error:[/bold error] Database connection '{name}' not found"
+            f"[bold error]Error: Database connection '{name}' not found[/bold error]"
         )
         logger.error("db.default.not_found", name=name)
         sys.exit(1)
 
     if config_manager.set_default_database(name):
-        console.print(f"[green]Successfully set '{name}' as default database[/green]")
+        console.print(
+            f"[success]Successfully set '{name}' as default database[/success]"
+        )
         logger.info("db.default.success", name=name)
     else:
         console.print(
-            f"[bold error]Error:[/bold error] Failed to set '{name}' as default"
+            f"[bold error]Error: Failed to set '{name}' as default[/bold error]"
         )
         logger.error("db.default.failed", name=name)
         sys.exit(1)
@@ -330,7 +476,7 @@ def test(
             db_config = config_manager.get_database(name)
             if not db_config:
                 console.print(
-                    f"[bold error]Error:[/bold error] Database connection '{name}' not found"
+                    f"[bold error]Error: Database connection '{name}' not found[/bold error]"
                 )
                 logger.error("db.test.not_found", name=name)
                 sys.exit(1)
@@ -338,7 +484,7 @@ def test(
             db_config = config_manager.get_default_database()
             if not db_config:
                 console.print(
-                    "[bold error]Error:[/bold error] No default database configured"
+                    "[bold error]Error: No default database configured[/bold error]"
                 )
                 console.print(
                     "Use 'sqlsaber db add <name>' to add a database connection"
@@ -350,14 +496,16 @@ def test(
 
         try:
             connection_string = db_config.to_connection_string()
-            db_conn = DatabaseConnection(connection_string)
+            db_conn = DatabaseConnection(
+                connection_string, excluded_schemas=db_config.exclude_schemas
+            )
 
             # Try to connect and run a simple query
             await db_conn.execute_query("SELECT 1 as test")
             await db_conn.close()
 
             console.print(
-                f"[green]✓ Connection to '{db_config.name}' successful[/green]"
+                f"[success]✓ Connection to '{db_config.name}' successful[/success]"
             )
             logger.info("db.test.success", name=db_config.name)
 
@@ -369,7 +517,7 @@ def test(
                 ),
                 error=str(e),
             )
-            console.print(f"[bold error]✗ Connection failed:[/bold error] {e}")
+            console.print(f"[bold error]✗ Connection failed: {e}[/bold error]")
             sys.exit(1)
 
     asyncio.run(test_connection())

sqlsaber/cli/display.py

@@ -406,9 +406,17 @@ class DisplayManager:
             for table_name, table_info in data.items():
                 self.console.print(f"\n[heading]Table: {table_name}[/heading]")
 
+                table_comment = table_info.get("comment")
+                if table_comment:
+                    self.console.print(f"[muted]Comment: {table_comment}[/muted]")
+
                 # Show columns
                 table_columns = table_info.get("columns", {})
                 if table_columns:
+                    include_column_comments = any(
+                        col_info.get("comment") for col_info in table_columns.values()
+                    )
+
                     # Create a table for columns
                     columns = [
                         {"name": "Column Name", "style": "column.name"},
@@ -416,6 +424,8 @@ class DisplayManager:
                         {"name": "Nullable", "style": "info"},
                         {"name": "Default", "style": "muted"},
                     ]
+                    if include_column_comments:
+                        columns.append({"name": "Comment", "style": "muted"})
                     col_table = self._create_table(columns, title="Columns")
 
                     for col_name, col_info in table_columns.items():
@@ -425,9 +435,15 @@ class DisplayManager:
                             if col_info.get("default")
                             else ""
                         )
-                        col_table.add_row(
-                            col_name, col_info.get("type", ""), nullable, default
-                        )
+                        row = [
+                            col_name,
+                            col_info.get("type", ""),
+                            nullable,
+                            default,
+                        ]
+                        if include_column_comments:
+                            row.append(col_info.get("comment") or "")
+                        col_table.add_row(*row)
 
                     self.console.print(col_table)
 

sqlsaber/cli/interactive.py

@@ -20,6 +20,7 @@ from sqlsaber.cli.completers import (
 )
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
+from sqlsaber.config.logging import get_logger
 from sqlsaber.database import (
     CSVConnection,
     DuckDBConnection,
@@ -30,7 +31,6 @@ from sqlsaber.database import (
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.theme.manager import get_theme_manager
 from sqlsaber.threads import ThreadStorage
-from sqlsaber.config.logging import get_logger
 
 if TYPE_CHECKING:
     from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
@@ -135,8 +135,10 @@ class InteractiveSession:
             )
 
         db_name = self.database_name or "Unknown"
+        model_name = self.sqlsaber_agent.agent.model.model_name
         self.console.print(
             f"[heading]\nConnected to {db_name} ({self._db_type_name()})[/heading]\n"
+            f"[heading]Model: {model_name}[/heading]\n"
         )
 
         if self._thread_id:
@@ -309,6 +311,8 @@ class InteractiveSession:
                         style=self.tm.pt_style(),
                     )
 
+                user_query = user_query.strip()
+
                 if not user_query:
                     continue
 
@@ -325,7 +329,7 @@ class InteractiveSession:
 
                 # Handle memory addition
                 if user_query.strip().startswith("#"):
-                    await self._handle_memory(user_query.strip()[1:].strip())
+                    await self._handle_memory(user_query[1:].strip())
                     continue
 
                 # Execute query with cancellation support

sqlsaber/cli/threads.py

@@ -229,7 +229,7 @@ def list_threads(
         logger.info("threads.cli.list.empty")
         return
     table = Table(title="Threads")
-    table.add_column("ID", style=tm.style("info"))
+    table.add_column("ID", style=tm.style("info"), no_wrap=True, min_width=36)
     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"))
@@ -318,7 +318,9 @@ def resume(
             )
             return
 
-        db_conn = DatabaseConnection(connection_string)
+        db_conn = DatabaseConnection(
+            connection_string, excluded_schemas=resolved.excluded_schemas
+        )
         try:
             sqlsaber_agent = SQLSaberAgent(db_conn, db_name)
             history = await store.get_thread_messages(thread_id)

sqlsaber/config/database.py

@@ -4,7 +4,7 @@ import json
 import os
 import platform
 import stat
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 from urllib.parse import quote_plus
@@ -29,6 +29,7 @@ class DatabaseConfig:
     ssl_cert: str | None = None
     ssl_key: str | None = None
     schema: str | None = None
+    exclude_schemas: list[str] = field(default_factory=list)
 
     def to_connection_string(self) -> str:
         """Convert config to database connection string."""
@@ -149,6 +150,7 @@ class DatabaseConfig:
             "ssl_cert": self.ssl_cert,
             "ssl_key": self.ssl_key,
             "schema": self.schema,
+            "exclude_schemas": self.exclude_schemas,
         }
 
     @classmethod
@@ -166,6 +168,7 @@ class DatabaseConfig:
             ssl_cert=data.get("ssl_cert"),
             ssl_key=data.get("ssl_key"),
             schema=data.get("schema"),
+            exclude_schemas=list(data.get("exclude_schemas", [])),
         )
 
 
@@ -246,6 +249,16 @@ class DatabaseConfigManager:
 
         self._save_config(config)
 
+    def update_database(self, db_config: DatabaseConfig) -> None:
+        """Update an existing database configuration."""
+        config = self._load_config()
+
+        if db_config.name not in config["connections"]:
+            raise ValueError(f"Database '{db_config.name}' does not exist")
+
+        config["connections"][db_config.name] = db_config.to_dict()
+        self._save_config(config)
+
     def get_database(self, name: str) -> DatabaseConfig | None:
         """Get a database configuration by name."""
         config = self._load_config()

sqlsaber/database/__init__.py

@@ -1,5 +1,7 @@
 """Database module for SQLSaber."""
 
+from collections.abc import Iterable
+
 from .base import (
     DEFAULT_QUERY_TIMEOUT,
     BaseDatabaseConnection,
@@ -18,23 +20,28 @@ from .schema import SchemaManager
 from .sqlite import SQLiteConnection, SQLiteSchemaIntrospector
 
 
-def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
+def DatabaseConnection(
+    connection_string: str, *, excluded_schemas: Iterable[str] | None = None
+) -> BaseDatabaseConnection:
     """Factory function to create appropriate database connection based on connection string."""
     if connection_string.startswith("postgresql://"):
-        return PostgreSQLConnection(connection_string)
+        conn = PostgreSQLConnection(connection_string)
     elif connection_string.startswith("mysql://"):
-        return MySQLConnection(connection_string)
+        conn = MySQLConnection(connection_string)
     elif connection_string.startswith("sqlite:///"):
-        return SQLiteConnection(connection_string)
+        conn = SQLiteConnection(connection_string)
     elif connection_string.startswith("duckdb://"):
-        return DuckDBConnection(connection_string)
+        conn = DuckDBConnection(connection_string)
     elif connection_string.startswith("csv:///"):
-        return CSVConnection(connection_string)
+        conn = CSVConnection(connection_string)
     else:
         raise ValueError(
             f"Unsupported database type in connection string: {connection_string}"
         )
 
+    conn.set_excluded_schemas(excluded_schemas)
+    return conn
+
 
 __all__ = [
     # Base classes and types