Fix Identity columns not being recognized during DML statements using…

… OUTPUT and WHERE clause (#3166)

Earlier when we executed DML queries that had OUTPUT clause along with a WHERE in the OUTPUT'ed query, we didnt expand the target list of this OUTPUT'ed query properly when the target list was not provided by user.
For example:
"UPDATE t1 SET TEXTVal = 'NewValue' OUTPUT Inserted.ID, Inserted.TEXTVal, Deleted.ID, Deleted.TEXTVal INTO t2 WHERE ID IN (1,2)" table t2's target list will be created during analyze phase and we were not ignoring the identity/computed columns while expanding the list.
With this commit we fix this by ignoring any identity/computed columns at the time of target-list fetch.

Tasks:
BABEL-5432

Signed-off-by: Kushaal Shroff <[email protected]>

contrib/babelfishpg_tsql/src/hooks.c

@@ -2378,7 +2378,13 @@ modify_insert_stmt(InsertStmt *stmt, Oid relid)
 
 		temp_col_list = NIL;
 
-		if (att->attnum > 0)
+		/*
+		 * We must skip the generated or identity columns while inserting.
+		 * Note that in case of IDENTITY_INSERT Explicit col-name must be specified
+		 * so stmt->col will be NOT NULL.
+		 */
+		if (att->attnum > 0
+			&& !(att->attgenerated || att->attidentity == ATTRIBUTE_IDENTITY_ALWAYS))
 		{
 			/*
  			* Do a deep copy of attname because tuple is a pointer 

test/JDBC/expected/output_with_where_tests.out

@@ -0,0 +1,89 @@
+-- INSERT with OUTPUT
+CREATE TABLE insert_table (ID INT IDENTITY(1,1), Name VARCHAR(50), Value INT);
+GO
+
+INSERT INTO insert_table (Name, Value) VALUES ('A', 10), ('B', 20), ('C', 30);
+GO
+~~ROW COUNT: 3~~
+
+
+
+DECLARE @ResultTable TABLE (NewID INT, NewName VARCHAR(50), NewValue INT);
+INSERT INTO insert_table (Name, Value)
+OUTPUT Inserted.ID, Inserted.Name, Inserted.Value INTO @ResultTable
+SELECT Name, Value
+FROM (VALUES ('D', 40), ('E', 50), ('F', 60), ('G', 70)) AS Source(Name, Value)
+WHERE Value > 45;
+GO
+~~ROW COUNT: 3~~
+
+
+DROP TABLE insert_table;
+GO
+
+-- DELETE with OUTPUT
+CREATE TABLE delete_table (id INT IDENTITY(1,1), name VARCHAR(50), Price DECIMAL(10,2));
+GO
+
+INSERT INTO delete_table (name, Price) VALUES ('Old Product 1', 99.99), ('Old Product 2', 149.99);
+GO
+~~ROW COUNT: 2~~
+
+
+CREATE TABLE #temp_delete_table (ID INT IDENTITY(1,1), DeletedID INT, DeletedName VARCHAR(50), DeletedPrice DECIMAL(10,2));
+GO
+
+DELETE FROM delete_table
+OUTPUT Deleted.id, Deleted.name, Deleted.Price INTO #temp_delete_table
+WHERE Price > 100;
+GO
+~~ROW COUNT: 1~~
+
+
+DROP TABLE delete_table;
+DROP TABLE #temp_delete_table;
+GO
+
+-- UPDATE with OUTPUT
+CREATE TABLE TestTable (ID INT, TEXTVal VARCHAR(100))
+GO
+
+INSERT INTO TestTable
+VALUES (1, 'John'),
+       (2, 'Jane'),
+       (3, 'Bob'),
+       (4, 'Alice');
+GO
+~~ROW COUNT: 4~~
+
+
+CREATE TABLE #TmpTable (
+    c1 INT IDENTITY,
+    ID_New INT,
+    TEXTVal_New VARCHAR(100),
+    ID_Old INT,
+    TEXTVal_Old VARCHAR(100)
+);
+GO
+
+UPDATE TestTable SET TEXTVal = 'NewValue'
+OUTPUT Inserted.ID, Inserted.TEXTVal, Deleted.ID, Deleted.TEXTVal INTO #TmpTable WHERE ID IN (1,2)
+GO
+~~ROW COUNT: 2~~
+
+
+-- additional tests to prove IDENTITY_INSERT is not affected by the fix:
+SET IDENTITY_INSERT #TmpTable ON
+GO
+
+UPDATE TestTable SET TEXTVal = 'NewValue'
+OUTPUT Inserted.ID, Inserted.TEXTVal, Deleted.ID, Deleted.TEXTVal INTO #TmpTable WHERE ID IN (1,2)
+GO
+~~ERROR (Code: 545)~~
+
+~~ERROR (Message: Explicit value must be specified for identity column in table '#tmptable' when IDENTITY_INSERT is set to ON)~~
+
+
+DROP TABLE TestTable;
+DROP TABLE #TmpTable;
+GO

test/JDBC/input/output_with_where_tests.sql

@@ -0,0 +1,73 @@
+-- INSERT with OUTPUT
+CREATE TABLE insert_table (ID INT IDENTITY(1,1), Name VARCHAR(50), Value INT);
+GO
+
+INSERT INTO insert_table (Name, Value) VALUES ('A', 10), ('B', 20), ('C', 30);
+GO
+
+DECLARE @ResultTable TABLE (NewID INT, NewName VARCHAR(50), NewValue INT);
+
+INSERT INTO insert_table (Name, Value)
+OUTPUT Inserted.ID, Inserted.Name, Inserted.Value INTO @ResultTable
+SELECT Name, Value
+FROM (VALUES ('D', 40), ('E', 50), ('F', 60), ('G', 70)) AS Source(Name, Value)
+WHERE Value > 45;
+GO
+
+DROP TABLE insert_table;
+GO
+
+-- DELETE with OUTPUT
+CREATE TABLE delete_table (id INT IDENTITY(1,1), name VARCHAR(50), Price DECIMAL(10,2));
+GO
+
+INSERT INTO delete_table (name, Price) VALUES ('Old Product 1', 99.99), ('Old Product 2', 149.99);
+GO
+
+CREATE TABLE #temp_delete_table (ID INT IDENTITY(1,1), DeletedID INT, DeletedName VARCHAR(50), DeletedPrice DECIMAL(10,2));
+GO
+
+DELETE FROM delete_table
+OUTPUT Deleted.id, Deleted.name, Deleted.Price INTO #temp_delete_table
+WHERE Price > 100;
+GO
+
+DROP TABLE delete_table;
+DROP TABLE #temp_delete_table;
+GO
+
+-- UPDATE with OUTPUT
+CREATE TABLE TestTable (ID INT, TEXTVal VARCHAR(100))
+GO
+
+INSERT INTO TestTable
+VALUES (1, 'John'),
+       (2, 'Jane'),
+       (3, 'Bob'),
+       (4, 'Alice');
+GO
+
+CREATE TABLE #TmpTable (
+    c1 INT IDENTITY,
+    ID_New INT,
+    TEXTVal_New VARCHAR(100),
+    ID_Old INT,
+    TEXTVal_Old VARCHAR(100)
+);
+GO
+
+UPDATE TestTable SET TEXTVal = 'NewValue'
+OUTPUT Inserted.ID, Inserted.TEXTVal, Deleted.ID, Deleted.TEXTVal INTO #TmpTable WHERE ID IN (1,2)
+GO
+
+-- additional tests to prove IDENTITY_INSERT is not affected by the fix:
+SET IDENTITY_INSERT #TmpTable ON
+GO
+
+UPDATE TestTable SET TEXTVal = 'NewValue'
+OUTPUT Inserted.ID, Inserted.TEXTVal, Deleted.ID, Deleted.TEXTVal INTO #TmpTable WHERE ID IN (1,2)
+GO
+
+DROP TABLE TestTable;
+DROP TABLE #TmpTable;
+GO