sqlsaber 0.21.0__tar.gz → 0.22.0__tar.gz

{sqlsaber-0.21.0 → sqlsaber-0.22.0}/PKG-INFO

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

{sqlsaber-0.21.0 → sqlsaber-0.22.0}/docs/src/content/docs/changelog.md

@@ -7,6 +7,16 @@ All notable changes to SQLsaber will be documented here.
 
 ### Unreleased
 
+### v0.22.0 - 2025-09-15
+
+#### Added
+
+- Query timeout protection to prevent runaway queries
+  - 30-second timeout applied to all database operations by default
+  - Both client-side and server-side timeout enforcement where supported (PostgreSQL, MySQL)
+  - Per-query timeout override parameter for edge cases
+  - Automatic rollback of transactions when queries timeout
+
 ### v0.21.0 - 2025-09-15
 
 #### Changed

{sqlsaber-0.21.0 → sqlsaber-0.22.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlsaber"
-version = "0.21.0"
+version = "0.22.0"
 description = "SQLsaber - Open-source agentic SQL assistant"
 readme = "README.md"
 requires-python = ">=3.12"

{sqlsaber-0.21.0 → sqlsaber-0.22.0}/src/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.22.0/tests/test_database/test_timeout.py

@@ -0,0 +1,124 @@
+"""Tests for query timeout functionality."""
+
+import asyncio
+
+import pytest
+
+from sqlsaber.database.connection import (
+    DEFAULT_QUERY_TIMEOUT,
+    DatabaseConnection,
+    QueryTimeoutError,
+    SQLiteConnection,
+)
+
+
+class TestQueryTimeout:
+    """Test query timeout functionality."""
+
+    def test_default_query_timeout_constant(self):
+        """Test that the default query timeout constant is defined."""
+        assert DEFAULT_QUERY_TIMEOUT == 30.0
+
+    @pytest.mark.asyncio
+    async def test_database_connection_creation(self):
+        """Test that DatabaseConnection can be created."""
+        connection_string = "sqlite:///:memory:"
+
+        conn = DatabaseConnection(connection_string)
+        assert conn.connection_string == connection_string
+
+        await conn.close()
+
+    @pytest.mark.asyncio
+    async def test_timeout_fallback_logic(self):
+        """Test timeout fallback: per-query > hardcoded default."""
+        connection_string = "sqlite:///:memory:"
+
+        # Test with no timeout - should use hardcoded default
+        conn1 = SQLiteConnection(connection_string)
+        result1 = await conn1.execute_query("SELECT 1 as test")
+        assert len(result1) == 1
+        await conn1.close()
+
+        # Test with per-query timeout override
+        conn2 = SQLiteConnection(connection_string)
+        result2 = await conn2.execute_query("SELECT 2 as test", timeout=60.0)
+        assert len(result2) == 1
+        await conn2.close()
+
+    @pytest.mark.asyncio
+    async def test_sqlite_timeout_functionality(self):
+        """Test timeout functionality with SQLite in-memory database."""
+        connection_string = "sqlite:///:memory:"
+        timeout = 1.0  # 1 second timeout
+
+        conn = SQLiteConnection(connection_string)
+
+        try:
+            # This query should complete quickly
+            result = await conn.execute_query("SELECT 1 as test", timeout=timeout)
+            assert len(result) == 1
+            assert result[0]["test"] == 1
+
+            # Test with no explicit timeout (uses default)
+            result2 = await conn.execute_query("SELECT 2 as test")
+            assert len(result2) == 1
+            assert result2[0]["test"] == 2
+
+        finally:
+            await conn.close()
+
+    @pytest.mark.asyncio
+    async def test_timeout_error_handling(self):
+        """Test that timeout errors are properly raised with asyncio.TimeoutError."""
+        connection_string = "sqlite:///:memory:"
+
+        conn = SQLiteConnection(connection_string)
+
+        try:
+            # Create a mock query that artificially delays to trigger timeout
+            async def mock_slow_execute_query(query, *args, timeout=None):
+                # Simulate a slow query by sleeping longer than timeout
+                await asyncio.sleep(0.1)  # 100ms, longer than 1ms timeout
+                return [{"result": "should not reach here"}]
+
+            # Replace the execute_query method temporarily
+            conn.execute_query = mock_slow_execute_query
+
+            # This should raise an asyncio.TimeoutError which gets converted to QueryTimeoutError
+            with pytest.raises(asyncio.TimeoutError):
+                await asyncio.wait_for(conn.execute_query("SELECT 1"), timeout=0.001)
+
+        finally:
+            await conn.close()
+
+    def test_query_timeout_error_message(self):
+        """Test QueryTimeoutError message formatting."""
+        timeout = 30.5
+        error = QueryTimeoutError(timeout)
+
+        assert str(error) == "Query exceeded timeout of 30.5s"
+        assert error.timeout == 30.5
+
+    @pytest.mark.asyncio
+    async def test_timeout_precedence(self):
+        """Test that query-level timeout overrides hardcoded default."""
+        connection_string = "sqlite:///:memory:"
+        query_timeout = 45.0
+
+        conn = SQLiteConnection(connection_string)
+
+        try:
+            # The query should use the query-specific timeout
+            # We can't easily test the actual timeout without creating a slow query,
+            # but we can verify the parameter is accepted
+            result = await conn.execute_query("SELECT 1 as test", timeout=query_timeout)
+            assert len(result) == 1
+            assert result[0]["test"] == 1
+
+        finally:
+            await conn.close()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

{sqlsaber-0.21.0 → sqlsaber-0.22.0}/uv.lock

@@ -1910,7 +1910,7 @@ wheels = [
 
 [[package]]
 name = "sqlsaber"
-version = "0.21.0"
+version = "0.22.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiomysql" },