Fix a corner case '=' not properly rewritten to AS (#3411)

Previously we rewrite "SELECT COL=expr" to "SELECT expr as "COL"", but
there's a corner case that for if condition, it didn't do the rewrite,
because different parser path from if condition to dml statement.

This commit fix this issue by adding the mutator logic to if statement
rewrite.

Task: BABEL-5320
Signed-off-by: Zhibai Song <[email protected]>

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -2936,34 +2936,54 @@ class tsqlBuilder : public tsqlCommonMutator
 		clear_rewritten_query_fragment();
 	}
 
-	PLtsql_expr *rewrite_if_condition(TSqlParser::Search_conditionContext *ctx)
+	PLtsql_expr *rewrite_if_condition(TSqlParser::Search_conditionContext *ctx, PLtsql_expr *expr)
 	{
-		PLtsql_expr *expr = makeTsqlExpr(ctx, false);
-		PLtsql_expr_query_mutator mutator(expr, ctx);
-		add_rewritten_query_fragment_to_mutator(&mutator);
-		mutator.run();
+		if (!expr)
+			expr = makeTsqlExpr(ctx, false);
+		if (statementMutator)
+		{
+			add_rewritten_query_fragment_to_mutator(statementMutator.get());
+			statementMutator->run();
+			statementMutator = nullptr;
+		}
+		else
+		{
+			PLtsql_expr_query_mutator mutator(expr, ctx);
+			add_rewritten_query_fragment_to_mutator(&mutator);
+			mutator.run();
+		}
 		clear_rewritten_query_fragment();
 
 		/* Now we can prepend SELECT to rewritten search_condition */
 		expr->query = strdup((std::string("SELECT ") + std::string(expr->query)).c_str());
 		return expr;
 	}
 
+	void enterSearch_condition(TSqlParser::Search_conditionContext *ctx) override
+	{
+		if (((TSqlParser::Cfl_statementContext *) ctx->parent->parent)->if_statement())
+		{
+			PLtsql_stmt_if *fragment = (PLtsql_stmt_if *) getPLtsql_fragment(ctx->parent->parent);
+			fragment->cond = makeTsqlExpr(ctx, false);
+			clear_rewritten_query_fragment();
+			statementMutator = std::make_unique<PLtsql_expr_query_mutator>(fragment->cond, ctx);
+		}
+	}
+
 	void exitSearch_condition(TSqlParser::Search_conditionContext *ctx) override
 	{
 		if (!ctx->parent || !ctx->parent->parent)
 			return;
 
 		if (((TSqlParser::Cfl_statementContext *) ctx->parent->parent)->if_statement())
 		{
-
 			PLtsql_stmt_if *fragment = (PLtsql_stmt_if *) getPLtsql_fragment(ctx->parent->parent);
-			fragment->cond = rewrite_if_condition(ctx);
+			fragment->cond = rewrite_if_condition(ctx, fragment->cond);
 		}
 		else if (((TSqlParser::Cfl_statementContext *) ctx->parent->parent)->while_statement())
 		{
 			PLtsql_stmt_while *fragment = (PLtsql_stmt_while *) getPLtsql_fragment(ctx->parent->parent);
-			fragment->cond = rewrite_if_condition(ctx);
+			fragment->cond = rewrite_if_condition(ctx, fragment->cond);
 		}
 	}
 };

test/JDBC/expected/BABEL-2514.out

@@ -197,3 +197,14 @@ int#!#int
 
 drop view if exists babel_2514_complex_view;
 go
+
+IF EXISTS 
+(
+SELECT tblnam = a.name FROM sysobjects a where a.name = 'tablename'
+) select 'Exists'; ELSE select 'Not Exists';
+GO
+~~START~~
+varchar
+Not Exists
+~~END~~
+

test/JDBC/input/BABEL-2514.sql

@@ -125,4 +125,10 @@ go
 select * from babel_2514_complex_view order by b;
 go
 drop view if exists babel_2514_complex_view;
-go
+go
+
+IF EXISTS 
+(
+SELECT tblnam = a.name FROM sysobjects a where a.name = 'tablename'
+) select 'Exists'; ELSE select 'Not Exists';
+GO