sqlsaber 0.20.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

sqlsaber/agents/pydantic_ai_agent.py

@@ -8,6 +8,7 @@ import httpx
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models.anthropic import AnthropicModel
 from pydantic_ai.models.google import GoogleModel
+from pydantic_ai.models.openai import OpenAIResponsesModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
 from pydantic_ai.providers.google import GoogleProvider
 
@@ -79,6 +80,10 @@ def build_sqlsaber_agent(
         provider_obj = AnthropicProvider(api_key="placeholder", http_client=http_client)
         model_obj = AnthropicModel(model_name_only, provider=provider_obj)
         agent = Agent(model_obj, name="sqlsaber")
+    elif provider == "openai":
+        # Use OpenAI Responses Model for structured output capabilities
+        model_obj = OpenAIResponsesModel(model_name_only)
+        agent = Agent(model_obj, name="sqlsaber")
     else:
         agent = Agent(cfg.model_name, name="sqlsaber")
 

sqlsaber/cli/display.py

@@ -1,12 +1,167 @@
-"""Display utilities for the CLI interface."""
+"""Display utilities for the CLI interface.
+
+All rendering occurs on the event loop thread.
+Streaming segments use Live Markdown; transient status and SQL blocks are also
+rendered with Live.
+"""
 
 import json
+from typing import Sequence, Type
 
-from rich.console import Console
-from rich.markdown import Markdown
+from pydantic_ai.messages import ModelResponsePart, TextPart
+from rich.columns import Columns
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.live import Live
+from rich.markdown import CodeBlock, Markdown
 from rich.panel import Panel
+from rich.spinner import Spinner
 from rich.syntax import Syntax
 from rich.table import Table
+from rich.text import Text
+
+
+class _SimpleCodeBlock(CodeBlock):
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        code = str(self.text).rstrip()
+        yield Syntax(
+            code,
+            self.lexer_name,
+            theme=self.theme,
+            background_color="default",
+            word_wrap=True,
+        )
+
+
+class LiveMarkdownRenderer:
+    """Handles Live markdown rendering with segment separation.
+
+    Supports different segment kinds: 'assistant', 'thinking', 'sql'.
+    Adds visible paragraph breaks between segments and renders code fences
+    with nicer formatting.
+    """
+
+    _patched_fences = False
+
+    def __init__(self, console: Console):
+        self.console = console
+        self._live: Live | None = None
+        self._status_live: Live | None = None
+        self._buffer: str = ""
+        self._current_kind: Type[ModelResponsePart] | None = None
+
+    def prepare_code_blocks(self) -> None:
+        """Patch rich Markdown fence rendering once for nicer code blocks."""
+        if LiveMarkdownRenderer._patched_fences:
+            return
+        # Guard with class check to avoid re-patching if already applied
+        if Markdown.elements.get("fence") is not _SimpleCodeBlock:
+            Markdown.elements["fence"] = _SimpleCodeBlock
+        LiveMarkdownRenderer._patched_fences = True
+
+    def ensure_segment(self, kind: Type[ModelResponsePart]) -> None:
+        """
+        Ensure a markdown Live segment is active for the given kind.
+
+        When switching kinds, end the previous segment and add a paragraph break.
+        """
+        # If a transient status is showing, clear it first (no paragraph break)
+        if self._status_live is not None:
+            self.end_status()
+        if self._live is not None and self._current_kind == kind:
+            return
+        if self._live is not None:
+            self.end()
+            self.paragraph_break()
+
+        self._start()
+        self._current_kind = kind
+
+    def append(self, text: str | None) -> None:
+        """Append text to the current markdown segment and refresh."""
+        if not text:
+            return
+        if self._live is None:
+            # default to assistant if no segment was ensured
+            self.ensure_segment(TextPart)
+
+        self._buffer += text
+        self._live.update(Markdown(self._buffer))
+
+    def end(self) -> None:
+        """Finalize and stop the current Live segment, if any."""
+        if self._live is None:
+            return
+        if self._buffer:
+            self._live.update(Markdown(self._buffer))
+        self._live.stop()
+        self._live = None
+        self._buffer = ""
+        self._current_kind = None
+
+    def end_if_active(self) -> None:
+        self.end()
+
+    def paragraph_break(self) -> None:
+        self.console.print()
+
+    def start_sql_block(self, sql: str) -> None:
+        """Render a SQL block using a transient Live markdown segment."""
+        if not sql or not isinstance(sql, str) or not sql.strip():
+            return
+        # Separate from surrounding content
+        self.end_if_active()
+        self.paragraph_break()
+        self._buffer = f"```sql\n{sql}\n```"
+        # Use context manager to auto-stop and persist final render
+        with Live(
+            Markdown(self._buffer),
+            console=self.console,
+            vertical_overflow="visible",
+            refresh_per_second=12,
+        ):
+            pass
+
+    def start_status(self, message: str = "Crunching data...") -> None:
+        """Show a transient status line with a spinner until streaming starts."""
+        if self._status_live is not None:
+            # Update existing status text
+            self._status_live.update(self._status_renderable(message))
+            return
+        live = Live(
+            self._status_renderable(message),
+            console=self.console,
+            transient=True,  # disappear when stopped
+            refresh_per_second=12,
+        )
+        self._status_live = live
+        live.start()
+
+    def end_status(self) -> None:
+        live = self._status_live
+        if live is None:
+            return
+        live.stop()
+        self._status_live = None
+
+    def _status_renderable(self, message: str):
+        spinner = Spinner("dots", style="yellow")
+        text = Text(f" {message}", style="yellow")
+        return Columns([spinner, text], expand=False)
+
+    def _start(self, initial_markdown: str = "") -> None:
+        if self._live is not None:
+            self.end()
+        self._buffer = initial_markdown or ""
+        live = Live(
+            Markdown(self._buffer),
+            console=self.console,
+            vertical_overflow="visible",
+            refresh_per_second=12,
+        )
+        self._live = live
+        live.start()
 
 
 class DisplayManager:
@@ -14,10 +169,11 @@ class DisplayManager:
 
     def __init__(self, console: Console):
         self.console = console
+        self.live = LiveMarkdownRenderer(console)
 
     def _create_table(
         self,
-        columns: list,
+        columns: Sequence[str | dict[str, str]],
         header_style: str = "bold blue",
         title: str | None = None,
     ) -> Table:
@@ -34,17 +190,24 @@ class DisplayManager:
 
     def show_tool_executing(self, tool_name: str, tool_input: dict):
         """Display tool execution details."""
-        self.console.print(f"\n[yellow]🔧 Using tool: {tool_name}[/yellow]")
+        # Normalized leading blank line before tool headers
+        self.show_newline()
         if tool_name == "list_tables":
-            self.console.print("[dim]  → Discovering available tables[/dim]")
+            self.console.print(
+                "[dim bold]:gear: Discovering available tables[/dim bold]"
+            )
         elif tool_name == "introspect_schema":
             pattern = tool_input.get("table_pattern", "all tables")
-            self.console.print(f"[dim]  → Examining schema for: {pattern}[/dim]")
+            self.console.print(
+                f"[dim bold]:gear: Examining schema for: {pattern}[/dim bold]"
+            )
         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("\n[bold green]Executing SQL:[/bold green]")
+            self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
             self.show_newline()
-            syntax = Syntax(query, "sql")
+            syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
             self.console.print(syntax)
 
     def show_text_stream(self, text: str):
@@ -99,10 +262,12 @@ class DisplayManager:
         """Display a newline for spacing."""
         self.console.print()
 
-    def show_table_list(self, tables_data: str):
+    def show_table_list(self, tables_data: str | dict):
         """Display the results from list_tables tool."""
         try:
-            data = json.loads(tables_data)
+            data = (
+                json.loads(tables_data) if isinstance(tables_data, str) else tables_data
+            )
 
             # Handle error case
             if "error" in data:
@@ -143,10 +308,12 @@ class DisplayManager:
         except Exception as e:
             self.show_error(f"Error displaying table list: {str(e)}")
 
-    def show_schema_info(self, schema_data: str):
+    def show_schema_info(self, schema_data: str | dict):
         """Display the results from introspect_schema tool."""
         try:
-            data = json.loads(schema_data)
+            data = (
+                json.loads(schema_data) if isinstance(schema_data, str) else schema_data
+            )
 
             # Handle error case
             if "error" in data:

sqlsaber/cli/interactive.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from pathlib import Path
+from textwrap import dedent
 
 import platformdirs
 from prompt_toolkit import PromptSession
@@ -10,6 +11,7 @@ from prompt_toolkit.patch_stdout import patch_stdout
 from prompt_toolkit.styles import Style
 from pydantic_ai import Agent
 from rich.console import Console
+from rich.markdown import Markdown
 from rich.panel import Panel
 
 from sqlsaber.cli.completers import (
@@ -102,14 +104,14 @@ class InteractiveSession:
                 )
             )
             self.console.print(
-                "\n",
-                "[dim] > Use '/clear' to reset conversation",
-                "[dim] > Use 'Ctrl+D', '/exit' or '/quit' to leave[/dim]",
-                "[dim] > Use 'Ctrl+C' to interrupt and return to prompt\n\n",
-                "[dim] > Start message with '#' to add something to agent's memory for this database",
-                "[dim] > Type '@' to get table name completions",
-                "[dim] > Press 'Esc-Enter' or 'Meta-Enter' to submit your question",
-                sep="\n",
+                Markdown(
+                    dedent("""
+                    - Use `/` for slash commands
+                    - Type `@` to get table name completions
+                    - Start message with `#` to add something to agent's memory
+                    - Use `Ctrl+C` to interrupt and `Ctrl+D` to exit
+                    """)
+                )
             )
 
         self.console.print(

sqlsaber/cli/streaming.py

@@ -1,7 +1,13 @@
-"""Streaming query handling for the CLI (pydantic-ai based)."""
+"""Streaming query handling for the CLI (pydantic-ai based).
+
+This module uses DisplayManager's LiveMarkdownRenderer to stream Markdown
+incrementally as the agent outputs tokens. Tool calls and results are
+rendered via DisplayManager helpers.
+"""
 
 import asyncio
 import json
+from functools import singledispatchmethod
 from typing import AsyncIterable
 
 from pydantic_ai import Agent, RunContext
@@ -22,56 +28,98 @@ from sqlsaber.cli.display import DisplayManager
 
 
 class StreamingQueryHandler:
-    """Handles streaming query execution and display using pydantic-ai events."""
+    """
+    Handles streaming query execution and display using pydantic-ai events.
+
+    Uses DisplayManager.live to render Markdown incrementally as text streams in.
+    """
 
     def __init__(self, console: Console):
         self.console = console
         self.display = DisplayManager(console)
 
-        self.status = self.console.status(
-            "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
-        )
-
     async def _event_stream_handler(
         self, ctx: RunContext, event_stream: AsyncIterable[AgentStreamEvent]
     ) -> None:
+        """
+        Handle pydantic-ai streaming events and update Live Markdown via DisplayManager.
+        """
+
         async for event in event_stream:
-            if isinstance(event, PartStartEvent):
-                if isinstance(event.part, (TextPart, ThinkingPart)):
-                    self.status.stop()
-                    self.display.show_text_stream(event.part.content)
-
-            elif isinstance(event, PartDeltaEvent):
-                if isinstance(event.delta, (TextPartDelta, ThinkingPartDelta)):
-                    delta = event.delta.content_delta or ""
-                    if delta:
-                        self.status.stop()
-                        self.display.show_text_stream(delta)
-
-            elif isinstance(event, FunctionToolCallEvent):
-                # Show tool execution start
-                self.status.stop()
-                args = event.part.args_as_dict()
-                self.display.show_newline()
-                self.display.show_tool_executing(event.part.tool_name, args)
-
-            elif isinstance(event, FunctionToolResultEvent):
-                self.status.stop()
-                # Route tool result to appropriate display
-                tool_name = event.result.tool_name
-                content = event.result.content
-                if tool_name == "list_tables":
-                    self.display.show_table_list(content)
-                elif tool_name == "introspect_schema":
-                    self.display.show_schema_info(content)
-                elif tool_name == "execute_sql":
+            await self.on_event(event, ctx)
+
+    # --- Event routing via singledispatchmethod ---------------------------------------
+    @singledispatchmethod
+    async def on_event(
+        self, event: AgentStreamEvent, ctx: RunContext
+    ) -> None:  # default
+        return
+
+    @on_event.register
+    async def _(self, event: PartStartEvent, ctx: RunContext) -> None:
+        if isinstance(event.part, TextPart):
+            self.display.live.ensure_segment(TextPart)
+            self.display.live.append(event.part.content)
+        elif isinstance(event.part, ThinkingPart):
+            self.display.live.ensure_segment(ThinkingPart)
+            self.display.live.append(event.part.content)
+
+    @on_event.register
+    async def _(self, event: PartDeltaEvent, ctx: RunContext) -> None:
+        d = event.delta
+        if isinstance(d, TextPartDelta):
+            delta = d.content_delta or ""
+            if delta:
+                self.display.live.ensure_segment(TextPart)
+                self.display.live.append(delta)
+        elif isinstance(d, ThinkingPartDelta):
+            delta = d.content_delta or ""
+            if delta:
+                self.display.live.ensure_segment(ThinkingPart)
+                self.display.live.append(delta)
+
+    @on_event.register
+    async def _(self, event: FunctionToolCallEvent, ctx: RunContext) -> None:
+        # Clear any status/markdown Live so tool output sits between
+        self.display.live.end_status()
+        self.display.live.end_if_active()
+        args = event.part.args_as_dict()
+
+        # Special handling: display SQL via Live as markdown code block
+        if event.part.tool_name == "execute_sql":
+            query = args.get("query") or ""
+            if isinstance(query, str) and query.strip():
+                self.display.live.start_sql_block(query)
+        else:
+            self.display.show_tool_executing(event.part.tool_name, args)
+
+    @on_event.register
+    async def _(self, event: FunctionToolResultEvent, ctx: RunContext) -> None:
+        # Route tool result to appropriate display
+        tool_name = event.result.tool_name
+        content = event.result.content
+        if tool_name == "list_tables":
+            self.display.show_table_list(content)
+        elif tool_name == "introspect_schema":
+            self.display.show_schema_info(content)
+        elif tool_name == "execute_sql":
+            data = {}
+            if isinstance(content, str):
+                try:
+                    data = json.loads(content)
+                except (json.JSONDecodeError, TypeError) as exc:
                     try:
-                        data = json.loads(content)
-                        if data.get("success") and data.get("results"):
-                            self.display.show_query_results(data["results"])  # type: ignore[arg-type]
-                    except json.JSONDecodeError:
-                        # If not JSON, ignore here
+                        self.console.log(f"Malformed execute_sql result: {exc}")
+                    except Exception:
                         pass
+            elif isinstance(content, dict):
+                data = content
+            if isinstance(data, dict) and data.get("success") and data.get("results"):
+                self.display.show_query_results(data["results"])  # type: ignore[arg-type]
+        # Add a blank line after tool output to separate from next segment
+        self.display.show_newline()
+        # Show status while agent sends a follow-up request to the model
+        self.display.live.start_status("Crunching data...")
 
     async def execute_streaming_query(
         self,
@@ -80,7 +128,8 @@ class StreamingQueryHandler:
         cancellation_token: asyncio.Event | None = None,
         message_history: list | None = None,
     ):
-        self.status.start()
+        # Prepare nicer code block rendering for Markdown
+        self.display.live.prepare_code_blocks()
         try:
             # If Anthropic OAuth, inject SQLsaber instructions before the first user prompt
             prepared_prompt: str | list[str] = user_query
@@ -104,30 +153,24 @@ class StreamingQueryHandler:
                     injected = "\n\n".join(parts)
                     prepared_prompt = [injected, user_query]
 
+            # Show a transient status until events start streaming
+            self.display.live.start_status("Crunching data...")
+
             # Run the agent with our event stream handler
             run = await agent.run(
                 prepared_prompt,
                 message_history=message_history,
                 event_stream_handler=self._event_stream_handler,
             )
-            # After the run completes, show the assistant's final text as markdown if available
-            try:
-                output = run.output
-                if isinstance(output, str) and output.strip():
-                    self.display.show_newline()
-                    self.display.show_markdown_response(
-                        [{"type": "text", "text": output}]
-                    )
-            except Exception as e:
-                self.display.show_error(str(e))
-                self.display.show_newline()
             return run
         except asyncio.CancelledError:
+            # Show interruption message outside of Live
             self.display.show_newline()
             self.console.print("[yellow]Query interrupted[/yellow]")
             return None
         finally:
+            # End any active status and live markdown segments
             try:
-                self.status.stop()
-            except Exception:
-                pass
+                self.display.live.end_status()
+            finally:
+                self.display.live.end_if_active()

sqlsaber/database/connection.py

@@ -1,5 +1,6 @@
 """Database connection management."""
 
+import asyncio
 import ssl
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -10,6 +11,17 @@ import aiomysql
 import aiosqlite
 import asyncpg
 
+# Default query timeout to prevent runaway queries
+DEFAULT_QUERY_TIMEOUT = 30.0  # seconds
+
+
+class QueryTimeoutError(RuntimeError):
+    """Exception raised when a query exceeds its timeout."""
+
+    def __init__(self, seconds: float):
+        self.timeout = seconds
+        super().__init__(f"Query exceeded timeout of {seconds}s")
+
 
 class BaseDatabaseConnection(ABC):
     """Abstract base class for database connections."""
@@ -29,11 +41,18 @@ class BaseDatabaseConnection(ABC):
         pass
 
     @abstractmethod
-    async def execute_query(self, query: str, *args) -> list[dict[str, Any]]:
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts.
 
         All queries run in a transaction that is rolled back at the end,
         ensuring no changes are persisted to the database.
+
+        Args:
+            query: SQL query to execute
+            *args: Query parameters
+            timeout: Query timeout in seconds (overrides default_timeout)
         """
         pass
 
@@ -111,21 +130,40 @@ class PostgreSQLConnection(BaseDatabaseConnection):
             await self._pool.close()
             self._pool = None
 
-    async def execute_query(self, query: str, *args) -> list[dict[str, Any]]:
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts.
 
         All queries run in a transaction that is rolled back at the end,
         ensuring no changes are persisted to the database.
         """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
         pool = await self.get_pool()
+
         async with pool.acquire() as conn:
             # Start a transaction that we'll always rollback
             transaction = conn.transaction()
             await transaction.start()
 
             try:
-                rows = await conn.fetch(query, *args)
+                # Set server-side timeout if specified
+                if effective_timeout:
+                    await conn.execute(
+                        f"SET LOCAL statement_timeout = {int(effective_timeout * 1000)}"
+                    )
+
+                # Execute query with client-side timeout
+                if effective_timeout:
+                    rows = await asyncio.wait_for(
+                        conn.fetch(query, *args), timeout=effective_timeout
+                    )
+                else:
+                    rows = await conn.fetch(query, *args)
+
                 return [dict(row) for row in rows]
+            except asyncio.TimeoutError as exc:
+                raise QueryTimeoutError(effective_timeout or 0) from exc
             finally:
                 # Always rollback to ensure no changes are committed
                 await transaction.rollback()
@@ -216,21 +254,44 @@ class MySQLConnection(BaseDatabaseConnection):
             await self._pool.wait_closed()
             self._pool = None
 
-    async def execute_query(self, query: str, *args) -> list[dict[str, Any]]:
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts.
 
         All queries run in a transaction that is rolled back at the end,
         ensuring no changes are persisted to the database.
         """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
         pool = await self.get_pool()
+
         async with pool.acquire() as conn:
             async with conn.cursor(aiomysql.DictCursor) as cursor:
                 # Start transaction
                 await conn.begin()
                 try:
-                    await cursor.execute(query, args if args else None)
-                    rows = await cursor.fetchall()
+                    # Set server-side timeout if specified
+                    if effective_timeout:
+                        await cursor.execute(
+                            f"SET SESSION MAX_EXECUTION_TIME = {int(effective_timeout * 1000)}"
+                        )
+
+                    # Execute query with client-side timeout
+                    if effective_timeout:
+                        await asyncio.wait_for(
+                            cursor.execute(query, args if args else None),
+                            timeout=effective_timeout,
+                        )
+                        rows = await asyncio.wait_for(
+                            cursor.fetchall(), timeout=effective_timeout
+                        )
+                    else:
+                        await cursor.execute(query, args if args else None)
+                        rows = await cursor.fetchall()
+
                     return [dict(row) for row in rows]
+                except asyncio.TimeoutError as exc:
+                    raise QueryTimeoutError(effective_timeout or 0) from exc
                 finally:
                     # Always rollback to ensure no changes are committed
                     await conn.rollback()
@@ -252,12 +313,16 @@ class SQLiteConnection(BaseDatabaseConnection):
         """SQLite connections are created per query, no persistent pool to close."""
         pass
 
-    async def execute_query(self, query: str, *args) -> list[dict[str, Any]]:
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts.
 
         All queries run in a transaction that is rolled back at the end,
         ensuring no changes are persisted to the database.
         """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
+
         async with aiosqlite.connect(self.database_path) as conn:
             # Enable row factory for dict-like access
             conn.row_factory = aiosqlite.Row
@@ -265,9 +330,22 @@ class SQLiteConnection(BaseDatabaseConnection):
             # Start transaction
             await conn.execute("BEGIN")
             try:
-                cursor = await conn.execute(query, args if args else ())
-                rows = await cursor.fetchall()
+                # Execute query with client-side timeout (SQLite has no server-side timeout)
+                if effective_timeout:
+                    cursor = await asyncio.wait_for(
+                        conn.execute(query, args if args else ()),
+                        timeout=effective_timeout,
+                    )
+                    rows = await asyncio.wait_for(
+                        cursor.fetchall(), timeout=effective_timeout
+                    )
+                else:
+                    cursor = await conn.execute(query, args if args else ())
+                    rows = await cursor.fetchall()
+
                 return [dict(row) for row in rows]
+            except asyncio.TimeoutError as exc:
+                raise QueryTimeoutError(effective_timeout or 0) from exc
             finally:
                 # Always rollback to ensure no changes are committed
                 await conn.rollback()
@@ -383,20 +461,35 @@ class CSVConnection(BaseDatabaseConnection):
         except Exception as e:
             raise ValueError(f"Error loading CSV file '{self.csv_path}': {str(e)}")
 
-    async def execute_query(self, query: str, *args) -> list[dict[str, Any]]:
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> list[dict[str, Any]]:
         """Execute a query and return results as list of dicts.
 
         All queries run in a transaction that is rolled back at the end,
         ensuring no changes are persisted to the database.
         """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
         conn = await self.get_pool()
 
         # Start transaction
         await conn.execute("BEGIN")
         try:
-            cursor = await conn.execute(query, args if args else ())
-            rows = await cursor.fetchall()
+            # Execute query with client-side timeout (CSV uses in-memory SQLite)
+            if effective_timeout:
+                cursor = await asyncio.wait_for(
+                    conn.execute(query, args if args else ()), timeout=effective_timeout
+                )
+                rows = await asyncio.wait_for(
+                    cursor.fetchall(), timeout=effective_timeout
+                )
+            else:
+                cursor = await conn.execute(query, args if args else ())
+                rows = await cursor.fetchall()
+
             return [dict(row) for row in rows]
+        except asyncio.TimeoutError as exc:
+            raise QueryTimeoutError(effective_timeout or 0) from exc
         finally:
             # Always rollback to ensure no changes are committed
             await conn.rollback()

{sqlsaber-0.20.0.dist-info → sqlsaber-0.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlsaber
-Version: 0.20.0
+Version: 0.22.0
 Summary: SQLsaber - Open-source agentic SQL assistant
 License-File: LICENSE
 Requires-Python: >=3.12

{sqlsaber-0.20.0.dist-info → sqlsaber-0.22.0.dist-info}/RECORD

@@ -3,17 +3,17 @@ sqlsaber/__main__.py,sha256=RIHxWeWh2QvLfah-2OkhI5IJxojWfy4fXpMnVEJYvxw,78
 sqlsaber/agents/__init__.py,sha256=i_MI2eWMQaVzGikKU71FPCmSQxNDKq36Imq1PrYoIPU,130
 sqlsaber/agents/base.py,sha256=7zOZTHKxUuU0uMc-NTaCkkBfDnU3jtwbT8_eP1ZtJ2k,2615
 sqlsaber/agents/mcp.py,sha256=GcJTx7YDYH6aaxIADEIxSgcWAdWakUx395JIzVnf17U,768
-sqlsaber/agents/pydantic_ai_agent.py,sha256=dGdsgyxCZvfK-v-MH8KimKOr-xb2aSfSWY8CMcOUCT8,6795
+sqlsaber/agents/pydantic_ai_agent.py,sha256=6RvG2O7G8P6NN9QaRXUodg5Q26QJ4ShGWoTGYbVQ5K4,7065
 sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
 sqlsaber/cli/auth.py,sha256=jTsRgbmlGPlASSuIKmdjjwfqtKvjfKd_cTYxX0-QqaQ,7400
 sqlsaber/cli/commands.py,sha256=mjLG9i1bXf0TEroxkIxq5O7Hhjufz3Ad72cyJz7vE1k,8128
 sqlsaber/cli/completers.py,sha256=HsUPjaZweLSeYCWkAcgMl8FylQ1xjWBWYTEL_9F6xfU,6430
 sqlsaber/cli/database.py,sha256=JKtHSN-BFzBa14REf0phFVQB7d67m1M5FFaD8N6DdrY,12966
-sqlsaber/cli/display.py,sha256=wa7BjTBwXwqLT145Q1AEL0C28pQJTrvDN10mnFMjqsg,8554
-sqlsaber/cli/interactive.py,sha256=suTZ-EvbaB21BsFsRc4MkjM89lZ2iJlYH4G1iYjW7PI,13213
+sqlsaber/cli/display.py,sha256=bul9Yzw8KFYkof-kDzeajpx2TtG9CjTaUiwWaTv95dQ,14293
+sqlsaber/cli/interactive.py,sha256=7uM4LoXbhPJr8o5yNjICSzL0uxZkp1psWrVq4G9V0OI,13118
 sqlsaber/cli/memory.py,sha256=OufHFJFwV0_GGn7LvKRTJikkWhV1IwNIUDOxFPHXOaQ,7794
 sqlsaber/cli/models.py,sha256=ZewtwGQwhd9b-yxBAPKePolvI1qQG-EkmeWAGMqtWNQ,8986
-sqlsaber/cli/streaming.py,sha256=WNqBYYbWtL5CNQkRg5YWhYpWKI8qz7JmqneB2DXTOHY,5259
+sqlsaber/cli/streaming.py,sha256=BeG7H38-I1n8b9R8XSBV-IqkxDRZhsWFW6sdvtbVi3o,6879
 sqlsaber/cli/threads.py,sha256=XUnLcCUe2wa_85IKdKmryqfiHTQu_IylET2Qo8oy1nk,11324
 sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
 sqlsaber/config/api_keys.py,sha256=RqWQCko1tY7sES7YOlexgBH5Hd5ne_kGXHdBDNqcV2U,3649
@@ -24,7 +24,7 @@ sqlsaber/config/oauth_tokens.py,sha256=C9z35hyx-PvSAYdC1LNf3rg9_wsEIY56hkEczelba
 sqlsaber/config/providers.py,sha256=JFjeJv1K5Q93zWSlWq3hAvgch1TlgoF0qFa0KJROkKY,2957
 sqlsaber/config/settings.py,sha256=vgb_RXaM-7DgbxYDmWNw1cSyMqwys4j3qNCvM4bljwI,5586
 sqlsaber/database/__init__.py,sha256=a_gtKRJnZVO8-fEZI7g3Z8YnGa6Nio-5Y50PgVp07ss,176
-sqlsaber/database/connection.py,sha256=kwx18bnwr4kyTUfQT0OW-DXzJUNWIQJP54spJBqU_48,15243
+sqlsaber/database/connection.py,sha256=1bDPEa6cmdh87gPfhNeBLpOdI0E2_2KlE74q_-4l_jI,18913
 sqlsaber/database/resolver.py,sha256=RPXF5EoKzvQDDLmPGNHYd2uG_oNICH8qvUjBp6iXmNY,3348
 sqlsaber/database/schema.py,sha256=r12qoN3tdtAXdO22EKlauAe7QwOm8lL2vTMM59XEMMY,26594
 sqlsaber/mcp/__init__.py,sha256=COdWq7wauPBp5Ew8tfZItFzbcLDSEkHBJSMhxzy8C9c,112
@@ -40,8 +40,8 @@ sqlsaber/tools/enums.py,sha256=CH32mL-0k9ZA18911xLpNtsgpV6tB85TktMj6uqGz54,411
 sqlsaber/tools/instructions.py,sha256=X-x8maVkkyi16b6Tl0hcAFgjiYceZaSwyWTfmrvx8U8,9024
 sqlsaber/tools/registry.py,sha256=HWOQMsNIdL4XZS6TeNUyrL-5KoSDH6PHsWd3X66o-18,3211
 sqlsaber/tools/sql_tools.py,sha256=hM6tKqW5MDhFUt6MesoqhTUqIpq_5baIIDoN1MjDCXY,9647
-sqlsaber-0.20.0.dist-info/METADATA,sha256=LJgfLWIWvI8ZgeLjZMLakKCmHn4EywL0MS9XeKLy63E,6178
-sqlsaber-0.20.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.20.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
-sqlsaber-0.20.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.20.0.dist-info/RECORD,,
+sqlsaber-0.22.0.dist-info/METADATA,sha256=T9TBoCGfPVrZKM-RnUVROqOKaBaU1KJdYKqh3a8Arr8,6178
+sqlsaber-0.22.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.22.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
+sqlsaber-0.22.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.22.0.dist-info/RECORD,,