From 0b401ee7b3b08eb84903ac283b65d6d071c77d06 Mon Sep 17 00:00:00 2001
From: Anju Bharti <66729219+anju15bharti@users.noreply.github.com>
Date: Wed, 31 Jul 2024 08:57:13 +0530
Subject: [PATCH] Block cross-db 'SELECT ... INTO table FROM ...' statement
 (#2744) (#2780)

Before this commit, Cross-db 'SELECT ... INTO table FROM ...' was allowed to be executed unintentionally since it was missed to be blocked when cross-db DMLs got supported. Due to this unblocked behaviour, It creates a table whose owner is current session's login since internally cross-db statement execution sets current role to current session's login.

This commit blocks cross-db 'SELECT ... INTO table FROM ...' statement by adding relevant check on ANTLR parser.

Task: BABEL-4934

Signed-off-by: Harsh Lunagariya lunharsh@amazon.com
---
 contrib/babelfishpg_tsql/src/tsqlIface.cpp |  16 ++
 test/JDBC/expected/BABEL-CROSS-DB.out      | 203 +++++++++++++++++++++
 test/JDBC/input/BABEL-CROSS-DB.mix         | 139 ++++++++++++++
 3 files changed, 358 insertions(+)

diff --git a/contrib/babelfishpg_tsql/src/tsqlIface.cpp b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
index 5b720dd746..760b1bd91a 100644
--- a/contrib/babelfishpg_tsql/src/tsqlIface.cpp
+++ b/contrib/babelfishpg_tsql/src/tsqlIface.cpp
@@ -1759,6 +1759,22 @@ class tsqlBuilder : public tsqlCommonMutator
 		if (is_schema_specified)
 			stmt->is_schema_specified = true;
 
+		if (is_cross_db)
+		{
+			if (ctx->select_statement_standalone() &&
+				ctx->select_statement_standalone()->select_statement() &&
+				ctx->select_statement_standalone()->select_statement()->query_expression() &&
+				ctx->select_statement_standalone()->select_statement()->query_expression()->query_specification() &&
+				ctx->select_statement_standalone()->select_statement()->query_expression()->query_specification()->INTO() &&
+				ctx->select_statement_standalone()->select_statement()->query_expression()->query_specification()->table_name())
+			{
+				throw PGErrorWrapperException(ERROR,
+						ERRCODE_FEATURE_NOT_SUPPORTED,
+						"cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.",
+						getLineAndPos(ctx->select_statement_standalone()));
+			}
+		}
+
 		if (is_compiling_create_function())
 		{
 			/* select without destination should be blocked. We can use already information about desitnation, which is already processed. */
diff --git a/test/JDBC/expected/BABEL-CROSS-DB.out b/test/JDBC/expected/BABEL-CROSS-DB.out
index 36c32036ab..d1d0ccf487 100644
--- a/test/JDBC/expected/BABEL-CROSS-DB.out
+++ b/test/JDBC/expected/BABEL-CROSS-DB.out
@@ -541,6 +541,209 @@ GO
 DROP DATABASE db2;
 GO
 
+-- BABEL-4934 Test blocking cross-db SELECT-INTO statement
+CREATE DATABASE db_4934_1;
+GO
+
+CREATE DATABASE db_4934_2;
+GO
+
+USE master;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_1;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_2;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_1;
+GO
+
+-- #4934.1 It should be blocked
+SELECT * INTO t222 FROM master.dbo.t1; 
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM master..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM db_4934_2..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM db_4934_2.dbo.t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM db_4934_2..t1, t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM t1, db_4934_1..t2;
+GO
+
+SELECT * INTO t222 FROM master..t1, db_4934_2..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM (SELECT * FROM master..t1) AS Subquery;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM (SELECT * FROM t1, master..t1) AS Subquery;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM (SELECT *, (SELECT * FROM master..t1) FROM t1) AS Subquery;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO t222 FROM (SELECT *, (SELECT * FROM t1) FROM master..t3) AS Subquery;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO master..t222 FROM db_4934_2..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO db_4934_1..t222 FROM db_4934_2..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+SELECT * INTO master..t222 FROM master..t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+-- #4934.2 Following statements will succeed across same database
+SELECT * INTO t3 FROM db_4934_1.dbo.t1;
+GO
+
+SELECT * INTO t4 FROM dbo.t1;
+GO
+
+SELECT * INTO t5 FROM db_4934_1..t1;
+GO
+
+SELECT * INTO t6 FROM db_4934_1..t1, db_4934_1..t2;
+GO
+
+SELECT * INTO t7 FROM (SELECT * FROM db_4934_1..t1) AS Subquery;
+GO
+
+SELECT * INTO t8 FROM (SELECT *, (SELECT * FROM db_4934_1..t2) FROM db_4934_1..t1) AS Subquery;
+GO
+
+-- validate the access
+SELECT * FROM t3, t4, t5, t6, t7, t8;
+GO
+~~START~~
+int#!#int#!#int#!#int#!#int#!#int#!#int#!#int
+~~END~~
+
+
+DROP TABLE t3, t4, t5, t6, t7, t8;
+GO
+
+-- #4934.3 Temporary table should not be blocked
+SELECT * INTO #t1 FROM db_4934_1.dbo.t1;
+GO
+
+SELECT * INTO #t2 FROM (SELECT * FROM db_4934_1.dbo.t1) AS Subquery;
+GO
+
+-- validate the access
+SELECT * FROM #t1, #t2;
+GO
+~~START~~
+int#!#int
+~~END~~
+
+
+DROP TABLE #t1, #t2;
+GO
+
+-- Even though this is same as above statement, this will still fail since
+-- internally it considers as cross-db statement (This behaviour is general to
+-- all of the applicable DMLs)
+SELECT * INTO tempdb..#t3 FROM db_4934_1.dbo.t1;
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: cross-db 'SELECT INTO' statement is not supported in Babelfish. As a workaround, consider running CREATE TABLE and INSERT-SELECT statements.)~~
+
+
+
+USE master;
+GO
+
+DROP TABLE t1, t2;
+GO
+
+DROP DATABASE db_4934_1;
+GO
+
+DROP DATABASE db_4934_2;
+GO
+
 -- psql
 ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'single-db';
 SELECT pg_reload_conf();
diff --git a/test/JDBC/input/BABEL-CROSS-DB.mix b/test/JDBC/input/BABEL-CROSS-DB.mix
index c86e7cd128..dd519efb70 100644
--- a/test/JDBC/input/BABEL-CROSS-DB.mix
+++ b/test/JDBC/input/BABEL-CROSS-DB.mix
@@ -359,6 +359,145 @@ GO
 DROP DATABASE db2;
 GO
 
+-- BABEL-4934 Test blocking cross-db SELECT-INTO statement
+CREATE DATABASE db_4934_1;
+GO
+
+CREATE DATABASE db_4934_2;
+GO
+
+USE master;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_1;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_2;
+GO
+
+CREATE TABLE t1(a int);
+GO
+
+CREATE TABLE t2(b int);
+GO
+
+USE db_4934_1;
+GO
+
+-- #4934.1 It should be blocked
+SELECT * INTO t222 FROM master.dbo.t1; 
+GO
+
+SELECT * INTO t222 FROM master..t1;
+GO
+
+SELECT * INTO t222 FROM db_4934_2..t1;
+GO
+
+SELECT * INTO t222 FROM db_4934_2.dbo.t1;
+GO
+
+SELECT * INTO t222 FROM db_4934_2..t1, t1;
+GO
+
+SELECT * INTO t222 FROM t1, db_4934_1..t2;
+GO
+
+SELECT * INTO t222 FROM master..t1, db_4934_2..t1;
+GO
+
+SELECT * INTO t222 FROM (SELECT * FROM master..t1) AS Subquery;
+GO
+
+SELECT * INTO t222 FROM (SELECT * FROM t1, master..t1) AS Subquery;
+GO
+
+SELECT * INTO t222 FROM (SELECT *, (SELECT * FROM master..t1) FROM t1) AS Subquery;
+GO
+
+SELECT * INTO t222 FROM (SELECT *, (SELECT * FROM t1) FROM master..t3) AS Subquery;
+GO
+
+SELECT * INTO master..t222 FROM db_4934_2..t1;
+GO
+
+SELECT * INTO db_4934_1..t222 FROM db_4934_2..t1;
+GO
+
+SELECT * INTO master..t222 FROM master..t1;
+GO
+
+-- #4934.2 Following statements will succeed across same database
+SELECT * INTO t3 FROM db_4934_1.dbo.t1;
+GO
+
+SELECT * INTO t4 FROM dbo.t1;
+GO
+
+SELECT * INTO t5 FROM db_4934_1..t1;
+GO
+
+SELECT * INTO t6 FROM db_4934_1..t1, db_4934_1..t2;
+GO
+
+SELECT * INTO t7 FROM (SELECT * FROM db_4934_1..t1) AS Subquery;
+GO
+
+SELECT * INTO t8 FROM (SELECT *, (SELECT * FROM db_4934_1..t2) FROM db_4934_1..t1) AS Subquery;
+GO
+
+-- validate the access
+SELECT * FROM t3, t4, t5, t6, t7, t8;
+GO
+
+DROP TABLE t3, t4, t5, t6, t7, t8;
+GO
+
+-- #4934.3 Temporary table should not be blocked
+SELECT * INTO #t1 FROM db_4934_1.dbo.t1;
+GO
+
+SELECT * INTO #t2 FROM (SELECT * FROM db_4934_1.dbo.t1) AS Subquery;
+GO
+
+-- validate the access
+SELECT * FROM #t1, #t2;
+GO
+
+DROP TABLE #t1, #t2;
+GO
+
+-- Even though this is same as above statement, this will still fail since
+-- internally it considers as cross-db statement (This behaviour is general to
+-- all of the applicable DMLs)
+SELECT * INTO tempdb..#t3 FROM db_4934_1.dbo.t1;
+GO
+
+
+USE master;
+GO
+
+DROP TABLE t1, t2;
+GO
+
+DROP DATABASE db_4934_1;
+GO
+
+DROP DATABASE db_4934_2;
+GO
+
 -- psql
 ALTER SYSTEM SET babelfishpg_tsql.migration_mode = 'single-db';
 SELECT pg_reload_conf();