diff --git a/datamind/capabilities/db/base.py b/datamind/capabilities/db/base.py index d0bfeab..0ebfc32 100644 --- a/datamind/capabilities/db/base.py +++ b/datamind/capabilities/db/base.py @@ -106,6 +106,8 @@ async def execute_readonly( ) -> QueryResult: if not sql or not sql.strip(): raise CapabilityError("db", "Empty SQL") + if row_limit < 1: + raise CapabilityError("db", "row_limit must be at least 1") if contains_multiple_statements(sql): raise MultiStatementSQLError( "multiple statements are not allowed (use a single SELECT)" diff --git a/datamind/config.py b/datamind/config.py index ea512e1..b251a26 100644 --- a/datamind/config.py +++ b/datamind/config.py @@ -112,7 +112,7 @@ class DBConfig(BaseModel): dialect: str = "sqlite" # sqlite | mysql | postgres | ... dsn: str | None = None # e.g. mysql+pymysql://user:pw@host/db read_only: bool = True - row_limit: int = 1000 + row_limit: int = Field(default=1000, ge=1) query_timeout_s: float = 10.0 diff --git a/datamind/tests/test_config.py b/datamind/tests/test_config.py index 8dda7ea..11743ed 100644 --- a/datamind/tests/test_config.py +++ b/datamind/tests/test_config.py @@ -14,7 +14,7 @@ import pytest from pydantic import ValidationError -from datamind.config import Settings +from datamind.config import DBConfig, Settings def test_nested_env_hydrates_required_llm(monkeypatch, tmp_path): @@ -101,3 +101,8 @@ def test_ensure_dirs_is_idempotent(monkeypatch, tmp_path): assert (tmp_path / "data" / "profiles" / "tp").is_dir() assert (tmp_path / "storage" / "tp").is_dir() + + +def test_db_config_rejects_nonpositive_row_limit(): + with pytest.raises(ValidationError): + DBConfig(row_limit=0) diff --git a/datamind/tests/test_db_safeguard.py b/datamind/tests/test_db_safeguard.py index bda0bd6..874043f 100644 --- a/datamind/tests/test_db_safeguard.py +++ b/datamind/tests/test_db_safeguard.py @@ -12,6 +12,8 @@ leading_verb, strip_comments, ) +from datamind.capabilities.db.providers.sqlite import SQLiteDialect +from datamind.core.errors import CapabilityError @pytest.mark.parametrize( @@ -67,3 +69,17 @@ def test_ensure_row_limit_strips_trailing_semicolon(): def test_strip_comments(): assert strip_comments("SELECT /* x */ 1") == "SELECT 1" assert strip_comments("-- line\nSELECT 1") == "\nSELECT 1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row_limit", [0, -1]) +async def test_execute_readonly_rejects_nonpositive_row_limit(tmp_path, row_limit): + dialect = SQLiteDialect() + engine = dialect.build_engine( + None, default_path=str(tmp_path / "db.sqlite") + ) + + with pytest.raises(CapabilityError, match="row_limit"): + await dialect.execute_readonly( + engine, "SELECT 1", row_limit=row_limit + )