Fix handling of UPDATE command when RETURNING clause updates declared…

… variables (#3028)

currently, Any variable being set as part of UPDATE...SET... command is being re-written using RETURNING clause. But there was an issue where any dynamic update on variable was not visible to subsequent statement execution. That issue is fixed with 6be12f1 and this commit aims to fix the same issue when dynamic variables are being updated as part of UPDATE...SET... command. In order to make variables updated dynamically, assignment statement is re-written using sys.pltsql_assign_var(dno, expr). For example,

SET @A+=expr, col=expr2

will be re-written to

SET col=expr2 ... RETURNING sys.pltsql_assign_var(dno, "@var" + cast((expr) as type))

Task: BABEL-5188
Signed-off-by: Dipesh Dhameliya <[email protected]>

.github/composite-actions/install-dependencies/action.yml

@@ -9,7 +9,7 @@ runs:
         curl https://packages.microsoft.com/keys/microsoft.asc | sudo apt-key add -
         curl https://packages.microsoft.com/config/ubuntu/20.04/prod.list | sudo tee /etc/apt/sources.list.d/msprod.list
         sudo apt-get update --fix-missing -y
-        sudo apt-get install uuid-dev openjdk-8-jre libicu-dev libxml2-dev openssl libssl-dev python3-dev libossp-uuid-dev libpq-dev cmake pkg-config g++ build-essential bison mssql-tools unixodbc-dev libsybdb5 freetds-dev freetds-common gdal-bin libgdal-dev libgeos-dev gdb
+        sudo apt-get install uuid-dev openjdk-8-jre libicu-dev libxml2-dev openssl libssl-dev python3-dev libossp-uuid-dev libpq-dev cmake pkg-config g++ build-essential bison mssql-tools unixodbc-dev libsybdb5 freetds-dev freetds-common gdal-bin libgdal-dev libgeos-dev gdb libkrb5-dev
         sudo apt install -y ccache
         sudo apt-get install lcov
         sudo /usr/sbin/update-ccache-symlinks

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--4.4.0--4.5.0.sql

@@ -54,6 +54,7 @@ BEGIN
 END;
 $$;
 
+
 /* Helper function to update local variables dynamically during execution */
 CREATE OR REPLACE FUNCTION sys.pltsql_assign_var(dno INT, val ANYELEMENT)
 RETURNS ANYELEMENT

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -7289,29 +7289,70 @@ void process_execsql_destination_update(TSqlParser::Update_statementContext *uct
 					appendStringInfo(&ds, ", ");
 				++returning_col_cnt;
 
-				if (elem->full_column_name())
+				if (elem->EQUAL(0) && elem->full_column_name())
 				{
 					/* "SET @a=col=expr" => "SET col=expr ... RETURNING sys.pltsql_assign_var(dno, cast(expr as type))" */
-					appendStringInfo(&ds, "sys.pltsql_assign_var(%d, cast(%s as %s))",
+					appendStringInfo(&ds, "sys.pltsql_assign_var(%d, %s)",
 										nse->itemno,
-										::getFullText(elem->full_column_name()).c_str(),
-										tsql_format_type_extended(var->datatype->typoid, var->datatype->atttypmod, FORMAT_TYPE_TYPEMOD_GIVEN));
+										rewrite_assignment_expression(var, elem->expression()));
 
 					removeTokenStringFromQuery(stmt->sqlstmt, elem->LOCAL_ID(), uctx);
 					removeTokenStringFromQuery(stmt->sqlstmt, elem->EQUAL(0), uctx);
 				}
-				else
+				else if(elem->EQUAL(0) && elem->expression())
 				{
 					/* "SET @a=expr, col=expr2" => "SET col=expr2 ... RETURNING sys.pltsql_assign_var(dno, cast(expr as type))" */
-					appendStringInfo(&ds, "sys.pltsql_assign_var(%d, cast(%s as %s))",
+					appendStringInfo(&ds, "sys.pltsql_assign_var(%d, %s)",
 										nse->itemno,
-										::getFullText(elem->expression()).c_str(),
-										tsql_format_type_extended(var->datatype->typoid, var->datatype->atttypmod, FORMAT_TYPE_TYPEMOD_GIVEN));
+										rewrite_assignment_expression(var, elem->expression()));
 
+					handle_local_ids_for_expression(elem->expression());
 					removeTokenStringFromQuery(stmt->sqlstmt, elem->LOCAL_ID(), uctx);
 					removeTokenStringFromQuery(stmt->sqlstmt, elem->EQUAL(0), uctx);
 					removeCtxStringFromQuery(stmt->sqlstmt, elem->expression(), uctx);
 				}
+				else if(elem->assignment_operator() && elem->expression())
+				{
+					/* "SET @a+=expr, col=expr2" => "SET col=expr2 ... RETURNING sys.pltsql_assign_var(dno, "@var" + cast((expr) as type))" */
+					tree::TerminalNode *anode = nullptr;
+
+					if (elem->assignment_operator())
+					{
+						if (elem->assignment_operator()->PLUS_ASSIGN())
+							anode = elem->assignment_operator()->PLUS_ASSIGN();
+						else if (elem->assignment_operator()->MINUS_ASSIGN())
+							anode = elem->assignment_operator()->MINUS_ASSIGN();
+						else if (elem->assignment_operator()->MULT_ASSIGN())
+							anode = elem->assignment_operator()->MULT_ASSIGN();
+						else if (elem->assignment_operator()->DIV_ASSIGN())
+							anode = elem->assignment_operator()->DIV_ASSIGN();
+						else if (elem->assignment_operator()->MOD_ASSIGN())
+							anode = elem->assignment_operator()->MOD_ASSIGN();
+						else if (elem->assignment_operator()->AND_ASSIGN())
+							anode = elem->assignment_operator()->AND_ASSIGN();
+						else if (elem->assignment_operator()->XOR_ASSIGN())
+							anode = elem->assignment_operator()->XOR_ASSIGN();
+						else if (elem->assignment_operator()->OR_ASSIGN())
+							anode = elem->assignment_operator()->OR_ASSIGN();
+						else
+							Assert(0);
+					}
+					appendStringInfo(&ds, "sys.pltsql_assign_var(%d, %s %s %s)",
+										nse->itemno,
+										delimitIfAtAtUserVarName(::getFullText(elem->LOCAL_ID())).c_str(),
+										rewrite_assign_operator(anode),
+										rewrite_assignment_expression(var, elem->expression()));
+
+					handle_local_ids_for_expression(elem->expression());
+					removeTokenStringFromQuery(stmt->sqlstmt, elem->LOCAL_ID(), uctx);
+					removeCtxStringFromQuery(stmt->sqlstmt, elem->assignment_operator(), uctx);
+					removeCtxStringFromQuery(stmt->sqlstmt, elem->expression(), uctx);
+				}
+				else
+				{
+					/* Syntax error should already be handled. */
+					Assert(0);
+				}
 
 				// Conceptually we have to remove any nearest COMMA.
 				// But code is little bit dirty to handle some corner cases (the first few elems are removed or the last few elems are removed)
@@ -7326,11 +7367,21 @@ void process_execsql_destination_update(TSqlParser::Update_statementContext *uct
 					/* remove prev COMMA by default */
 					removeTokenStringFromQuery(stmt->sqlstmt, uctx->COMMA(i-1), uctx);
 				}
+
 			}
 			else
 				comma_carry_over = false;
 		}
 
+		/*
+		 * Should remove semi colon before appending RETURNING clause. Otherwise syntax
+		 * error will be thrown.
+		 */
+		if (uctx->SEMI())
+		{
+			removeTokenStringFromQuery(stmt->sqlstmt, uctx->SEMI(), uctx);
+		}
+
 		pltsql_adddatum((PLtsql_datum *) target);
 
 		stmt->target = (PLtsql_variable *)target;