Skip to content

Commit

Permalink
ALTER PROCEDURE throwing unexpected state error (#2697)
Browse files Browse the repository at this point in the history 
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 <owjco@amazon.com>
  • Loading branch information
@Jakeowen1
Jakeowen1 committed Jul 2, 2024
1 parent c569e7e commit fff46f2
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 89 deletions.
187 changes: 99 additions & 88 deletions contrib/babelfishpg_tsql/src/pl_handler.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2379,104 +2379,115 @@ 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;
bool with_recompile = false;

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;
bool with_recompile = false;

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.
*/
origname_location = intVal((Node *) defel->arg);
location_cell = option;
pfree(defel);
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);
}
else if (strcmp(defel->defname, "recompile") == 0)
{
/*
* ALTER PROCEDURE ... WITH RECOMPILE
* Record RECOMPILE in catalog
*/
with_recompile = true;
}
}
else if (strcmp(defel->defname, "recompile") == 0)
{
/*
* ALTER PROCEDURE ... WITH RECOMPILE
* Record RECOMPILE in catalog
*/
with_recompile = true;

/* 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, with_recompile);
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);
}
/* 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, with_recompile);
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:
{
Expand Down
38 changes: 38 additions & 0 deletions test/JDBC/expected/alter-procedure-schema.out
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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~~

30 changes: 29 additions & 1 deletion test/JDBC/input/alter/alter-procedure-schema.mix
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit fff46f2

Please sign in to comment.