refactor(FIR-49930): decoupling query preprocessing logic (#467)

Author: ptiurin

src/firebolt/async_db/cursor.py

@@ -25,7 +25,6 @@
     UPDATE_ENDPOINT_HEADER,
     UPDATE_PARAMETERS_HEADER,
     CursorState,
-    ParameterStyle,
 )
 from firebolt.common.cursor.base_cursor import (
     BaseCursor,
@@ -38,6 +37,10 @@
     check_not_closed,
     check_query_executed,
 )
+from firebolt.common.cursor.statement_planners import (
+    ExecutionPlan,
+    StatementPlannerFactory,
+)
 from firebolt.common.row_set.asynchronous.base import BaseAsyncRowSet
 from firebolt.common.row_set.asynchronous.in_memory import InMemoryAsyncRowSet
 from firebolt.common.row_set.asynchronous.streaming import StreamingAsyncRowSet
@@ -222,117 +225,84 @@ async def _do_execute(
         from firebolt.async_db import paramstyle
 
         try:
-            parameter_style = ParameterStyle(paramstyle)
-        except ValueError:
-            raise ProgrammingError(f"Unsupported paramstyle: {paramstyle}")
-        try:
-            if parameter_style == ParameterStyle.FB_NUMERIC:
-                await self._execute_fb_numeric(
-                    raw_query, parameters, timeout, async_execution, streaming
-                )
-            else:
-                queries: List[Union[SetParameter, str]] = (
-                    [raw_query]
-                    if skip_parsing
-                    else self._formatter.split_format_sql(raw_query, parameters)
-                )
-                timeout_controller = TimeoutController(timeout)
-                if len(queries) > 1 and async_execution:
-                    raise FireboltError(
-                        "Server side async does not support multi-statement queries"
-                    )
-                for query in queries:
-                    await self._execute_single_query(
-                        query, timeout_controller, async_execution, streaming
-                    )
+            statement_planner = StatementPlannerFactory.create_planner(
+                paramstyle, self._formatter
+            )
+
+            plan = statement_planner.create_execution_plan(
+                raw_query, parameters, skip_parsing, async_execution, streaming
+            )
+            await self._execute_plan(plan, timeout)
             self._state = CursorState.DONE
         except Exception:
             self._state = CursorState.ERROR
             raise
 
-    async def _execute_fb_numeric(
+    async def _execute_plan(
         self,
-        query: str,
-        parameters: Sequence[Sequence[ParameterType]],
+        plan: ExecutionPlan,
         timeout: Optional[float],
-        async_execution: bool,
-        streaming: bool,
     ) -> None:
-        Cursor._log_query(query)
+        """Execute an execution plan."""
         timeout_controller = TimeoutController(timeout)
-        timeout_controller.raise_if_timeout()
-        query_params = self._build_fb_numeric_query_params(
-            parameters, streaming, async_execution
-        )
-        resp = await self._api_request(
-            query,
-            query_params,
-            timeout=timeout_controller.remaining(),
-        )
-        await self._raise_if_error(resp)
-        if async_execution:
-            await resp.aread()
-            self._parse_async_response(resp)
-        else:
-            await self._parse_response_headers(resp.headers)
-            await self._append_row_set_from_response(resp)
+
+        for query in plan.queries:
+            if isinstance(query, SetParameter):
+                if plan.async_execution:
+                    raise FireboltError(
+                        "Server side async does not support set statements, "
+                        "please use execute to set this parameter"
+                    )
+                await self._validate_set_parameter(
+                    query, timeout_controller.remaining()
+                )
+            else:
+                # Regular query execution
+                await self._execute_single_query(
+                    query,
+                    plan.query_params,
+                    timeout_controller,
+                    plan.async_execution,
+                    plan.streaming,
+                )
 
     async def _execute_single_query(
         self,
-        query: Union[SetParameter, str],
+        query: str,
+        query_params: Optional[Dict[str, Any]],
         timeout_controller: TimeoutController,
         async_execution: bool,
         streaming: bool,
     ) -> None:
+        """Execute a single query."""
         start_time = time.time()
         Cursor._log_query(query)
         timeout_controller.raise_if_timeout()
 
-        if isinstance(query, SetParameter):
-            if async_execution:
-                raise FireboltError(
-                    "Server side async does not support set statements, "
-                    "please use execute to set this parameter"
-                )
-            await self._validate_set_parameter(query, timeout_controller.remaining())
-        else:
-            await self._handle_query_execution(
-                query, timeout_controller, async_execution, streaming
-            )
-
-        if not async_execution:
-            logger.info(
-                f"Query fetched {self.rowcount} rows in"
-                f" {time.time() - start_time} seconds."
-            )
-        else:
-            logger.info("Query submitted for async execution.")
+        final_params = query_params or {}
 
-    async def _handle_query_execution(
-        self,
-        query: str,
-        timeout_controller: TimeoutController,
-        async_execution: bool,
-        streaming: bool,
-    ) -> None:
-        query_params: Dict[str, Any] = {
-            "output_format": self._get_output_format(streaming)
-        }
-        if async_execution:
-            query_params["async"] = True
         resp = await self._api_request(
             query,
-            query_params,
+            final_params,
             timeout=timeout_controller.remaining(),
         )
         await self._raise_if_error(resp)
+
         if async_execution:
             await resp.aread()
             self._parse_async_response(resp)
         else:
             await self._parse_response_headers(resp.headers)
             await self._append_row_set_from_response(resp)
 
+        if not async_execution:
+            logger.info(
+                f"Query fetched {self.rowcount} rows in"
+                f" {time.time() - start_time} seconds."
+            )
+        else:
+            logger.info("Query submitted for async execution.")
+
     async def use_database(self, database: str, cache: bool = True) -> None:
         """Switch the current database context with caching."""
         if cache:

src/firebolt/common/cursor/base_cursor.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
-import json
 import logging
 import re
 from types import TracebackType
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 from httpx import URL, Response
 
 from firebolt.client.auth.base import Auth
 from firebolt.client.client import AsyncClient, Client
-from firebolt.common._types import ParameterType, RawColType, SetParameter
+from firebolt.common._types import RawColType, SetParameter
 from firebolt.common.constants import (
     DISALLOWED_PARAMETER_LIST,
     IMMUTABLE_PARAMETER_LIST,
@@ -238,42 +237,6 @@ def _log_query(query: Union[str, SetParameter]) -> None:
         ):
             logger.debug(f"Running query: {query}")
 
-    def _build_fb_numeric_query_params(
-        self,
-        parameters: Sequence[Sequence[ParameterType]],
-        streaming: bool,
-        async_execution: bool,
-    ) -> Dict[str, Any]:
-        """
-        Build query parameters dictionary for fb_numeric paramstyle.
-
-        Args:
-            parameters: A sequence of parameter sequences. For fb_numeric,
-                       only the first parameter sequence is used.
-            streaming: Whether streaming is enabled
-            async_execution: Whether async execution is enabled
-
-        Returns:
-            Dictionary of query parameters to send with the request
-        """
-        param_list = parameters[0] if parameters else []
-        query_parameters = [
-            {
-                "name": f"${i+1}",
-                "value": self._formatter.convert_parameter_for_serialization(value),
-            }
-            for i, value in enumerate(param_list)
-        ]
-
-        query_params: Dict[str, Any] = {
-            "output_format": self._get_output_format(streaming),
-        }
-        if query_parameters:
-            query_params["query_parameters"] = json.dumps(query_parameters)
-        if async_execution:
-            query_params["async"] = True
-        return query_params
-
     @property
     def engine_name(self) -> str:
         """