additional tests for UPDATE/DELETE of table variable in function (#3070)

Additional tests and test cases for detecting allowed UPDATE/DELETE on a table variables inside a T-SQL function, for cases where the FROM-clause contains an ANSI join. Issues Resolved: BABEL-5357 Signed-off-by: Rob Verschoor [email protected]
babelfish-for-postgresql · Nov 4, 2024 · 9eccba1 · 9eccba1
1 parent 5444c49
commit 9eccba1
Show file tree

Hide file tree

Showing 5 changed files with 355 additions and 26 deletions.
diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
@@ -1985,37 +1985,17 @@ class tsqlBuilder : public tsqlCommonMutator
 				{
 					if (ddl_object && ddl_object->full_object_name())
 					{
-						line_and_pos = getLineAndPos(ddl_object);
 						// DML target can be an alias: verify that the alias is for a table variable
-						bool aliasIsTabVar = false;
 						if (table_sources)
 						{
-							std::string dmlTarget = getFullText(ddl_object->full_object_name());
-							line_and_pos = getLineAndPos(ddl_object->full_object_name());
-							std::vector<TSqlParser::Table_source_itemContext *> table_source_item;
-							if (ctx->update_statement())
-								table_source_item = table_sources->table_source_item();
-							else
-								table_source_item = table_sources->table_source_item();
-
-							for (size_t i=0; i<table_source_item.size(); ++i)
+							line_and_pos = getLineAndPos(ddl_object);
+							// DML target can be an alias: verify that the alias is for a table variable
+							if (table_sources)
 							{
-								// Find FROM-clause alias matching the DML target
-								if (table_source_item[i]->as_table_alias().size() == 0)
-									continue;  // No alias found in this FROM-clause item
-
-								std::string alias = getFullText(table_source_item[i]->as_table_alias()[0]->table_alias());
-								if (pg_strcasecmp(alias.c_str(), dmlTarget.c_str()) != 0)
-									continue;  // This alias is not matching the DML target
-
-								if (table_source_item[i]->local_id())
-								{
-									// Table variable, with alias matching the DML target
-									aliasIsTabVar = true;
-									break;
-								}
+								line_and_pos = getLineAndPos(ddl_object->full_object_name());								
+								std::string dmlTarget = getFullText(ddl_object->full_object_name());
+								dmlTargetAllowed = dmlTargetIsTabvar(table_sources->table_source_item(), dmlTarget);
 							}
-							dmlTargetAllowed = aliasIsTabVar;
 						}
 					}
 				}
@@ -2065,6 +2045,43 @@ class tsqlBuilder : public tsqlCommonMutator
 		clear_tables_info();
 	}
 
+	bool dmlTargetIsTabvar(std::vector<TSqlParser::Table_source_itemContext *> table_source_item, std::string dmlTarget)
+	{
+		// Potential optimization: as soon as we find the alias matching the DML target, and it's not a table variable, 
+		// we can stop searching through the rest of the table_source_items. 
+		// But some quick testing did not show a measurable difference so this may not be worth doing.
+
+		// Search through table_source_item, including nested ones, to find the alias and whether it is for a table variable
+		for (size_t i=0; i<table_source_item.size(); ++i)
+		{
+			// First recurse deep until the lowest level of nested table_source_item, if any
+			if (table_source_item[i]->table_source_item().size() > 0)
+			{
+				if (dmlTargetIsTabvar(table_source_item[i]->table_source_item(), dmlTarget))
+					return true;
+			}
+
+			// At this point, there is no deeper nested table_source_item so we can look for the alias
+
+			// Find alias matching the DML target
+			if (table_source_item[i]->as_table_alias().size() == 0)
+				continue;  // No alias found in this table_source_item
+
+			std::string alias = getFullText(table_source_item[i]->as_table_alias()[0]->table_alias());
+			if (pg_strcasecmp(alias.c_str(), dmlTarget.c_str()) != 0)
+				continue;  // This alias is not matching the DML target
+
+			if (table_source_item[i]->local_id())
+			{
+				// Table variable, with alias matching the DML target
+				return true;
+			}
+		}
+
+		// We did not find the target alias, or it was not a table variable
+		return false;
+	}
+
 	void exitSelect_statement(TSqlParser::Select_statementContext *selectCtx) override
 	{
 		if (statementMutator)

diff --git a/test/JDBC/expected/tabvar_in_function-vu-cleanup.out b/test/JDBC/expected/tabvar_in_function-vu-cleanup.out
@@ -12,5 +12,17 @@ drop function f4_tabvar_in_function_del
 go
 drop function f4_tabvar_in_function_upd
 go
+drop function f6_tabvar_in_function_del
+go
+drop function f6_tabvar_in_function_upd
+go
+drop function f6a_tabvar_in_function_del
+go
+drop function f6a_tabvar_in_function_upd
+go
+drop function f7_tabvar_in_function_del
+go
+drop function f7_tabvar_in_function_upd
+go
 drop table t1_tabvar_in_function_ins
 go
diff --git a/test/JDBC/expected/tabvar_in_function-vu-verify.out b/test/JDBC/expected/tabvar_in_function-vu-verify.out
@@ -195,3 +195,170 @@ go
 ~~ERROR (Message: 'DELETE' cannot be used within a function)~~
 
 
+create function f6_tabvar_in_function_upd(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update t set a = t.a + 1 from @tv as t join t1_tabvar_in_function_ins as tv on tv.a = t.a
+    return
+end
+go
+select * from dbo.f6_tabvar_in_function_upd(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f6_tabvar_in_function_del(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update @tv set a = a + 1
+    insert @tv values(@p1)
+    delete t from @tv as t inner join t1_tabvar_in_function_ins as tv on tv.a = t.a
+    return
+end
+go
+select * from dbo.f6_tabvar_in_function_del(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f6a_tabvar_in_function_upd(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update t set a = t.a + 1 from t1_tabvar_in_function_ins as tv join @tv as t on tv.a = t.a
+    return
+end
+go
+select * from dbo.f6a_tabvar_in_function_upd(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f6a_tabvar_in_function_del(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update @tv set a = a + 1
+    insert @tv values(@p1)
+    delete t from t1_tabvar_in_function_ins as tv join @tv as t on tv.a = t.a
+    return
+end
+go
+select * from dbo.f6a_tabvar_in_function_del(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f7_tabvar_in_function_upd(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update t set a = t.a + 1 from t1_tabvar_in_function_ins as t1 left join t1_tabvar_in_function_ins as tv on tv.a = t1.a right join @tv t on tv.a = t.a 
+    return
+end
+go
+select * from dbo.f7_tabvar_in_function_upd(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f7_tabvar_in_function_del(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update @tv set a = a + 1
+    insert @tv values(@p1)
+    delete t from t1_tabvar_in_function_ins as t1 join t1_tabvar_in_function_ins as tv on tv.a = t1.a join @tv t on tv.a = t.a 
+    return
+end
+go
+select * from dbo.f7_tabvar_in_function_del(123)
+go
+~~START~~
+int
+124
+~~END~~
+
+
+create function f8_tabvar_in_function_upd(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update t set a = t.a + 1 from @tv as tv cross join t1_tabvar_in_function_ins as t
+    return
+end
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: 'UPDATE' cannot be used within a function)~~
+
+
+create function f8_tabvar_in_function_del(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update @tv set a = a + 1
+    insert @tv values(@p1)
+    delete t from @tv as tv inner join t1_tabvar_in_function_ins as t on tv.a = t.a
+    return
+end
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: 'DELETE' cannot be used within a function)~~
+
+
+create function f9_tabvar_in_function_upd(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update t set a = t.a + 1 from @tv as t1 full outer join t1_tabvar_in_function_ins as tv on tv.a = t1.a cross join t1_tabvar_in_function_ins t
+    return
+end
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: 'UPDATE' cannot be used within a function)~~
+
+
+create function f9_tabvar_in_function_del(@p1 int)
+returns @tv table(a int)
+as
+begin
+    insert @tv values(@p1)
+    update @tv set a = a + 1
+    insert @tv values(@p1)
+    delete t from @tv as t1 right join t1_tabvar_in_function_ins as tv on tv.a = t1.a full outer join t1_tabvar_in_function_ins t on tv.a = t.a 
+    return
+end
+go
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: 'DELETE' cannot be used within a function)~~
+
diff --git a/test/JDBC/input/tabvar_in_function-vu-cleanup.sql b/test/JDBC/input/tabvar_in_function-vu-cleanup.sql
@@ -12,5 +12,17 @@ drop function f4_tabvar_in_function_del
 go
 drop function f4_tabvar_in_function_upd
 go
+drop function f6_tabvar_in_function_del
+go
+drop function f6_tabvar_in_function_upd
+go
+drop function f6a_tabvar_in_function_del
+go
+drop function f6a_tabvar_in_function_upd
+go
+drop function f7_tabvar_in_function_del
+go
+drop function f7_tabvar_in_function_upd
+go
 drop table t1_tabvar_in_function_ins
 go