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

sqlsaber/cli/display.py

@@ -93,12 +93,15 @@ class LiveMarkdownRenderer:
         """Finalize and stop the current Live segment, if any."""
         if self._live is None:
             return
-        if self._buffer:
-            self._live.update(Markdown(self._buffer))
+        # Persist the *final* render exactly once, then shut Live down.
+        buf = self._buffer
         self._live.stop()
         self._live = None
         self._buffer = ""
         self._current_kind = None
+        # Print the complete markdown to scroll-back for permanent reference
+        if buf:
+            self.console.print(Markdown(buf))
 
     def end_if_active(self) -> None:
         self.end()
@@ -154,10 +157,12 @@ class LiveMarkdownRenderer:
         if self._live is not None:
             self.end()
         self._buffer = initial_markdown or ""
+        # NOTE: Use transient=True so the live widget disappears on exit,
+        # giving a clean transition to the final printed result.
         live = Live(
             Markdown(self._buffer),
             console=self.console,
-            vertical_overflow="visible",
+            transient=True,
             refresh_per_second=12,
         )
         self._live = live
@@ -251,6 +256,15 @@ class DisplayManager:
         """Display error message."""
         self.console.print(f"\n[bold red]Error:[/bold red] {error_message}")
 
+    def show_sql_error(self, error_message: str, suggestions: list[str] | None = None):
+        """Display SQL-specific error with optional suggestions."""
+        self.show_newline()
+        self.console.print(f"[bold red]SQL error:[/bold red] {error_message}")
+        if suggestions:
+            self.console.print("[yellow]Hints:[/yellow]")
+            for suggestion in suggestions:
+                self.console.print(f"  • {suggestion}")
+
     def show_processing(self, message: str):
         """Display processing message."""
         self.console.print()  # Add newline

sqlsaber/cli/streaming.py

@@ -114,8 +114,14 @@ class StreamingQueryHandler:
                         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]
+
+            if isinstance(data, dict):
+                if data.get("success") and data.get("results"):
+                    self.display.show_query_results(data["results"])  # type: ignore[arg-type]
+                elif "error" in data:
+                    self.display.show_sql_error(
+                        data.get("error"), data.get("suggestions")
+                    )
         # 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

sqlsaber/cli/threads.py

@@ -131,6 +131,10 @@ def _render_transcript(
                             and data.get("results")
                         ):
                             dm.show_query_results(data["results"])  # type: ignore[arg-type]
+                        elif isinstance(data, dict) and "error" in data:
+                            dm.show_sql_error(
+                                data.get("error"), data.get("suggestions")
+                            )
                         else:
                             console.print(
                                 Panel.fit(

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.21.0.dist-info → sqlsaber-0.23.0.dist-info}/METADATA

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

{sqlsaber-0.21.0.dist-info → sqlsaber-0.23.0.dist-info}/RECORD

@@ -9,12 +9,12 @@ 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=bul9Yzw8KFYkof-kDzeajpx2TtG9CjTaUiwWaTv95dQ,14293
+sqlsaber/cli/display.py,sha256=9uXg0GqFXGAwBj1O7-i6aLPQkGq-iURrsKHk6PWTq1E,15025
 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=BeG7H38-I1n8b9R8XSBV-IqkxDRZhsWFW6sdvtbVi3o,6879
-sqlsaber/cli/threads.py,sha256=XUnLcCUe2wa_85IKdKmryqfiHTQu_IylET2Qo8oy1nk,11324
+sqlsaber/cli/streaming.py,sha256=Eo5CNUgDGY1WYP90jwDA2aY7RefN-TfcStA6NyjUQTY,7076
+sqlsaber/cli/threads.py,sha256=HJ6v9wEv1as21B7IJglYs3q6LH7Plv2oheLbM5YEQQA,11549
 sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
 sqlsaber/config/api_keys.py,sha256=RqWQCko1tY7sES7YOlexgBH5Hd5ne_kGXHdBDNqcV2U,3649
 sqlsaber/config/auth.py,sha256=b5qB2h1doXyO9Bn8z0CcL8LAR2jF431gGXBGKLgTmtQ,2756
@@ -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.21.0.dist-info/METADATA,sha256=7NgOXIfUrri2-SHemXf6TVRLBqj2KuzweGafg-NPZuQ,6178
-sqlsaber-0.21.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.21.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
-sqlsaber-0.21.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.21.0.dist-info/RECORD,,
+sqlsaber-0.23.0.dist-info/METADATA,sha256=olVI7W6UB-F6oZUqOjTWhJAvYtDbXXO0xoeOVYYmpfM,6178
+sqlsaber-0.23.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.23.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
+sqlsaber-0.23.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.23.0.dist-info/RECORD,,