Merge remote-tracking branch 'upstream/BABEL_4_X_DEV' into jira-babel…

…-5155

contrib/babelfishpg_tds/src/backend/tds/tds_srv.c

@@ -203,6 +203,7 @@ pe_tds_init(void)
 	pltsql_plugin_handler_ptr->set_context_info = &set_tds_context_info;
 	pltsql_plugin_handler_ptr->get_datum_from_byte_ptr = &TdsBytePtrToDatum;
 	pltsql_plugin_handler_ptr->get_datum_from_date_time_struct = &TdsDateTimeTypeToDatum;
+	pltsql_plugin_handler_ptr->set_reset_tds_connection_flag = &SetResetTDSConnectionFlag;
 
 	invalidate_stat_table_hook = invalidate_stat_table;
 	guc_newval_hook = TdsSetGucStatVariable;

contrib/babelfishpg_tds/src/backend/tds/tdsprotocol.c

@@ -73,6 +73,7 @@ typedef ResetConnectionData *ResetConnection;
 TdsRequestCtrlData *TdsRequestCtrl = NULL;
 
 ResetConnection resetCon = NULL;
+static bool resetTdsConnectionFlag = false;
 
 /* Local functions */
 static void ResetTDSConnection(void);
@@ -154,13 +155,28 @@ ResetTDSConnection(void)
 	TdsResetCache();
 	TdsResponseReset();
 	TdsResetBcpOffset();
-	SetConfigOption("default_transaction_isolation", isolationOld,
-					PGC_BACKEND, PGC_S_CLIENT);
+	/* Retore previous isolation level when not called by sys.sp_reset_connection */
+	if (!resetTdsConnectionFlag)
+	{
+		SetConfigOption("default_transaction_isolation", isolationOld,
+						PGC_BACKEND, PGC_S_CLIENT);
+	}
 
 	tvp_lookup_list = NIL;
 
-	/* send an environement change token */
-	TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL);
+	/* send an environement change token is its not called via sys.sp_reset_connection procedure */
+	if (!resetTdsConnectionFlag)
+	{
+		TdsSendEnvChange(TDS_ENVID_RESETCON, NULL, NULL);
+	}
+}
+
+/*
+ * SetResetTDSConnectionFlag - Sets reset tds connection flag
+ */
+void SetResetTDSConnectionFlag()
+{
+	resetTdsConnectionFlag = true;	
 }
 
 /*
@@ -674,6 +690,12 @@ TdsSocketBackend(void)
 						/* Ready to fetch the next request */
 						TdsRequestCtrl->phase = TDS_REQUEST_PHASE_FETCH;
 
+						if (resetTdsConnectionFlag)
+						{
+							ResetTDSConnection();
+							resetTdsConnectionFlag = false;
+						}
+
 						break;
 					}
 				case TDS_REQUEST_PHASE_ERROR:

contrib/babelfishpg_tds/src/include/tds_protocol.h

@@ -77,5 +77,6 @@ typedef struct
 } TdsRequestCtrlData;
 
 extern TdsRequestCtrlData *TdsRequestCtrl;
+extern void SetResetTDSConnectionFlag(void);
 
 #endif							/* TDS_PROTOCOL_H */

contrib/babelfishpg_tsql/sql/sys_procedures.sql

@@ -329,3 +329,7 @@ GRANT EXECUTE ON PROCEDURE sys.sp_dropextendedproperty TO PUBLIC;
 CREATE OR REPLACE PROCEDURE sys.sp_enum_oledb_providers()
 AS 'babelfishpg_tsql', 'sp_enum_oledb_providers_internal' LANGUAGE C;
 GRANT EXECUTE on PROCEDURE sys.sp_enum_oledb_providers() TO PUBLIC;
+
+CREATE OR REPLACE PROCEDURE sys.sp_reset_connection()
+AS 'babelfishpg_tsql', 'sp_reset_connection_internal' LANGUAGE C;
+GRANT EXECUTE ON PROCEDURE sys.sp_reset_connection() TO PUBLIC;

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--4.3.0--4.4.0.sql

@@ -1910,6 +1910,10 @@ CALL sys.babelfish_drop_deprecated_object('function', 'sys', 'sp_pkeys_internal_
 -- Please have this be one of the last statements executed in this upgrade script.
 DROP PROCEDURE sys.babelfish_drop_deprecated_object(varchar, varchar, varchar);
 
+CREATE OR REPLACE PROCEDURE sys.sp_reset_connection()
+AS 'babelfishpg_tsql', 'sp_reset_connection_internal' LANGUAGE C;
+GRANT EXECUTE ON PROCEDURE sys.sp_reset_connection() TO PUBLIC;
+
 -- After upgrade, always run analyze for all babelfish catalogs.
 CALL sys.analyze_babelfish_catalogs();
 

contrib/babelfishpg_tsql/src/dbcmds.c

@@ -88,12 +88,16 @@ gen_createdb_subcmds(const char *dbname, const char *owner)
 	const char     *db_owner;
 	const char     *guest;
 	const char     *guest_schema;
+	Oid       	owner_oid;
+	bool     	owner_is_sa;
 
 	schema = get_dbo_schema_name(dbname);
 	dbo = get_dbo_role_name(dbname);
 	db_owner = get_db_owner_name(dbname);
 	guest = get_guest_role_name(dbname);
 	guest_schema = get_guest_schema_name(dbname);
+	owner_oid = get_role_oid(owner, true);
+	owner_is_sa = role_is_sa(owner_oid);
 
 	/*
 	 * To avoid SQL injection, we generate statement parsetree with dummy
@@ -104,7 +108,10 @@ gen_createdb_subcmds(const char *dbname, const char *owner)
 	appendStringInfo(&query, "CREATE ROLE dummy CREATEROLE INHERIT; ");
 	appendStringInfo(&query, "CREATE ROLE dummy INHERIT CREATEROLE ROLE sysadmin IN ROLE dummy; ");
 	appendStringInfo(&query, "GRANT CREATE, CONNECT, TEMPORARY ON DATABASE dummy TO dummy; ");
-	appendStringInfo(&query, "GRANT dummy TO dummy; ");
+
+	/* Only grant dbo to owner if owner is not master user  */
+	if (!owner_is_sa)
+		appendStringInfo(&query, "GRANT dummy TO dummy; ");
 
 	if (guest)
 	{
@@ -125,9 +132,19 @@ gen_createdb_subcmds(const char *dbname, const char *owner)
 	res = raw_parser(query.data, RAW_PARSE_DEFAULT);
 
 	if (guest)
-		expected_stmt_num = list_length(logins) > 0 ? 10 : 9;
+	{
+		if (!owner_is_sa)
+			expected_stmt_num = list_length(logins) > 0 ? 10 : 9;	
+		else
+			expected_stmt_num = list_length(logins) > 0 ? 9 : 8;
+	}
 	else
-		expected_stmt_num = 7;
+	{
+		expected_stmt_num = 6;
+
+		if (!owner_is_sa)
+			expected_stmt_num++;
+	}
 
 	if (list_length(res) != expected_stmt_num)
 		ereport(ERROR,
@@ -145,11 +162,13 @@ gen_createdb_subcmds(const char *dbname, const char *owner)
 	stmt = parsetree_nth_stmt(res, i++);
 	update_GrantStmt(stmt, get_database_name(MyDatabaseId), NULL, dbo, NULL);
 
-	/* Grant dbo role to owner */
-	stmt = parsetree_nth_stmt(res, i++);
-
-	update_GrantRoleStmt(stmt, list_make1(make_accesspriv_node(dbo)),
-						 list_make1(make_rolespec_node(owner)));
+	if (!owner_is_sa)
+	{
+		/* Grant dbo role to owner */
+		stmt = parsetree_nth_stmt(res, i++);
+		update_GrantRoleStmt(stmt, list_make1(make_accesspriv_node(dbo)),
+							list_make1(make_rolespec_node(owner)));
+	}
 
 	if (guest)
 	{

contrib/babelfishpg_tsql/src/pltsql.h

@@ -1797,6 +1797,8 @@ typedef struct PLtsql_protocol_plugin
 
 	char	   *(*get_physical_schema_name) (char *db_name, const char *schema_name);
 
+	void		(*set_reset_tds_connection_flag) ();
+
 	/* Session level GUCs */
 	bool		quoted_identifier;
 	bool		arithabort;
@@ -1810,7 +1812,6 @@ typedef struct PLtsql_protocol_plugin
 	int			datefirst;
 	int			lock_timeout;
 	const char *language;
-
 } PLtsql_protocol_plugin;
 
 /*

contrib/babelfishpg_tsql/src/procedures.c

@@ -78,6 +78,7 @@ PG_FUNCTION_INFO_V1(sp_babelfish_volatility);
 PG_FUNCTION_INFO_V1(sp_rename_internal);
 PG_FUNCTION_INFO_V1(sp_execute_postgresql);
 PG_FUNCTION_INFO_V1(sp_enum_oledb_providers_internal);
+PG_FUNCTION_INFO_V1(sp_reset_connection_internal);
 PG_FUNCTION_INFO_V1(sp_renamedb_internal);
 
 extern void delete_cached_batch(int handle);
@@ -4198,3 +4199,14 @@ sp_enum_oledb_providers_internal(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+Datum
+sp_reset_connection_internal(PG_FUNCTION_ARGS)
+{
+	if (*pltsql_protocol_plugin_ptr) 
+	{
+		(*pltsql_protocol_plugin_ptr)->set_reset_tds_connection_flag();
+	}
+
+	PG_RETURN_VOID();
+}

contrib/babelfishpg_tsql/src/rolecmds.c

@@ -588,6 +588,14 @@ grant_revoke_dbo_to_login(const char* login, const char* db_name, bool is_grant)
 	PlannedStmt *wrapper;
 
 	const char *dbo_role_name = get_dbo_role_name(db_name);
+
+	/*
+	 * If login i.e old_owner/new_owner is master user 
+	 * then skip grant/revoke dbo to login
+	 * since it will always be the member of sysadmin.
+	 */
+	if (role_is_sa(get_role_oid(login, true)))
+		return;
 
 	initStringInfo(&query);
 

test/JDBC/expected/Test-sp_reset_connection.out

@@ -0,0 +1,109 @@
+-- 1. Test resets GUC variables
+SET lock_timeout 0;
+GO
+SELECT @@lock_timeout;
+GO
+~~START~~
+int
+0
+~~END~~
+
+EXEC sys.sp_reset_connection
+-- TODO: GUC is not resetting
+SELECT @@lock_timeout;
+GO
+~~START~~
+int
+0
+~~END~~
+
+
+-- 2. Test open transactions are aborted on reset
+DROP TABLE IF EXISTS sp_reset_connection_test_table;
+CREATE TABLE sp_reset_connection_test_table(id int);
+BEGIN TRANSACTION
+INSERT INTO sp_reset_connection_test_table VALUES(1)
+GO
+~~ROW COUNT: 1~~
+
+EXEC sys.sp_reset_connection
+GO
+COMMIT TRANSACTION
+GO
+~~ERROR (Code: 3902)~~
+
+~~ERROR (Message: COMMIT can only be used in transaction blocks)~~
+
+SELECT * FROM sp_reset_connection_test_table
+GO
+~~START~~
+int
+~~END~~
+
+
+-- 3. Test temp tables are deleted on reset
+CREATE TABLE #babel_temp_table (ID INT identity(1,1), Data INT)
+INSERT INTO #babel_temp_table (Data) VALUES (100), (200), (300)
+GO
+~~ROW COUNT: 3~~
+
+SELECT * from #babel_temp_table
+GO
+~~START~~
+int#!#int
+1#!#100
+2#!#200
+3#!#300
+~~END~~
+
+EXEC sys.sp_reset_connection
+GO
+SELECT * from #babel_temp_table
+GO
+~~ERROR (Code: 33557097)~~
+
+~~ERROR (Message: relation "#babel_temp_table" does not exist)~~
+
+
+-- 4. Test isolation level is reset
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+1
+~~END~~
+
+EXEC sys.sp_reset_connection
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+2
+~~END~~
+
+
+-- 5. Test sp_reset_connection called with sp_prepexec
+SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+1
+~~END~~
+
+DECLARE @handle int;
+EXEC SP_PREPARE @handle output, NULL, N'exec sys.sp_reset_connection'
+EXEC SP_EXECUTE @handle
+GO
+GO
+select transaction_isolation_level from sys.dm_exec_sessions where session_id=@@SPID
+GO
+~~START~~
+smallint
+2
+~~END~~
+