fix object resolution inside pltsql routines

Signed-off-by: Tanzeel Khan <[email protected]>

contrib/babelfishpg_tsql/sql/sys_functions.sql

@@ -2403,15 +2403,15 @@ LANGUAGE C PARALLEL SAFE IMMUTABLE;
 
 CREATE OR REPLACE FUNCTION sys.db_id() RETURNS SMALLINT
 AS 'babelfishpg_tsql', 'babelfish_db_id'
-LANGUAGE C PARALLEL SAFE IMMUTABLE;
+LANGUAGE C PARALLEL SAFE STABLE;
 
 CREATE OR REPLACE FUNCTION sys.db_name(int) RETURNS sys.nvarchar(128)
 AS 'babelfishpg_tsql', 'babelfish_db_name'
 LANGUAGE C PARALLEL SAFE IMMUTABLE;
 
 CREATE OR REPLACE FUNCTION sys.db_name() RETURNS sys.nvarchar(128)
 AS 'babelfishpg_tsql', 'babelfish_db_name'
-LANGUAGE C PARALLEL SAFE IMMUTABLE;
+LANGUAGE C PARALLEL SAFE STABLE;
 
 CREATE OR REPLACE FUNCTION sys.exp(IN arg DOUBLE PRECISION)
 RETURNS DOUBLE PRECISION

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--5.0.0--5.1.0.sql

@@ -39,6 +39,13 @@ $$;
  * So make sure that any SQL statement (DDL/DML) being added here can be executed multiple times without affecting
  * final behaviour.
  */
+CREATE OR REPLACE FUNCTION sys.db_id() RETURNS SMALLINT
+AS 'babelfishpg_tsql', 'babelfish_db_id'
+LANGUAGE C PARALLEL SAFE STABLE;
+
+CREATE OR REPLACE FUNCTION sys.db_name() RETURNS sys.nvarchar(128)
+AS 'babelfishpg_tsql', 'babelfish_db_name'
+LANGUAGE C PARALLEL SAFE STABLE;
 
 -- After upgrade, always run analyze for all babelfish catalogs.
 CALL sys.analyze_babelfish_catalogs();

contrib/babelfishpg_tsql/src/catalog.c

@@ -1001,47 +1001,33 @@ get_authid_user_ext_physical_name(const char *db_name, const char *login)
 char *
 get_authid_user_ext_schema_name(const char *db_name, const char *user)
 {
-	Relation	bbf_authid_user_ext_rel;
 	HeapTuple	tuple_user_ext;
-	ScanKeyData key[2];
-	TableScanDesc scan;
 	char	   *schema_name = NULL;
-	NameData   *user_name;
 
 	if (!db_name || !user)
 		return NULL;
 
-	bbf_authid_user_ext_rel = table_open(get_authid_user_ext_oid(),
-										 RowExclusiveLock);
+	tuple_user_ext = SearchSysCache1(AUTHIDUSEREXTROLENAME, CStringGetDatum(user));
 
-	user_name = (NameData *) palloc0(NAMEDATALEN);
-	snprintf(user_name->data, NAMEDATALEN, "%s", user);
-	ScanKeyInit(&key[0],
-				Anum_bbf_authid_user_ext_rolname,
-				BTEqualStrategyNumber, F_NAMEEQ,
-				NameGetDatum(user_name));
-	ScanKeyInit(&key[1],
-				Anum_bbf_authid_user_ext_database_name,
-				BTEqualStrategyNumber, F_TEXTEQ,
-				CStringGetTextDatum(db_name));
-
-	scan = table_beginscan_catalog(bbf_authid_user_ext_rel, 2, key);
-
-	tuple_user_ext = heap_getnext(scan, ForwardScanDirection);
 	if (HeapTupleIsValid(tuple_user_ext))
 	{
-		Datum		datum;
-		bool		is_null;
-
-		datum = heap_getattr(tuple_user_ext,
-							 Anum_bbf_authid_user_ext_default_schema_name,
-							 bbf_authid_user_ext_rel->rd_att,
-							 &is_null);
-		schema_name = pstrdup(TextDatumGetCString(datum));
-	}
+		Datum datum = SysCacheGetAttrNotNull(AUTHIDUSEREXTROLENAME, tuple_user_ext,
+											 Anum_bbf_authid_user_ext_database_name);
+		char  *db_name_cstring = TextDatumGetCString(datum);
 
-	table_endscan(scan);
-	table_close(bbf_authid_user_ext_rel, RowExclusiveLock);
+		if (strcmp(db_name_cstring, db_name) == 0)
+		{
+			Datum 	schema_datum = SysCacheGetAttrNotNull(AUTHIDUSEREXTROLENAME, tuple_user_ext,
+											   Anum_bbf_authid_user_ext_default_schema_name);
+			char  	*default_schema_name = TextDatumGetCString(schema_datum);
+
+			if (strlen(default_schema_name))
+				schema_name = default_schema_name;
+		}
+
+		pfree(db_name_cstring);
+		ReleaseSysCache(tuple_user_ext);
+	}
 
 	return schema_name;
 }

contrib/babelfishpg_tsql/src/hooks.c

@@ -188,7 +188,6 @@ static bool pltsql_bbfViewHasInsteadofTrigger(Relation view, CmdType event);
 static bool plsql_TriggerRecursiveCheck(ResultRelInfo *resultRelInfo);
 static bool bbf_check_rowcount_hook(int es_processed);
 
-static char *get_local_schema_for_bbf_functions(Oid proc_nsp_oid);
 extern bool called_from_tsql_insert_exec();
 extern bool called_for_tsql_itvf_func();
 static void is_function_pg_stat_valid(FunctionCallInfo fcinfo,
@@ -290,7 +289,6 @@ static table_variable_satisfies_update_hook_type prev_table_variable_satisfies_u
 static table_variable_satisfies_vacuum_hook_type prev_table_variable_satisfies_vacuum = NULL;
 static table_variable_satisfies_vacuum_horizon_hook_type prev_table_variable_satisfies_vacuum_horizon = NULL;
 static drop_relation_refcnt_hook_type prev_drop_relation_refcnt_hook = NULL;
-static set_local_schema_for_func_hook_type prev_set_local_schema_for_func_hook = NULL;
 static bbf_get_sysadmin_oid_hook_type prev_bbf_get_sysadmin_oid_hook = NULL;
 static get_bbf_admin_oid_hook_type prev_get_bbf_admin_oid_hook = NULL;
 static transform_pivot_clause_hook_type pre_transform_pivot_clause_hook = NULL;
@@ -473,9 +471,6 @@ InstallExtendedHooks(void)
 	prev_drop_relation_refcnt_hook = drop_relation_refcnt_hook;
 	drop_relation_refcnt_hook = pltsql_drop_relation_refcnt_hook;
 
-	prev_set_local_schema_for_func_hook = set_local_schema_for_func_hook;
-	set_local_schema_for_func_hook = get_local_schema_for_bbf_functions;
-
 	prev_bbf_get_sysadmin_oid_hook = bbf_get_sysadmin_oid_hook;
 	bbf_get_sysadmin_oid_hook = get_sysadmin_oid;
 
@@ -586,7 +581,6 @@ UninstallExtendedHooks(void)
 	IsToastRelationHook = PrevIsToastRelationHook;
 	IsToastClassHook = PrevIsToastClassHook;
 	drop_relation_refcnt_hook = prev_drop_relation_refcnt_hook;
-	set_local_schema_for_func_hook = prev_set_local_schema_for_func_hook;
 	bbf_get_sysadmin_oid_hook = prev_bbf_get_sysadmin_oid_hook;
 	get_bbf_admin_oid_hook = prev_get_bbf_admin_oid_hook;
 	transform_pivot_clause_hook = pre_transform_pivot_clause_hook;
@@ -5093,37 +5087,6 @@ sort_nulls_first(SortGroupClause * sortcl, bool reverse)
 	}
 }
 
-
-static char *
-get_local_schema_for_bbf_functions(Oid proc_nsp_oid)
-{
-	HeapTuple 	 	tuple;
-	char 			*func_schema_name = NULL,
-					*new_search_path = NULL;
-	char  			*func_dbo_schema;
-	const char		*cur_dbname = get_cur_db_name();
-
-	tuple = SearchSysCache1(NAMESPACEOID,
-						ObjectIdGetDatum(proc_nsp_oid));
-	if(HeapTupleIsValid(tuple))
-	{
-		func_schema_name = NameStr(((Form_pg_namespace) GETSTRUCT(tuple))->nspname);
-		func_dbo_schema = get_dbo_schema_name(cur_dbname);
-
-		if(strcmp(func_schema_name, func_dbo_schema) != 0
-			&& strcmp(func_schema_name, "sys") != 0)
-			new_search_path = psprintf("%s, %s, \"$user\", sys, pg_catalog",
-										quote_identifier(func_schema_name),
-										quote_identifier(func_dbo_schema));
-
-		ReleaseSysCache(tuple);
-
-		pfree(func_dbo_schema);
-	}
-
-	return new_search_path;
-}
-
 static ResTarget *
 make_restarget_from_cstr_list(List * l)
 {

contrib/babelfishpg_tsql/src/iterative_exec.c

@@ -26,6 +26,7 @@ static void restore_ctx_full(PLtsql_execstate *estate);
 static ErrorData *restore_ctx_partial1(PLtsql_execstate *estate);
 static void restore_ctx_partial2(PLtsql_execstate *estate);
 static void set_exec_error_data(char *procedure, int number, int severity, int state, bool rethrow);
+static void set_search_path_for_pltsql_stmt(PLtsql_stmt *stmt);
 static void reset_exec_error_data(PLtsql_execstate *estate);
 static void assert_equal_estate_err(PLtsql_estate_err *err1, PLtsql_estate_err *err2);
 static void read_raiserror_params(PLtsql_execstate *estate, List *params, int paramno,
@@ -46,6 +47,9 @@ extern PLtsql_estate_err *pltsql_clone_estate_err(PLtsql_estate_err *err);
 extern void prepare_format_string(StringInfo buf, char *msg_string, int nargs,
 								  Datum *args, Oid *argtypes, bool *argisnull);
 
+char	   *pltsql_search_path = NULL;
+char	   *current_db_search_path = NULL;
+
 static int
 exec_stmt_goto(PLtsql_execstate *estate, PLtsql_stmt_goto *stmt)
 {
@@ -645,6 +649,8 @@ dispatch_stmt(PLtsql_execstate *estate, PLtsql_stmt *stmt)
 	/* reset number of tuple processed in previous command */
 	estate->eval_processed = 0;
 
+	set_search_path_for_pltsql_stmt(stmt);
+
 	switch (stmt->cmd_type)
 	{
 		case PLTSQL_STMT_ASSIGN:
@@ -2184,3 +2190,78 @@ send_env_change_token_on_txn_abort(void)
 	if (*pltsql_protocol_plugin_ptr && (*pltsql_protocol_plugin_ptr)->send_env_change_binary)
 		((*pltsql_protocol_plugin_ptr)->send_env_change_binary) (ENVCHANGE_ROLLBACKTXN, NULL, 0, &txnId, sizeof(uint64_t));
 }
+
+/*
+ * We will keep verifying the search path is correct
+ * before executing any pltsql statement
+ */
+static void
+set_search_path_for_pltsql_stmt(PLtsql_stmt *stmt)
+{
+	char		*cur_dbname = get_cur_db_name();
+	char 		*new_search_path = NULL;
+
+	if (!IS_TDS_CONN())
+	{
+		pfree(cur_dbname);
+		return;
+	}
+
+	if(stmt->cmd_type == PLTSQL_STMT_EXECSQL)
+	{
+		PLtsql_stmt_execsql *execsqlstmt = (PLtsql_stmt_execsql *)stmt;
+		if (execsqlstmt->is_create_view && execsqlstmt->schema_name != NULL && strcmp(execsqlstmt->schema_name, "sys") != 0
+			&& strcmp(execsqlstmt->schema_name, "pg_catalog") != 0)
+		{
+			char *physical_schema = get_physical_schema_name(cur_dbname, execsqlstmt->schema_name);
+			char *dbo_schema = get_dbo_schema_name(cur_dbname);
+
+			new_search_path = psprintf(PLTSQL_SEARCH_PATH_BUFFER, quote_identifier(physical_schema), quote_identifier(dbo_schema));
+
+			pfree(physical_schema);
+			pfree(dbo_schema);
+		}
+		else if (execsqlstmt->is_ddl)
+		{
+			char *dbo_schema = get_dbo_schema_name(cur_dbname);
+			char *schema = get_authid_user_ext_schema_name(cur_dbname, GetUserNameFromId(GetUserId(), false));
+			char *physical_schema = NULL;
+
+			if (schema)
+			{
+				physical_schema = get_physical_schema_name(cur_dbname, schema);
+				pfree(schema);
+			}
+			new_search_path = psprintf(PLTSQL_SEARCH_PATH_BUFFER,
+							physical_schema ? quote_identifier(physical_schema) : quote_identifier(dbo_schema),
+							quote_identifier(dbo_schema));
+		}
+	}
+
+	PG_TRY();
+	{
+		pltsql_check_search_path = false;
+
+		/* Override case */
+		if (new_search_path)
+		{
+			SetConfigOption("search_path", new_search_path,
+							PGC_SUSET, PGC_S_SESSION);
+			pfree(new_search_path);
+		}
+		/* default search inside pltsql routines */
+		else if (pltsql_search_path)
+			SetConfigOption("search_path", pltsql_search_path,
+							PGC_SUSET, PGC_S_SESSION);
+		/* current active logical database search path */
+		else if (current_db_search_path)
+			SetConfigOption("search_path", current_db_search_path,
+							PGC_SUSET, PGC_S_SESSION);
+		}
+	PG_FINALLY();
+	{
+		pltsql_check_search_path = true;
+		pfree(cur_dbname);
+	}
+	PG_END_TRY();
+}

contrib/babelfishpg_tsql/src/pl_comp.c

@@ -42,6 +42,7 @@
 #include "pltsql.h"
 #include "pltsql-2.h"
 #include "analyzer.h"
+#include "catalog.h"
 #include "codegen.h"
 #include "iterative_exec.h"
 #include "multidb.h"
@@ -147,6 +148,7 @@ static void delete_function(PLtsql_function *func);
 
 extern Portal ActivePortal;
 extern bool pltsql_function_parse_error_transpose(const char *prosrc);
+static char *get_local_schema_for_bbf_functions(Oid proc_nsp_oid, int16 dbid);
 
 /* ----------
  * pltsql_compile		Make an execution tree for a PL/tsql function.
@@ -967,6 +969,18 @@ do_compile(FunctionCallInfo fcinfo,
 	for (i = 0; i < function->fn_nargs; i++)
 		function->fn_argvarnos[i] = in_arg_varnos[i];
 
+	/* store the logical db which contains this function */
+	if (IS_TDS_CONN() && procStruct->provolatile != PROVOLATILE_IMMUTABLE)
+	{
+		function->fn_dbid = get_dbid_from_physical_schema_name(get_namespace_name(procStruct->pronamespace), true);
+		function->fn_search_path = get_local_schema_for_bbf_functions(procStruct->pronamespace, function->fn_dbid);
+	}
+	else
+	{
+		function->fn_dbid = InvalidDbid;
+		function->fn_search_path = NULL;
+	}
+
 	pltsql_finish_datums(function);
 
 	/* Debug dump for completed functions */
@@ -1112,6 +1126,8 @@ pltsql_compile_inline(char *proc_source, InlineCodeBlockArgs *args)
 	function->fn_retbyval = true;
 	function->fn_rettyplen = sizeof(int32);
 	function->fn_tupdesc = NULL;
+	function->fn_dbid = InvalidDbid;
+	function->fn_search_path = NULL;
 
 	/*
 	 * Remember if function is STABLE/IMMUTABLE.  XXX would it be better to
@@ -3252,3 +3268,34 @@ reset_cached_batch()
 		delete_cached_batch(cur_handle_id--);
 	cur_handle_id = 1;
 }
+
+static char *
+get_local_schema_for_bbf_functions(Oid proc_nsp_oid, int16 dbid)
+{
+	HeapTuple 	 	tuple;
+	char 			*func_schema_name = NULL,
+					*new_search_path = NULL;
+	char  			*func_dbo_schema;
+
+	if (!IS_TDS_CONN() || !DbidIsValid(dbid) || proc_nsp_oid == get_sys_namespace_oid())
+		return NULL;
+
+	tuple = SearchSysCache1(NAMESPACEOID,
+						ObjectIdGetDatum(proc_nsp_oid));
+	if(HeapTupleIsValid(tuple))
+	{
+		func_schema_name = get_namespace_name(proc_nsp_oid);
+		func_dbo_schema = get_dbo_schema_name(get_db_name(dbid));
+
+		new_search_path = psprintf(PLTSQL_SEARCH_PATH_BUFFER,
+									quote_identifier(func_schema_name),
+									quote_identifier(func_dbo_schema));
+
+		ReleaseSysCache(tuple);
+
+		pfree(func_schema_name);
+		pfree(func_dbo_schema);
+	}
+
+	return new_search_path;
+}