ALTER PROCEDURE throwing unexpected state error (#2725)

This change prevents certain alter procedure calls in PSQL dialect from attempting to use the TSQL alter procedure implementation by adding a dialect check. Procedures created in PSQL and altered with security definer would throw the following error ERROR: StartTransactionCommand: unexpected state STARTED Task: BABEL-5074 Signed-off-by: Jake Owen <[email protected]>
babelfish-for-postgresql · Jul 9, 2024 · d33f8b4 · d33f8b4
1 parent 4ea4dd5
commit d33f8b4
Show file tree

Hide file tree

Showing 3 changed files with 157 additions and 80 deletions.
diff --git a/contrib/babelfishpg_tsql/src/pl_handler.c b/contrib/babelfishpg_tsql/src/pl_handler.c
@@ -2314,95 +2314,106 @@ bbf_ProcessUtility(PlannedStmt *pstmt,
 	{
 		case T_AlterFunctionStmt:
 			{
-				/*
-				* For ALTER PROC, we will:
-				* 1. Save important pg_proc metadata from the current proc (oid, proacl)
-				* 2. drop the current proc
-				* 3. create the new proc
-				* 4. update the pg_proc entry for the new proc with metadata from the old proc
-				* 5. update the babelfish_function_ext entry for the existing proc with new metadata based on the new proc
-				*/
-				AlterFunctionStmt *stmt = (AlterFunctionStmt *) parsetree;
-				bool 				isCompleteQuery = (context != PROCESS_UTILITY_SUBCOMMAND);
-				bool 				needCleanup;
-				Oid					oldoid;
-				Acl					*proacl;
-				bool				isSameProc;
-				ObjectAddress 		address;
-				CreateFunctionStmt	*cfs;
-				ListCell 			*option, *location_cell = NULL;
-				int 				origname_location = -1;
-
-				if (stmt->objtype != OBJECT_PROCEDURE)
-					break;
+				if (sql_dialect == SQL_DIALECT_TSQL)
+				{
+				       /*
+					* For ALTER PROC, we will:
+					* 1. Save important pg_proc metadata from the current proc (oid, proacl)
+					* 2. drop the current proc
+					* 3. create the new proc
+					* 4. update the pg_proc entry for the new proc with metadata from the old proc
+					* 5. update the babelfish_function_ext entry for the existing proc with new metadata based on the new proc
+					*/
+					AlterFunctionStmt *stmt = (AlterFunctionStmt *) parsetree;
+					bool 				isCompleteQuery = (context != PROCESS_UTILITY_SUBCOMMAND);
+					bool 				needCleanup;
+					Oid					oldoid;
+					Acl					*proacl;
+					bool				isSameProc;
+					ObjectAddress 		address;
+					CreateFunctionStmt	*cfs;
+					ListCell 			*option, *location_cell = NULL;
+					int 				origname_location = -1;
+
+					if (!IS_TDS_CLIENT())
+					{
+						ereport(ERROR,
+								(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+								errmsg("TSQL ALTER PROCEDURE is not supported from PostgreSQL endpoint.")));
+					}
 
-				/* All event trigger calls are done only when isCompleteQuery is true */
-				needCleanup = isCompleteQuery && EventTriggerBeginCompleteQuery();
+					if (stmt->objtype != OBJECT_PROCEDURE)
+						break;
 
-				/* PG_TRY block is to ensure we call EventTriggerEndCompleteQuery */
-				PG_TRY();
-				{
-					StartTransactionCommand();
-					if (isCompleteQuery)
-						EventTriggerDDLCommandStart(parsetree);
+					/* All event trigger calls are done only when isCompleteQuery is true */
+					needCleanup = isCompleteQuery && EventTriggerBeginCompleteQuery();
 
-					foreach (option, stmt->actions)
+					/* PG_TRY block is to ensure we call EventTriggerEndCompleteQuery */
+					PG_TRY();
 					{
-						DefElem *defel = (DefElem *) lfirst(option);
-						if (strcmp(defel->defname, "location") == 0)
+						StartTransactionCommand();
+						if (isCompleteQuery)
+							EventTriggerDDLCommandStart(parsetree);
+
+						foreach (option, stmt->actions)
 						{
-							/*
-							* location is an implicit option in tsql dialect,
-							* we use this mechanism to store location of function
-							* name so that we can extract original input function
-							* name from queryString.
+							DefElem *defel = (DefElem *) lfirst(option);
+							if (strcmp(defel->defname, "location") == 0)
+							{
+							       /*
+								* location is an implicit option in tsql dialect,
+								* we use this mechanism to store location of function
+								* name so that we can extract original input function
+								* name from queryString.
+								*/
+								origname_location = intVal((Node *) defel->arg);
+								location_cell = option;
+								pfree(defel);
+							}
+						}
+
+						/* delete location cell if it exists as it is for internal use only */
+						if (location_cell)
+							stmt->actions = list_delete_cell(stmt->actions, location_cell);
+
+						/* make a CreateFunctionStmt to pass into CreateFunction() */
+						cfs = makeNode(CreateFunctionStmt);
+						cfs->is_procedure = true;
+						cfs->replace = true;
+						cfs->funcname = stmt->func->objname;
+						cfs->parameters = stmt->func->objfuncargs;
+						cfs->returnType = NULL;
+						cfs->options = stmt->actions;
+
+						pltsql_proc_get_oid_proname_proacl(stmt, pstate, &oldoid, &proacl, &isSameProc);
+						if (!isSameProc) /* i.e. different signature */
+							RemoveFunctionById(oldoid);
+						address = CreateFunction(pstate, cfs); /* if this is the same proc, will just update the existing one */
+						pg_proc_update_oid_acl(address, oldoid, proacl);
+						/* Update function/procedure related metadata in babelfish catalog */
+						pltsql_store_func_default_positions(address, cfs->parameters, queryString, origname_location);
+						if (!isSameProc) {
+						       /*
+							* When the signatures differ we need to manually update the 'function_args' column in 
+							* the 'bbf_schema_permissions' catalog
 							*/
-							origname_location = intVal((Node *) defel->arg);
-							location_cell = option;
-							pfree(defel);
+							alter_bbf_schema_permissions_catalog(stmt->func, cfs->parameters, OBJECT_PROCEDURE, oldoid);
 						}
+						/* Clean up table entries for the create function statement */
+						deleteDependencyRecordsFor(DefaultAclRelationId, address.objectId, false);
+						deleteDependencyRecordsFor(ProcedureRelationId, address.objectId, false);
+						deleteSharedDependencyRecordsFor(ProcedureRelationId, address.objectId, 0);
+						CommitTransactionCommand();
 					}
-
-					/* delete location cell if it exists as it is for internal use only */
-					if (location_cell)
-						stmt->actions = list_delete_cell(stmt->actions, location_cell);
-
-					/* make a CreateFunctionStmt to pass into CreateFunction() */
-					cfs = makeNode(CreateFunctionStmt);
-					cfs->is_procedure = true;
-					cfs->replace = true;
-					cfs->funcname = stmt->func->objname;
-					cfs->parameters = stmt->func->objfuncargs;
-					cfs->returnType = NULL;
-					cfs->options = stmt->actions;
-
-					pltsql_proc_get_oid_proname_proacl(stmt, pstate, &oldoid, &proacl, &isSameProc);
-					if (!isSameProc) /* i.e. different signature */
-						RemoveFunctionById(oldoid);
-					address = CreateFunction(pstate, cfs); /* if this is the same proc, will just update the existing one */
-					pg_proc_update_oid_acl(address, oldoid, proacl);
-					/* Update function/procedure related metadata in babelfish catalog */
-					pltsql_store_func_default_positions(address, cfs->parameters, queryString, origname_location);
-					if (!isSameProc) {
-						/*
-						 * When the signatures differ we need to manually update the 'function_args' column in 
-						 * the 'bbf_schema_permissions' catalog
-						 */
-						alter_bbf_schema_permissions_catalog(stmt->func, cfs->parameters, OBJECT_PROCEDURE, oldoid);
+					PG_FINALLY();
+					{
+						if (needCleanup)
+							EventTriggerEndCompleteQuery();
 					}
-					/* Clean up table entries for the create function statement */
-					deleteDependencyRecordsFor(DefaultAclRelationId, address.objectId, false);
-					deleteDependencyRecordsFor(ProcedureRelationId, address.objectId, false);
-					deleteSharedDependencyRecordsFor(ProcedureRelationId, address.objectId, 0);
-					CommitTransactionCommand();
-				}
-				PG_FINALLY();
-				{
-					if (needCleanup)
-						EventTriggerEndCompleteQuery();
+					PG_END_TRY();
+					return;
 				}
-				PG_END_TRY();
-				return;
+				break;
 			}
 		case T_AlterTableStmt:
 			{

diff --git a/test/JDBC/expected/alter-procedure-schema.out b/test/JDBC/expected/alter-procedure-schema.out
@@ -425,3 +425,41 @@ void
 -- tsql 
 drop login alter_proc_l4;
 go
+
+-- psql currentSchema=master_dbo,public
+-- Test psql procedures altered with security definer do not throw StartTransactionCommand: unexpected state STARTED error 
+-- Test alter procedure using tsql dialect in PSQL port throws error
+CREATE PROCEDURE master_dbo.p1() AS $$ BEGIN SELECT 1; END $$ LANGUAGE plpgsql;
+go
+
+alter procedure master_dbo.p1 security definer;
+go
+
+drop procedure master_dbo.p1;
+go
+
+set babelfishpg_tsql.sql_dialect = "tsql";
+GO
+
+CREATE PROCEDURE tsqlp1 as select 1
+go
+
+ALTER PROCEDURE tsqlp1 as select 2
+go
+~~ERROR (Code: 0)~~
+
+~~ERROR (Message: ERROR: TSQL ALTER PROCEDURE is not supported from PostgreSQL endpoint.
+    Server SQLState: 0A000)~~
+
+
+drop procedure tsqlp1
+go
+
+-- Set dialect back to postgres
+select set_config('babelfishpg_tsql.sql_dialect', 'postgres', null);
+GO
+~~START~~
+text
+postgres
+~~END~~
+
diff --git a/test/JDBC/input/alter/alter-procedure-schema.mix b/test/JDBC/input/alter/alter-procedure-schema.mix
@@ -306,4 +306,32 @@ go
 
 -- tsql 
 drop login alter_proc_l4;
-go
+go
+
+-- psql currentSchema=master_dbo,public
+-- Test psql procedures altered with security definer do not throw StartTransactionCommand: unexpected state STARTED error 
+-- Test alter procedure using tsql dialect in PSQL port throws error
+CREATE PROCEDURE master_dbo.p1() AS $$ BEGIN SELECT 1; END $$ LANGUAGE plpgsql;
+go
+
+alter procedure master_dbo.p1 security definer;
+go
+
+drop procedure master_dbo.p1;
+go
+
+set babelfishpg_tsql.sql_dialect = "tsql";
+GO
+
+CREATE PROCEDURE tsqlp1 as select 1
+go
+
+ALTER PROCEDURE tsqlp1 as select 2
+go
+
+drop procedure tsqlp1
+go
+
+-- Set dialect back to postgres
+select set_config('babelfishpg_tsql.sql_dialect', 'postgres', null);
+GO