refactor: simplify bulk_insert validation and reuse existing execution logic

- Replace sqlparse with simple string checks in _validate_bulk_insert_query
- Add extra_params parameter to _build_fb_numeric_query_params for extensibility
- Refactor _executemany_bulk_insert to preprocess and delegate to existing methods
- Use TimeoutController for consistent timeout handling
- Parametrize integration tests to avoid duplication

Addresses PR feedback from ptiurin on #463

Co-Authored-By: petro.tiurin@firebolt.io <petro.tiurin@firebolt.io>

Author: devin-ai-integration[bot]

src/firebolt/async_db/cursor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import logging
 import time
 import warnings
@@ -462,20 +461,16 @@ async def executemany(
 
     def _validate_bulk_insert_query(self, query: str) -> None:
         """Validate that query is an INSERT statement for bulk_insert."""
-        statements = parse_sql(query)
-        if not statements:
-            raise ProgrammingError("bulk_insert requires a valid INSERT statement")
+        query_normalized = query.lstrip().lower()
 
-        if len(statements) > 1:
+        if not query_normalized.startswith("insert"):
             raise ProgrammingError(
-                "bulk_insert does not support multi-statement queries"
+                "bulk_insert is only supported for INSERT statements"
             )
 
-        statement_type = statements[0].get_type()
-        if statement_type != "INSERT":
+        if ";" in query.strip().rstrip(";"):
             raise ProgrammingError(
-                f"bulk_insert is only supported for INSERT statements. "
-                f"Got {statement_type} statement"
+                "bulk_insert does not support multi-statement queries"
             )
 
     async def _executemany_bulk_insert(
@@ -497,37 +492,32 @@ async def _executemany_bulk_insert(
         except ValueError:
             raise ProgrammingError(f"Unsupported paramstyle: {paramstyle}")
 
+        concatenated_query = "; ".join([query] * len(parameters_seq))
+
         await self._close_rowset_and_reset()
         self._row_set = InMemoryAsyncRowSet()
 
         try:
             if parameter_style == ParameterStyle.FB_NUMERIC:
-                concatenated_query = "; ".join([query] * len(parameters_seq))
-
                 flattened_params: List[ParameterType] = []
                 for param_set in parameters_seq:
                     flattened_params.extend(param_set)
 
-                query_parameters = [
-                    {
-                        "name": f"${i+1}",
-                        "value": self._formatter.convert_parameter_for_serialization(
-                            value
-                        ),
-                    }
-                    for i, value in enumerate(flattened_params)
-                ]
-
-                query_params: Dict[str, Any] = {
-                    "output_format": self._get_output_format(False),
-                    "merge_prepared_statement_batches": "true",
-                }
-                if query_parameters:
-                    query_params["query_parameters"] = json.dumps(query_parameters)
-
                 Cursor._log_query(concatenated_query)
+                timeout_controller = TimeoutController(timeout_seconds)
+                timeout_controller.raise_if_timeout()
+
+                query_params = self._build_fb_numeric_query_params(
+                    [flattened_params],
+                    streaming=False,
+                    async_execution=False,
+                    extra_params={"merge_prepared_statement_batches": "true"},
+                )
+
                 resp = await self._api_request(
-                    concatenated_query, query_params, timeout=timeout_seconds
+                    concatenated_query,
+                    query_params,
+                    timeout=timeout_controller.remaining(),
                 )
                 await self._raise_if_error(resp)
                 await self._parse_response_headers(resp.headers)

src/firebolt/common/cursor/base_cursor.py

@@ -243,6 +243,7 @@ def _build_fb_numeric_query_params(
         parameters: Sequence[Sequence[ParameterType]],
         streaming: bool,
         async_execution: bool,
+        extra_params: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """
         Build query parameters dictionary for fb_numeric paramstyle.
@@ -252,6 +253,7 @@ def _build_fb_numeric_query_params(
                        only the first parameter sequence is used.
             streaming: Whether streaming is enabled
             async_execution: Whether async execution is enabled
+            extra_params: Optional additional query parameters to include
 
         Returns:
             Dictionary of query parameters to send with the request
@@ -272,6 +274,8 @@ def _build_fb_numeric_query_params(
             query_params["query_parameters"] = json.dumps(query_parameters)
         if async_execution:
             query_params["async"] = True
+        if extra_params:
+            query_params.update(extra_params)
         return query_params
 
     @property

src/firebolt/db/cursor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import logging
 import time
 from abc import ABCMeta, abstractmethod
@@ -464,20 +463,16 @@ def executemany(
 
     def _validate_bulk_insert_query(self, query: str) -> None:
         """Validate that query is an INSERT statement for bulk_insert."""
-        statements = parse_sql(query)
-        if not statements:
-            raise ProgrammingError("bulk_insert requires a valid INSERT statement")
+        query_normalized = query.lstrip().lower()
 
-        if len(statements) > 1:
+        if not query_normalized.startswith("insert"):
             raise ProgrammingError(
-                "bulk_insert does not support multi-statement queries"
+                "bulk_insert is only supported for INSERT statements"
             )
 
-        statement_type = statements[0].get_type()
-        if statement_type != "INSERT":
+        if ";" in query.strip().rstrip(";"):
             raise ProgrammingError(
-                f"bulk_insert is only supported for INSERT statements. "
-                f"Got {statement_type} statement"
+                "bulk_insert does not support multi-statement queries"
             )
 
     def _executemany_bulk_insert(
@@ -499,37 +494,32 @@ def _executemany_bulk_insert(
         except ValueError:
             raise ProgrammingError(f"Unsupported paramstyle: {paramstyle}")
 
+        concatenated_query = "; ".join([query] * len(parameters_seq))
+
         self._close_rowset_and_reset()
         self._row_set = InMemoryRowSet()
 
         try:
             if parameter_style == ParameterStyle.FB_NUMERIC:
-                concatenated_query = "; ".join([query] * len(parameters_seq))
-
                 flattened_params: List[ParameterType] = []
                 for param_set in parameters_seq:
                     flattened_params.extend(param_set)
 
-                query_parameters = [
-                    {
-                        "name": f"${i+1}",
-                        "value": self._formatter.convert_parameter_for_serialization(
-                            value
-                        ),
-                    }
-                    for i, value in enumerate(flattened_params)
-                ]
-
-                query_params: Dict[str, Any] = {
-                    "output_format": self._get_output_format(False),
-                    "merge_prepared_statement_batches": "true",
-                }
-                if query_parameters:
-                    query_params["query_parameters"] = json.dumps(query_parameters)
-
                 Cursor._log_query(concatenated_query)
+                timeout_controller = TimeoutController(timeout_seconds)
+                timeout_controller.raise_if_timeout()
+
+                query_params = self._build_fb_numeric_query_params(
+                    [flattened_params],
+                    streaming=False,
+                    async_execution=False,
+                    extra_params={"merge_prepared_statement_batches": "true"},
+                )
+
                 resp = self._api_request(
-                    concatenated_query, query_params, timeout=timeout_seconds
+                    concatenated_query,
+                    query_params,
+                    timeout=timeout_controller.remaining(),
                 )
                 self._raise_if_error(resp)
                 self._parse_response_headers(resp.headers)

tests/integration/dbapi/async/V2/test_queries_async.py

@@ -282,56 +282,45 @@ async def test_parameterized_query_with_special_chars(connection: Connection) ->
         ], "Invalid data in table after parameterized insert"
 
 
-async def test_executemany_bulk_insert(connection: Connection) -> None:
+@mark.parametrize("paramstyle", ["qmark", "fb_numeric"])
+async def test_executemany_bulk_insert(connection: Connection, paramstyle: str) -> None:
     """executemany with bulk_insert=True inserts data correctly."""
-    async with connection.cursor() as c:
-        await c.execute('DROP TABLE IF EXISTS "test_bulk_insert_async"')
-        await c.execute(
-            'CREATE FACT TABLE "test_bulk_insert_async"(id int, name string) primary index id'
-        )
+    import firebolt.async_db as db_module
 
-        import firebolt.async_db as db_module
+    original_paramstyle = db_module.paramstyle
 
-        original_paramstyle = db_module.paramstyle
-        db_module.paramstyle = "qmark"
+    try:
+        db_module.paramstyle = paramstyle
 
-        try:
-            await c.executemany(
-                'INSERT INTO "test_bulk_insert_async" VALUES (?, ?)',
-                [(1, "alice"), (2, "bob"), (3, "charlie")],
-                bulk_insert=True,
+        async with connection.cursor() as c:
+            await c.execute('DROP TABLE IF EXISTS "test_bulk_insert_async"')
+            await c.execute(
+                'CREATE FACT TABLE "test_bulk_insert_async"(id int, name string) primary index id'
             )
 
+            if paramstyle == "qmark":
+                await c.executemany(
+                    'INSERT INTO "test_bulk_insert_async" VALUES (?, ?)',
+                    [(1, "alice"), (2, "bob"), (3, "charlie")],
+                    bulk_insert=True,
+                )
+            else:
+                await c.executemany(
+                    'INSERT INTO "test_bulk_insert_async" VALUES ($1, $2)',
+                    [(1, "alice"), (2, "bob"), (3, "charlie")],
+                    bulk_insert=True,
+                )
+
             await c.execute('SELECT * FROM "test_bulk_insert_async" ORDER BY id')
             data = await c.fetchall()
             assert len(data) == 3
             assert data[0] == [1, "alice"]
             assert data[1] == [2, "bob"]
             assert data[2] == [3, "charlie"]
-        finally:
-            db_module.paramstyle = original_paramstyle
-
-        await c.execute('DELETE FROM "test_bulk_insert_async"')
-
-        db_module.paramstyle = "fb_numeric"
-
-        try:
-            await c.executemany(
-                'INSERT INTO "test_bulk_insert_async" VALUES ($1, $2)',
-                [(4, "david"), (5, "eve"), (6, "frank")],
-                bulk_insert=True,
-            )
-
-            await c.execute('SELECT * FROM "test_bulk_insert_async" ORDER BY id')
-            data = await c.fetchall()
-            assert len(data) == 3
-            assert data[0] == [4, "david"]
-            assert data[1] == [5, "eve"]
-            assert data[2] == [6, "frank"]
-        finally:
-            db_module.paramstyle = original_paramstyle
 
-        await c.execute('DROP TABLE "test_bulk_insert_async"')
+            await c.execute('DROP TABLE "test_bulk_insert_async"')
+    finally:
+        db_module.paramstyle = original_paramstyle
 
 
 async def test_multi_statement_query(connection: Connection) -> None:

tests/integration/dbapi/sync/V2/test_queries.py

@@ -283,56 +283,45 @@ def test_empty_query(c: Cursor, query: str, params: tuple) -> None:
         )
 
 
-def test_executemany_bulk_insert(connection: Connection) -> None:
+@mark.parametrize("paramstyle", ["qmark", "fb_numeric"])
+def test_executemany_bulk_insert(connection: Connection, paramstyle: str) -> None:
     """executemany with bulk_insert=True inserts data correctly."""
-    with connection.cursor() as c:
-        c.execute('DROP TABLE IF EXISTS "test_bulk_insert"')
-        c.execute(
-            'CREATE FACT TABLE "test_bulk_insert"(id int, name string) primary index id'
-        )
+    import firebolt.db as db_module
 
-        import firebolt.db as db_module
+    original_paramstyle = db_module.paramstyle
 
-        original_paramstyle = db_module.paramstyle
-        db_module.paramstyle = "qmark"
+    try:
+        db_module.paramstyle = paramstyle
 
-        try:
-            c.executemany(
-                'INSERT INTO "test_bulk_insert" VALUES (?, ?)',
-                [(1, "alice"), (2, "bob"), (3, "charlie")],
-                bulk_insert=True,
+        with connection.cursor() as c:
+            c.execute('DROP TABLE IF EXISTS "test_bulk_insert"')
+            c.execute(
+                'CREATE FACT TABLE "test_bulk_insert"(id int, name string) primary index id'
             )
 
+            if paramstyle == "qmark":
+                c.executemany(
+                    'INSERT INTO "test_bulk_insert" VALUES (?, ?)',
+                    [(1, "alice"), (2, "bob"), (3, "charlie")],
+                    bulk_insert=True,
+                )
+            else:
+                c.executemany(
+                    'INSERT INTO "test_bulk_insert" VALUES ($1, $2)',
+                    [(1, "alice"), (2, "bob"), (3, "charlie")],
+                    bulk_insert=True,
+                )
+
             c.execute('SELECT * FROM "test_bulk_insert" ORDER BY id')
             data = c.fetchall()
             assert len(data) == 3
             assert data[0] == [1, "alice"]
             assert data[1] == [2, "bob"]
             assert data[2] == [3, "charlie"]
-        finally:
-            db_module.paramstyle = original_paramstyle
-
-        c.execute('DELETE FROM "test_bulk_insert"')
-
-        db_module.paramstyle = "fb_numeric"
-
-        try:
-            c.executemany(
-                'INSERT INTO "test_bulk_insert" VALUES ($1, $2)',
-                [(4, "david"), (5, "eve"), (6, "frank")],
-                bulk_insert=True,
-            )
-
-            c.execute('SELECT * FROM "test_bulk_insert" ORDER BY id')
-            data = c.fetchall()
-            assert len(data) == 3
-            assert data[0] == [4, "david"]
-            assert data[1] == [5, "eve"]
-            assert data[2] == [6, "frank"]
-        finally:
-            db_module.paramstyle = original_paramstyle
 
-        c.execute('DROP TABLE "test_bulk_insert"')
+            c.execute('DROP TABLE "test_bulk_insert"')
+    finally:
+        db_module.paramstyle = original_paramstyle
 
 
 def test_multi_statement_query(connection: Connection) -> None: