add integration tests

Author: ptiurin

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

@@ -4,6 +4,7 @@
 from random import randint
 from typing import Callable, List, Tuple
 
+import trio
 from pytest import mark, raises
 
 import firebolt.async_db
@@ -283,34 +284,73 @@ 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, fb_numeric_paramstyle: None
+@mark.parametrize(
+    "paramstyle,query,test_data",
+    [
+        (
+            "fb_numeric",
+            'INSERT INTO "test_tbl" VALUES ($1, $2)',
+            [(1, "alice"), (2, "bob"), (3, "charlie")],
+        ),
+        (
+            "qmark",
+            'INSERT INTO "test_tbl" VALUES (?, ?)',
+            [(4, "david"), (5, "eve"), (6, "frank")],
+        ),
+    ],
+)
+async def test_executemany_bulk_insert_paramstyles(
+    connection: Connection,
+    paramstyle: str,
+    query: str,
+    test_data: List[Tuple],
+    create_drop_test_table_setup_teardown_async: Callable,
 ) -> None:
-    """executemany with bulk_insert=True inserts data correctly."""
+    """executemany with bulk_insert=True works correctly for both paramstyles."""
+    # Set the paramstyle for this test
+    original_paramstyle = firebolt.async_db.paramstyle
+    firebolt.async_db.paramstyle = paramstyle
+    # Generate a unique label for this test execution
+    unique_label = f"test_bulk_insert_async_{paramstyle}_{randint(100000, 999999)}"
+    table_name = "test_tbl"
+
     try:
-        firebolt.async_db.paramstyle = "fb_numeric"
+        c = connection.cursor()
 
-        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'
-            )
+        # Can't do this for fb_numeric yet - FIR-49970
+        if paramstyle != "fb_numeric":
+            await c.execute(f"SET query_label = '{unique_label}'")
 
-            await c.executemany(
-                'INSERT INTO "test_bulk_insert_async" VALUES ($1, $2)',
-                [(1, "alice"), (2, "bob"), (3, "charlie")],
-                bulk_insert=True,
-            )
+        # Execute bulk insert
+        await c.executemany(
+            query,
+            test_data,
+            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"]
+        # Verify the data was inserted correctly
+        await c.execute(f'SELECT * FROM "{table_name}" ORDER BY id')
+        data = await c.fetchall()
+        assert len(data) == len(test_data)
+        for i, (expected_id, expected_name) in enumerate(test_data):
+            assert data[i] == [expected_id, expected_name]
+
+        # Verify that only one INSERT query was executed with our unique label
+        # Can't do this for fb_numeric yet - FIR-49970
+        if paramstyle != "fb_numeric":
+            # Wait a moment to ensure query history is updated
+            await trio.sleep(10)
+            await c.execute(
+                "SELECT COUNT(*) FROM information_schema.engine_query_history "
+                f"WHERE query_label = '{unique_label}' AND query_text LIKE 'INSERT INTO%'"
+                " AND status = 'ENDED_SUCCESSFULLY'"
+            )
+            query_count = (await c.fetchone())[0]
+            assert (
+                query_count == 1
+            ), f"Expected 1 INSERT query with label '{unique_label}', but found {query_count}"
     finally:
-        async with connection.cursor() as c:
-            await c.execute('DROP TABLE IF EXISTS "test_bulk_insert_async"')
+        firebolt.async_db.paramstyle = original_paramstyle
 
 
 async def test_multi_statement_query(connection: Connection) -> None:

tests/integration/dbapi/conftest.py

@@ -12,9 +12,7 @@
 
 LOGGER = getLogger(__name__)
 
-CREATE_TEST_TABLE = (
-    'CREATE DIMENSION TABLE IF NOT EXISTS "test_tbl" (id int, name string)'
-)
+CREATE_TEST_TABLE = 'CREATE TABLE IF NOT EXISTS "test_tbl" (id int, name string)'
 DROP_TEST_TABLE = 'DROP TABLE IF EXISTS "test_tbl" CASCADE'
 
 LONG_SELECT_DEFAULT_V1 = 250000000000
@@ -59,7 +57,7 @@ def create_drop_test_table_setup_teardown(connection: Connection) -> None:
 
 @fixture
 async def create_drop_test_table_setup_teardown_async(connection: Connection) -> None:
-    with connection.cursor() as c:
+    async with connection.cursor() as c:
         await c.execute(CREATE_TEST_TABLE)
         yield c
         await c.execute(DROP_TEST_TABLE)

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

@@ -1,4 +1,5 @@
 import math
+import time
 from datetime import date, datetime
 from decimal import Decimal
 from random import randint
@@ -7,6 +8,7 @@
 
 from pytest import mark, raises
 
+import firebolt.db
 from firebolt.client.auth import Auth
 from firebolt.common._types import ColType
 from firebolt.common.row_set.types import Column
@@ -283,31 +285,73 @@ def test_empty_query(c: Cursor, query: str, params: tuple) -> None:
         )
 
 
-def test_executemany_bulk_insert(
-    connection: Connection, fb_numeric_paramstyle: None
+@mark.parametrize(
+    "paramstyle,query,test_data",
+    [
+        (
+            "fb_numeric",
+            'INSERT INTO "test_tbl" VALUES ($1, $2)',
+            [(1, "alice"), (2, "bob"), (3, "charlie")],
+        ),
+        (
+            "qmark",
+            'INSERT INTO "test_tbl" VALUES (?, ?)',
+            [(4, "david"), (5, "eve"), (6, "frank")],
+        ),
+    ],
+)
+def test_executemany_bulk_insert_paramstyles(
+    connection: Connection,
+    paramstyle: str,
+    query: str,
+    test_data: List[Tuple],
+    create_drop_test_table_setup_teardown: Callable,
 ) -> None:
-    """executemany with bulk_insert=True inserts data correctly."""
+    """executemany with bulk_insert=True works correctly for both paramstyles."""
+    # Set the paramstyle for this test
+    original_paramstyle = firebolt.db.paramstyle
+    firebolt.db.paramstyle = paramstyle
+    # Generate a unique label for this test execution
+    unique_label = f"test_bulk_insert_{paramstyle}_{randint(100000, 999999)}"
+    table_name = "test_tbl"
+
     try:
-        with connection.cursor() as c:
-            c.execute('DROP TABLE IF EXISTS "test_bulk_insert"')
+        c = connection.cursor()
+
+        # Can't do this for fb_numeric yet - FIR-49970
+        if paramstyle != "fb_numeric":
+            c.execute(f"SET query_label = '{unique_label}'")
+
+        # Execute bulk insert
+        c.executemany(
+            query,
+            test_data,
+            bulk_insert=True,
+        )
+
+        # Verify the data was inserted correctly
+        c.execute(f'SELECT * FROM "{table_name}" ORDER BY id')
+        data = c.fetchall()
+        assert len(data) == len(test_data)
+        for i, (expected_id, expected_name) in enumerate(test_data):
+            assert data[i] == [expected_id, expected_name]
+
+        # Verify that only one INSERT query was executed with our unique label
+        # Can't do this for fb_numeric yet - FIR-49970
+        if paramstyle != "fb_numeric":
+            # Wait a moment to ensure query history is updated
+            time.sleep(10)
             c.execute(
-                'CREATE FACT TABLE "test_bulk_insert"(id int, name string) primary index id'
-            )
-            c.executemany(
-                'INSERT INTO "test_bulk_insert" VALUES ($1, $2)',
-                [(1, "alice"), (2, "bob"), (3, "charlie")],
-                bulk_insert=True,
+                "SELECT COUNT(*) FROM information_schema.engine_query_history "
+                f"WHERE query_label = '{unique_label}' AND query_text LIKE 'INSERT INTO%'"
+                " AND status = 'ENDED_SUCCESSFULLY'"
             )
-
-            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"]
+            query_count = c.fetchone()[0]
+            assert (
+                query_count == 1
+            ), f"Expected 1 INSERT query with label '{unique_label}', but found {query_count}"
     finally:
-        with connection.cursor() as c:
-            c.execute('DROP TABLE IF EXISTS "test_bulk_insert"')
+        firebolt.db.paramstyle = original_paramstyle
 
 
 def test_multi_statement_query(connection: Connection) -> None: