Support expressions in SELECT...OFFSET...FETCH clauses (#2298)

In SELECT...OFFSET...FETCH, T-SQL allows expressions for the expression in OFFSET and FETCH. In Babelfish any expression involving an operator or function call requires that the expression is enclosed in brackets.
This fix add those brackets as part of ANTLR rewriting.

Signed-off-by: Rob Verschoor [email protected]
Task: BABEL-1174

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -197,6 +197,7 @@ static void handleBitNotOperator(TSqlParser::Unary_op_exprContext *ctx);
 static void handleBitOperators(TSqlParser::Plus_minus_bit_exprContext *ctx);
 static void handleModuloOperator(TSqlParser::Mult_div_percent_exprContext *ctx);
 static void handleAtAtVarInPredicate(TSqlParser::PredicateContext *ctx);
+static void handleOrderByOffsetFetch(TSqlParser::Order_by_clauseContext *ctx);
 
 /*
  * Structure / Utility function for general purpose of query string modification
@@ -659,9 +660,14 @@ void PLtsql_expr_query_mutator::run()
 		const std::u32string& orig_text = utf8_to_utf32(entry.second.first.c_str());
 		const std::u32string& repl_text = utf8_to_utf32(entry.second.second.c_str());
 		if (isSelectFragment) offset += fragment_SELECT_prefix.length(); // because this is an expression prefixed with 'SELECT '
-
+					
 		if (orig_text.length() == 0 || orig_text.c_str(), query.substr(offset, orig_text.length()) == orig_text) // local_id maybe already deleted in some cases such as select-assignment. check here if it still exists)
 		{
+			// Note: the test below does not work, and has never worked, because size_t will not be negative, 
+			// and the result of the subtraction is also of type size_t.
+			// This test has been in the code since day 1. 
+			// When making the test work, some test cases will start failing as they run into this condition 
+			// (test table_variable_xact_errors and two variants). Therefore, not touching the test for now.
 			if (offset - cursor < 0)
 				throw PGErrorWrapperException(ERROR, ERRCODE_INTERNAL_ERROR, "can't mutate an internal query. might be due to multiple mutations on the same position", 0, 0);
 			if (offset - cursor > 0) // if offset==cursor, no need to copy
@@ -2281,7 +2287,12 @@ class tsqlBuilder : public tsqlCommonMutator
 			}
 		}
 	}
-
+
+	void exitOrder_by_clause(TSqlParser::Order_by_clauseContext *ctx) override
+	{
+		handleOrderByOffsetFetch(ctx);
+	}
+
 	// NB: the following are copied in tsqlMutator
 	void exitColumn_def_table_constraints(TSqlParser::Column_def_table_constraintsContext *ctx)
 	{
@@ -8542,3 +8553,30 @@ void handleAtAtVarInPredicate(TSqlParser::PredicateContext *ctx)
 	}
 	return;
 }
+
+static void
+handleOrderByOffsetFetch(TSqlParser::Order_by_clauseContext *ctx)
+{
+	// Add brackets around the expressions for OFFSET..ROWS and FETCH..ROWS
+
+	if (ctx->offset_exp) 
+	{
+		// Do not rewrite the entire expression since that will break the logic in the mutator when there is something inside the
+		// expression that also needs rewriting (like a local variable @p which needs to be rewritten as "@p").
+		// Instead, insert an opening and closing bracket in the right places.
+		// Also, do not add a rewrite at the start position of the expression since there may be an '@' for a local var 
+		// at that position and the rewrite to double-quote the variable will be lost as a result.
+		rewritten_query_fragment.emplace(std::make_pair((ctx->offset_exp->start->getStartIndex() - 1), std::make_pair("", " (")));
+		rewritten_query_fragment.emplace(std::make_pair((ctx->offset_exp->stop->getStopIndex() + 1), std::make_pair("", ") ")));
+	}
+
+	if (ctx->fetch_exp) 
+	{
+		// See comment for offset_exp above.
+		rewritten_query_fragment.emplace(std::make_pair((ctx->fetch_exp->start->getStartIndex() - 1), std::make_pair("", " (")));
+		rewritten_query_fragment.emplace(std::make_pair((ctx->fetch_exp->stop->getStopIndex() + 1), std::make_pair("", ") ")));
+	}
+
+	return;
+}
+

test/JDBC/expected/order_by_offset_fetch_rows-before-15_6-or-16_2-vu-cleanup.out

@@ -0,0 +1,8 @@
+DROP PROC p1_upgr_order_by_offset_fetch
+go
+DROP FUNCTION f1_upgr_order_by_offset_fetch
+go
+DROP VIEW v1_upgr_order_by_offset_fetch 
+go
+DROP table t1_upgr_order_by_offset_fetch
+go

test/JDBC/expected/order_by_offset_fetch_rows-before-15_6-or-16_2-vu-prepare.out

@@ -0,0 +1,8 @@
+create table t1_upgr_order_by_offset_fetch(a int, b int)
+go
+CREATE PROC p1_upgr_order_by_offset_fetch @p int=1,@q int=1  AS SELECT * FROM t1_upgr_order_by_offset_fetch ORDER BY b OFFSET @p*@q ROWS FETCH NEXT square(@q)+1 ROWS ONLY
+go
+CREATE FUNCTION f1_upgr_order_by_offset_fetch(@p int, @q int) returns int as begin declare @v int SELECT @v=count(a) FROM t1_upgr_order_by_offset_fetch group by b ORDER BY b OFFSET @p*@q ROWS FETCH NEXT square(@q)+1 ROWS ONLY return @v end
+go
+CREATE VIEW v1_upgr_order_by_offset_fetch as select * FROM t1_upgr_order_by_offset_fetch ORDER BY b OFFSET 5 ROWS FETCH NEXT 3 rows only 
+go

test/JDBC/expected/order_by_offset_fetch_rows-before-15_6-or-16_2-vu-verify.out

@@ -0,0 +1,73 @@
+insert t1_upgr_order_by_offset_fetch select generate_series, 0 from generate_series(1,100)
+go
+~~ROW COUNT: 100~~
+
+update t1_upgr_order_by_offset_fetch set b=a
+go
+~~ROW COUNT: 100~~
+
+exec p1_upgr_order_by_offset_fetch 1, 3
+go
+~~START~~
+int#!#int
+4#!#4
+5#!#5
+6#!#6
+7#!#7
+8#!#8
+9#!#9
+10#!#10
+11#!#11
+12#!#12
+13#!#13
+~~END~~
+
+exec p1_upgr_order_by_offset_fetch 2, 3
+go
+~~START~~
+int#!#int
+7#!#7
+8#!#8
+9#!#9
+10#!#10
+11#!#11
+12#!#12
+13#!#13
+14#!#14
+15#!#15
+16#!#16
+~~END~~
+
+p1_upgr_order_by_offset_fetch 3, 3
+go
+~~START~~
+int#!#int
+10#!#10
+11#!#11
+12#!#12
+13#!#13
+14#!#14
+15#!#15
+16#!#16
+17#!#17
+18#!#18
+19#!#19
+~~END~~
+
+select dbo.f1_upgr_order_by_offset_fetch(1,1)
+go
+~~START~~
+int
+1
+~~END~~
+
+select * from v1_upgr_order_by_offset_fetch
+go
+~~START~~
+int#!#int
+6#!#6
+7#!#7
+8#!#8
+~~END~~
+
+

test/JDBC/expected/order_by_offset_fetch_rows-vu-cleanup.out

@@ -0,0 +1,8 @@
+DROP PROC p1_order_by_offset_fetch
+go
+DROP FUNCTION f1_order_by_offset_fetch
+go
+DROP VIEW v1_order_by_offset_fetch 
+go
+DROP table t1_order_by_offset_fetch
+go

test/JDBC/expected/order_by_offset_fetch_rows-vu-prepare.out

@@ -0,0 +1,16 @@
+create table t1_order_by_offset_fetch(a int, b int)
+go
+insert t1_order_by_offset_fetch select generate_series, 0 from generate_series(1,100)
+go
+~~ROW COUNT: 100~~
+
+update t1_order_by_offset_fetch set b=a
+go
+~~ROW COUNT: 100~~
+
+CREATE PROC p1_order_by_offset_fetch @p int=1,@q int=1  AS SELECT * FROM t1_order_by_offset_fetch ORDER BY b OFFSET @p*@q ROWS FETCH NEXT square(@q)+1 ROWS ONLY
+go
+CREATE FUNCTION f1_order_by_offset_fetch(@p int, @q int) returns int as begin declare @v int SELECT @v=count(a) FROM t1_order_by_offset_fetch group by b ORDER BY b OFFSET @p*@q ROWS FETCH NEXT square(@q)+1 ROWS ONLY return @v end
+go
+CREATE VIEW v1_order_by_offset_fetch as select * FROM t1_order_by_offset_fetch ORDER BY b OFFSET 5 ROWS FETCH NEXT 3 rows only 
+go