Merge branch 'main' into feat-remove-parameters

Author: ptiurin

docsrc/Connecting_and_queries.rst

@@ -196,14 +196,14 @@ To get started, follow the steps below:
         ) as connection:
             # Create a cursor
             cursor = connection.cursor()
-        
+
             # Execute a simple test query
             cursor.execute("SELECT 1")
 
 .. note::
 
-    Firebolt Core is assumed to be running locally on the default port (3473). For instructions 
-    on how to run Firebolt Core locally using Docker, refer to the 
+    Firebolt Core is assumed to be running locally on the default port (3473). For instructions
+    on how to run Firebolt Core locally using Docker, refer to the
     `official docs <https://docs.firebolt.io/firebolt-core/firebolt-core-get-started>`_.
 
 
@@ -404,7 +404,7 @@ parameters equal in length to the number of placeholders in the statement.
         "INSERT INTO test_table2 VALUES ($1, $2, $3)",
         (2, "world", "2018-01-02"),
     )
-    
+
     # paramstyle only needs to be set once, it will be used for all subsequent queries
 
     cursor.execute(
@@ -437,6 +437,58 @@ as values in the second argument.
     cursor.close()
 
 
+Bulk insert for improved performance
+--------------------------------------
+
+For inserting large amounts of data more efficiently, you can use the ``bulk_insert`` parameter
+with ``executemany()``. This concatenates multiple INSERT statements into a single batch request,
+which can significantly improve performance when inserting many rows.
+
+**Note:** The ``bulk_insert`` parameter only works with INSERT statements and supports both
+``fb_numeric`` and ``qmark`` parameter styles. Using it with other statement types will
+raise an error.
+
+**Example with QMARK parameter style (default):**
+
+::
+
+    # Using the default qmark parameter style
+    cursor.executemany(
+        "INSERT INTO test_table VALUES (?, ?, ?)",
+        (
+            (1, "apple", "2019-01-01"),
+            (2, "banana", "2020-01-01"),
+            (3, "carrot", "2021-01-01"),
+            (4, "donut", "2022-01-01"),
+            (5, "eggplant", "2023-01-01")
+        ),
+        bulk_insert=True  # Enable bulk insert for better performance
+    )
+
+**Example with FB_NUMERIC parameter style:**
+
+::
+
+    import firebolt.db
+    # Set paramstyle to "fb_numeric" for server-side parameter substitution
+    firebolt.db.paramstyle = "fb_numeric"
+
+    cursor.executemany(
+        "INSERT INTO test_table VALUES ($1, $2, $3)",
+        (
+            (1, "apple", "2019-01-01"),
+            (2, "banana", "2020-01-01"),
+            (3, "carrot", "2021-01-01"),
+            (4, "donut", "2022-01-01"),
+            (5, "eggplant", "2023-01-01")
+        ),
+        bulk_insert=True  # Enable bulk insert for better performance
+    )
+
+When ``bulk_insert=True``, the SDK concatenates all INSERT statements into a single batch
+and sends them to the server for optimized batch processing.
+
+
 Setting session parameters
 --------------------------------------
 
@@ -731,7 +783,7 @@ of execute_async is -1, which is the rowcount for queries where it's not applica
     cursor.execute_async("INSERT INTO my_table VALUES (5, 'egg', '2022-01-01')")
     token = cursor.async_query_token
 
-Trying to access `async_query_token` before calling `execute_async` will raise an exception. 
+Trying to access `async_query_token` before calling `execute_async` will raise an exception.
 
 .. note::
     Multiple-statement queries are not supported for asynchronous queries. However, you can run each statement
@@ -746,9 +798,9 @@ Monitoring the query status
 To check the async query status you need to retrieve the token of the query. The token is a unique
 identifier for the query and can be used to fetch the query status. You can store this token
 outside of the current process and use it later to check the query status. :ref:`Connection <firebolt.db:Connection>` object
-has two methods to check the query status: :py:meth:`firebolt.db.connection.Connection.is_async_query_running` and 
-:py:meth:`firebolt.db.connection.Connection.is_async_query_successful`.`is_async_query_running` will return True 
-if the query is still running, and False otherwise. `is_async_query_successful` will return True if the query 
+has two methods to check the query status: :py:meth:`firebolt.db.connection.Connection.is_async_query_running` and
+:py:meth:`firebolt.db.connection.Connection.is_async_query_successful`.`is_async_query_running` will return True
+if the query is still running, and False otherwise. `is_async_query_successful` will return True if the query
 has finished successfully, None if query is still running and False if the query has failed.
 
 ::
@@ -779,7 +831,7 @@ will send a cancel request to the server and the query will be stopped.
 
     token = cursor.async_query_token
     connection.cancel_async_query(token)
-    
+
     # Verify that the query was cancelled
     running = connection.is_async_query_running(token)
     print(running) # False

src/firebolt/async_db/cursor.py

@@ -26,7 +26,6 @@
     UPDATE_ENDPOINT_HEADER,
     UPDATE_PARAMETERS_HEADER,
     CursorState,
-    ParameterStyle,
 )
 from firebolt.common.cursor.base_cursor import (
     BaseCursor,
@@ -40,6 +39,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
@@ -221,124 +224,98 @@ async def _do_execute(
         timeout: Optional[float] = None,
         async_execution: bool = False,
         streaming: bool = False,
+        bulk_insert: bool = False,
     ) -> None:
         await self._close_rowset_and_reset()
         self._row_set = StreamingAsyncRowSet() if streaming else InMemoryAsyncRowSet()
+
         # Import paramstyle from module level
         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,
+                bulk_insert,
+            )
+            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
-            )
+        final_params = query_params or {}
 
-        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 _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:
@@ -421,6 +398,7 @@ async def executemany(
         query: str,
         parameters_seq: Sequence[Sequence[ParameterType]],
         timeout_seconds: Optional[float] = None,
+        bulk_insert: bool = False,
     ) -> Union[int, str]:
         """Prepare and execute a database query.
 
@@ -438,6 +416,9 @@ async def executemany(
                 `SET param=value` statement before it. All parameters are stored in
                 cursor object until it's closed. They can also be removed with
                 `flush_parameters` method call.
+            Bulk insert: When bulk_insert=True, multiple INSERT queries are
+                concatenated and sent as a single batch for improved performance.
+                Only supported for INSERT statements.
 
         Args:
             query (str): SQL query to execute.
@@ -446,11 +427,15 @@ async def executemany(
                query with actual values from each set in a sequence. Resulting queries
                for each subset are executed sequentially.
             timeout_seconds (Optional[float]): Query execution timeout in seconds.
+            bulk_insert (bool): When True, concatenates multiple INSERT queries
+               into a single batch request. Only supported for INSERT statements.
 
         Returns:
             int: Query row count.
         """
-        await self._do_execute(query, parameters_seq, timeout=timeout_seconds)
+        await self._do_execute(
+            query, parameters_seq, timeout=timeout_seconds, bulk_insert=bulk_insert
+        )
         return self.rowcount
 
     @check_not_closed