From 05e3fa47711b563a3eb7b51e7131f707123d79b8 Mon Sep 17 00:00:00 2001 From: Zhibai Song Date: Wed, 15 Jan 2025 15:50:57 -0800 Subject: [PATCH] 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 --- contrib/babelfishpg_tsql/src/tsqlIface.cpp | 36 +++++++++++++++++----- test/JDBC/expected/BABEL-2514.out | 11 +++++++ test/JDBC/input/BABEL-2514.sql | 8 ++++- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp index 3be79c3866..51a03fc443 100644 --- a/contrib/babelfishpg_tsql/src/tsqlIface.cpp +++ b/contrib/babelfishpg_tsql/src/tsqlIface.cpp @@ -2936,12 +2936,22 @@ 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 */ @@ -2949,6 +2959,17 @@ class tsqlBuilder : public tsqlCommonMutator 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(fragment->cond, ctx); + } + } + void exitSearch_condition(TSqlParser::Search_conditionContext *ctx) override { if (!ctx->parent || !ctx->parent->parent) @@ -2956,14 +2977,13 @@ class tsqlBuilder : public tsqlCommonMutator 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); } } }; diff --git a/test/JDBC/expected/BABEL-2514.out b/test/JDBC/expected/BABEL-2514.out index 3940f553e6..3d638679f5 100644 --- a/test/JDBC/expected/BABEL-2514.out +++ b/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~~ + diff --git a/test/JDBC/input/BABEL-2514.sql b/test/JDBC/input/BABEL-2514.sql index fa240b5ed7..ccc48e2613 100644 --- a/test/JDBC/input/BABEL-2514.sql +++ b/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 \ No newline at end of file +go + +IF EXISTS +( +SELECT tblnam = a.name FROM sysobjects a where a.name = 'tablename' +) select 'Exists'; ELSE select 'Not Exists'; +GO