sqlsaber 0.1.0__py3-none-any.whl

sqlsaber/cli/commands.py

@@ -0,0 +1,132 @@
+"""CLI command definitions and handlers."""
+
+import asyncio
+from typing import Optional
+
+import typer
+from rich.console import Console
+
+from sqlsaber.agents.anthropic import AnthropicSQLAgent
+from sqlsaber.cli.database import create_db_app
+from sqlsaber.cli.interactive import InteractiveSession
+from sqlsaber.cli.memory import create_memory_app
+from sqlsaber.cli.models import create_models_app
+from sqlsaber.cli.streaming import StreamingQueryHandler
+from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.database.connection import DatabaseConnection
+
+app = typer.Typer(
+    name="sqlsaber",
+    help="SQLSaber - Use the agent Luke!\n\nSQL assistant for your database",
+    add_completion=True,
+)
+
+
+console = Console()
+config_manager = DatabaseConfigManager()
+
+
+@app.callback()
+def main_callback(
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """
+    Query your database using natural language.
+
+    Examples:
+        sb query                             # Start interactive mode
+        sb query "show me all users"         # Run a single query with default database
+        sb query -d mydb "show me users"     # Run a query with specific database
+    """
+    pass
+
+
+@app.command()
+def query(
+    query_text: Optional[str] = typer.Argument(
+        None,
+        help="SQL query in natural language (if not provided, starts interactive mode)",
+    ),
+    database: Optional[str] = typer.Option(
+        None,
+        "--database",
+        "-d",
+        help="Database connection name (uses default if not specified)",
+    ),
+):
+    """Run a query against the database or start interactive mode."""
+
+    async def run_session():
+        # Get database configuration
+        if database:
+            db_config = config_manager.get_database(database)
+            if not db_config:
+                console.print(
+                    f"[bold red]Error:[/bold red] Database connection '{database}' not found."
+                )
+                console.print("Use 'sqlsaber db list' to see available connections.")
+                raise typer.Exit(1)
+        else:
+            db_config = config_manager.get_default_database()
+            if not db_config:
+                console.print(
+                    "[bold red]Error:[/bold red] No database connections configured."
+                )
+                console.print(
+                    "Use 'sqlsaber db add <name>' to add a database connection."
+                )
+                raise typer.Exit(1)
+
+        # Create database connection
+        try:
+            connection_string = db_config.to_connection_string()
+            db_conn = DatabaseConnection(connection_string)
+        except Exception as e:
+            console.print(
+                f"[bold red]Error creating database connection:[/bold red] {e}"
+            )
+            raise typer.Exit(1)
+
+        # Create agent instance with database name for memory context
+        agent = AnthropicSQLAgent(db_conn, db_config.name)
+
+        try:
+            if query_text:
+                # Single query mode with streaming
+                streaming_handler = StreamingQueryHandler(console)
+                await streaming_handler.execute_streaming_query(query_text, agent)
+            else:
+                # Interactive mode
+                session = InteractiveSession(console, agent)
+                await session.run()
+
+        finally:
+            # Clean up
+            await db_conn.close()
+            console.print("\n[green]Goodbye![/green]")
+
+    # Run the async function
+    asyncio.run(run_session())
+
+
+# Add database management commands after main callback is defined
+db_app = create_db_app()
+app.add_typer(db_app, name="db")
+
+# Add memory management commands
+memory_app = create_memory_app()
+app.add_typer(memory_app, name="memory")
+
+# Add model management commands
+models_app = create_models_app()
+app.add_typer(models_app, name="models")
+
+
+def main():
+    """Entry point for the CLI application."""
+    app()

sqlsaber/cli/database.py

@@ -0,0 +1,275 @@
+"""Database management CLI commands."""
+
+import asyncio
+import getpass
+from pathlib import Path
+from typing import Optional
+
+import questionary
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from sqlsaber.config.database import DatabaseConfig, DatabaseConfigManager
+from sqlsaber.database.connection import DatabaseConnection
+
+# Global instances for CLI commands
+console = Console()
+config_manager = DatabaseConfigManager()
+
+# Create the database management CLI app
+db_app = typer.Typer(
+    name="db",
+    help="Manage database connections",
+    add_completion=True,
+)
+
+
+@db_app.command("add")
+def add_database(
+    name: str = typer.Argument(..., help="Name for the database connection"),
+    type: str = typer.Option(
+        "postgresql",
+        "--type",
+        "-t",
+        help="Database type (postgresql, mysql, sqlite)",
+    ),
+    host: Optional[str] = typer.Option(None, "--host", "-h", help="Database host"),
+    port: Optional[int] = typer.Option(None, "--port", "-p", help="Database port"),
+    database: Optional[str] = typer.Option(
+        None, "--database", "--db", help="Database name"
+    ),
+    username: Optional[str] = typer.Option(None, "--username", "-u", help="Username"),
+    interactive: bool = typer.Option(
+        True, "--interactive/--no-interactive", help="Use interactive mode"
+    ),
+):
+    """Add a new database connection."""
+
+    if interactive:
+        # Interactive mode - prompt for all required fields
+        console.print(f"[bold]Adding database connection: {name}[/bold]")
+
+        # Database type
+        if not type or type == "postgresql":
+            type = questionary.select(
+                "Database type:",
+                choices=["postgresql", "mysql", "sqlite"],
+                default="postgresql",
+            ).ask()
+
+        if type == "sqlite":
+            # SQLite only needs database path
+            database = database or questionary.path("Database file path:").ask()
+            database = str(Path(database).expanduser())
+            host = "localhost"
+            port = 0
+            username = "sqlite"
+            password = ""
+        else:
+            # PostgreSQL/MySQL need connection details
+            host = host or questionary.text("Host:", default="localhost").ask()
+
+            default_port = 5432 if type == "postgresql" else 3306
+            port = port or int(
+                questionary.text("Port:", default=str(default_port)).ask()
+            )
+
+            database = database or questionary.text("Database name:").ask()
+            username = username or questionary.text("Username:").ask()
+
+            # Ask for password
+            password = getpass.getpass("Password (stored in your OS keychain): ")
+    else:
+        # Non-interactive mode - use provided values or defaults
+        if type == "sqlite":
+            if not database:
+                console.print(
+                    "[bold red]Error:[/bold red] Database file path is required for SQLite"
+                )
+                raise typer.Exit(1)
+            host = "localhost"
+            port = 0
+            username = "sqlite"
+            password = ""
+        else:
+            if not all([host, database, username]):
+                console.print(
+                    "[bold red]Error:[/bold red] Host, database, and username are required"
+                )
+                raise typer.Exit(1)
+
+            if port is None:
+                port = 5432 if type == "postgresql" else 3306
+
+            password = (
+                getpass.getpass("Password (stored in your OS keychain): ")
+                if questionary.confirm("Enter password?").ask()
+                else ""
+            )
+
+    # Create database config
+    # At this point, all required values should be set
+    assert database is not None, "Database should be set by now"
+    if type != "sqlite":
+        assert host is not None, "Host should be set by now"
+        assert port is not None, "Port should be set by now"
+        assert username is not None, "Username should be set by now"
+
+    db_config = DatabaseConfig(
+        name=name,
+        type=type,
+        host=host,
+        port=port,
+        database=database,
+        username=username,
+    )
+
+    try:
+        # Add the configuration
+        config_manager.add_database(db_config, password if password else None)
+        console.print(f"[green]Successfully added database connection '{name}'[/green]")
+
+        # Set as default if it's the first one
+        if len(config_manager.list_databases()) == 1:
+            console.print(f"[blue]Set '{name}' as default database[/blue]")
+
+    except Exception as e:
+        console.print(f"[bold red]Error adding database:[/bold red] {e}")
+        raise typer.Exit(1)
+
+
+@db_app.command("list")
+def list_databases():
+    """List all configured database connections."""
+    databases = config_manager.list_databases()
+    default_name = config_manager.get_default_name()
+
+    if not databases:
+        console.print("[yellow]No database connections configured[/yellow]")
+        console.print("Use 'sqlsaber db add <name>' to add a database connection")
+        return
+
+    table = Table(title="Database Connections")
+    table.add_column("Name", style="cyan")
+    table.add_column("Type", style="magenta")
+    table.add_column("Host", style="green")
+    table.add_column("Port", style="yellow")
+    table.add_column("Database", style="blue")
+    table.add_column("Username", style="white")
+    table.add_column("Default", style="bold red")
+
+    for db in databases:
+        is_default = "✓" if db.name == default_name else ""
+        table.add_row(
+            db.name,
+            db.type,
+            db.host,
+            str(db.port) if db.port else "",
+            db.database,
+            db.username,
+            is_default,
+        )
+
+    console.print(table)
+
+
+@db_app.command("remove")
+def remove_database(
+    name: str = typer.Argument(..., help="Name of the database connection to remove"),
+):
+    """Remove a database connection."""
+    if not config_manager.get_database(name):
+        console.print(
+            f"[bold red]Error:[/bold red] Database connection '{name}' not found"
+        )
+        raise typer.Exit(1)
+
+    if questionary.confirm(
+        f"Are you sure you want to remove database connection '{name}'?"
+    ).ask():
+        if config_manager.remove_database(name):
+            console.print(
+                f"[green]Successfully removed database connection '{name}'[/green]"
+            )
+        else:
+            console.print(
+                f"[bold red]Error:[/bold red] Failed to remove database connection '{name}'"
+            )
+            raise typer.Exit(1)
+    else:
+        console.print("Operation cancelled")
+
+
+@db_app.command("set-default")
+def set_default_database(
+    name: str = typer.Argument(
+        ..., help="Name of the database connection to set as default"
+    ),
+):
+    """Set the default database connection."""
+    if not config_manager.get_database(name):
+        console.print(
+            f"[bold red]Error:[/bold red] Database connection '{name}' not found"
+        )
+        raise typer.Exit(1)
+
+    if config_manager.set_default_database(name):
+        console.print(f"[green]Successfully set '{name}' as default database[/green]")
+    else:
+        console.print(f"[bold red]Error:[/bold red] Failed to set '{name}' as default")
+        raise typer.Exit(1)
+
+
+@db_app.command("test")
+def test_database(
+    name: Optional[str] = typer.Argument(
+        None,
+        help="Name of the database connection to test (uses default if not specified)",
+    ),
+):
+    """Test a database connection."""
+
+    async def test_connection():
+        if name:
+            db_config = config_manager.get_database(name)
+            if not db_config:
+                console.print(
+                    f"[bold red]Error:[/bold red] Database connection '{name}' not found"
+                )
+                raise typer.Exit(1)
+        else:
+            db_config = config_manager.get_default_database()
+            if not db_config:
+                console.print(
+                    "[bold red]Error:[/bold red] No default database configured"
+                )
+                console.print(
+                    "Use 'sqlsaber db add <name>' to add a database connection"
+                )
+                raise typer.Exit(1)
+
+        console.print(f"[blue]Testing connection to '{db_config.name}'...[/blue]")
+
+        try:
+            connection_string = db_config.to_connection_string()
+            db_conn = DatabaseConnection(connection_string)
+
+            # 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]"
+            )
+
+        except Exception as e:
+            console.print(f"[bold red]✗ Connection failed:[/bold red] {e}")
+            raise typer.Exit(1)
+
+    asyncio.run(test_connection())
+
+
+def create_db_app() -> typer.Typer:
+    """Return the database management CLI app."""
+    return db_app

sqlsaber/cli/display.py

@@ -0,0 +1,207 @@
+"""Display utilities for the CLI interface."""
+
+import json
+from typing import Optional
+
+from rich.console import Console
+from rich.syntax import Syntax
+from rich.table import Table
+
+
+class DisplayManager:
+    """Manages display formatting and output for the CLI."""
+
+    def __init__(self, console: Console):
+        self.console = console
+
+    def _create_table(
+        self,
+        columns: list,
+        header_style: str = "bold blue",
+        title: Optional[str] = None,
+    ) -> Table:
+        """Create a Rich table with specified columns."""
+        table = Table(show_header=True, header_style=header_style, title=title)
+        for col in columns:
+            if isinstance(col, dict):
+                table.add_column(
+                    col["name"], style=col.get("style"), justify=col.get("justify")
+                )
+            else:
+                table.add_column(col)
+        return table
+
+    def show_tool_started(self, tool_name: str):
+        """Display tool started message."""
+        self.console.print(f"\n[yellow]🔧 Using tool: {tool_name}[/yellow]")
+
+    def show_tool_executing(self, tool_name: str, tool_input: dict):
+        """Display tool execution details."""
+        if tool_name == "list_tables":
+            self.console.print("[dim]  → Discovering available tables[/dim]")
+        elif tool_name == "introspect_schema":
+            pattern = tool_input.get("table_pattern", "all tables")
+            self.console.print(f"[dim]  → Examining schema for: {pattern}[/dim]")
+        elif tool_name == "execute_sql":
+            query = tool_input.get("query", "")
+            self.console.print("\n[bold green]Executing SQL:[/bold green]")
+            syntax = Syntax(query, "sql")
+            self.console.print(syntax)
+
+    def show_text_stream(self, text: str):
+        """Display streaming text."""
+        self.console.print(text, end="", markup=False)
+
+    def show_query_results(self, results: list):
+        """Display query results in a formatted table."""
+        if not results:
+            return
+
+        self.console.print(
+            f"\n[bold magenta]Results ({len(results)} rows):[/bold magenta]"
+        )
+
+        # Create table with columns from first result
+        columns = list(results[0].keys())
+        table = self._create_table(columns)
+
+        # Add rows (show first 20 rows)
+        for row in results[:20]:
+            table.add_row(*[str(row[key]) for key in columns])
+
+        self.console.print(table)
+
+        if len(results) > 20:
+            self.console.print(
+                f"[yellow]... and {len(results) - 20} more rows[/yellow]"
+            )
+
+    def show_error(self, error_message: str):
+        """Display error message."""
+        self.console.print(f"\n[bold red]Error:[/bold red] {error_message}")
+
+    def show_processing(self, message: str):
+        """Display processing message."""
+        self.console.print()  # Add newline
+        return self.console.status(
+            f"[yellow]{message}[/yellow]", spinner="bouncingBall"
+        )
+
+    def show_newline(self):
+        """Display a newline for spacing."""
+        self.console.print()
+
+    def show_table_list(self, tables_data: str):
+        """Display the results from list_tables tool."""
+        try:
+            data = json.loads(tables_data)
+
+            # Handle error case
+            if "error" in data:
+                self.show_error(data["error"])
+                return
+
+            tables = data.get("tables", [])
+            total_tables = data.get("total_tables", 0)
+
+            if not tables:
+                self.console.print("[yellow]No tables found in the database.[/yellow]")
+                return
+
+            self.console.print(
+                f"\n[bold green]Database Tables ({total_tables} total):[/bold green]"
+            )
+
+            # Create a rich table for displaying table information
+            columns = [
+                {"name": "Schema", "style": "cyan"},
+                {"name": "Table Name", "style": "white"},
+                {"name": "Type", "style": "yellow"},
+                {"name": "Row Count", "justify": "right", "style": "magenta"},
+            ]
+            table = self._create_table(columns)
+
+            # Add rows
+            for table_info in tables:
+                schema = table_info.get("schema", "")
+                name = table_info.get("name", "")
+                table_type = table_info.get("type", "")
+                row_count = table_info.get("row_count", 0)
+
+                # Format row count with commas for readability
+                formatted_count = f"{row_count:,}" if row_count else "0"
+
+                table.add_row(schema, name, table_type, formatted_count)
+
+            self.console.print(table)
+
+        except json.JSONDecodeError:
+            self.show_error("Failed to parse table list data")
+        except Exception as e:
+            self.show_error(f"Error displaying table list: {str(e)}")
+
+    def show_schema_info(self, schema_data: str):
+        """Display the results from introspect_schema tool."""
+        try:
+            data = json.loads(schema_data)
+
+            # Handle error case
+            if "error" in data:
+                self.show_error(data["error"])
+                return
+
+            if not data:
+                self.console.print("[yellow]No schema information found.[/yellow]")
+                return
+
+            self.console.print(
+                f"\n[bold green]Schema Information ({len(data)} tables):[/bold green]"
+            )
+
+            # Display each table's schema
+            for table_name, table_info in data.items():
+                self.console.print(f"\n[bold cyan]Table: {table_name}[/bold cyan]")
+
+                # Show columns
+                table_columns = table_info.get("columns", {})
+                if table_columns:
+                    # Create a table for columns
+                    columns = [
+                        {"name": "Column Name", "style": "white"},
+                        {"name": "Type", "style": "yellow"},
+                        {"name": "Nullable", "style": "cyan"},
+                        {"name": "Default", "style": "dim"},
+                    ]
+                    col_table = self._create_table(columns, title="Columns")
+
+                    for col_name, col_info in table_columns.items():
+                        nullable = "✓" if col_info.get("nullable", False) else "✗"
+                        default = (
+                            str(col_info.get("default", ""))
+                            if col_info.get("default")
+                            else ""
+                        )
+                        col_table.add_row(
+                            col_name, col_info.get("type", ""), nullable, default
+                        )
+
+                    self.console.print(col_table)
+
+                # Show primary keys
+                primary_keys = table_info.get("primary_keys", [])
+                if primary_keys:
+                    self.console.print(
+                        f"[bold yellow]Primary Keys:[/bold yellow] {', '.join(primary_keys)}"
+                    )
+
+                # Show foreign keys
+                foreign_keys = table_info.get("foreign_keys", [])
+                if foreign_keys:
+                    self.console.print("[bold magenta]Foreign Keys:[/bold magenta]")
+                    for fk in foreign_keys:
+                        self.console.print(f"  • {fk}")
+
+        except json.JSONDecodeError:
+            self.show_error("Failed to parse schema data")
+        except Exception as e:
+            self.show_error(f"Error displaying schema information: {str(e)}")