fix(connection, cursor): raise DataError when bind parameter limit is exceeded (#165) (#187)

* raise DataError in Connection execute
* rethrow error in Cursor insert_data_bulk
* add integration and unit test for cursor
* rectify integration test

Co-authored-by: Soksamnang Lim <soklim@amazon.com>

Author: soksamnanglim

redshift_connector/core.py

@@ -36,6 +36,7 @@
     ArrayContentNotHomogenousError,
     ArrayContentNotSupportedError,
     DatabaseError,
+    DataError,
     Error,
     IntegrityError,
     InterfaceError,
@@ -1710,6 +1711,9 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
             #   Int32 - The OID of the parameter data type.
             val: typing.Union[bytes, bytearray] = bytearray(statement_name_bin)
             typing.cast(bytearray, val).extend(statement.encode(_client_encoding) + NULL_BYTE)
+            if len(params) > 32767:
+                raise DataError("Prepared statement exceeds bind parameter limit 32767. {} bind parameters were "
+                                "provided. Please retry with fewer bind parameters.".format(len(params)))
             typing.cast(bytearray, val).extend(h_pack(len(params)))
             for oid, fc, send_func in params:  # type: ignore
                 # Parse message doesn't seem to handle the -1 type_oid for NULL

redshift_connector/cursor.py

@@ -340,7 +340,7 @@ def insert_data_bulk(
                     self.execute(insert_stmt, values_list)
 
         except Exception as e:
-            raise InterfaceError(e)
+            raise e
         finally:
             # reset paramstyle to it's original value
             self.paramstyle = orig_paramstyle

test/integration/test_cursor.py

@@ -4,7 +4,7 @@
 import pytest  # type: ignore
 
 import redshift_connector
-from redshift_connector import InterfaceError
+from redshift_connector import InterfaceError, DataError
 
 
 @pytest.mark.parametrize("col_name", (("apples", "apples"), ("author‎ ", "author\u200e")))
@@ -66,3 +66,130 @@ def test_insert_data_invalid_column_raises(mocked_csv, db_kwargs):
                     delimiter=",",
                     batch_size=3,
                 )
+
+
+# max binding parameters for a prepared statement
+max_params = 32767
+
+def test_insert_data_raises_too_many_params(db_kwargs):
+    prepared_stmt = (
+            "INSERT INTO githubissue165 (col1) VALUES " + "(%s), " * max_params + "(%s);"
+    )
+    params = [1 for _ in range(max_params + 1)]
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("create temporary table githubissue165 (col1 int)")
+
+            with pytest.raises(
+                    DataError,
+                    match=f"Prepared statement exceeds bind parameter limit 32767. {32768} bind parameters were "
+                          f"provided.",
+            ):
+                cursor.execute(prepared_stmt, params)
+
+
+def test_insert_data_raises_no_exception(db_kwargs):
+    prepared_stmt_32767 = (
+            "INSERT INTO githubissue165 (col1) VALUES "
+            + "(%s), " * (max_params - 1)
+            + "(%s);"
+    )
+    params_32767 = [1 for _ in range(max_params)]
+
+    prepared_stmt_32766 = (
+            "INSERT INTO githubissue165 (col1) VALUES "
+            + "(%s), " * (max_params - 2)
+            + "(%s);"
+    )
+    params_32766 = [1 for _ in range(max_params - 1)]
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("create temporary table githubissue165 (col1 int)")
+            try:
+                cursor.execute(prepared_stmt_32767, params_32767)
+            except Exception as e:
+                assert (
+                    False
+                ), f"'execute' with {max_params} bind parameters raised an exception {e}"
+            try:
+                cursor.execute(prepared_stmt_32766, params_32766)
+            except Exception as e:
+                assert (
+                    False
+                ), f"'execute' with {max_params - 1} bind parameters raised an exception {e}"
+
+
+indices, names = (
+    [0],
+    ["col1"],
+)
+
+
+@patch("builtins.open", new_callable=mock_open)
+def test_insert_data_bulk_raises_too_many_params(mocked_csv, db_kwargs):
+    csv_str = "\col1\n" + "1\n" * max_params + "1"  # 32768 rows
+    mocked_csv.side_effect = [StringIO(csv_str)]
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("create temporary table githubissue165 (col1 int)")
+            with pytest.raises(
+                    DataError,
+                    match="Prepared statement exceeds bind parameter limit 32767.",
+            ):
+                cursor.insert_data_bulk(
+                    filename="mocked_csv",
+                    table_name="githubissue165",
+                    parameter_indices=indices,
+                    column_names=["col1"],
+                    delimiter=",",
+                    batch_size=max_params + 1,
+                )
+
+
+@patch("builtins.open", new_callable=mock_open)
+def test_insert_data_bulk_raises_no_exception_32766(mocked_csv_32766, db_kwargs):
+    csv_str_32766 = "\col1\n" + "1\n" * (max_params - 2) + "1"
+    mocked_csv_32766.side_effect = [StringIO(csv_str_32766)]
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("create temporary table githubissue165 (col1 int)")
+            try:
+                cursor.insert_data_bulk(
+                    filename="mocked_csv_32766",
+                    table_name="githubissue165",
+                    parameter_indices=indices,
+                    column_names=["col1"],
+                    delimiter=",",
+                    batch_size=max_params - 1,
+                )
+            except Exception as e:
+                assert (
+                    False
+                ), f"'insert_data_bulk' with {max_params - 1} bind parameters raised an exception {e}"
+
+
+@patch("builtins.open", new_callable=mock_open)
+def test_insert_data_bulk_raises_no_exception_32767(mocked_csv_32767, db_kwargs):
+    csv_str_32767 = "\col1\n" + "1\n" * (max_params - 1) + "1"
+    mocked_csv_32767.side_effect = [StringIO(csv_str_32767)]
+
+    with redshift_connector.connect(**db_kwargs) as conn:
+        with conn.cursor() as cursor:
+            cursor.execute("create temporary table githubissue165 (col1 int)")
+            try:
+                cursor.insert_data_bulk(
+                    filename="mocked_csv_32767",
+                    table_name="githubissue165",
+                    parameter_indices=indices,
+                    column_names=["col1"],
+                    delimiter=",",
+                    batch_size=max_params,
+                )
+            except Exception as e:
+                assert (
+                    False
+                ), f"'insert_data_bulk' with {max_params} bind parameters raised an exception {e}"

test/unit/test_cursor.py

@@ -6,7 +6,7 @@
 
 import pytest  # type: ignore
 
-from redshift_connector import Connection, Cursor, InterfaceError
+from redshift_connector import Connection, Cursor, InterfaceError, DataError
 
 IS_SINGLE_DATABASE_METADATA_TOGGLE: typing.List[bool] = [True, False]
 
@@ -406,3 +406,59 @@ def test_insert_data_uses_batch_size(mocked_csv, batch_size, mocker):
             actual_insert_stmts_executed += 1
 
     assert actual_insert_stmts_executed == ceil(3 / batch_size)
+
+max_params = 32767
+
+@patch("builtins.open", new_callable=mock_open)
+def test_insert_data_bulk_raises_too_many_parameters(mocked_csv, mocker):
+    # mock fetchone to return "True" to ensure the table_name and column_name
+    # validation steps pass
+    mocker.patch("redshift_connector.Cursor.fetchone", return_value=[1])
+
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+
+    # mock out the connection to raise DataError.
+    mock_cursor._c = Mock()
+    mocker.patch.object(mock_cursor._c, "execute", side_effect=DataError("Prepared statement exceeds bind parameter "
+                                                                         "limit 32767."))
+    mock_cursor.paramstyle = "mocked"
+
+    indexes, names = (
+        [0],
+        ["col1"],
+    )
+
+    csv_str = "\col1\n" + "1\n" * max_params + "1"  # 32768 rows
+    mocked_csv.side_effect = [StringIO(csv_str)]
+
+    with pytest.raises(
+        DataError, match="Prepared statement exceeds bind parameter limit 32767."
+    ):
+        mock_cursor.insert_data_bulk(
+            filename="mocked_csv",
+            table_name="githubissue165",
+            parameter_indices=indexes,
+            column_names=["col1"],
+            delimiter=",",
+            batch_size=max_params + 1,
+        )
+
+
+@patch("builtins.open", new_callable=mock_open)
+def test_insert_data_raises_too_many_parameters(mocker):
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+
+    # mock out the connection to raise DataError.
+    mock_cursor._c = Mock()
+    mock_cursor._c.execute.side_effect = DataError(
+        "Prepared statement exceeds bind " "parameter limit 32767."
+    )
+    mock_cursor.paramstyle = "mocked"
+
+    prepared_stmt = "INSERT INTO githubissue165 (col1) VALUES " + "(%s), " * max_params + "(%s);"
+    params = [1 for _ in range(max_params + 1)]
+
+    with pytest.raises(
+        DataError, match="Prepared statement exceeds bind parameter limit 32767."
+    ):
+        mock_cursor.execute(prepared_stmt, params)