sqlsaber 0.23.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

sqlsaber/agents/base.py

@@ -8,6 +8,7 @@ from typing import Any, AsyncIterator
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
     CSVConnection,
+    DuckDBConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -51,7 +52,9 @@ class BaseSQLAgent(ABC):
         elif isinstance(self.db, SQLiteConnection):
             return "SQLite"
         elif isinstance(self.db, CSVConnection):
-            return "SQLite"  # we convert csv to in-memory sqlite
+            return "DuckDB"
+        elif isinstance(self.db, DuckDBConnection):
+            return "DuckDB"
         else:
             return "database"  # Fallback
 

sqlsaber/agents/pydantic_ai_agent.py

@@ -17,6 +17,7 @@ from sqlsaber.config.settings import Config
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
     CSVConnection,
+    DuckDBConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -169,7 +170,9 @@ def _get_database_type_name(db: BaseDatabaseConnection) -> str:
         return "MySQL"
     elif isinstance(db, SQLiteConnection):
         return "SQLite"
+    elif isinstance(db, DuckDBConnection):
+        return "DuckDB"
     elif isinstance(db, CSVConnection):
-        return "SQLite"
+        return "DuckDB"
     else:
         return "database"

sqlsaber/cli/commands.py

@@ -46,7 +46,7 @@ def meta_handler(
         str | None,
         cyclopts.Parameter(
             ["--database", "-d"],
-            help="Database connection name, file path (CSV/SQLite), or connection string (postgresql://, mysql://) (uses default if not specified)",
+            help="Database connection name, file path (CSV/SQLite/DuckDB), or connection string (postgresql://, mysql://, duckdb://) (uses default if not specified)",
         ),
     ] = None,
 ):
@@ -59,8 +59,10 @@ def meta_handler(
         saber -d mydb "show me users"          # Run a query with specific database
         saber -d data.csv "show me users"      # Run a query with ad-hoc CSV file
         saber -d data.db "show me users"       # Run a query with ad-hoc SQLite file
+        saber -d data.duckdb "show me users"   # Run a query with ad-hoc DuckDB file
         saber -d "postgresql://user:pass@host:5432/db" "show users"  # PostgreSQL connection string
         saber -d "mysql://user:pass@host:3306/db" "show users"       # MySQL connection string
+        saber -d "duckdb:///data.duckdb" "show users"                 # DuckDB connection string
         echo "show me all users" | saber       # Read query from stdin
         cat query.txt | saber                  # Read query from file via stdin
     """
@@ -80,7 +82,7 @@ def query(
         str | None,
         cyclopts.Parameter(
             ["--database", "-d"],
-            help="Database connection name, file path (CSV/SQLite), or connection string (postgresql://, mysql://) (uses default if not specified)",
+            help="Database connection name, file path (CSV/SQLite/DuckDB), or connection string (postgresql://, mysql://, duckdb://) (uses default if not specified)",
         ),
     ] = None,
 ):
@@ -97,8 +99,10 @@ def query(
         saber "show me all users"         # Run a single query
         saber -d data.csv "show users"    # Run a query with ad-hoc CSV file
         saber -d data.db "show users"     # Run a query with ad-hoc SQLite file
+        saber -d data.duckdb "show users" # Run a query with ad-hoc DuckDB file
         saber -d "postgresql://user:pass@host:5432/db" "show users"  # PostgreSQL connection string
         saber -d "mysql://user:pass@host:3306/db" "show users"       # MySQL connection string
+        saber -d "duckdb:///data.duckdb" "show users"                 # DuckDB connection string
         echo "show me all users" | saber  # Read query from stdin
     """
 
@@ -111,6 +115,7 @@ def query(
         from sqlsaber.database.connection import (
             CSVConnection,
             DatabaseConnection,
+            DuckDBConnection,
             MySQLConnection,
             PostgreSQLConnection,
             SQLiteConnection,
@@ -149,15 +154,18 @@ def query(
                 # Single query mode with streaming
                 streaming_handler = StreamingQueryHandler(console)
                 # Compute DB type for the greeting line
-                db_type = (
-                    "PostgreSQL"
-                    if isinstance(db_conn, PostgreSQLConnection)
-                    else "MySQL"
-                    if isinstance(db_conn, MySQLConnection)
-                    else "SQLite"
-                    if isinstance(db_conn, (SQLiteConnection, CSVConnection))
-                    else "database"
-                )
+                if isinstance(db_conn, PostgreSQLConnection):
+                    db_type = "PostgreSQL"
+                elif isinstance(db_conn, MySQLConnection):
+                    db_type = "MySQL"
+                elif isinstance(db_conn, DuckDBConnection):
+                    db_type = "DuckDB"
+                elif isinstance(db_conn, SQLiteConnection):
+                    db_type = "SQLite"
+                elif isinstance(db_conn, CSVConnection):
+                    db_type = "DuckDB"
+                else:
+                    db_type = "database"
                 console.print(
                     f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
                 )

sqlsaber/cli/database.py

@@ -31,7 +31,7 @@ def add(
         str,
         cyclopts.Parameter(
             ["--type", "-t"],
-            help="Database type (postgresql, mysql, sqlite)",
+            help="Database type (postgresql, mysql, sqlite, duckdb)",
         ),
     ] = "postgresql",
     host: Annotated[
@@ -87,17 +87,17 @@ def add(
         if not type or type == "postgresql":
             type = questionary.select(
                 "Database type:",
-                choices=["postgresql", "mysql", "sqlite"],
+                choices=["postgresql", "mysql", "sqlite", "duckdb"],
                 default="postgresql",
             ).ask()
 
-        if type == "sqlite":
-            # SQLite only needs database path
+        if type in {"sqlite", "duckdb"}:
+            # SQLite/DuckDB only need database file path
             database = database or questionary.path("Database file path:").ask()
             database = str(Path(database).expanduser().resolve())
             host = "localhost"
             port = 0
-            username = "sqlite"
+            username = type
             password = ""
         else:
             # PostgreSQL/MySQL need connection details
@@ -182,6 +182,17 @@ def add(
             port = 0
             username = "sqlite"
             password = ""
+        elif type == "duckdb":
+            if not database:
+                console.print(
+                    "[bold red]Error:[/bold red] Database file path is required for DuckDB"
+                )
+                sys.exit(1)
+            database = str(Path(database).expanduser().resolve())
+            host = "localhost"
+            port = 0
+            username = "duckdb"
+            password = ""
         else:
             if not all([host, database, username]):
                 console.print(
@@ -264,7 +275,7 @@ def list():
             if db.ssl_ca or db.ssl_cert:
                 ssl_status += " (certs)"
         else:
-            ssl_status = "disabled" if db.type != "sqlite" else "N/A"
+            ssl_status = "disabled" if db.type not in {"sqlite", "duckdb"} else "N/A"
 
         table.add_row(
             db.name,

sqlsaber/cli/display.py

@@ -198,22 +198,32 @@ class DisplayManager:
         # Normalized leading blank line before tool headers
         self.show_newline()
         if tool_name == "list_tables":
-            self.console.print(
-                "[dim bold]:gear: Discovering available tables[/dim bold]"
-            )
+            if self.console.is_terminal:
+                self.console.print(
+                    "[dim bold]:gear: Discovering available tables[/dim bold]"
+                )
+            else:
+                self.console.print("**Discovering available tables**\n")
         elif tool_name == "introspect_schema":
             pattern = tool_input.get("table_pattern", "all tables")
-            self.console.print(
-                f"[dim bold]:gear: Examining schema for: {pattern}[/dim bold]"
-            )
+            if self.console.is_terminal:
+                self.console.print(
+                    f"[dim bold]:gear: Examining schema for: {pattern}[/dim bold]"
+                )
+            else:
+                self.console.print(f"**Examining schema for:** {pattern}\n")
         elif tool_name == "execute_sql":
             # For streaming, we render SQL via LiveMarkdownRenderer; keep Syntax
             # rendering for threads show/resume. Controlled by include_sql flag.
             query = tool_input.get("query", "")
-            self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
-            self.show_newline()
-            syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
-            self.console.print(syntax)
+            if self.console.is_terminal:
+                self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
+                self.show_newline()
+                syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
+                self.console.print(syntax)
+            else:
+                self.console.print("**Executing SQL:**\n")
+                self.console.print(f"```sql\n{query}\n```\n")
 
     def show_text_stream(self, text: str):
         """Display streaming text."""
@@ -225,9 +235,12 @@ class DisplayManager:
         if not results:
             return
 
-        self.console.print(
-            f"\n[bold magenta]Results ({len(results)} rows):[/bold magenta]"
-        )
+        if self.console.is_terminal:
+            self.console.print(
+                f"\n[bold magenta]Results ({len(results)} rows):[/bold magenta]"
+            )
+        else:
+            self.console.print(f"\n**Results ({len(results)} rows):**\n")
 
         # Create table with columns from first result
         all_columns = list(results[0].keys())
@@ -235,9 +248,14 @@ class DisplayManager:
 
         # Show warning if columns were truncated
         if len(all_columns) > 15:
-            self.console.print(
-                f"[yellow]Note: Showing first 15 of {len(all_columns)} columns[/yellow]"
-            )
+            if self.console.is_terminal:
+                self.console.print(
+                    f"[yellow]Note: Showing first 15 of {len(all_columns)} columns[/yellow]"
+                )
+            else:
+                self.console.print(
+                    f"*Note: Showing first 15 of {len(all_columns)} columns*\n"
+                )
 
         table = self._create_table(display_columns)
 
@@ -248,9 +266,14 @@ class DisplayManager:
         self.console.print(table)
 
         if len(results) > 20:
-            self.console.print(
-                f"[yellow]... and {len(results) - 20} more rows[/yellow]"
-            )
+            if self.console.is_terminal:
+                self.console.print(
+                    f"[yellow]... and {len(results) - 20} more rows[/yellow]"
+                )
+            else:
+                self.console.print(
+                    f"*... and {len(results) - 20} more rows*\n"
+                )
 
     def show_error(self, error_message: str):
         """Display error message."""
@@ -385,6 +408,13 @@ class DisplayManager:
                     for fk in foreign_keys:
                         self.console.print(f"  • {fk}")
 
+                # Show indexes
+                indexes = table_info.get("indexes", [])
+                if indexes:
+                    self.console.print("[bold blue]Indexes:[/bold blue]")
+                    for idx in indexes:
+                        self.console.print(f"  • {idx}")
+
         except json.JSONDecodeError:
             self.show_error("Failed to parse schema data")
         except Exception as e:

sqlsaber/cli/interactive.py

@@ -23,6 +23,7 @@ from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
 from sqlsaber.database.connection import (
     CSVConnection,
+    DuckDBConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -85,8 +86,12 @@ class InteractiveSession:
             if isinstance(self.db_conn, PostgreSQLConnection)
             else "MySQL"
             if isinstance(self.db_conn, MySQLConnection)
+            else "DuckDB"
+            if isinstance(self.db_conn, DuckDBConnection)
+            else "DuckDB"
+            if isinstance(self.db_conn, CSVConnection)
             else "SQLite"
-            if isinstance(self.db_conn, (SQLiteConnection, CSVConnection))
+            if isinstance(self.db_conn, SQLiteConnection)
             else "database"
         )
 

sqlsaber/cli/threads.py

@@ -38,6 +38,8 @@ def _render_transcript(
     from sqlsaber.cli.display import DisplayManager
 
     dm = DisplayManager(console)
+    # Check if output is being redirected (for clean markdown export)
+    is_redirected = not console.is_terminal
 
     # Locate indices of user prompts
     user_indices: list[int] = []
@@ -78,11 +80,17 @@ def _render_transcript(
                                 parts.append(str(seg))
                     text = "\n".join([s for s in parts if s]) or None
                 if text:
-                    console.print(
-                        Panel.fit(Markdown(text), title="User", border_style="cyan")
-                    )
+                    if is_redirected:
+                        console.print(f"**User:**\n\n{text}\n")
+                    else:
+                        console.print(
+                            Panel.fit(Markdown(text), title="User", border_style="cyan")
+                        )
                     return
-        console.print(Panel.fit("(no content)", title="User", border_style="cyan"))
+        if is_redirected:
+            console.print("**User:** (no content)\n")
+        else:
+            console.print(Panel.fit("(no content)", title="User", border_style="cyan"))
 
     def _render_response(message: ModelMessage) -> None:
         for part in getattr(message, "parts", []):
@@ -90,11 +98,14 @@ def _render_transcript(
             if kind == "text":
                 text = getattr(part, "content", "")
                 if isinstance(text, str) and text.strip():
-                    console.print(
-                        Panel.fit(
-                            Markdown(text), title="Assistant", border_style="green"
+                    if is_redirected:
+                        console.print(f"**Assistant:**\n\n{text}\n")
+                    else:
+                        console.print(
+                            Panel.fit(
+                                Markdown(text), title="Assistant", border_style="green"
+                            )
                         )
-                    )
             elif kind in ("tool-call", "builtin-tool-call"):
                 name = getattr(part, "tool_name", "tool")
                 args = getattr(part, "args", None)
@@ -135,6 +146,20 @@ def _render_transcript(
                             dm.show_sql_error(
                                 data.get("error"), data.get("suggestions")
                             )
+                        else:
+                            if is_redirected:
+                                console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                            else:
+                                console.print(
+                                    Panel.fit(
+                                        content_str,
+                                        title=f"Tool result: {name}",
+                                        border_style="yellow",
+                                    )
+                                )
+                    except Exception:
+                        if is_redirected:
+                            console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
                         else:
                             console.print(
                                 Panel.fit(
@@ -143,7 +168,10 @@ def _render_transcript(
                                     border_style="yellow",
                                 )
                             )
-                    except Exception:
+                else:
+                    if is_redirected:
+                        console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                    else:
                         console.print(
                             Panel.fit(
                                 content_str,
@@ -151,14 +179,6 @@ def _render_transcript(
                                 border_style="yellow",
                             )
                         )
-                else:
-                    console.print(
-                        Panel.fit(
-                            content_str,
-                            title=f"Tool result: {name}",
-                            border_style="yellow",
-                        )
-                    )
         # Thinking parts omitted
 
     for start_idx, end_idx in slices or [(0, len(all_msgs))]:
@@ -270,7 +290,10 @@ def resume(
         try:
             agent = build_sqlsaber_agent(db_conn, db_name)
             history = await store.get_thread_messages(thread_id)
-            console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
+            if console.is_terminal:
+                console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
+            else:
+                console.print(f"# Thread: {thread.id}\n")
             _render_transcript(console, history, None)
             session = InteractiveSession(
                 console=console,

sqlsaber/config/database.py

@@ -18,7 +18,7 @@ class DatabaseConfig:
     """Database connection configuration."""
 
     name: str
-    type: str  # postgresql, mysql, sqlite, csv
+    type: str  # postgresql, mysql, sqlite, duckdb, csv
     host: str | None
     port: int | None
     database: str
@@ -90,6 +90,8 @@ class DatabaseConfig:
 
         elif self.type == "sqlite":
             return f"sqlite:///{self.database}"
+        elif self.type == "duckdb":
+            return f"duckdb:///{self.database}"
         elif self.type == "csv":
             # For CSV files, database field contains the file path
             base_url = f"csv:///{self.database}"