Fix multi-row TVP SELECT returning last row (#4817)

Bug introduced in c497926 uses currentColumn which is being reset to 0 and creating an ordering issue in insert stmt for the TVP datatype. With this we want to fix the param ordering for the insertion.

Issues Resolved
BABEL-6511

Sign-off-by: Kushaal Shroff <kushaal@amazon.com>

Author: KushaalShroff

contrib/babelfishpg_tds/src/backend/tds/tdstypeio.c

@@ -44,6 +44,7 @@
 #include "src/include/tds_typeio.h"
 #include "src/include/err_handler.h"
 #include "src/include/tds_instr.h"
+#include "common/int.h"
 
 #include "tds_data_map.c"		/* include tables that used to initialize
 								 * hashmaps */
@@ -2120,26 +2121,14 @@ FetchTvpTypeOid(const ParameterToken token, char *tvpName)
 	char	   *query;
 
 	if ((rc = SPI_connect()) < 0)
-	{
-		/* Reset dialect. */
-		set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
-						  GUC_CONTEXT_CONFIG,
-						  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 		elog(ERROR, "SPI_connect() failed in TDS Listener "
 			 "with return code %d", rc);
-	}
 
 	query = psprintf("SELECT '%s'::regtype::oid", tvpName);
 
 	rc = SPI_execute(query, false, 1);
 	if (rc != SPI_OK_SELECT)
-	{
-		/* Reset dialect. */
-		set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
-						  GUC_CONTEXT_CONFIG,
-						  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
-		elog(ERROR, "Failed to insert in the underlying table for table-valued parameter: %d", rc);
-	}
+		elog(ERROR, "Failed to fetch TVP type OID: %d", rc);
 	tupdesc = SPI_tuptable->tupdesc;
 	row = SPI_tuptable->vals[0];
 
@@ -2217,6 +2206,9 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 					  GUC_CONTEXT_CONFIG,
 					  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 
+	PG_TRY();
+	{
+
 	if (!xactStarted)
 		StartTransactionCommand();
 	PushActiveSnapshot(GetTransactionSnapshot());
@@ -2238,10 +2230,6 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 
 	if (rc != SPI_OK_UTILITY)
 	{
-		/* Reset dialect. */
-		set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
-						  GUC_CONTEXT_CONFIG,
-						  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 		elog(ERROR, "Failed to create the underlying table for table-valued parameter: %d", rc);
 	}
 
@@ -2252,10 +2240,25 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 
 	{
 		char	   *src;
-		int			nargs = token->tvpInfo->colCount * token->tvpInfo->rowCount;
-		Datum	   *values = palloc(nargs * sizeof(Datum));
-		char	   *nulls = palloc(nargs * sizeof(char));
-		Oid		   *argtypes = palloc(nargs * sizeof(Datum));
+		int			nargs;
+		Datum	   *values;
+		char	   *nulls;
+		Oid		   *argtypes;
+		int			paramIndex = 0;
+
+		/*
+		 * pg_mul_s32_overflow guards colCount * rowCount, ensuring nargs
+		 * fits in int32. Since nargs <= INT32_MAX, the subsequent
+		 * palloc_array allocations cannot overflow on 64-bit systems.
+		 */
+		if (pg_mul_s32_overflow(token->tvpInfo->colCount, token->tvpInfo->rowCount, &nargs))
+			ereport(ERROR,
+					(errcode(ERRCODE_PROTOCOL_VIOLATION),
+					 errmsg("TVP row/column count too large")));
+
+		values = palloc_array(Datum, nargs);
+		nulls = palloc_array(char, nargs);
+		argtypes = palloc_array(Oid, nargs);
 
 		query = " ";
 
@@ -2274,67 +2277,67 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 			{
 				temp = &(row->columnValues[currentColumn]);
 				tempFuncInfo = TdsLookupTypeFunctionsByTdsId(colMetaData[currentColumn].columnTdsType, colMetaData[currentColumn].maxLen);
-				GetPgOid(argtypes[currentColumn], tempFuncInfo);
-				if (row->isNull[currentColumn] == 'n')
-					nulls[currentColumn] = row->isNull[currentColumn];
-				else
+				GetPgOid(argtypes[paramIndex], tempFuncInfo);
+
+				if (row->isNull[currentColumn] != 'n')
 					switch (colMetaData[currentColumn].columnTdsType)
 					{
 						case TDS_TYPE_CHAR:
 						case TDS_TYPE_VARCHAR:
-							values[currentColumn] = TdsTypeVarcharToDatum(temp, colMetaData[currentColumn].encoding, colMetaData[currentColumn].columnTdsType);
+							values[paramIndex] = TdsTypeVarcharToDatum(temp, colMetaData[currentColumn].encoding, colMetaData[currentColumn].columnTdsType);
 							break;
 						case TDS_TYPE_NCHAR:
 						case TDS_TYPE_NVARCHAR:
-							values[currentColumn] = TdsTypeNCharToDatum(temp);
+							values[paramIndex] = TdsTypeNCharToDatum(temp);
 							break;
 						case TDS_TYPE_INTEGER:
 						case TDS_TYPE_BIT:
-							values[currentColumn] = TdsTypeIntegerToDatum(temp, colMetaData[currentColumn].maxLen);
+							values[paramIndex] = TdsTypeIntegerToDatum(temp, colMetaData[currentColumn].maxLen);
 							break;
 						case TDS_TYPE_FLOAT:
-							values[currentColumn] = TdsTypeFloatToDatum(temp, colMetaData[currentColumn].maxLen);
+							values[paramIndex] = TdsTypeFloatToDatum(temp, colMetaData[currentColumn].maxLen);
 							break;
 						case TDS_TYPE_NUMERICN:
 						case TDS_TYPE_DECIMALN:
-							values[currentColumn] = TdsTypeNumericToDatum(temp, colMetaData[currentColumn].scale);
+							values[paramIndex] = TdsTypeNumericToDatum(temp, colMetaData[currentColumn].scale);
 							break;
 						case TDS_TYPE_VARBINARY:
 						case TDS_TYPE_BINARY:
-							values[currentColumn] = TdsTypeVarbinaryToDatum(temp);
-							argtypes[currentColumn] = tempFuncInfo->ttmtypeid;
+							values[paramIndex] = TdsTypeVarbinaryToDatum(temp);
+							argtypes[paramIndex] = tempFuncInfo->ttmtypeid;
 							break;
 						case TDS_TYPE_DATE:
-							values[currentColumn] = TdsTypeDateToDatum(temp);
+							values[paramIndex] = TdsTypeDateToDatum(temp);
 							break;
 						case TDS_TYPE_TIME:
-							values[currentColumn] = TdsTypeTimeToDatum(temp, colMetaData[currentColumn].scale, temp->len);
+							values[paramIndex] = TdsTypeTimeToDatum(temp, colMetaData[currentColumn].scale, temp->len);
 							break;
 						case TDS_TYPE_DATETIMEOFFSET:
-							values[currentColumn] = TdsTypeDatetimeoffsetToDatum(temp, colMetaData[currentColumn].scale, temp->len);
+							values[paramIndex] = TdsTypeDatetimeoffsetToDatum(temp, colMetaData[currentColumn].scale, temp->len);
 							break;
 						case TDS_TYPE_DATETIME2:
-							values[currentColumn] = TdsTypeDatetime2ToDatum(temp, colMetaData[currentColumn].scale, temp->len);
+							values[paramIndex] = TdsTypeDatetime2ToDatum(temp, colMetaData[currentColumn].scale, temp->len);
 							break;
 						case TDS_TYPE_DATETIMEN:
-							values[currentColumn] = TdsTypeDatetimeToDatum(temp);
+							values[paramIndex] = TdsTypeDatetimeToDatum(temp);
 							break;
 						case TDS_TYPE_MONEYN:
-							values[currentColumn] = TdsTypeMoneyToDatum(temp);
+							values[paramIndex] = TdsTypeMoneyToDatum(temp);
 							break;
 						case TDS_TYPE_XML:
-							values[currentColumn] = TdsTypeXMLToDatum(temp);
+							values[paramIndex] = TdsTypeXMLToDatum(temp);
 							break;
 						case TDS_TYPE_UNIQUEIDENTIFIER:
-							values[currentColumn] = TdsTypeUIDToDatum(temp);
+							values[paramIndex] = TdsTypeUIDToDatum(temp);
 							break;
 						case TDS_TYPE_SQLVARIANT:
-							values[currentColumn] = TdsTypeSqlVariantToDatum(temp);
+							values[paramIndex] = TdsTypeSqlVariantToDatum(temp);
 							break;
 					}
 				/* Build a string for bind parameters. */
-				currentQuery = psprintf("%s,$%d", currentQuery, currentColumn + 1);
-				nulls[currentColumn] = row->isNull[currentColumn];
+				currentQuery = psprintf("%s,$%d", currentQuery, paramIndex + 1);
+				nulls[paramIndex] = row->isNull[currentColumn];
+				paramIndex++;
 				currentColumn++;
 			}
 			row = row->nextRow;
@@ -2354,14 +2357,8 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 
 			src = psprintf("Insert into %s values %s", finalTableName, query);
 			if ((rc = SPI_connect()) < 0)
-			{
-				/* Reset dialect. */
-				set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
-								  GUC_CONTEXT_CONFIG,
-								  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 				elog(ERROR, "SPI_connect() failed in TDS Listener "
 					 "with return code %d", rc);
-			}
 
 			rc = SPI_execute_with_args(src,
 									   nargs, argtypes,
@@ -2370,10 +2367,6 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 
 			if (rc != SPI_OK_INSERT)
 			{
-				/* Reset dialect. */
-				set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
-								  GUC_CONTEXT_CONFIG,
-								  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 				elog(ERROR, "Failed to insert in the underlying table for table-valued parameter: %d", rc);
 			}
 
@@ -2384,10 +2377,15 @@ TdsRecvTypeTable(const char *message, const ParameterToken token)
 		if (!xactStarted)
 			CommitTransactionCommand();
 
+	}
+	}
+	PG_FINALLY();
+	{
 		set_config_option("babelfishpg_tsql.sql_dialect", "tsql",
 						  GUC_CONTEXT_CONFIG,
 						  PGC_S_SESSION, GUC_ACTION_SAVE, true, 0, false);
 	}
+	PG_END_TRY();
 
 	/* Free all the pointers. */
 	while (token->tvpInfo->rowData)

test/JDBC/expected/TestTvp.out

@@ -51,3 +51,16 @@ int
 drop procedure table_variable_vu_proc1;
 drop type table_variable_vu_type;
 drop type tableType;
+
+-- Multi-row TVP test
+create type tvp_multirow_test_type as table (a int, b smallint, c bigint, d tinyint, e bit, f char(10), g nchar(10), h varchar(10), i nvarchar(10), l varbinary(10), m binary(10), n date, o datetime, p money, q uniqueidentifier, r float, s real, t numeric(4,3), u decimal(5,3), v time, w datetime2)
+prepst#!#Select * from ? order by a#!#tvp|-|tvp_multirow_test_type|-|utils/tvp-dotnet-multirow.csv|-|utils/tvp-dotnet-multirow.csv
+~~START~~
+int#!#smallint#!#bigint#!#tinyint#!#bit#!#char#!#nchar#!#varchar#!#nvarchar#!#varbinary#!#binary#!#date#!#datetime#!#money#!#uniqueidentifier#!#float#!#real#!#numeric#!#numeric#!#time#!#datetime2
+<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>#!#<NULL>
+1#!#1#!#100#!#10#!#1#!#hello     #!#world     #!#foo#!#bar#!#0A1B#!#0C0D0000000000000000#!#2023-01-15#!#2023-01-15 08:30:00.0#!#100.5000#!#CE8AF10A-2709-43B0-9E4E-A02753929D17#!#3.14#!#2.71#!#1.234#!#12.345#!#08:30:00.0000000#!#2023-01-15 08:30:00.0000000
+2#!#<NULL>#!#200#!#<NULL>#!#0#!#<NULL>#!#test      #!#abc#!#<NULL>#!#<NULL>#!#1A2B0000000000000000#!#2024-06-20#!#<NULL>#!#<NULL>#!#A1B2C3D4-E5F6-7890-ABCD-EF1234567890#!#<NULL>#!#1.41#!#<NULL>#!#99.999#!#12:00:00.0000000#!#2024-06-20 12:00:00.0000000
+3#!#3#!#<NULL>#!#30#!#<NULL>#!#data      #!#<NULL>#!#xyz#!#nval#!#FFEE#!#<NULL>#!#<NULL>#!#2025-12-31 23:59:59.0#!#999.9900#!#<NULL>#!#2.718#!#<NULL>#!#9.876#!#<NULL>#!#<NULL>#!#2025-12-31 23:59:59.0000000
+~~END~~
+
+drop type tvp_multirow_test_type;

test/JDBC/input/datatypes/TestTvp.txt

@@ -14,3 +14,8 @@ declare @var1 table_variable_vu_type insert into @var1 values ('1', 2, 3, 4) exe
 drop procedure table_variable_vu_proc1;
 drop type table_variable_vu_type;
 drop type tableType;
+
+-- Multi-row TVP test
+create type tvp_multirow_test_type as table (a int, b smallint, c bigint, d tinyint, e bit, f char(10), g nchar(10), h varchar(10), i nvarchar(10), l varbinary(10), m binary(10), n date, o datetime, p money, q uniqueidentifier, r float, s real, t numeric(4,3), u decimal(5,3), v time, w datetime2)
+prepst#!#Select * from @a order by a#!#tvp|-|tvp_multirow_test_type|-|utils/tvp-dotnet-multirow.csv|-|utils/tvp-dotnet-multirow.csv
+drop type tvp_multirow_test_type;

test/JDBC/src/main/java/com/sqlsamples/CompareResults.java

@@ -290,7 +290,19 @@ static Object parse_data(String result, String datatype, Logger logger) {
                     || datatype.equalsIgnoreCase("varbinary")
                     || datatype.equalsIgnoreCase("timestamp")
                     || datatype.equalsIgnoreCase("udt")) {
-                return result;
+                /* convert hex strings to byte[] for binary/varbinary types */
+                int len = result.length();
+                if (len % 2 != 0)
+                    throw new IllegalArgumentException("Hex string has odd length: " + result);
+                byte[] bytes = new byte[len / 2];
+                for (int idx = 0; idx < len; idx += 2) {
+                    int hi = Character.digit(result.charAt(idx), 16);
+                    int lo = Character.digit(result.charAt(idx + 1), 16);
+                    if (hi < 0 || lo < 0)
+                        throw new IllegalArgumentException("Invalid hex character in: " + result);
+                    bytes[idx / 2] = (byte) ((hi << 4) + lo);
+                }
+                return bytes;
             } else if (datatype.equalsIgnoreCase("decimal")
                     || datatype.equalsIgnoreCase("money")
                     || datatype.equalsIgnoreCase("smallmoney")

test/JDBC/utils/tvp-dotnet-multirow.csv

@@ -0,0 +1,5 @@
+a-int,b-smallint,c-bigint,d-tinyint,e-bit,f-char-10,g-nchar-10,h-varchar-10,i-nvarchar-10,l-varbinary-10,m-binary-10,n-date,o-datetime,p-money,q-uniqueidentifier,r-float,s-real,t-numeric-4-3,u-decimal-5-3,v-time,w-datetime2
+1,1,100,10,TRUE,hello,world,foo,bar,0A1B,0C0D,2023-01-15,2023-01-15 08:30:00,100.50,ce8af10a-2709-43b0-9e4e-a02753929d17,3.14,2.71,1.234,12.345,08:30:00,2023-01-15 08:30:00
+2,<NULL>,200,<NULL>,FALSE,<NULL>,test,abc,<NULL>,<NULL>,1A2B,2024-06-20,<NULL>,<NULL>,a1b2c3d4-e5f6-7890-abcd-ef1234567890,<NULL>,1.41,<NULL>,99.999,12:00:00,2024-06-20 12:00:00
+3,3,<NULL>,30,<NULL>,data,<NULL>,xyz,nval,FFEE,<NULL>,<NULL>,2025-12-31 23:59:59,999.99,<NULL>,2.718,<NULL>,9.876,<NULL>,<NULL>,2025-12-31 23:59:59
+<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>,<NULL>