Skip to content

Better error messages on SQL client #93

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lotif merged 3 commits into from
Mar 26, 2026
Merged

Better error messages on SQL client #93

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7f6bb1c
Raising error on internal method instead of returning false because i…
lotif Mar 26, 2026
5be4339
Updating model name as gemini 3 has been deprecated
lotif Mar 26, 2026
525842e
Fixing tests
lotif Mar 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ GOOGLE_API_KEY="..." # Used by google-adk

# Model selection (see https://ai.google.dev/gemini-api/docs/models)
# Stable: gemini-2.5-pro, gemini-2.5-flash, gemini-2.5-flash-lite
# Preview: gemini-3-pro-preview, gemini-3-flash-preview
# Preview: gemini-3.1-pro-preview, gemini-3.1-flash-preview
DEFAULT_PLANNER_MODEL="gemini-2.5-pro"
DEFAULT_WORKER_MODEL="gemini-2.5-flash"
DEFAULT_EVALUATOR_MODEL="gemini-2.5-pro"
Expand Down
5 changes: 3 additions & 2 deletions aieng-eval-agents/aieng/agent_evals/tools/sql_database.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,9 +166,9 @@ def _is_safe_readonly_query(self, query: str) -> bool:

return is_safe
except Exception as e:
logger.error("SQL Parsing Error: %s", e)
logger.exception("SQL Parsing Error: %s", e)
# If we can't parse it, we don't run it.
return False
raise e
Copy link
Copy Markdown
Collaborator

@fcogidi fcogidi Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The caller of this method actually raises a PermissionError:

if not self._is_safe_readonly_query(query):
    raise PermissionError("Security Violation: Query contains prohibited WRITE operations.")

If the error is raised here, then the other one is not required - the caller doesn't need the if statement anymore.

Copy link
Copy Markdown
Collaborator Author

@lotif lotif Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not exactly, those are 2 different error cases.

The caller will raise the PermissionError when this function returns False because it has detected an unsafe statement in the query, for example if it has a DELETE statement.

The function will raise an exception if it encounters an error during the check for unsafe statements, for example if the SQL query has a syntax issue (this is how I found this problem).

It is important that the agent receives two different messages for those two cases so they can fix their SQL statements appropriately.

Copy link
Copy Markdown
Collaborator

@fcogidi fcogidi Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ahh, right! My mistake


def get_schema_info(self, table_names: Optional[list[str]] = None) -> str:
"""Return schema for specific tables/views or all if None.
Expand Down Expand Up @@ -296,6 +296,7 @@ def execute(self, query: str) -> str:
return "\n".join(output)

except Exception as e:
logger.exception("Error executing query: %s", e)
error_msg = str(e)
return f"Query Error: {error_msg}"

Expand Down
13 changes: 6 additions & 7 deletions aieng-eval-agents/tests/aieng/agent_evals/tools/test_sql_database.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import pytest
import sqlglot
from aieng.agent_evals.tools import ReadOnlySqlDatabase, ReadOnlySqlPolicy, sql_database
from sqlalchemy import create_engine

Expand Down Expand Up @@ -154,13 +155,11 @@ def test_is_safe_readonly_query_blocks_forbidden_nodes_even_if_root_is_allowed(s
assert db._is_safe_readonly_query("DELETE FROM customers") is False


def test_is_safe_readonly_query_returns_false_on_parse_errors(
sqlite_db_uri: str, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Return False (fail closed) if parsing raises."""
db = ReadOnlySqlDatabase(connection_uri=sqlite_db_uri)
monkeypatch.setattr("aieng.agent_evals.tools.sql_database.sqlglot.parse", lambda _query: 1 / 0)
assert db._is_safe_readonly_query("SELECT 1") is False
def test_is_safe_readonly_query_returns_exception_on_parse_errors(sqlite_db_uri: str) -> None:
"""Return exception if parsing raises an exception."""
with pytest.raises(sqlglot.errors.ParseError):
db = ReadOnlySqlDatabase(connection_uri=sqlite_db_uri)
assert db._is_safe_readonly_query("SERECT 1") is False


def test_get_schema_info_includes_tables_and_views(default_db: ReadOnlySqlDatabase) -> None:
Expand Down