refactor(iql): errors definition (#73)

src/dbally/iql/_exceptions.py

@@ -1,32 +1,71 @@
 import ast
-from typing import Optional, Union
+from typing import List, Optional
 
 from dbally.exceptions import DbAllyError
 
 
 class IQLError(DbAllyError):
     """Base exception for all IQL parsing related exceptions."""
 
-    def __init__(self, message: str, node: Union[ast.stmt, ast.expr], source: str) -> None:
-        message = message + ": " + source[node.col_offset : node.end_col_offset]
-
+    def __init__(self, message: str, source: str) -> None:
         super().__init__(message)
-        self.node = node
         self.source = source
 
 
-class IQLArgumentParsingError(IQLError):
+class IQLSyntaxError(IQLError):
+    """Raised when IQL syntax is invalid."""
+
+    def __init__(self, source: str) -> None:
+        message = f"Syntax error in: {source}"
+        super().__init__(message, source)
+
+
+class IQLEmptyExpressionError(IQLError):
+    """Raised when IQL expression is empty."""
+
+    def __init__(self, source: str) -> None:
+        message = "Empty IQL expression"
+        super().__init__(message, source)
+
+
+class IQLMultipleExpressionsError(IQLError):
+    """Raised when IQL contains multiple expressions."""
+
+    def __init__(self, nodes: List[ast.stmt], source: str) -> None:
+        message = "Multiple expressions or statements in IQL are not supported"
+        super().__init__(message, source)
+        self.nodes = nodes
+
+
+class IQLExpressionError(IQLError):
+    """Raised when IQL expression is invalid."""
+
+    def __init__(self, message: str, node: ast.expr, source: str) -> None:
+        message = f"{message}: {source[node.col_offset : node.end_col_offset]}"
+        super().__init__(message, source)
+        self.node = node
+
+
+class IQLNoExpressionError(IQLExpressionError):
+    """Raised when IQL expression is not found."""
+
+    def __init__(self, node: ast.stmt, source: str) -> None:
+        message = "No expression found in IQL"
+        super().__init__(message, node, source)
+
+
+class IQLArgumentParsingError(IQLExpressionError):
     """Raised when an argument cannot be parsed into a valid IQL."""
 
-    def __init__(self, node: Union[ast.stmt, ast.expr], source: str) -> None:
+    def __init__(self, node: ast.expr, source: str) -> None:
         message = "Not a valid IQL argument"
         super().__init__(message, node, source)
 
 
-class IQLUnsupportedSyntaxError(IQLError):
+class IQLUnsupportedSyntaxError(IQLExpressionError):
     """Raised when trying to parse an unsupported syntax."""
 
-    def __init__(self, node: Union[ast.stmt, ast.expr], source: str, context: Optional[str] = None) -> None:
+    def __init__(self, node: ast.expr, source: str, context: Optional[str] = None) -> None:
         node_name = node.__class__.__name__
 
         message = f"{node_name} syntax is not supported in IQL"
@@ -37,13 +76,21 @@ def __init__(self, node: Union[ast.stmt, ast.expr], source: str, context: Option
         super().__init__(message, node, source)
 
 
-class IQLFunctionNotExists(IQLError):
+class IQLFunctionNotExists(IQLExpressionError):
     """Raised when IQL contains function call to a function that not exists."""
 
     def __init__(self, node: ast.Name, source: str) -> None:
         message = f"Function {node.id} not exists"
         super().__init__(message, node, source)
 
 
-class IQLArgumentValidationError(IQLError):
+class IQLIncorrectNumberArgumentsError(IQLExpressionError):
+    """Raised when IQL contains too many arguments for a function."""
+
+    def __init__(self, node: ast.Call, source: str) -> None:
+        message = f"The method {node.func.id} has incorrect number of arguments"
+        super().__init__(message, node, source)
+
+
+class IQLArgumentValidationError(IQLExpressionError):
     """Raised when argument is not valid for a given method."""

src/dbally/iql/_processor.py

@@ -6,8 +6,12 @@
 from dbally.iql._exceptions import (
     IQLArgumentParsingError,
     IQLArgumentValidationError,
-    IQLError,
+    IQLEmptyExpressionError,
     IQLFunctionNotExists,
+    IQLIncorrectNumberArgumentsError,
+    IQLMultipleExpressionsError,
+    IQLNoExpressionError,
+    IQLSyntaxError,
     IQLUnsupportedSyntaxError,
 )
 from dbally.iql._type_validators import validate_arg_type
@@ -33,21 +37,28 @@ async def process(self) -> syntax.Node:
         Process IQL string to root IQL.Node.
 
         Returns:
-            IQL.Node which is root of the tree representing IQL query.
+            IQL node which is root of the tree representing IQL query.
 
         Raises:
-             IQLError: if parsing fails.
+            IQLError: If parsing fails.
         """
         self.source = self._to_lower_except_in_quotes(self.source, ["AND", "OR", "NOT"])
 
-        ast_tree = ast.parse(self.source)
-        first_element = ast_tree.body[0]
+        try:
+            ast_tree = ast.parse(self.source)
+        except (SyntaxError, ValueError) as exc:
+            raise IQLSyntaxError(self.source) from exc
 
-        if not isinstance(first_element, ast.Expr):
-            raise IQLError("Not a valid IQL expression", first_element, self.source)
+        if not ast_tree.body:
+            raise IQLEmptyExpressionError(self.source)
 
-        root = await self._parse_node(first_element.value)
-        return root
+        if len(ast_tree.body) > 1:
+            raise IQLMultipleExpressionsError(ast_tree.body, self.source)
+
+        if not isinstance(ast_tree.body[0], ast.Expr):
+            raise IQLNoExpressionError(ast_tree.body[0], self.source)
+
+        return await self._parse_node(ast_tree.body[0].value)
 
     async def _parse_node(self, node: Union[ast.expr, ast.Expr]) -> syntax.Node:
         if isinstance(node, ast.BoolOp):
@@ -82,7 +93,7 @@ async def _parse_call(self, node: ast.Call) -> syntax.FunctionCall:
         args = []
 
         if len(func_def.parameters) != len(node.args):
-            raise ValueError(f"The method {func.id} has incorrect number of arguments")
+            raise IQLIncorrectNumberArgumentsError(node, self.source)
 
         for arg, arg_def in zip(node.args, func_def.parameters):
             arg_value = self._parse_arg(arg)

src/dbally/iql/_query.py

@@ -33,11 +33,15 @@ async def parse(
         Parse IQL string to IQLQuery object.
 
         Args:
-            source: IQL string that needs to be parsed
-            allowed_functions: list of IQL functions that are allowed for this query
-            event_tracker: EventTracker object to track events
+            source: IQL string that needs to be parsed.
+            allowed_functions: List of IQL functions that are allowed for this query.
+            event_tracker: EventTracker object to track events.
+
         Returns:
-             IQLQuery object
+            IQLQuery object.
+
+        Raises:
+            IQLError: If parsing fails.
         """
         root = await IQLProcessor(source, allowed_functions, event_tracker=event_tracker).process()
         return cls(root=root, source=source)

src/dbally/iql_generator/iql_generator.py

@@ -56,6 +56,9 @@ async def generate_iql(
 
         Returns:
             Generated IQL query.
+
+        Raises:
+            IQLError: If IQL generation fails after all retries.
         """
         prompt_format = IQLGenerationPromptFormat(
             question=question,
@@ -64,7 +67,7 @@ async def generate_iql(
         )
         formatted_prompt = self._prompt_template.format_prompt(prompt_format)
 
-        for _ in range(n_retries + 1):
+        for retry in range(n_retries + 1):
             try:
                 response = await self._llm.generate_text(
                     prompt=formatted_prompt,
@@ -80,5 +83,7 @@ async def generate_iql(
                     event_tracker=event_tracker,
                 )
             except IQLError as exc:
+                if retry == n_retries:
+                    raise exc
                 formatted_prompt = formatted_prompt.add_assistant_message(response)
                 formatted_prompt = formatted_prompt.add_user_message(ERROR_MESSAGE.format(error=exc))

src/dbally/views/structured.py

@@ -55,6 +55,9 @@ async def ask(
 
         Returns:
             The result of the query.
+
+        Raises:
+            IQLError: If the generated IQL query is not valid.
         """
         iql_generator = self.get_iql_generator(llm)
 

tests/integration/test_llm_options.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, call
+from unittest.mock import ANY, AsyncMock, call, patch
 
 import pytest
 
@@ -24,11 +24,12 @@ async def test_llm_options_propagation():
     collection.add(MockView1)
     collection.add(MockView2)
 
-    await collection.ask(
-        question="Mock question",
-        return_natural_response=True,
-        llm_options=custom_options,
-    )
+    with patch("dbally.iql.IQLQuery.parse", AsyncMock()):
+        await collection.ask(
+            question="Mock question",
+            return_natural_response=True,
+            llm_options=custom_options,
+        )
 
     assert llm.client.call.call_count == 3
 

tests/unit/iql/test_iql_parser.py

@@ -4,7 +4,15 @@
 import pytest
 
 from dbally.iql import IQLArgumentParsingError, IQLQuery, IQLUnsupportedSyntaxError, syntax
-from dbally.iql._exceptions import IQLArgumentValidationError, IQLFunctionNotExists
+from dbally.iql._exceptions import (
+    IQLArgumentValidationError,
+    IQLEmptyExpressionError,
+    IQLFunctionNotExists,
+    IQLIncorrectNumberArgumentsError,
+    IQLMultipleExpressionsError,
+    IQLNoExpressionError,
+    IQLSyntaxError,
+)
 from dbally.iql._processor import IQLProcessor
 from dbally.views.exposed_functions import ExposedFunction, MethodParamWithTyping
 
@@ -68,6 +76,78 @@ async def test_iql_parser_arg_error():
     assert exc_info.match(re.escape("Not a valid IQL argument: lambda x: x + 1"))
 
 
+async def test_iql_parser_syntax_error():
+    with pytest.raises(IQLSyntaxError) as exc_info:
+        await IQLQuery.parse(
+            "filter_by_age(",
+            allowed_functions=[
+                ExposedFunction(
+                    name="filter_by_age",
+                    description="",
+                    parameters=[
+                        MethodParamWithTyping(name="age", type=int),
+                    ],
+                ),
+            ],
+        )
+
+    assert exc_info.match(re.escape("Syntax error in: filter_by_age("))
+
+
+async def test_iql_parser_multiple_expression_error():
+    with pytest.raises(IQLMultipleExpressionsError) as exc_info:
+        await IQLQuery.parse(
+            "filter_by_age\nfilter_by_age",
+            allowed_functions=[
+                ExposedFunction(
+                    name="filter_by_age",
+                    description="",
+                    parameters=[
+                        MethodParamWithTyping(name="age", type=int),
+                    ],
+                ),
+            ],
+        )
+
+    assert exc_info.match(re.escape("Multiple expressions or statements in IQL are not supported"))
+
+
+async def test_iql_parser_empty_expression_error():
+    with pytest.raises(IQLEmptyExpressionError) as exc_info:
+        await IQLQuery.parse(
+            "",
+            allowed_functions=[
+                ExposedFunction(
+                    name="filter_by_age",
+                    description="",
+                    parameters=[
+                        MethodParamWithTyping(name="age", type=int),
+                    ],
+                ),
+            ],
+        )
+
+    assert exc_info.match(re.escape("Empty IQL expression"))
+
+
+async def test_iql_parser_no_expression_error():
+    with pytest.raises(IQLNoExpressionError) as exc_info:
+        await IQLQuery.parse(
+            "import filter_by_age",
+            allowed_functions=[
+                ExposedFunction(
+                    name="filter_by_age",
+                    description="",
+                    parameters=[
+                        MethodParamWithTyping(name="age", type=int),
+                    ],
+                ),
+            ],
+        )
+
+    assert exc_info.match(re.escape("No expression found in IQL: import filter_by_age"))
+
+
 async def test_iql_parser_unsupported_syntax_error():
     with pytest.raises(IQLUnsupportedSyntaxError) as exc_info:
         await IQLQuery.parse(
@@ -104,6 +184,26 @@ async def test_iql_parser_method_not_exists():
     assert exc_info.match(re.escape("Function filter_by_how_old_somebody_is not exists: filter_by_how_old_somebody_is"))
 
 
+async def test_iql_parser_incorrect_number_of_arguments_fail():
+    with pytest.raises(IQLIncorrectNumberArgumentsError) as exc_info:
+        await IQLQuery.parse(
+            "filter_by_age('too old', 40)",
+            allowed_functions=[
+                ExposedFunction(
+                    name="filter_by_age",
+                    description="",
+                    parameters=[
+                        MethodParamWithTyping(name="age", type=int),
+                    ],
+                ),
+            ],
+        )
+
+    assert exc_info.match(
+        re.escape("The method filter_by_age has incorrect number of arguments: filter_by_age('too old', 40)")
+    )
+
+
 async def test_iql_parser_argument_validation_fail():
     with pytest.raises(IQLArgumentValidationError) as exc_info:
         await IQLQuery.parse(